chore(a2a-server): refactor a2a-server src directory (#7593)

This commit is contained in:
Sam McCauley
2025-09-03 10:24:48 -04:00
committed by GitHub
parent abddd2b6ee
commit 2782af3f57
17 changed files with 253 additions and 229 deletions

View File

@@ -11,7 +11,7 @@
"type": "module",
"main": "dist/server.js",
"scripts": {
"start": "node dist/src/server.js",
"start": "node dist/src/http/server.js",
"build": "node ../../scripts/build_package.js",
"lint": "eslint . --ext .ts,.tsx",
"format": "prettier --write .",

View File

@@ -4,10 +4,7 @@
* 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 { Message, Task as SDKTask } from '@a2a-js/sdk';
import type {
TaskStore,
AgentExecutor,
@@ -15,8 +12,6 @@ import type {
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,
@@ -24,18 +19,16 @@ import type {
} 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 { 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 { 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 }>();
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.
@@ -77,48 +70,10 @@ class TaskWrapper {
}
}
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 {
export class CoderAgentExecutor implements AgentExecutor {
private tasks: Map<string, TaskWrapper> = new Map();
// Track tasks with an active execution loop.
private executingTasks = new Set<string>();
@@ -639,147 +594,3 @@ class CoderAgentExecutor implements AgentExecutor {
}
}
}
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);
}
}

View File

@@ -7,7 +7,7 @@
import { describe, it, expect, vi } from 'vitest';
import { Task } from './task.js';
import type { Config, ToolCallRequestInfo } from '@google/gemini-cli-core';
import { createMockConfig } from './testing_utils.js';
import { createMockConfig } from '../utils/testing_utils.js';
import type { ExecutionEventBus } from '@a2a-js/sdk/server';
describe('Task', () => {

View File

@@ -37,10 +37,10 @@ import type {
Artifact,
} from '@a2a-js/sdk';
import { v4 as uuidv4 } from 'uuid';
import { logger } from './logger.js';
import { logger } from '../utils/logger.js';
import * as fs from 'node:fs';
import { CoderAgentEvent } from './types.js';
import { CoderAgentEvent } from '../types.js';
import type {
CoderAgentMessage,
StateChange,
@@ -49,7 +49,7 @@ import type {
TaskMetadata,
Thought,
ThoughtSummary,
} from './types.js';
} from '../types.js';
import type { PartUnion, Part as genAiPart } from '@google/genai';
export class Task {

View File

@@ -22,10 +22,10 @@ import {
DEFAULT_GEMINI_MODEL,
} from '@google/gemini-cli-core';
import { logger } from './logger.js';
import { logger } from '../utils/logger.js';
import type { Settings } from './settings.js';
import type { Extension } from './extension.js';
import { type AgentSettings, CoderAgentEvent } from './types.js';
import { type AgentSettings, CoderAgentEvent } from '../types.js';
export async function loadConfig(
settings: Settings,

View File

@@ -10,7 +10,7 @@ import type { MCPServerConfig } from '@google/gemini-cli-core';
import * as fs from 'node:fs';
import * as path from 'node:path';
import * as os from 'node:os';
import { logger } from './logger.js';
import { logger } from '../utils/logger.js';
export const EXTENSIONS_DIRECTORY_NAME = path.join('.gemini', 'extensions');
export const EXTENSIONS_CONFIG_FILENAME = 'gemini-extension.json';

View File

@@ -27,13 +27,13 @@ import {
it,
vi,
} from 'vitest';
import { createApp } from './agent.js';
import { createApp } from './app.js';
import {
assertUniqueFinalEventIsLast,
assertTaskCreationAndWorkingStatus,
createStreamMessageRequest,
createMockConfig,
} from './testing_utils.js';
} from '../utils/testing_utils.js';
import { MockTool } from '@google/gemini-cli-core';
const mockToolConfirmationFn = async () =>
@@ -57,15 +57,15 @@ const streamToSSEEvents = (
// Mock the logger to avoid polluting test output
// Comment out to debug tests
vi.mock('./logger.js', () => ({
vi.mock('../utils/logger.js', () => ({
logger: { info: vi.fn(), warn: vi.fn(), error: vi.fn() },
}));
let config: Config;
const getToolRegistrySpy = vi.fn().mockReturnValue(ApprovalMode.DEFAULT);
const getApprovalModeSpy = vi.fn();
vi.mock('./config.js', async () => {
const actual = await vi.importActual('./config.js');
vi.mock('../config/config.js', async () => {
const actual = await vi.importActual('../config/config.js');
return {
...actual,
loadConfig: vi.fn().mockImplementation(async () => {

View File

@@ -0,0 +1,200 @@
/**
* @license
* Copyright 2025 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import express from 'express';
import type { AgentCard } from '@a2a-js/sdk';
import type { TaskStore } 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 { v4 as uuidv4 } from 'uuid';
import { logger } from '../utils/logger.js';
import type { AgentSettings } from '../types.js';
import { GCSTaskStore, NoOpTaskStore } from '../persistence/gcs.js';
import { CoderAgentExecutor } from '../agent/executor.js';
import { requestStorage } from './requestStorage.js';
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,
};
export function updateCoderAgentCardUrl(port: number) {
coderAgentCard.url = `http://localhost:${port}/`;
}
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);
}
}
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);
}
}

View File

@@ -7,22 +7,22 @@
import { describe, it, expect, beforeAll, afterAll, vi } from 'vitest';
import request from 'supertest';
import type express from 'express';
import { createApp, updateCoderAgentCardUrl } from './agent.js';
import { createApp, updateCoderAgentCardUrl } from './app.js';
import * as fs from 'node:fs';
import * as path from 'node:path';
import * as os from 'node:os';
import type { Server } from 'node:http';
import type { TaskMetadata } from './types.js';
import type { TaskMetadata } from '../types.js';
import type { AddressInfo } from 'node:net';
// Mock the logger to avoid polluting test output
// Comment out to help debug
vi.mock('./logger.js', () => ({
vi.mock('../utils/logger.js', () => ({
logger: { info: vi.fn(), warn: vi.fn(), error: vi.fn() },
}));
// Mock Task.create to avoid its complex setup
vi.mock('./task.js', () => {
vi.mock('../agent/task.js', () => {
class MockTask {
id: string;
contextId: string;

View File

@@ -0,0 +1,10 @@
/**
* @license
* Copyright 2025 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import type express from 'express';
import { AsyncLocalStorage } from 'node:async_hooks';
export const requestStorage = new AsyncLocalStorage<{ req: express.Request }>();

View File

@@ -7,8 +7,8 @@
import * as url from 'node:url';
import * as path from 'node:path';
import { logger } from './logger.js';
import { main } from './agent.js';
import { logger } from '../utils/logger.js';
import { main } from './app.js';
// Check if the module is the main script being run. path.resolve() creates a
// canonical, absolute path, which avoids cross-platform issues.

View File

@@ -4,5 +4,6 @@
* SPDX-License-Identifier: Apache-2.0
*/
export * from './agent.js';
export * from './agent/executor.js';
export * from './http/app.js';
export * from './types.js';

View File

@@ -16,9 +16,9 @@ import type { Mocked, MockedClass, Mock } from 'vitest';
import { describe, it, expect, beforeEach, vi } from 'vitest';
import { GCSTaskStore, NoOpTaskStore } from './gcs.js';
import { logger } from './logger.js';
import * as configModule from './config.js';
import * as metadataModule from './metadata_types.js';
import { logger } from '../utils/logger.js';
import * as configModule from '../config/config.js';
import * as metadataModule from '../metadata_types.js';
// Mock dependencies
vi.mock('@google-cloud/storage');
@@ -42,7 +42,7 @@ vi.mock('node:fs', async () => {
vi.mock('tar');
vi.mock('zlib');
vi.mock('uuid');
vi.mock('./logger', () => ({
vi.mock('../utils/logger.js', () => ({
logger: {
info: vi.fn(),
warn: vi.fn(),
@@ -50,8 +50,10 @@ vi.mock('./logger', () => ({
debug: vi.fn(),
},
}));
vi.mock('./config');
vi.mock('./metadata_types');
vi.mock('../config/config.js', () => ({
setTargetDir: vi.fn(),
}));
vi.mock('../metadata_types');
vi.mock('node:stream/promises', () => ({
pipeline: vi.fn(),
}));

View File

@@ -13,12 +13,12 @@ import { tmpdir } from 'node:os';
import { join } from 'node:path';
import type { Task as SDKTask } from '@a2a-js/sdk';
import type { TaskStore } from '@a2a-js/sdk/server';
import { logger } from './logger.js';
import { setTargetDir } from './config.js';
import { logger } from '../utils/logger.js';
import { setTargetDir } from '../config/config.js';
import {
getPersistedState,
type PersistedTaskMetadata,
} from './metadata_types.js';
} from '../metadata_types.js';
import { v4 as uuidv4 } from 'uuid';
type ObjectType = 'metadata' | 'workspace';