mirror of
https://github.com/google-gemini/gemini-cli.git
synced 2026-04-14 15:21:01 -07:00
Add a2a-server package to gemini-cli (#6597)
This commit is contained in:
Adam Weidman
785
packages/a2a-server/src/agent.ts
Normal file
785
packages/a2a-server/src/agent.ts
Normal file
@@ -0,0 +1,785 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2025 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import express from 'express';
|
||||
import { AsyncLocalStorage } from 'node:async_hooks';
|
||||
|
||||
import type { Message, Task as SDKTask, AgentCard } from '@a2a-js/sdk';
|
||||
import type {
|
||||
TaskStore,
|
||||
AgentExecutor,
|
||||
AgentExecutionEvent,
|
||||
RequestContext,
|
||||
ExecutionEventBus,
|
||||
} from '@a2a-js/sdk/server';
|
||||
import { DefaultRequestHandler, InMemoryTaskStore } from '@a2a-js/sdk/server';
|
||||
import { A2AExpressApp } from '@a2a-js/sdk/server/express'; // Import server components
|
||||
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 './logger.js';
|
||||
import type { StateChange, AgentSettings } from './types.js';
|
||||
import { CoderAgentEvent } from './types.js';
|
||||
import { loadConfig, loadEnvironment, setTargetDir } from './config.js';
|
||||
import { loadSettings } from './settings.js';
|
||||
import { loadExtensions } from './extension.js';
|
||||
import { Task } from './task.js';
|
||||
import { GCSTaskStore, NoOpTaskStore } from './gcs.js';
|
||||
import type { PersistedStateMetadata } from './metadata_types.js';
|
||||
import { getPersistedState, setPersistedState } from './metadata_types.js';
|
||||
|
||||
const requestStorage = new AsyncLocalStorage<{ req: express.Request }>();
|
||||
|
||||
/**
|
||||
* 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;
|
||||
}
|
||||
}
|
||||
|
||||
const coderAgentCard: AgentCard = {
|
||||
name: 'Gemini SDLC Agent',
|
||||
description:
|
||||
'An agent that generates code based on natural language instructions and streams file outputs.',
|
||||
url: 'http://localhost:41242/',
|
||||
provider: {
|
||||
organization: 'Google',
|
||||
url: 'https://google.com',
|
||||
},
|
||||
protocolVersion: '0.3.0',
|
||||
version: '0.0.2', // Incremented version
|
||||
capabilities: {
|
||||
streaming: true,
|
||||
pushNotifications: false,
|
||||
stateTransitionHistory: true,
|
||||
},
|
||||
securitySchemes: undefined,
|
||||
security: undefined,
|
||||
defaultInputModes: ['text'],
|
||||
defaultOutputModes: ['text'],
|
||||
skills: [
|
||||
{
|
||||
id: 'code_generation',
|
||||
name: 'Code Generation',
|
||||
description:
|
||||
'Generates code snippets or complete files based on user requests, streaming the results.',
|
||||
tags: ['code', 'development', 'programming'],
|
||||
examples: [
|
||||
'Write a python function to calculate fibonacci numbers.',
|
||||
'Create an HTML file with a basic button that alerts "Hello!" when clicked.',
|
||||
],
|
||||
inputModes: ['text'],
|
||||
outputModes: ['text'],
|
||||
},
|
||||
],
|
||||
supportsAuthenticatedExtendedCard: false,
|
||||
};
|
||||
|
||||
/**
|
||||
* CoderAgentExecutor implements the agent's core logic for code generation.
|
||||
*/
|
||||
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,
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
export function updateCoderAgentCardUrl(port: number) {
|
||||
coderAgentCard.url = `http://localhost:${port}/`;
|
||||
}
|
||||
|
||||
export async function main() {
|
||||
try {
|
||||
const expressApp = await createApp();
|
||||
const port = process.env['CODER_AGENT_PORT'] || 0;
|
||||
|
||||
const server = expressApp.listen(port, () => {
|
||||
const address = server.address();
|
||||
let actualPort;
|
||||
if (process.env['CODER_AGENT_PORT']) {
|
||||
actualPort = process.env['CODER_AGENT_PORT'];
|
||||
} else if (address && typeof address !== 'string') {
|
||||
actualPort = address.port;
|
||||
} else {
|
||||
throw new Error('[Core Agent] Could not find port number.');
|
||||
}
|
||||
updateCoderAgentCardUrl(Number(actualPort));
|
||||
logger.info(
|
||||
`[CoreAgent] Agent Server started on http://localhost:${actualPort}`,
|
||||
);
|
||||
logger.info(
|
||||
`[CoreAgent] Agent Card: http://localhost:${actualPort}/.well-known/agent-card.json`,
|
||||
);
|
||||
logger.info('[CoreAgent] Press Ctrl+C to stop the server');
|
||||
});
|
||||
} catch (error) {
|
||||
logger.error('[CoreAgent] Error during startup:', error);
|
||||
process.exit(1);
|
||||
}
|
||||
}
|
||||
|
||||
export async function createApp() {
|
||||
try {
|
||||
// loadEnvironment() is called within getConfig now
|
||||
const bucketName = process.env['GCS_BUCKET_NAME'];
|
||||
let taskStoreForExecutor: TaskStore;
|
||||
let taskStoreForHandler: TaskStore;
|
||||
|
||||
if (bucketName) {
|
||||
logger.info(`Using GCSTaskStore with bucket: ${bucketName}`);
|
||||
const gcsTaskStore = new GCSTaskStore(bucketName);
|
||||
taskStoreForExecutor = gcsTaskStore;
|
||||
taskStoreForHandler = new NoOpTaskStore(gcsTaskStore);
|
||||
} else {
|
||||
logger.info('Using InMemoryTaskStore');
|
||||
const inMemoryTaskStore = new InMemoryTaskStore();
|
||||
taskStoreForExecutor = inMemoryTaskStore;
|
||||
taskStoreForHandler = inMemoryTaskStore;
|
||||
}
|
||||
|
||||
const agentExecutor = new CoderAgentExecutor(taskStoreForExecutor);
|
||||
|
||||
const requestHandler = new DefaultRequestHandler(
|
||||
coderAgentCard,
|
||||
taskStoreForHandler,
|
||||
agentExecutor,
|
||||
);
|
||||
|
||||
let expressApp = express();
|
||||
expressApp.use((req, res, next) => {
|
||||
requestStorage.run({ req }, next);
|
||||
});
|
||||
|
||||
const appBuilder = new A2AExpressApp(requestHandler);
|
||||
expressApp = appBuilder.setupRoutes(expressApp, '');
|
||||
expressApp.use(express.json());
|
||||
|
||||
expressApp.post('/tasks', async (req, res) => {
|
||||
try {
|
||||
const taskId = uuidv4();
|
||||
const agentSettings = req.body.agentSettings as
|
||||
| AgentSettings
|
||||
| undefined;
|
||||
const contextId = req.body.contextId || uuidv4();
|
||||
const wrapper = await agentExecutor.createTask(
|
||||
taskId,
|
||||
contextId,
|
||||
agentSettings,
|
||||
);
|
||||
await taskStoreForExecutor.save(wrapper.toSDKTask());
|
||||
res.status(201).json(wrapper.id);
|
||||
} catch (error) {
|
||||
logger.error('[CoreAgent] Error creating task:', error);
|
||||
const errorMessage =
|
||||
error instanceof Error
|
||||
? error.message
|
||||
: 'Unknown error creating task';
|
||||
res.status(500).send({ error: errorMessage });
|
||||
}
|
||||
});
|
||||
|
||||
expressApp.get('/tasks/metadata', async (req, res) => {
|
||||
// This endpoint is only meaningful if the task store is in-memory.
|
||||
if (!(taskStoreForExecutor instanceof InMemoryTaskStore)) {
|
||||
res.status(501).send({
|
||||
error:
|
||||
'Listing all task metadata is only supported when using InMemoryTaskStore.',
|
||||
});
|
||||
}
|
||||
try {
|
||||
const wrappers = agentExecutor.getAllTasks();
|
||||
if (wrappers && wrappers.length > 0) {
|
||||
const tasksMetadata = await Promise.all(
|
||||
wrappers.map((wrapper) => wrapper.task.getMetadata()),
|
||||
);
|
||||
res.status(200).json(tasksMetadata);
|
||||
} else {
|
||||
res.status(204).send();
|
||||
}
|
||||
} catch (error) {
|
||||
logger.error('[CoreAgent] Error getting all task metadata:', error);
|
||||
const errorMessage =
|
||||
error instanceof Error
|
||||
? error.message
|
||||
: 'Unknown error getting task metadata';
|
||||
res.status(500).send({ error: errorMessage });
|
||||
}
|
||||
});
|
||||
|
||||
expressApp.get('/tasks/:taskId/metadata', async (req, res) => {
|
||||
const taskId = req.params.taskId;
|
||||
let wrapper = agentExecutor.getTask(taskId);
|
||||
if (!wrapper) {
|
||||
const sdkTask = await taskStoreForExecutor.load(taskId);
|
||||
if (sdkTask) {
|
||||
wrapper = await agentExecutor.reconstruct(sdkTask);
|
||||
}
|
||||
}
|
||||
if (!wrapper) {
|
||||
res.status(404).send({ error: 'Task not found' });
|
||||
return;
|
||||
}
|
||||
res.json({ metadata: await wrapper.task.getMetadata() });
|
||||
});
|
||||
return expressApp;
|
||||
} catch (error) {
|
||||
logger.error('[CoreAgent] Error during startup:', error);
|
||||
process.exit(1);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user