mirror of
https://github.com/google-gemini/gemini-cli.git
synced 2026-07-09 01:27:41 -07:00
chore(a2a-server): refactor a2a-server src directory (#7593)
This commit is contained in:
Sam McCauley
@@ -0,0 +1,596 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2025 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import type { Message, Task as SDKTask } from '@a2a-js/sdk';
|
||||
import type {
|
||||
TaskStore,
|
||||
AgentExecutor,
|
||||
AgentExecutionEvent,
|
||||
RequestContext,
|
||||
ExecutionEventBus,
|
||||
} from '@a2a-js/sdk/server';
|
||||
import type {
|
||||
ToolCallRequestInfo,
|
||||
ServerGeminiToolCallRequestEvent,
|
||||
Config,
|
||||
} from '@google/gemini-cli-core';
|
||||
import { GeminiEventType } from '@google/gemini-cli-core';
|
||||
import { v4 as uuidv4 } from 'uuid';
|
||||
import { logger } from '../utils/logger.js';
|
||||
import type { StateChange, AgentSettings } from '../types.js';
|
||||
import { CoderAgentEvent } from '../types.js';
|
||||
import { loadConfig, loadEnvironment, setTargetDir } from '../config/config.js';
|
||||
import { loadSettings } from '../config/settings.js';
|
||||
import { loadExtensions } from '../config/extension.js';
|
||||
import { Task } from './task.js';
|
||||
import type { PersistedStateMetadata } from '../metadata_types.js';
|
||||
import { getPersistedState, setPersistedState } from '../metadata_types.js';
|
||||
import { requestStorage } from '../http/requestStorage.js';
|
||||
|
||||
/**
|
||||
* Provides a wrapper for Task. Passes data from Task to SDKTask.
|
||||
* The idea is to use this class inside CoderAgentExecutor to replace Task.
|
||||
*/
|
||||
class TaskWrapper {
|
||||
task: Task;
|
||||
agentSettings: AgentSettings;
|
||||
|
||||
constructor(task: Task, agentSettings: AgentSettings) {
|
||||
this.task = task;
|
||||
this.agentSettings = agentSettings;
|
||||
}
|
||||
|
||||
get id() {
|
||||
return this.task.id;
|
||||
}
|
||||
|
||||
toSDKTask(): SDKTask {
|
||||
const persistedState: PersistedStateMetadata = {
|
||||
_agentSettings: this.agentSettings,
|
||||
_taskState: this.task.taskState,
|
||||
};
|
||||
|
||||
const sdkTask: SDKTask = {
|
||||
id: this.task.id,
|
||||
contextId: this.task.contextId,
|
||||
kind: 'task',
|
||||
status: {
|
||||
state: this.task.taskState,
|
||||
timestamp: new Date().toISOString(),
|
||||
},
|
||||
metadata: setPersistedState({}, persistedState),
|
||||
history: [],
|
||||
artifacts: [],
|
||||
};
|
||||
sdkTask.metadata!['_contextId'] = this.task.contextId;
|
||||
return sdkTask;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* CoderAgentExecutor implements the agent's core logic for code generation.
|
||||
*/
|
||||
export class CoderAgentExecutor implements AgentExecutor {
|
||||
private tasks: Map<string, TaskWrapper> = new Map();
|
||||
// Track tasks with an active execution loop.
|
||||
private executingTasks = new Set<string>();
|
||||
|
||||
constructor(private taskStore?: TaskStore) {}
|
||||
|
||||
private async getConfig(
|
||||
agentSettings: AgentSettings,
|
||||
taskId: string,
|
||||
): Promise<Config> {
|
||||
const workspaceRoot = setTargetDir(agentSettings);
|
||||
loadEnvironment(); // Will override any global env with workspace envs
|
||||
const settings = loadSettings(workspaceRoot);
|
||||
const extensions = loadExtensions(workspaceRoot);
|
||||
return await loadConfig(settings, extensions, taskId);
|
||||
}
|
||||
|
||||
/**
|
||||
* Reconstructs TaskWrapper from SDKTask.
|
||||
*/
|
||||
async reconstruct(
|
||||
sdkTask: SDKTask,
|
||||
eventBus?: ExecutionEventBus,
|
||||
): Promise<TaskWrapper> {
|
||||
const metadata = sdkTask.metadata || {};
|
||||
const persistedState = getPersistedState(metadata);
|
||||
|
||||
if (!persistedState) {
|
||||
throw new Error(
|
||||
`Cannot reconstruct task ${sdkTask.id}: missing persisted state in metadata.`,
|
||||
);
|
||||
}
|
||||
|
||||
const agentSettings = persistedState._agentSettings;
|
||||
const config = await this.getConfig(agentSettings, sdkTask.id);
|
||||
const contextId =
|
||||
(metadata['_contextId'] as string) || (sdkTask.contextId as string);
|
||||
const runtimeTask = await Task.create(
|
||||
sdkTask.id,
|
||||
contextId,
|
||||
config,
|
||||
eventBus,
|
||||
);
|
||||
runtimeTask.taskState = persistedState._taskState;
|
||||
await runtimeTask.geminiClient.initialize(
|
||||
runtimeTask.config.getContentGeneratorConfig(),
|
||||
);
|
||||
|
||||
const wrapper = new TaskWrapper(runtimeTask, agentSettings);
|
||||
this.tasks.set(sdkTask.id, wrapper);
|
||||
logger.info(`Task ${sdkTask.id} reconstructed from store.`);
|
||||
return wrapper;
|
||||
}
|
||||
|
||||
async createTask(
|
||||
taskId: string,
|
||||
contextId: string,
|
||||
agentSettingsInput?: AgentSettings,
|
||||
eventBus?: ExecutionEventBus,
|
||||
): Promise<TaskWrapper> {
|
||||
const agentSettings = agentSettingsInput || ({} as AgentSettings);
|
||||
const config = await this.getConfig(agentSettings, taskId);
|
||||
const runtimeTask = await Task.create(taskId, contextId, config, eventBus);
|
||||
await runtimeTask.geminiClient.initialize(
|
||||
runtimeTask.config.getContentGeneratorConfig(),
|
||||
);
|
||||
|
||||
const wrapper = new TaskWrapper(runtimeTask, agentSettings);
|
||||
this.tasks.set(taskId, wrapper);
|
||||
logger.info(`New task ${taskId} created.`);
|
||||
return wrapper;
|
||||
}
|
||||
|
||||
getTask(taskId: string): TaskWrapper | undefined {
|
||||
return this.tasks.get(taskId);
|
||||
}
|
||||
|
||||
getAllTasks(): TaskWrapper[] {
|
||||
return Array.from(this.tasks.values());
|
||||
}
|
||||
|
||||
cancelTask = async (
|
||||
taskId: string,
|
||||
eventBus: ExecutionEventBus,
|
||||
): Promise<void> => {
|
||||
logger.info(
|
||||
`[CoderAgentExecutor] Received cancel request for task ${taskId}`,
|
||||
);
|
||||
const wrapper = this.tasks.get(taskId);
|
||||
|
||||
if (!wrapper) {
|
||||
logger.warn(
|
||||
`[CoderAgentExecutor] Task ${taskId} not found for cancellation.`,
|
||||
);
|
||||
eventBus.publish({
|
||||
kind: 'status-update',
|
||||
taskId,
|
||||
contextId: uuidv4(),
|
||||
status: {
|
||||
state: 'failed',
|
||||
message: {
|
||||
kind: 'message',
|
||||
role: 'agent',
|
||||
parts: [{ kind: 'text', text: `Task ${taskId} not found.` }],
|
||||
messageId: uuidv4(),
|
||||
taskId,
|
||||
},
|
||||
},
|
||||
final: true,
|
||||
});
|
||||
return;
|
||||
}
|
||||
|
||||
const { task } = wrapper;
|
||||
|
||||
if (task.taskState === 'canceled' || task.taskState === 'failed') {
|
||||
logger.info(
|
||||
`[CoderAgentExecutor] Task ${taskId} is already in a final state: ${task.taskState}. No action needed for cancellation.`,
|
||||
);
|
||||
eventBus.publish({
|
||||
kind: 'status-update',
|
||||
taskId,
|
||||
contextId: task.contextId,
|
||||
status: {
|
||||
state: task.taskState,
|
||||
message: {
|
||||
kind: 'message',
|
||||
role: 'agent',
|
||||
parts: [
|
||||
{
|
||||
kind: 'text',
|
||||
text: `Task ${taskId} is already ${task.taskState}.`,
|
||||
},
|
||||
],
|
||||
messageId: uuidv4(),
|
||||
taskId,
|
||||
},
|
||||
},
|
||||
final: true,
|
||||
});
|
||||
return;
|
||||
}
|
||||
|
||||
try {
|
||||
logger.info(
|
||||
`[CoderAgentExecutor] Initiating cancellation for task ${taskId}.`,
|
||||
);
|
||||
task.cancelPendingTools('Task canceled by user request.');
|
||||
|
||||
const stateChange: StateChange = {
|
||||
kind: CoderAgentEvent.StateChangeEvent,
|
||||
};
|
||||
task.setTaskStateAndPublishUpdate(
|
||||
'canceled',
|
||||
stateChange,
|
||||
'Task canceled by user request.',
|
||||
undefined,
|
||||
true,
|
||||
);
|
||||
logger.info(
|
||||
`[CoderAgentExecutor] Task ${taskId} cancellation processed. Saving state.`,
|
||||
);
|
||||
await this.taskStore?.save(wrapper.toSDKTask());
|
||||
logger.info(`[CoderAgentExecutor] Task ${taskId} state CANCELED saved.`);
|
||||
} catch (error) {
|
||||
const errorMessage =
|
||||
error instanceof Error ? error.message : 'Unknown error';
|
||||
logger.error(
|
||||
`[CoderAgentExecutor] Error during task cancellation for ${taskId}: ${errorMessage}`,
|
||||
error,
|
||||
);
|
||||
eventBus.publish({
|
||||
kind: 'status-update',
|
||||
taskId,
|
||||
contextId: task.contextId,
|
||||
status: {
|
||||
state: 'failed',
|
||||
message: {
|
||||
kind: 'message',
|
||||
role: 'agent',
|
||||
parts: [
|
||||
{
|
||||
kind: 'text',
|
||||
text: `Failed to process cancellation for task ${taskId}: ${errorMessage}`,
|
||||
},
|
||||
],
|
||||
messageId: uuidv4(),
|
||||
taskId,
|
||||
},
|
||||
},
|
||||
final: true,
|
||||
});
|
||||
}
|
||||
};
|
||||
|
||||
async execute(
|
||||
requestContext: RequestContext,
|
||||
eventBus: ExecutionEventBus,
|
||||
): Promise<void> {
|
||||
const userMessage = requestContext.userMessage as Message;
|
||||
const sdkTask = requestContext.task as SDKTask | undefined;
|
||||
|
||||
const taskId = sdkTask?.id || userMessage.taskId || uuidv4();
|
||||
const contextId =
|
||||
userMessage.contextId ||
|
||||
sdkTask?.contextId ||
|
||||
sdkTask?.metadata?.['_contextId'] ||
|
||||
uuidv4();
|
||||
|
||||
logger.info(
|
||||
`[CoderAgentExecutor] Executing for taskId: ${taskId}, contextId: ${contextId}`,
|
||||
);
|
||||
logger.info(
|
||||
`[CoderAgentExecutor] userMessage: ${JSON.stringify(userMessage)}`,
|
||||
);
|
||||
eventBus.on('event', (event: AgentExecutionEvent) =>
|
||||
logger.info('[EventBus event]: ', event),
|
||||
);
|
||||
|
||||
const store = requestStorage.getStore();
|
||||
if (!store) {
|
||||
logger.error(
|
||||
'[CoderAgentExecutor] Could not get request from async local storage. Cancellation on socket close will not be handled for this request.',
|
||||
);
|
||||
}
|
||||
|
||||
const abortController = new AbortController();
|
||||
const abortSignal = abortController.signal;
|
||||
|
||||
if (store) {
|
||||
// Grab the raw socket from the request object
|
||||
const socket = store.req.socket;
|
||||
const onClientEnd = () => {
|
||||
logger.info(
|
||||
`[CoderAgentExecutor] Client socket closed for task ${taskId}. Cancelling execution.`,
|
||||
);
|
||||
if (!abortController.signal.aborted) {
|
||||
abortController.abort();
|
||||
}
|
||||
// Clean up the listener to prevent memory leaks
|
||||
socket.removeListener('close', onClientEnd);
|
||||
};
|
||||
|
||||
// Listen on the socket's 'end' event (remote closed the connection)
|
||||
socket.on('end', onClientEnd);
|
||||
|
||||
// It's also good practice to remove the listener if the task completes successfully
|
||||
abortSignal.addEventListener('abort', () => {
|
||||
socket.removeListener('end', onClientEnd);
|
||||
});
|
||||
logger.info(
|
||||
`[CoderAgentExecutor] Socket close handler set up for task ${taskId}.`,
|
||||
);
|
||||
}
|
||||
|
||||
let wrapper: TaskWrapper | undefined = this.tasks.get(taskId);
|
||||
|
||||
if (wrapper) {
|
||||
wrapper.task.eventBus = eventBus;
|
||||
logger.info(`[CoderAgentExecutor] Task ${taskId} found in memory cache.`);
|
||||
} else if (sdkTask) {
|
||||
logger.info(
|
||||
`[CoderAgentExecutor] Task ${taskId} found in TaskStore. Reconstructing...`,
|
||||
);
|
||||
try {
|
||||
wrapper = await this.reconstruct(sdkTask, eventBus);
|
||||
} catch (e) {
|
||||
logger.error(
|
||||
`[CoderAgentExecutor] Failed to hydrate task ${taskId}:`,
|
||||
e,
|
||||
);
|
||||
const stateChange: StateChange = {
|
||||
kind: CoderAgentEvent.StateChangeEvent,
|
||||
};
|
||||
eventBus.publish({
|
||||
kind: 'status-update',
|
||||
taskId,
|
||||
contextId: sdkTask.contextId,
|
||||
status: {
|
||||
state: 'failed',
|
||||
message: {
|
||||
kind: 'message',
|
||||
role: 'agent',
|
||||
parts: [
|
||||
{
|
||||
kind: 'text',
|
||||
text: 'Internal error: Task state lost or corrupted.',
|
||||
},
|
||||
],
|
||||
messageId: uuidv4(),
|
||||
taskId,
|
||||
contextId: sdkTask.contextId,
|
||||
} as Message,
|
||||
},
|
||||
final: true,
|
||||
metadata: { coderAgent: stateChange },
|
||||
});
|
||||
return;
|
||||
}
|
||||
} else {
|
||||
logger.info(`[CoderAgentExecutor] Creating new task ${taskId}.`);
|
||||
const agentSettings = userMessage.metadata?.[
|
||||
'coderAgent'
|
||||
] as AgentSettings;
|
||||
wrapper = await this.createTask(
|
||||
taskId,
|
||||
contextId as string,
|
||||
agentSettings,
|
||||
eventBus,
|
||||
);
|
||||
const newTaskSDK = wrapper.toSDKTask();
|
||||
eventBus.publish({
|
||||
...newTaskSDK,
|
||||
kind: 'task',
|
||||
status: { state: 'submitted', timestamp: new Date().toISOString() },
|
||||
history: [userMessage],
|
||||
});
|
||||
try {
|
||||
await this.taskStore?.save(newTaskSDK);
|
||||
logger.info(`[CoderAgentExecutor] New task ${taskId} saved to store.`);
|
||||
} catch (saveError) {
|
||||
logger.error(
|
||||
`[CoderAgentExecutor] Failed to save new task ${taskId} to store:`,
|
||||
saveError,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
if (!wrapper) {
|
||||
logger.error(
|
||||
`[CoderAgentExecutor] Task ${taskId} is unexpectedly undefined after load/create.`,
|
||||
);
|
||||
return;
|
||||
}
|
||||
|
||||
const currentTask = wrapper.task;
|
||||
|
||||
if (['canceled', 'failed', 'completed'].includes(currentTask.taskState)) {
|
||||
logger.warn(
|
||||
`[CoderAgentExecutor] Attempted to execute task ${taskId} which is already in state ${currentTask.taskState}. Ignoring.`,
|
||||
);
|
||||
return;
|
||||
}
|
||||
|
||||
if (this.executingTasks.has(taskId)) {
|
||||
logger.info(
|
||||
`[CoderAgentExecutor] Task ${taskId} has a pending execution. Processing message and yielding.`,
|
||||
);
|
||||
currentTask.eventBus = eventBus;
|
||||
for await (const _ of currentTask.acceptUserMessage(
|
||||
requestContext,
|
||||
abortController.signal,
|
||||
)) {
|
||||
logger.info(
|
||||
`[CoderAgentExecutor] Processing user message ${userMessage.messageId} in secondary execution loop for task ${taskId}.`,
|
||||
);
|
||||
}
|
||||
// End this execution-- the original/source will be resumed.
|
||||
return;
|
||||
}
|
||||
|
||||
logger.info(
|
||||
`[CoderAgentExecutor] Starting main execution for message ${userMessage.messageId} for task ${taskId}.`,
|
||||
);
|
||||
this.executingTasks.add(taskId);
|
||||
|
||||
try {
|
||||
let agentTurnActive = true;
|
||||
logger.info(`[CoderAgentExecutor] Task ${taskId}: Processing user turn.`);
|
||||
let agentEvents = currentTask.acceptUserMessage(
|
||||
requestContext,
|
||||
abortSignal,
|
||||
);
|
||||
|
||||
while (agentTurnActive) {
|
||||
logger.info(
|
||||
`[CoderAgentExecutor] Task ${taskId}: Processing agent turn (LLM stream).`,
|
||||
);
|
||||
const toolCallRequests: ToolCallRequestInfo[] = [];
|
||||
for await (const event of agentEvents) {
|
||||
if (abortSignal.aborted) {
|
||||
logger.warn(
|
||||
`[CoderAgentExecutor] Task ${taskId}: Abort signal received during agent event processing.`,
|
||||
);
|
||||
throw new Error('Execution aborted');
|
||||
}
|
||||
if (event.type === GeminiEventType.ToolCallRequest) {
|
||||
toolCallRequests.push(
|
||||
(event as ServerGeminiToolCallRequestEvent).value,
|
||||
);
|
||||
continue;
|
||||
}
|
||||
await currentTask.acceptAgentMessage(event);
|
||||
}
|
||||
|
||||
if (abortSignal.aborted) throw new Error('Execution aborted');
|
||||
|
||||
if (toolCallRequests.length > 0) {
|
||||
logger.info(
|
||||
`[CoderAgentExecutor] Task ${taskId}: Found ${toolCallRequests.length} tool call requests. Scheduling as a batch.`,
|
||||
);
|
||||
await currentTask.scheduleToolCalls(toolCallRequests, abortSignal);
|
||||
}
|
||||
|
||||
logger.info(
|
||||
`[CoderAgentExecutor] Task ${taskId}: Waiting for pending tools if any.`,
|
||||
);
|
||||
await currentTask.waitForPendingTools();
|
||||
logger.info(
|
||||
`[CoderAgentExecutor] Task ${taskId}: All pending tools completed or none were pending.`,
|
||||
);
|
||||
|
||||
if (abortSignal.aborted) throw new Error('Execution aborted');
|
||||
|
||||
const completedTools = currentTask.getAndClearCompletedTools();
|
||||
|
||||
if (completedTools.length > 0) {
|
||||
// If all completed tool calls were canceled, manually add them to history and set state to input-required, final:true
|
||||
if (completedTools.every((tool) => tool.status === 'cancelled')) {
|
||||
logger.info(
|
||||
`[CoderAgentExecutor] Task ${taskId}: All tool calls were cancelled. Updating history and ending agent turn.`,
|
||||
);
|
||||
currentTask.addToolResponsesToHistory(completedTools);
|
||||
agentTurnActive = false;
|
||||
const stateChange: StateChange = {
|
||||
kind: CoderAgentEvent.StateChangeEvent,
|
||||
};
|
||||
currentTask.setTaskStateAndPublishUpdate(
|
||||
'input-required',
|
||||
stateChange,
|
||||
undefined,
|
||||
undefined,
|
||||
true,
|
||||
);
|
||||
} else {
|
||||
logger.info(
|
||||
`[CoderAgentExecutor] Task ${taskId}: Found ${completedTools.length} completed tool calls. Sending results back to LLM.`,
|
||||
);
|
||||
|
||||
agentEvents = currentTask.sendCompletedToolsToLlm(
|
||||
completedTools,
|
||||
abortSignal,
|
||||
);
|
||||
// Continue the loop to process the LLM response to the tool results.
|
||||
}
|
||||
} else {
|
||||
logger.info(
|
||||
`[CoderAgentExecutor] Task ${taskId}: No more tool calls to process. Ending agent turn.`,
|
||||
);
|
||||
agentTurnActive = false;
|
||||
}
|
||||
}
|
||||
|
||||
logger.info(
|
||||
`[CoderAgentExecutor] Task ${taskId}: Agent turn finished, setting to input-required.`,
|
||||
);
|
||||
const stateChange: StateChange = {
|
||||
kind: CoderAgentEvent.StateChangeEvent,
|
||||
};
|
||||
currentTask.setTaskStateAndPublishUpdate(
|
||||
'input-required',
|
||||
stateChange,
|
||||
undefined,
|
||||
undefined,
|
||||
true,
|
||||
);
|
||||
} catch (error) {
|
||||
if (abortSignal.aborted) {
|
||||
logger.warn(`[CoderAgentExecutor] Task ${taskId} execution aborted.`);
|
||||
currentTask.cancelPendingTools('Execution aborted');
|
||||
if (
|
||||
currentTask.taskState !== 'canceled' &&
|
||||
currentTask.taskState !== 'failed'
|
||||
) {
|
||||
currentTask.setTaskStateAndPublishUpdate(
|
||||
'input-required',
|
||||
{ kind: CoderAgentEvent.StateChangeEvent },
|
||||
'Execution aborted by client.',
|
||||
undefined,
|
||||
true,
|
||||
);
|
||||
}
|
||||
} else {
|
||||
const errorMessage =
|
||||
error instanceof Error ? error.message : 'Agent execution error';
|
||||
logger.error(
|
||||
`[CoderAgentExecutor] Error executing agent for task ${taskId}:`,
|
||||
error,
|
||||
);
|
||||
currentTask.cancelPendingTools(errorMessage);
|
||||
if (currentTask.taskState !== 'failed') {
|
||||
const stateChange: StateChange = {
|
||||
kind: CoderAgentEvent.StateChangeEvent,
|
||||
};
|
||||
currentTask.setTaskStateAndPublishUpdate(
|
||||
'failed',
|
||||
stateChange,
|
||||
errorMessage,
|
||||
undefined,
|
||||
true,
|
||||
);
|
||||
}
|
||||
}
|
||||
} finally {
|
||||
this.executingTasks.delete(taskId);
|
||||
logger.info(
|
||||
`[CoderAgentExecutor] Saving final state for task ${taskId}.`,
|
||||
);
|
||||
try {
|
||||
await this.taskStore?.save(wrapper.toSDKTask());
|
||||
logger.info(`[CoderAgentExecutor] Task ${taskId} state saved.`);
|
||||
} catch (saveError) {
|
||||
logger.error(
|
||||
`[CoderAgentExecutor] Failed to save task ${taskId} state in finally block:`,
|
||||
saveError,
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,59 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2025 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { describe, it, expect, vi } from 'vitest';
|
||||
import { Task } from './task.js';
|
||||
import type { Config, ToolCallRequestInfo } from '@google/gemini-cli-core';
|
||||
import { createMockConfig } from '../utils/testing_utils.js';
|
||||
import type { ExecutionEventBus } from '@a2a-js/sdk/server';
|
||||
|
||||
describe('Task', () => {
|
||||
it('scheduleToolCalls should not modify the input requests array', async () => {
|
||||
const mockConfig = createMockConfig();
|
||||
|
||||
const mockEventBus: ExecutionEventBus = {
|
||||
publish: vi.fn(),
|
||||
on: vi.fn(),
|
||||
off: vi.fn(),
|
||||
once: vi.fn(),
|
||||
removeAllListeners: vi.fn(),
|
||||
finished: vi.fn(),
|
||||
};
|
||||
|
||||
// The Task constructor is private. We'll bypass it for this unit test.
|
||||
// @ts-expect-error - Calling private constructor for test purposes.
|
||||
const task = new Task(
|
||||
'task-id',
|
||||
'context-id',
|
||||
mockConfig as Config,
|
||||
mockEventBus,
|
||||
);
|
||||
|
||||
task['setTaskStateAndPublishUpdate'] = vi.fn();
|
||||
task['getProposedContent'] = vi.fn().mockResolvedValue('new content');
|
||||
|
||||
const requests: ToolCallRequestInfo[] = [
|
||||
{
|
||||
callId: '1',
|
||||
name: 'replace',
|
||||
args: {
|
||||
file_path: 'test.txt',
|
||||
old_string: 'old',
|
||||
new_string: 'new',
|
||||
},
|
||||
isClientInitiated: false,
|
||||
prompt_id: 'prompt-id-1',
|
||||
},
|
||||
];
|
||||
|
||||
const originalRequests = JSON.parse(JSON.stringify(requests));
|
||||
const abortController = new AbortController();
|
||||
|
||||
await task.scheduleToolCalls(requests, abortController.signal);
|
||||
|
||||
expect(requests).toEqual(originalRequests);
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,936 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2025 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import {
|
||||
CoreToolScheduler,
|
||||
GeminiClient,
|
||||
GeminiEventType,
|
||||
ToolConfirmationOutcome,
|
||||
ApprovalMode,
|
||||
getAllMCPServerStatuses,
|
||||
MCPServerStatus,
|
||||
isNodeError,
|
||||
parseAndFormatApiError,
|
||||
} from '@google/gemini-cli-core';
|
||||
import type {
|
||||
ToolConfirmationPayload,
|
||||
CompletedToolCall,
|
||||
ToolCall,
|
||||
ToolCallRequestInfo,
|
||||
ServerGeminiErrorEvent,
|
||||
ServerGeminiStreamEvent,
|
||||
ToolCallConfirmationDetails,
|
||||
Config,
|
||||
UserTierId,
|
||||
} from '@google/gemini-cli-core';
|
||||
import type { RequestContext } from '@a2a-js/sdk/server';
|
||||
import { type ExecutionEventBus } from '@a2a-js/sdk/server';
|
||||
import type {
|
||||
TaskStatusUpdateEvent,
|
||||
TaskArtifactUpdateEvent,
|
||||
TaskState,
|
||||
Message,
|
||||
Part,
|
||||
Artifact,
|
||||
} from '@a2a-js/sdk';
|
||||
import { v4 as uuidv4 } from 'uuid';
|
||||
import { logger } from '../utils/logger.js';
|
||||
import * as fs from 'node:fs';
|
||||
|
||||
import { CoderAgentEvent } from '../types.js';
|
||||
import type {
|
||||
CoderAgentMessage,
|
||||
StateChange,
|
||||
ToolCallUpdate,
|
||||
TextContent,
|
||||
TaskMetadata,
|
||||
Thought,
|
||||
ThoughtSummary,
|
||||
} from '../types.js';
|
||||
import type { PartUnion, Part as genAiPart } from '@google/genai';
|
||||
|
||||
export class Task {
|
||||
id: string;
|
||||
contextId: string;
|
||||
scheduler: CoreToolScheduler;
|
||||
config: Config;
|
||||
geminiClient: GeminiClient;
|
||||
pendingToolConfirmationDetails: Map<string, ToolCallConfirmationDetails>;
|
||||
taskState: TaskState;
|
||||
eventBus?: ExecutionEventBus;
|
||||
completedToolCalls: CompletedToolCall[];
|
||||
skipFinalTrueAfterInlineEdit = false;
|
||||
|
||||
// For tool waiting logic
|
||||
private pendingToolCalls: Map<string, string> = new Map(); //toolCallId --> status
|
||||
private toolCompletionPromise?: Promise<void>;
|
||||
private toolCompletionNotifier?: {
|
||||
resolve: () => void;
|
||||
reject: (reason?: Error) => void;
|
||||
};
|
||||
|
||||
private constructor(
|
||||
id: string,
|
||||
contextId: string,
|
||||
config: Config,
|
||||
eventBus?: ExecutionEventBus,
|
||||
) {
|
||||
this.id = id;
|
||||
this.contextId = contextId;
|
||||
this.config = config;
|
||||
this.scheduler = this.createScheduler();
|
||||
this.geminiClient = new GeminiClient(this.config);
|
||||
this.pendingToolConfirmationDetails = new Map();
|
||||
this.taskState = 'submitted';
|
||||
this.eventBus = eventBus;
|
||||
this.completedToolCalls = [];
|
||||
this._resetToolCompletionPromise();
|
||||
this.config.setFlashFallbackHandler(
|
||||
async (currentModel: string, fallbackModel: string): Promise<boolean> => {
|
||||
config.setModel(fallbackModel); // gemini-cli-core sets to DEFAULT_GEMINI_FLASH_MODEL
|
||||
// Switch model for future use but return false to stop current retry
|
||||
return false;
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
static async create(
|
||||
id: string,
|
||||
contextId: string,
|
||||
config: Config,
|
||||
eventBus?: ExecutionEventBus,
|
||||
): Promise<Task> {
|
||||
return new Task(id, contextId, config, eventBus);
|
||||
}
|
||||
|
||||
// Note: `getAllMCPServerStatuses` retrieves the status of all MCP servers for the entire
|
||||
// process. This is not scoped to the individual task but reflects the global connection
|
||||
// state managed within the @gemini-cli/core module.
|
||||
async getMetadata(): Promise<TaskMetadata> {
|
||||
const toolRegistry = await this.config.getToolRegistry();
|
||||
const mcpServers = this.config.getMcpServers() || {};
|
||||
const serverStatuses = getAllMCPServerStatuses();
|
||||
const servers = Object.keys(mcpServers).map((serverName) => ({
|
||||
name: serverName,
|
||||
status: serverStatuses.get(serverName) || MCPServerStatus.DISCONNECTED,
|
||||
tools: toolRegistry.getToolsByServer(serverName).map((tool) => ({
|
||||
name: tool.name,
|
||||
description: tool.description,
|
||||
parameterSchema: tool.schema.parameters,
|
||||
})),
|
||||
}));
|
||||
|
||||
const availableTools = toolRegistry.getAllTools().map((tool) => ({
|
||||
name: tool.name,
|
||||
description: tool.description,
|
||||
parameterSchema: tool.schema.parameters,
|
||||
}));
|
||||
|
||||
const metadata: TaskMetadata = {
|
||||
id: this.id,
|
||||
contextId: this.contextId,
|
||||
taskState: this.taskState,
|
||||
model: this.config.getContentGeneratorConfig().model,
|
||||
mcpServers: servers,
|
||||
availableTools,
|
||||
};
|
||||
return metadata;
|
||||
}
|
||||
|
||||
private _resetToolCompletionPromise(): void {
|
||||
this.toolCompletionPromise = new Promise((resolve, reject) => {
|
||||
this.toolCompletionNotifier = { resolve, reject };
|
||||
});
|
||||
// If there are no pending calls when reset, resolve immediately.
|
||||
if (this.pendingToolCalls.size === 0 && this.toolCompletionNotifier) {
|
||||
this.toolCompletionNotifier.resolve();
|
||||
}
|
||||
}
|
||||
|
||||
private _registerToolCall(toolCallId: string, status: string): void {
|
||||
const wasEmpty = this.pendingToolCalls.size === 0;
|
||||
this.pendingToolCalls.set(toolCallId, status);
|
||||
if (wasEmpty) {
|
||||
this._resetToolCompletionPromise();
|
||||
}
|
||||
logger.info(
|
||||
`[Task] Registered tool call: ${toolCallId}. Pending: ${this.pendingToolCalls.size}`,
|
||||
);
|
||||
}
|
||||
|
||||
private _resolveToolCall(toolCallId: string): void {
|
||||
if (this.pendingToolCalls.has(toolCallId)) {
|
||||
this.pendingToolCalls.delete(toolCallId);
|
||||
logger.info(
|
||||
`[Task] Resolved tool call: ${toolCallId}. Pending: ${this.pendingToolCalls.size}`,
|
||||
);
|
||||
if (this.pendingToolCalls.size === 0 && this.toolCompletionNotifier) {
|
||||
this.toolCompletionNotifier.resolve();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async waitForPendingTools(): Promise<void> {
|
||||
if (this.pendingToolCalls.size === 0) {
|
||||
return Promise.resolve();
|
||||
}
|
||||
logger.info(
|
||||
`[Task] Waiting for ${this.pendingToolCalls.size} pending tool(s)...`,
|
||||
);
|
||||
return this.toolCompletionPromise;
|
||||
}
|
||||
|
||||
cancelPendingTools(reason: string): void {
|
||||
if (this.pendingToolCalls.size > 0) {
|
||||
logger.info(
|
||||
`[Task] Cancelling all ${this.pendingToolCalls.size} pending tool calls. Reason: ${reason}`,
|
||||
);
|
||||
}
|
||||
if (this.toolCompletionNotifier) {
|
||||
this.toolCompletionNotifier.reject(new Error(reason));
|
||||
}
|
||||
this.pendingToolCalls.clear();
|
||||
// Reset the promise for any future operations, ensuring it's in a clean state.
|
||||
this._resetToolCompletionPromise();
|
||||
}
|
||||
|
||||
private _createTextMessage(
|
||||
text: string,
|
||||
role: 'agent' | 'user' = 'agent',
|
||||
): Message {
|
||||
return {
|
||||
kind: 'message',
|
||||
role,
|
||||
parts: [{ kind: 'text', text }],
|
||||
messageId: uuidv4(),
|
||||
taskId: this.id,
|
||||
contextId: this.contextId,
|
||||
};
|
||||
}
|
||||
|
||||
private _createStatusUpdateEvent(
|
||||
stateToReport: TaskState,
|
||||
coderAgentMessage: CoderAgentMessage,
|
||||
message?: Message,
|
||||
final = false,
|
||||
timestamp?: string,
|
||||
metadataError?: string,
|
||||
): TaskStatusUpdateEvent {
|
||||
const metadata: {
|
||||
coderAgent: CoderAgentMessage;
|
||||
model: string;
|
||||
userTier?: UserTierId;
|
||||
error?: string;
|
||||
} = {
|
||||
coderAgent: coderAgentMessage,
|
||||
model: this.config.getModel(),
|
||||
userTier: this.geminiClient.getUserTier(),
|
||||
};
|
||||
|
||||
if (metadataError) {
|
||||
metadata.error = metadataError;
|
||||
}
|
||||
|
||||
return {
|
||||
kind: 'status-update',
|
||||
taskId: this.id,
|
||||
contextId: this.contextId,
|
||||
status: {
|
||||
state: stateToReport,
|
||||
message, // Shorthand property
|
||||
timestamp: timestamp || new Date().toISOString(),
|
||||
},
|
||||
final,
|
||||
metadata,
|
||||
};
|
||||
}
|
||||
|
||||
setTaskStateAndPublishUpdate(
|
||||
newState: TaskState,
|
||||
coderAgentMessage: CoderAgentMessage,
|
||||
messageText?: string,
|
||||
messageParts?: Part[], // For more complex messages
|
||||
final = false,
|
||||
metadataError?: string,
|
||||
): void {
|
||||
this.taskState = newState;
|
||||
let message: Message | undefined;
|
||||
|
||||
if (messageText) {
|
||||
message = this._createTextMessage(messageText);
|
||||
} else if (messageParts) {
|
||||
message = {
|
||||
kind: 'message',
|
||||
role: 'agent',
|
||||
parts: messageParts,
|
||||
messageId: uuidv4(),
|
||||
taskId: this.id,
|
||||
contextId: this.contextId,
|
||||
};
|
||||
}
|
||||
|
||||
const event = this._createStatusUpdateEvent(
|
||||
this.taskState,
|
||||
coderAgentMessage,
|
||||
message,
|
||||
final,
|
||||
undefined,
|
||||
metadataError,
|
||||
);
|
||||
this.eventBus?.publish(event);
|
||||
}
|
||||
|
||||
private _schedulerOutputUpdate(
|
||||
toolCallId: string,
|
||||
outputChunk: string,
|
||||
): void {
|
||||
logger.info(
|
||||
'[Task] Scheduler output update for tool call ' +
|
||||
toolCallId +
|
||||
': ' +
|
||||
outputChunk,
|
||||
);
|
||||
const artifact: Artifact = {
|
||||
artifactId: `tool-${toolCallId}-output`,
|
||||
parts: [
|
||||
{
|
||||
kind: 'text',
|
||||
text: outputChunk,
|
||||
} as Part,
|
||||
],
|
||||
};
|
||||
const artifactEvent: TaskArtifactUpdateEvent = {
|
||||
kind: 'artifact-update',
|
||||
taskId: this.id,
|
||||
contextId: this.contextId,
|
||||
artifact,
|
||||
append: true,
|
||||
lastChunk: false,
|
||||
};
|
||||
this.eventBus?.publish(artifactEvent);
|
||||
}
|
||||
|
||||
private async _schedulerAllToolCallsComplete(
|
||||
completedToolCalls: CompletedToolCall[],
|
||||
): Promise<void> {
|
||||
logger.info(
|
||||
'[Task] All tool calls completed by scheduler (batch):',
|
||||
completedToolCalls.map((tc) => tc.request.callId),
|
||||
);
|
||||
this.completedToolCalls.push(...completedToolCalls);
|
||||
completedToolCalls.forEach((tc) => {
|
||||
this._resolveToolCall(tc.request.callId);
|
||||
});
|
||||
}
|
||||
|
||||
private _schedulerToolCallsUpdate(toolCalls: ToolCall[]): void {
|
||||
logger.info(
|
||||
'[Task] Scheduler tool calls updated:',
|
||||
toolCalls.map((tc) => `${tc.request.callId} (${tc.status})`),
|
||||
);
|
||||
|
||||
// Update state and send continuous, non-final updates
|
||||
toolCalls.forEach((tc) => {
|
||||
const previousStatus = this.pendingToolCalls.get(tc.request.callId);
|
||||
const hasChanged = previousStatus !== tc.status;
|
||||
|
||||
// Resolve tool call if it has reached a terminal state
|
||||
if (['success', 'error', 'cancelled'].includes(tc.status)) {
|
||||
this._resolveToolCall(tc.request.callId);
|
||||
} else {
|
||||
// This will update the map
|
||||
this._registerToolCall(tc.request.callId, tc.status);
|
||||
}
|
||||
|
||||
if (tc.status === 'awaiting_approval' && tc.confirmationDetails) {
|
||||
this.pendingToolConfirmationDetails.set(
|
||||
tc.request.callId,
|
||||
tc.confirmationDetails,
|
||||
);
|
||||
}
|
||||
|
||||
// Only send an update if the status has actually changed.
|
||||
if (hasChanged) {
|
||||
const message = this.toolStatusMessage(tc, this.id, this.contextId);
|
||||
const coderAgentMessage: CoderAgentMessage =
|
||||
tc.status === 'awaiting_approval'
|
||||
? { kind: CoderAgentEvent.ToolCallConfirmationEvent }
|
||||
: { kind: CoderAgentEvent.ToolCallUpdateEvent };
|
||||
|
||||
const event = this._createStatusUpdateEvent(
|
||||
this.taskState,
|
||||
coderAgentMessage,
|
||||
message,
|
||||
false, // Always false for these continuous updates
|
||||
);
|
||||
this.eventBus?.publish(event);
|
||||
}
|
||||
});
|
||||
|
||||
if (this.config.getApprovalMode() === ApprovalMode.YOLO) {
|
||||
logger.info('[Task] YOLO mode enabled. Auto-approving all tool calls.');
|
||||
toolCalls.forEach((tc: ToolCall) => {
|
||||
if (tc.status === 'awaiting_approval' && tc.confirmationDetails) {
|
||||
tc.confirmationDetails.onConfirm(ToolConfirmationOutcome.ProceedOnce);
|
||||
this.pendingToolConfirmationDetails.delete(tc.request.callId);
|
||||
}
|
||||
});
|
||||
return;
|
||||
}
|
||||
|
||||
const allPendingStatuses = Array.from(this.pendingToolCalls.values());
|
||||
const isAwaitingApproval = allPendingStatuses.some(
|
||||
(status) => status === 'awaiting_approval',
|
||||
);
|
||||
const allPendingAreStable = allPendingStatuses.every(
|
||||
(status) =>
|
||||
status === 'awaiting_approval' ||
|
||||
status === 'success' ||
|
||||
status === 'error' ||
|
||||
status === 'cancelled',
|
||||
);
|
||||
|
||||
// 1. Are any pending tool calls awaiting_approval
|
||||
// 2. Are all pending tool calls in a stable state (i.e. not in validing or executing)
|
||||
// 3. After an inline edit, the edited tool call will send awaiting_approval THEN scheduled. We wait for the next update in this case.
|
||||
if (
|
||||
isAwaitingApproval &&
|
||||
allPendingAreStable &&
|
||||
!this.skipFinalTrueAfterInlineEdit
|
||||
) {
|
||||
this.skipFinalTrueAfterInlineEdit = false;
|
||||
|
||||
// We don't need to send another message, just a final status update.
|
||||
this.setTaskStateAndPublishUpdate(
|
||||
'input-required',
|
||||
{ kind: CoderAgentEvent.StateChangeEvent },
|
||||
undefined,
|
||||
undefined,
|
||||
/*final*/ true,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
private createScheduler(): CoreToolScheduler {
|
||||
const scheduler = new CoreToolScheduler({
|
||||
outputUpdateHandler: this._schedulerOutputUpdate.bind(this),
|
||||
onAllToolCallsComplete: this._schedulerAllToolCallsComplete.bind(this),
|
||||
onToolCallsUpdate: this._schedulerToolCallsUpdate.bind(this),
|
||||
getPreferredEditor: () => 'vscode',
|
||||
config: this.config,
|
||||
onEditorClose: () => {},
|
||||
});
|
||||
return scheduler;
|
||||
}
|
||||
|
||||
private toolStatusMessage(
|
||||
tc: ToolCall,
|
||||
taskId: string,
|
||||
contextId: string,
|
||||
): Message {
|
||||
const messageParts: Part[] = [];
|
||||
|
||||
// Create a serializable version of the ToolCall (pick necesssary
|
||||
// properties/avoic methods causing circular reference errors)
|
||||
const serializableToolCall: { [key: string]: unknown } = {
|
||||
request: tc.request,
|
||||
status: tc.status,
|
||||
};
|
||||
|
||||
// For WaitingToolCall type
|
||||
if ('confirmationDetails' in tc) {
|
||||
serializableToolCall['confirmationDetails'] = tc.confirmationDetails;
|
||||
}
|
||||
|
||||
if (tc.tool) {
|
||||
serializableToolCall['tool'] = {
|
||||
name: tc.tool.name,
|
||||
displayName: tc.tool.displayName,
|
||||
description: tc.tool.description,
|
||||
kind: tc.tool.kind,
|
||||
isOutputMarkdown: tc.tool.isOutputMarkdown,
|
||||
canUpdateOutput: tc.tool.canUpdateOutput,
|
||||
schema: tc.tool.schema,
|
||||
parameterSchema: tc.tool.parameterSchema,
|
||||
};
|
||||
}
|
||||
|
||||
messageParts.push({
|
||||
kind: 'data',
|
||||
data: serializableToolCall as ToolCall,
|
||||
} as Part);
|
||||
|
||||
return {
|
||||
kind: 'message',
|
||||
role: 'agent',
|
||||
parts: messageParts,
|
||||
messageId: uuidv4(),
|
||||
taskId,
|
||||
contextId,
|
||||
};
|
||||
}
|
||||
|
||||
private async getProposedContent(
|
||||
file_path: string,
|
||||
old_string: string,
|
||||
new_string: string,
|
||||
): Promise<string> {
|
||||
try {
|
||||
const currentContent = fs.readFileSync(file_path, 'utf8');
|
||||
return this._applyReplacement(
|
||||
currentContent,
|
||||
old_string,
|
||||
new_string,
|
||||
old_string === '' && currentContent === '',
|
||||
);
|
||||
} catch (err) {
|
||||
if (!isNodeError(err) || err.code !== 'ENOENT') throw err;
|
||||
return '';
|
||||
}
|
||||
}
|
||||
|
||||
private _applyReplacement(
|
||||
currentContent: string | null,
|
||||
oldString: string,
|
||||
newString: string,
|
||||
isNewFile: boolean,
|
||||
): string {
|
||||
if (isNewFile) {
|
||||
return newString;
|
||||
}
|
||||
if (currentContent === null) {
|
||||
// Should not happen if not a new file, but defensively return empty or newString if oldString is also empty
|
||||
return oldString === '' ? newString : '';
|
||||
}
|
||||
// If oldString is empty and it's not a new file, do not modify the content.
|
||||
if (oldString === '' && !isNewFile) {
|
||||
return currentContent;
|
||||
}
|
||||
return currentContent.replaceAll(oldString, newString);
|
||||
}
|
||||
|
||||
async scheduleToolCalls(
|
||||
requests: ToolCallRequestInfo[],
|
||||
abortSignal: AbortSignal,
|
||||
): Promise<void> {
|
||||
if (requests.length === 0) {
|
||||
return;
|
||||
}
|
||||
|
||||
const updatedRequests = await Promise.all(
|
||||
requests.map(async (request) => {
|
||||
if (
|
||||
request.name === 'replace' &&
|
||||
request.args &&
|
||||
!request.args['newContent'] &&
|
||||
request.args['file_path'] &&
|
||||
request.args['old_string'] &&
|
||||
request.args['new_string']
|
||||
) {
|
||||
const newContent = await this.getProposedContent(
|
||||
request.args['file_path'] as string,
|
||||
request.args['old_string'] as string,
|
||||
request.args['new_string'] as string,
|
||||
);
|
||||
return { ...request, args: { ...request.args, newContent } };
|
||||
}
|
||||
return request;
|
||||
}),
|
||||
);
|
||||
|
||||
logger.info(
|
||||
`[Task] Scheduling batch of ${updatedRequests.length} tool calls.`,
|
||||
);
|
||||
const stateChange: StateChange = {
|
||||
kind: CoderAgentEvent.StateChangeEvent,
|
||||
};
|
||||
this.setTaskStateAndPublishUpdate('working', stateChange);
|
||||
|
||||
await this.scheduler.schedule(updatedRequests, abortSignal);
|
||||
}
|
||||
|
||||
async acceptAgentMessage(event: ServerGeminiStreamEvent): Promise<void> {
|
||||
const stateChange: StateChange = {
|
||||
kind: CoderAgentEvent.StateChangeEvent,
|
||||
};
|
||||
switch (event.type) {
|
||||
case GeminiEventType.Content:
|
||||
logger.info('[Task] Sending agent message content...');
|
||||
this._sendTextContent(event.value);
|
||||
break;
|
||||
case GeminiEventType.ToolCallRequest:
|
||||
// This is now handled by the agent loop, which collects all requests
|
||||
// and calls scheduleToolCalls once.
|
||||
logger.warn(
|
||||
'[Task] A single tool call request was passed to acceptAgentMessage. This should be handled in a batch by the agent. Ignoring.',
|
||||
);
|
||||
break;
|
||||
case GeminiEventType.ToolCallResponse:
|
||||
// This event type from ServerGeminiStreamEvent might be for when LLM *generates* a tool response part.
|
||||
// The actual execution result comes via user message.
|
||||
logger.info(
|
||||
'[Task] Received tool call response from LLM (part of generation):',
|
||||
event.value,
|
||||
);
|
||||
break;
|
||||
case GeminiEventType.ToolCallConfirmation:
|
||||
// This is when LLM requests confirmation, not when user provides it.
|
||||
logger.info(
|
||||
'[Task] Received tool call confirmation request from LLM:',
|
||||
event.value.request.callId,
|
||||
);
|
||||
this.pendingToolConfirmationDetails.set(
|
||||
event.value.request.callId,
|
||||
event.value.details,
|
||||
);
|
||||
// This will be handled by the scheduler and _schedulerToolCallsUpdate will set InputRequired if needed.
|
||||
// No direct state change here, scheduler drives it.
|
||||
break;
|
||||
case GeminiEventType.UserCancelled:
|
||||
logger.info('[Task] Received user cancelled event from LLM stream.');
|
||||
this.cancelPendingTools('User cancelled via LLM stream event');
|
||||
this.setTaskStateAndPublishUpdate(
|
||||
'input-required',
|
||||
stateChange,
|
||||
'Task cancelled by user',
|
||||
undefined,
|
||||
true,
|
||||
);
|
||||
break;
|
||||
case GeminiEventType.Thought:
|
||||
logger.info('[Task] Sending agent thought...');
|
||||
this._sendThought(event.value);
|
||||
break;
|
||||
case GeminiEventType.ChatCompressed:
|
||||
break;
|
||||
case GeminiEventType.Finished:
|
||||
logger.info(`[Task ${this.id}] Agent finished its turn.`);
|
||||
break;
|
||||
case GeminiEventType.Error:
|
||||
default: {
|
||||
// Block scope for lexical declaration
|
||||
const errorEvent = event as ServerGeminiErrorEvent; // Type assertion
|
||||
const errorMessage =
|
||||
errorEvent.value?.error.message ?? 'Unknown error from LLM stream';
|
||||
logger.error(
|
||||
'[Task] Received error event from LLM stream:',
|
||||
errorMessage,
|
||||
);
|
||||
|
||||
let errMessage = 'Unknown error from LLM stream';
|
||||
if (errorEvent.value) {
|
||||
errMessage = parseAndFormatApiError(errorEvent.value);
|
||||
}
|
||||
this.cancelPendingTools(`LLM stream error: ${errorMessage}`);
|
||||
this.setTaskStateAndPublishUpdate(
|
||||
this.taskState,
|
||||
stateChange,
|
||||
`Agent Error, unknown agent message: ${errorMessage}`,
|
||||
undefined,
|
||||
false,
|
||||
errMessage,
|
||||
);
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private async _handleToolConfirmationPart(part: Part): Promise<boolean> {
|
||||
if (
|
||||
part.kind !== 'data' ||
|
||||
!part.data ||
|
||||
typeof part.data['callId'] !== 'string' ||
|
||||
typeof part.data['outcome'] !== 'string'
|
||||
) {
|
||||
return false;
|
||||
}
|
||||
|
||||
const callId = part.data['callId'] as string;
|
||||
const outcomeString = part.data['outcome'] as string;
|
||||
let confirmationOutcome: ToolConfirmationOutcome | undefined;
|
||||
|
||||
if (outcomeString === 'proceed_once') {
|
||||
confirmationOutcome = ToolConfirmationOutcome.ProceedOnce;
|
||||
} else if (outcomeString === 'cancel') {
|
||||
confirmationOutcome = ToolConfirmationOutcome.Cancel;
|
||||
} else if (outcomeString === 'proceed_always') {
|
||||
confirmationOutcome = ToolConfirmationOutcome.ProceedAlways;
|
||||
} else if (outcomeString === 'proceed_always_server') {
|
||||
confirmationOutcome = ToolConfirmationOutcome.ProceedAlwaysServer;
|
||||
} else if (outcomeString === 'proceed_always_tool') {
|
||||
confirmationOutcome = ToolConfirmationOutcome.ProceedAlwaysTool;
|
||||
} else if (outcomeString === 'modify_with_editor') {
|
||||
confirmationOutcome = ToolConfirmationOutcome.ModifyWithEditor;
|
||||
} else {
|
||||
logger.warn(
|
||||
`[Task] Unknown tool confirmation outcome: "${outcomeString}" for callId: ${callId}`,
|
||||
);
|
||||
return false;
|
||||
}
|
||||
|
||||
const confirmationDetails = this.pendingToolConfirmationDetails.get(callId);
|
||||
|
||||
if (!confirmationDetails) {
|
||||
logger.warn(
|
||||
`[Task] Received tool confirmation for unknown or already processed callId: ${callId}`,
|
||||
);
|
||||
return false;
|
||||
}
|
||||
|
||||
logger.info(
|
||||
`[Task] Handling tool confirmation for callId: ${callId} with outcome: ${outcomeString}`,
|
||||
);
|
||||
try {
|
||||
// Temporarily unset GCP environment variables so they do not leak into
|
||||
// tool calls.
|
||||
const gcpProject = process.env['GOOGLE_CLOUD_PROJECT'];
|
||||
const gcpCreds = process.env['GOOGLE_APPLICATION_CREDENTIALS'];
|
||||
try {
|
||||
delete process.env['GOOGLE_CLOUD_PROJECT'];
|
||||
delete process.env['GOOGLE_APPLICATION_CREDENTIALS'];
|
||||
|
||||
// This will trigger the scheduler to continue or cancel the specific tool.
|
||||
// The scheduler's onToolCallsUpdate will then reflect the new state (e.g., executing or cancelled).
|
||||
|
||||
// If `edit` tool call, pass updated payload if presesent
|
||||
if (confirmationDetails.type === 'edit') {
|
||||
const payload = part.data['newContent']
|
||||
? ({
|
||||
newContent: part.data['newContent'] as string,
|
||||
} as ToolConfirmationPayload)
|
||||
: undefined;
|
||||
this.skipFinalTrueAfterInlineEdit = !!payload;
|
||||
await confirmationDetails.onConfirm(confirmationOutcome, payload);
|
||||
} else {
|
||||
await confirmationDetails.onConfirm(confirmationOutcome);
|
||||
}
|
||||
} finally {
|
||||
if (gcpProject) {
|
||||
process.env['GOOGLE_CLOUD_PROJECT'] = gcpProject;
|
||||
}
|
||||
if (gcpCreds) {
|
||||
process.env['GOOGLE_APPLICATION_CREDENTIALS'] = gcpCreds;
|
||||
}
|
||||
}
|
||||
|
||||
// Do not delete if modifying, a subsequent tool confirmation for the same
|
||||
// callId will be passed with ProceedOnce/Cancel/etc
|
||||
// Note !== ToolConfirmationOutcome.ModifyWithEditor does not work!
|
||||
if (confirmationOutcome !== 'modify_with_editor') {
|
||||
this.pendingToolConfirmationDetails.delete(callId);
|
||||
}
|
||||
|
||||
// If outcome is Cancel, scheduler should update status to 'cancelled', which then resolves the tool.
|
||||
// If ProceedOnce, scheduler updates to 'executing', then eventually 'success'/'error', which resolves.
|
||||
return true;
|
||||
} catch (error) {
|
||||
logger.error(
|
||||
`[Task] Error during tool confirmation for callId ${callId}:`,
|
||||
error,
|
||||
);
|
||||
// If confirming fails, we should probably mark this tool as failed
|
||||
this._resolveToolCall(callId); // Resolve it as it won't proceed.
|
||||
const errorMessageText =
|
||||
error instanceof Error
|
||||
? error.message
|
||||
: `Error processing tool confirmation for ${callId}`;
|
||||
const message = this._createTextMessage(errorMessageText);
|
||||
const toolCallUpdate: ToolCallUpdate = {
|
||||
kind: CoderAgentEvent.ToolCallUpdateEvent,
|
||||
};
|
||||
const event = this._createStatusUpdateEvent(
|
||||
this.taskState,
|
||||
toolCallUpdate,
|
||||
message,
|
||||
false,
|
||||
);
|
||||
this.eventBus?.publish(event);
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
getAndClearCompletedTools(): CompletedToolCall[] {
|
||||
const tools = [...this.completedToolCalls];
|
||||
this.completedToolCalls = [];
|
||||
return tools;
|
||||
}
|
||||
|
||||
addToolResponsesToHistory(completedTools: CompletedToolCall[]): void {
|
||||
logger.info(
|
||||
`[Task] Adding ${completedTools.length} tool responses to history without generating a new response.`,
|
||||
);
|
||||
const responsesToAdd = completedTools.flatMap(
|
||||
(toolCall) => toolCall.response.responseParts,
|
||||
);
|
||||
|
||||
for (const response of responsesToAdd) {
|
||||
let parts: genAiPart[];
|
||||
if (Array.isArray(response)) {
|
||||
parts = response;
|
||||
} else if (typeof response === 'string') {
|
||||
parts = [{ text: response }];
|
||||
} else {
|
||||
parts = [response];
|
||||
}
|
||||
this.geminiClient.addHistory({
|
||||
role: 'user',
|
||||
parts,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
async *sendCompletedToolsToLlm(
|
||||
completedToolCalls: CompletedToolCall[],
|
||||
aborted: AbortSignal,
|
||||
): AsyncGenerator<ServerGeminiStreamEvent> {
|
||||
if (completedToolCalls.length === 0) {
|
||||
yield* (async function* () {})(); // Yield nothing
|
||||
return;
|
||||
}
|
||||
|
||||
const llmParts: PartUnion[] = [];
|
||||
logger.info(
|
||||
`[Task] Feeding ${completedToolCalls.length} tool responses to LLM.`,
|
||||
);
|
||||
for (const completedToolCall of completedToolCalls) {
|
||||
logger.info(
|
||||
`[Task] Adding tool response for "${completedToolCall.request.name}" (callId: ${completedToolCall.request.callId}) to LLM input.`,
|
||||
);
|
||||
const responseParts = completedToolCall.response.responseParts;
|
||||
if (Array.isArray(responseParts)) {
|
||||
llmParts.push(...responseParts);
|
||||
} else {
|
||||
llmParts.push(responseParts);
|
||||
}
|
||||
}
|
||||
|
||||
logger.info('[Task] Sending new parts to agent.');
|
||||
const stateChange: StateChange = {
|
||||
kind: CoderAgentEvent.StateChangeEvent,
|
||||
};
|
||||
// Set task state to working as we are about to call LLM
|
||||
this.setTaskStateAndPublishUpdate('working', stateChange);
|
||||
// TODO: Determine what it mean to have, then add a prompt ID.
|
||||
yield* this.geminiClient.sendMessageStream(
|
||||
llmParts,
|
||||
aborted,
|
||||
/*prompt_id*/ '',
|
||||
);
|
||||
}
|
||||
|
||||
async *acceptUserMessage(
|
||||
requestContext: RequestContext,
|
||||
aborted: AbortSignal,
|
||||
): AsyncGenerator<ServerGeminiStreamEvent> {
|
||||
const userMessage = requestContext.userMessage;
|
||||
const llmParts: PartUnion[] = [];
|
||||
let anyConfirmationHandled = false;
|
||||
let hasContentForLlm = false;
|
||||
|
||||
for (const part of userMessage.parts) {
|
||||
const confirmationHandled = await this._handleToolConfirmationPart(part);
|
||||
if (confirmationHandled) {
|
||||
anyConfirmationHandled = true;
|
||||
// If a confirmation was handled, the scheduler will now run the tool (or cancel it).
|
||||
// We don't send anything to the LLM for this part.
|
||||
// The subsequent tool execution will eventually lead to resolveToolCall.
|
||||
continue;
|
||||
}
|
||||
|
||||
if (part.kind === 'text') {
|
||||
llmParts.push({ text: part.text });
|
||||
hasContentForLlm = true;
|
||||
}
|
||||
}
|
||||
|
||||
if (hasContentForLlm) {
|
||||
logger.info('[Task] Sending new parts to LLM.');
|
||||
const stateChange: StateChange = {
|
||||
kind: CoderAgentEvent.StateChangeEvent,
|
||||
};
|
||||
// Set task state to working as we are about to call LLM
|
||||
this.setTaskStateAndPublishUpdate('working', stateChange);
|
||||
// TODO: Determine what it mean to have, then add a prompt ID.
|
||||
yield* this.geminiClient.sendMessageStream(
|
||||
llmParts,
|
||||
aborted,
|
||||
/*prompt_id*/ '',
|
||||
);
|
||||
} else if (anyConfirmationHandled) {
|
||||
logger.info(
|
||||
'[Task] User message only contained tool confirmations. Scheduler is active. No new input for LLM this turn.',
|
||||
);
|
||||
// Ensure task state reflects that scheduler might be working due to confirmation.
|
||||
// If scheduler is active, it will emit its own status updates.
|
||||
// If all pending tools were just confirmed, waitForPendingTools will handle the wait.
|
||||
// If some tools are still pending approval, scheduler would have set InputRequired.
|
||||
// If not, and no new text, we are just waiting.
|
||||
if (
|
||||
this.pendingToolCalls.size > 0 &&
|
||||
this.taskState !== 'input-required'
|
||||
) {
|
||||
const stateChange: StateChange = {
|
||||
kind: CoderAgentEvent.StateChangeEvent,
|
||||
};
|
||||
this.setTaskStateAndPublishUpdate('working', stateChange); // Reflect potential background activity
|
||||
}
|
||||
yield* (async function* () {})(); // Yield nothing
|
||||
} else {
|
||||
logger.info(
|
||||
'[Task] No relevant parts in user message for LLM interaction or tool confirmation.',
|
||||
);
|
||||
// If there's no new text and no confirmations, and no pending tools,
|
||||
// it implies we might need to signal input required if nothing else is happening.
|
||||
// However, the agent.ts will make this determination after waitForPendingTools.
|
||||
yield* (async function* () {})(); // Yield nothing
|
||||
}
|
||||
}
|
||||
|
||||
_sendTextContent(content: string): void {
|
||||
if (content === '') {
|
||||
return;
|
||||
}
|
||||
logger.info('[Task] Sending text content to event bus.');
|
||||
const message = this._createTextMessage(content);
|
||||
const textContent: TextContent = {
|
||||
kind: CoderAgentEvent.TextContentEvent,
|
||||
};
|
||||
this.eventBus?.publish(
|
||||
this._createStatusUpdateEvent(
|
||||
this.taskState,
|
||||
textContent,
|
||||
message,
|
||||
false,
|
||||
),
|
||||
);
|
||||
}
|
||||
|
||||
_sendThought(content: ThoughtSummary): void {
|
||||
if (!content.subject && !content.description) {
|
||||
return;
|
||||
}
|
||||
logger.info('[Task] Sending thought to event bus.');
|
||||
const message: Message = {
|
||||
kind: 'message',
|
||||
role: 'agent',
|
||||
parts: [
|
||||
{
|
||||
kind: 'data',
|
||||
data: content,
|
||||
} as Part,
|
||||
],
|
||||
messageId: uuidv4(),
|
||||
taskId: this.id,
|
||||
contextId: this.contextId,
|
||||
};
|
||||
const thought: Thought = {
|
||||
kind: CoderAgentEvent.ThoughtEvent,
|
||||
};
|
||||
this.eventBus?.publish(
|
||||
this._createStatusUpdateEvent(this.taskState, thought, message, false),
|
||||
);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user