mirror of
https://github.com/google-gemini/gemini-cli.git
synced 2026-06-11 20:07:00 -07:00
chore(a2a-server): refactor a2a-server src directory (#7593)
This commit is contained in:
Sam McCauley
@@ -0,0 +1,625 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2025 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import type { Config } from '@google/gemini-cli-core';
|
||||
import {
|
||||
GeminiEventType,
|
||||
ApprovalMode,
|
||||
type ToolCallConfirmationDetails,
|
||||
} from '@google/gemini-cli-core';
|
||||
import type {
|
||||
TaskStatusUpdateEvent,
|
||||
SendStreamingMessageSuccessResponse,
|
||||
} from '@a2a-js/sdk';
|
||||
import type express from 'express';
|
||||
import type { Server } from 'node:http';
|
||||
import request from 'supertest';
|
||||
import {
|
||||
afterAll,
|
||||
afterEach,
|
||||
beforeEach,
|
||||
beforeAll,
|
||||
describe,
|
||||
expect,
|
||||
it,
|
||||
vi,
|
||||
} from 'vitest';
|
||||
import { createApp } from './app.js';
|
||||
import {
|
||||
assertUniqueFinalEventIsLast,
|
||||
assertTaskCreationAndWorkingStatus,
|
||||
createStreamMessageRequest,
|
||||
createMockConfig,
|
||||
} from '../utils/testing_utils.js';
|
||||
import { MockTool } from '@google/gemini-cli-core';
|
||||
|
||||
const mockToolConfirmationFn = async () =>
|
||||
({}) as unknown as ToolCallConfirmationDetails;
|
||||
|
||||
const streamToSSEEvents = (
|
||||
stream: string,
|
||||
): SendStreamingMessageSuccessResponse[] =>
|
||||
stream
|
||||
.split('\n\n')
|
||||
.filter(Boolean) // Remove empty strings from trailing newlines
|
||||
.map((chunk) => {
|
||||
const dataLine = chunk
|
||||
.split('\n')
|
||||
.find((line) => line.startsWith('data: '));
|
||||
if (!dataLine) {
|
||||
throw new Error(`Invalid SSE chunk found: "${chunk}"`);
|
||||
}
|
||||
return JSON.parse(dataLine.substring(6));
|
||||
});
|
||||
|
||||
// Mock the logger to avoid polluting test output
|
||||
// Comment out to debug tests
|
||||
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/config.js', async () => {
|
||||
const actual = await vi.importActual('../config/config.js');
|
||||
return {
|
||||
...actual,
|
||||
loadConfig: vi.fn().mockImplementation(async () => {
|
||||
const mockConfig = createMockConfig({
|
||||
getToolRegistry: getToolRegistrySpy,
|
||||
getApprovalMode: getApprovalModeSpy,
|
||||
});
|
||||
config = mockConfig as Config;
|
||||
return config;
|
||||
}),
|
||||
};
|
||||
});
|
||||
|
||||
// Mock the GeminiClient to avoid actual API calls
|
||||
const sendMessageStreamSpy = vi.fn();
|
||||
vi.mock('@google/gemini-cli-core', async () => {
|
||||
const actual = await vi.importActual('@google/gemini-cli-core');
|
||||
return {
|
||||
...actual,
|
||||
GeminiClient: vi.fn().mockImplementation(() => ({
|
||||
sendMessageStream: sendMessageStreamSpy,
|
||||
getUserTier: vi.fn().mockReturnValue('free'),
|
||||
initialize: vi.fn(),
|
||||
})),
|
||||
};
|
||||
});
|
||||
|
||||
describe('E2E Tests', () => {
|
||||
let app: express.Express;
|
||||
let server: Server;
|
||||
|
||||
beforeAll(async () => {
|
||||
app = await createApp();
|
||||
server = app.listen(0); // Listen on a random available port
|
||||
});
|
||||
|
||||
beforeEach(() => {
|
||||
getApprovalModeSpy.mockReturnValue(ApprovalMode.DEFAULT);
|
||||
});
|
||||
|
||||
afterAll(
|
||||
() =>
|
||||
new Promise<void>((resolve) => {
|
||||
server.close(() => {
|
||||
resolve();
|
||||
});
|
||||
}),
|
||||
);
|
||||
|
||||
afterEach(() => {
|
||||
vi.clearAllMocks();
|
||||
});
|
||||
|
||||
it('should create a new task and stream status updates (text-content) via POST /', async () => {
|
||||
sendMessageStreamSpy.mockImplementation(async function* () {
|
||||
yield* [{ type: 'content', value: 'Hello how are you?' }];
|
||||
});
|
||||
|
||||
const agent = request.agent(app);
|
||||
const res = await agent
|
||||
.post('/')
|
||||
.send(createStreamMessageRequest('hello', 'a2a-test-message'))
|
||||
.set('Content-Type', 'application/json')
|
||||
.expect(200);
|
||||
|
||||
const events = streamToSSEEvents(res.text);
|
||||
|
||||
assertTaskCreationAndWorkingStatus(events);
|
||||
|
||||
// Status update: text-content
|
||||
const textContentEvent = events[2].result as TaskStatusUpdateEvent;
|
||||
expect(textContentEvent.kind).toBe('status-update');
|
||||
expect(textContentEvent.status.state).toBe('working');
|
||||
expect(textContentEvent.metadata?.['coderAgent']).toMatchObject({
|
||||
kind: 'text-content',
|
||||
});
|
||||
expect(textContentEvent.status.message?.parts).toMatchObject([
|
||||
{ kind: 'text', text: 'Hello how are you?' },
|
||||
]);
|
||||
|
||||
// Status update: input-required (final)
|
||||
const finalEvent = events[3].result as TaskStatusUpdateEvent;
|
||||
expect(finalEvent.kind).toBe('status-update');
|
||||
expect(finalEvent.status?.state).toBe('input-required');
|
||||
expect(finalEvent.final).toBe(true);
|
||||
|
||||
assertUniqueFinalEventIsLast(events);
|
||||
expect(events.length).toBe(4);
|
||||
});
|
||||
|
||||
it('should create a new task, schedule a tool call, and wait for approval', async () => {
|
||||
// First call yields the tool request
|
||||
sendMessageStreamSpy.mockImplementationOnce(async function* () {
|
||||
yield* [
|
||||
{
|
||||
type: GeminiEventType.ToolCallRequest,
|
||||
value: {
|
||||
callId: 'test-call-id',
|
||||
name: 'test-tool',
|
||||
args: {},
|
||||
},
|
||||
},
|
||||
];
|
||||
});
|
||||
// Subsequent calls yield nothing
|
||||
sendMessageStreamSpy.mockImplementation(async function* () {
|
||||
yield* [];
|
||||
});
|
||||
|
||||
const mockTool = new MockTool({
|
||||
name: 'test-tool',
|
||||
shouldConfirmExecute: vi.fn(mockToolConfirmationFn),
|
||||
});
|
||||
|
||||
getToolRegistrySpy.mockReturnValue({
|
||||
getAllTools: vi.fn().mockReturnValue([mockTool]),
|
||||
getToolsByServer: vi.fn().mockReturnValue([]),
|
||||
getTool: vi.fn().mockReturnValue(mockTool),
|
||||
});
|
||||
|
||||
const agent = request.agent(app);
|
||||
const res = await agent
|
||||
.post('/')
|
||||
.send(createStreamMessageRequest('run a tool', 'a2a-tool-test-message'))
|
||||
.set('Content-Type', 'application/json')
|
||||
.expect(200);
|
||||
|
||||
const events = streamToSSEEvents(res.text);
|
||||
assertTaskCreationAndWorkingStatus(events);
|
||||
|
||||
// Status update: working
|
||||
const workingEvent2 = events[2].result as TaskStatusUpdateEvent;
|
||||
expect(workingEvent2.kind).toBe('status-update');
|
||||
expect(workingEvent2.status.state).toBe('working');
|
||||
expect(workingEvent2.metadata?.['coderAgent']).toMatchObject({
|
||||
kind: 'state-change',
|
||||
});
|
||||
|
||||
// Status update: tool-call-update
|
||||
const toolCallUpdateEvent = events[3].result as TaskStatusUpdateEvent;
|
||||
expect(toolCallUpdateEvent.kind).toBe('status-update');
|
||||
expect(toolCallUpdateEvent.status.state).toBe('working');
|
||||
expect(toolCallUpdateEvent.metadata?.['coderAgent']).toMatchObject({
|
||||
kind: 'tool-call-update',
|
||||
});
|
||||
expect(toolCallUpdateEvent.status.message?.parts).toMatchObject([
|
||||
{
|
||||
data: {
|
||||
status: 'validating',
|
||||
request: { callId: 'test-call-id' },
|
||||
},
|
||||
},
|
||||
]);
|
||||
|
||||
// State update: awaiting_approval update
|
||||
const toolCallConfirmationEvent = events[4].result as TaskStatusUpdateEvent;
|
||||
expect(toolCallConfirmationEvent.kind).toBe('status-update');
|
||||
expect(toolCallConfirmationEvent.metadata?.['coderAgent']).toMatchObject({
|
||||
kind: 'tool-call-confirmation',
|
||||
});
|
||||
expect(toolCallConfirmationEvent.status.message?.parts).toMatchObject([
|
||||
{
|
||||
data: {
|
||||
status: 'awaiting_approval',
|
||||
request: { callId: 'test-call-id' },
|
||||
},
|
||||
},
|
||||
]);
|
||||
expect(toolCallConfirmationEvent.status?.state).toBe('working');
|
||||
|
||||
assertUniqueFinalEventIsLast(events);
|
||||
expect(events.length).toBe(6);
|
||||
});
|
||||
|
||||
it('should handle multiple tool calls in a single turn', async () => {
|
||||
// First call yields the tool request
|
||||
sendMessageStreamSpy.mockImplementationOnce(async function* () {
|
||||
yield* [
|
||||
{
|
||||
type: GeminiEventType.ToolCallRequest,
|
||||
value: {
|
||||
callId: 'test-call-id-1',
|
||||
name: 'test-tool-1',
|
||||
args: {},
|
||||
},
|
||||
},
|
||||
{
|
||||
type: GeminiEventType.ToolCallRequest,
|
||||
value: {
|
||||
callId: 'test-call-id-2',
|
||||
name: 'test-tool-2',
|
||||
args: {},
|
||||
},
|
||||
},
|
||||
];
|
||||
});
|
||||
// Subsequent calls yield nothing
|
||||
sendMessageStreamSpy.mockImplementation(async function* () {
|
||||
yield* [];
|
||||
});
|
||||
|
||||
const mockTool1 = new MockTool({
|
||||
name: 'test-tool-1',
|
||||
displayName: 'Test Tool 1',
|
||||
shouldConfirmExecute: vi.fn(mockToolConfirmationFn),
|
||||
});
|
||||
const mockTool2 = new MockTool({
|
||||
name: 'test-tool-2',
|
||||
displayName: 'Test Tool 2',
|
||||
shouldConfirmExecute: vi.fn(mockToolConfirmationFn),
|
||||
});
|
||||
|
||||
getToolRegistrySpy.mockReturnValue({
|
||||
getAllTools: vi.fn().mockReturnValue([mockTool1, mockTool2]),
|
||||
getToolsByServer: vi.fn().mockReturnValue([]),
|
||||
getTool: vi.fn().mockImplementation((name: string) => {
|
||||
if (name === 'test-tool-1') return mockTool1;
|
||||
if (name === 'test-tool-2') return mockTool2;
|
||||
return undefined;
|
||||
}),
|
||||
});
|
||||
|
||||
const agent = request.agent(app);
|
||||
const res = await agent
|
||||
.post('/')
|
||||
.send(
|
||||
createStreamMessageRequest(
|
||||
'run two tools',
|
||||
'a2a-multi-tool-test-message',
|
||||
),
|
||||
)
|
||||
.set('Content-Type', 'application/json')
|
||||
.expect(200);
|
||||
|
||||
const events = streamToSSEEvents(res.text);
|
||||
assertTaskCreationAndWorkingStatus(events);
|
||||
|
||||
// Second working update
|
||||
const workingEvent = events[2].result as TaskStatusUpdateEvent;
|
||||
expect(workingEvent.kind).toBe('status-update');
|
||||
expect(workingEvent.status.state).toBe('working');
|
||||
|
||||
// State Update: Validate each tool call
|
||||
const toolCallValidateEvent1 = events[3].result as TaskStatusUpdateEvent;
|
||||
expect(toolCallValidateEvent1.metadata?.['coderAgent']).toMatchObject({
|
||||
kind: 'tool-call-update',
|
||||
});
|
||||
expect(toolCallValidateEvent1.status.message?.parts).toMatchObject([
|
||||
{
|
||||
data: {
|
||||
status: 'validating',
|
||||
request: { callId: 'test-call-id-1' },
|
||||
},
|
||||
},
|
||||
]);
|
||||
const toolCallValidateEvent2 = events[4].result as TaskStatusUpdateEvent;
|
||||
expect(toolCallValidateEvent2.metadata?.['coderAgent']).toMatchObject({
|
||||
kind: 'tool-call-update',
|
||||
});
|
||||
expect(toolCallValidateEvent2.status.message?.parts).toMatchObject([
|
||||
{
|
||||
data: {
|
||||
status: 'validating',
|
||||
request: { callId: 'test-call-id-2' },
|
||||
},
|
||||
},
|
||||
]);
|
||||
|
||||
// State Update: Set each tool call to awaiting
|
||||
const toolCallAwaitEvent1 = events[5].result as TaskStatusUpdateEvent;
|
||||
expect(toolCallAwaitEvent1.metadata?.['coderAgent']).toMatchObject({
|
||||
kind: 'tool-call-confirmation',
|
||||
});
|
||||
expect(toolCallAwaitEvent1.status.message?.parts).toMatchObject([
|
||||
{
|
||||
data: {
|
||||
status: 'awaiting_approval',
|
||||
request: { callId: 'test-call-id-1' },
|
||||
},
|
||||
},
|
||||
]);
|
||||
const toolCallAwaitEvent2 = events[6].result as TaskStatusUpdateEvent;
|
||||
expect(toolCallAwaitEvent2.metadata?.['coderAgent']).toMatchObject({
|
||||
kind: 'tool-call-confirmation',
|
||||
});
|
||||
expect(toolCallAwaitEvent2.status.message?.parts).toMatchObject([
|
||||
{
|
||||
data: {
|
||||
status: 'awaiting_approval',
|
||||
request: { callId: 'test-call-id-2' },
|
||||
},
|
||||
},
|
||||
]);
|
||||
|
||||
assertUniqueFinalEventIsLast(events);
|
||||
expect(events.length).toBe(8);
|
||||
});
|
||||
|
||||
it('should handle tool calls that do not require approval', async () => {
|
||||
// First call yields the tool request
|
||||
sendMessageStreamSpy.mockImplementationOnce(async function* () {
|
||||
yield* [
|
||||
{
|
||||
type: GeminiEventType.ToolCallRequest,
|
||||
value: {
|
||||
callId: 'test-call-id-no-approval',
|
||||
name: 'test-tool-no-approval',
|
||||
args: {},
|
||||
},
|
||||
},
|
||||
];
|
||||
});
|
||||
// Second call, after the tool runs, yields the final text
|
||||
sendMessageStreamSpy.mockImplementationOnce(async function* () {
|
||||
yield* [{ type: 'content', value: 'Tool executed successfully.' }];
|
||||
});
|
||||
|
||||
const mockTool = new MockTool({
|
||||
name: 'test-tool-no-approval',
|
||||
displayName: 'Test Tool No Approval',
|
||||
execute: vi.fn().mockResolvedValue({
|
||||
llmContent: 'Tool executed successfully.',
|
||||
returnDisplay: 'Tool executed successfully.',
|
||||
}),
|
||||
});
|
||||
|
||||
getToolRegistrySpy.mockReturnValue({
|
||||
getAllTools: vi.fn().mockReturnValue([mockTool]),
|
||||
getToolsByServer: vi.fn().mockReturnValue([]),
|
||||
getTool: vi.fn().mockReturnValue(mockTool),
|
||||
});
|
||||
|
||||
const agent = request.agent(app);
|
||||
const res = await agent
|
||||
.post('/')
|
||||
.send(
|
||||
createStreamMessageRequest(
|
||||
'run a tool without approval',
|
||||
'a2a-no-approval-test-message',
|
||||
),
|
||||
)
|
||||
.set('Content-Type', 'application/json')
|
||||
.expect(200);
|
||||
|
||||
const events = streamToSSEEvents(res.text);
|
||||
assertTaskCreationAndWorkingStatus(events);
|
||||
|
||||
// Status update: working
|
||||
const workingEvent2 = events[2].result as TaskStatusUpdateEvent;
|
||||
expect(workingEvent2.kind).toBe('status-update');
|
||||
expect(workingEvent2.status.state).toBe('working');
|
||||
|
||||
// Status update: tool-call-update (validating)
|
||||
const validatingEvent = events[3].result as TaskStatusUpdateEvent;
|
||||
expect(validatingEvent.metadata?.['coderAgent']).toMatchObject({
|
||||
kind: 'tool-call-update',
|
||||
});
|
||||
expect(validatingEvent.status.message?.parts).toMatchObject([
|
||||
{
|
||||
data: {
|
||||
status: 'validating',
|
||||
request: { callId: 'test-call-id-no-approval' },
|
||||
},
|
||||
},
|
||||
]);
|
||||
|
||||
// Status update: tool-call-update (scheduled)
|
||||
const scheduledEvent = events[4].result as TaskStatusUpdateEvent;
|
||||
expect(scheduledEvent.metadata?.['coderAgent']).toMatchObject({
|
||||
kind: 'tool-call-update',
|
||||
});
|
||||
expect(scheduledEvent.status.message?.parts).toMatchObject([
|
||||
{
|
||||
data: {
|
||||
status: 'scheduled',
|
||||
request: { callId: 'test-call-id-no-approval' },
|
||||
},
|
||||
},
|
||||
]);
|
||||
|
||||
// Status update: tool-call-update (executing)
|
||||
const executingEvent = events[5].result as TaskStatusUpdateEvent;
|
||||
expect(executingEvent.metadata?.['coderAgent']).toMatchObject({
|
||||
kind: 'tool-call-update',
|
||||
});
|
||||
expect(executingEvent.status.message?.parts).toMatchObject([
|
||||
{
|
||||
data: {
|
||||
status: 'executing',
|
||||
request: { callId: 'test-call-id-no-approval' },
|
||||
},
|
||||
},
|
||||
]);
|
||||
|
||||
// Status update: tool-call-update (success)
|
||||
const successEvent = events[6].result as TaskStatusUpdateEvent;
|
||||
expect(successEvent.metadata?.['coderAgent']).toMatchObject({
|
||||
kind: 'tool-call-update',
|
||||
});
|
||||
expect(successEvent.status.message?.parts).toMatchObject([
|
||||
{
|
||||
data: {
|
||||
status: 'success',
|
||||
request: { callId: 'test-call-id-no-approval' },
|
||||
},
|
||||
},
|
||||
]);
|
||||
|
||||
// Status update: working (before sending tool result to LLM)
|
||||
const workingEvent3 = events[7].result as TaskStatusUpdateEvent;
|
||||
expect(workingEvent3.kind).toBe('status-update');
|
||||
expect(workingEvent3.status.state).toBe('working');
|
||||
|
||||
// Status update: text-content (final LLM response)
|
||||
const textContentEvent = events[8].result as TaskStatusUpdateEvent;
|
||||
expect(textContentEvent.metadata?.['coderAgent']).toMatchObject({
|
||||
kind: 'text-content',
|
||||
});
|
||||
expect(textContentEvent.status.message?.parts).toMatchObject([
|
||||
{ text: 'Tool executed successfully.' },
|
||||
]);
|
||||
|
||||
assertUniqueFinalEventIsLast(events);
|
||||
expect(events.length).toBe(10);
|
||||
});
|
||||
|
||||
it('should bypass tool approval in YOLO mode', async () => {
|
||||
// First call yields the tool request
|
||||
sendMessageStreamSpy.mockImplementationOnce(async function* () {
|
||||
yield* [
|
||||
{
|
||||
type: GeminiEventType.ToolCallRequest,
|
||||
value: {
|
||||
callId: 'test-call-id-yolo',
|
||||
name: 'test-tool-yolo',
|
||||
args: {},
|
||||
},
|
||||
},
|
||||
];
|
||||
});
|
||||
// Second call, after the tool runs, yields the final text
|
||||
sendMessageStreamSpy.mockImplementationOnce(async function* () {
|
||||
yield* [{ type: 'content', value: 'Tool executed successfully.' }];
|
||||
});
|
||||
|
||||
// Set approval mode to yolo
|
||||
getApprovalModeSpy.mockReturnValue(ApprovalMode.YOLO);
|
||||
|
||||
const mockTool = new MockTool({
|
||||
name: 'test-tool-yolo',
|
||||
displayName: 'Test Tool YOLO',
|
||||
execute: vi.fn().mockResolvedValue({
|
||||
llmContent: 'Tool executed successfully.',
|
||||
returnDisplay: 'Tool executed successfully.',
|
||||
}),
|
||||
});
|
||||
|
||||
getToolRegistrySpy.mockReturnValue({
|
||||
getAllTools: vi.fn().mockReturnValue([mockTool]),
|
||||
getToolsByServer: vi.fn().mockReturnValue([]),
|
||||
getTool: vi.fn().mockReturnValue(mockTool),
|
||||
});
|
||||
|
||||
const agent = request.agent(app);
|
||||
const res = await agent
|
||||
.post('/')
|
||||
.send(
|
||||
createStreamMessageRequest(
|
||||
'run a tool in yolo mode',
|
||||
'a2a-yolo-mode-test-message',
|
||||
),
|
||||
)
|
||||
.set('Content-Type', 'application/json')
|
||||
.expect(200);
|
||||
|
||||
const events = streamToSSEEvents(res.text);
|
||||
assertTaskCreationAndWorkingStatus(events);
|
||||
|
||||
// Status update: working
|
||||
const workingEvent2 = events[2].result as TaskStatusUpdateEvent;
|
||||
expect(workingEvent2.kind).toBe('status-update');
|
||||
expect(workingEvent2.status.state).toBe('working');
|
||||
|
||||
// Status update: tool-call-update (validating)
|
||||
const validatingEvent = events[3].result as TaskStatusUpdateEvent;
|
||||
expect(validatingEvent.metadata?.['coderAgent']).toMatchObject({
|
||||
kind: 'tool-call-update',
|
||||
});
|
||||
expect(validatingEvent.status.message?.parts).toMatchObject([
|
||||
{
|
||||
data: {
|
||||
status: 'validating',
|
||||
request: { callId: 'test-call-id-yolo' },
|
||||
},
|
||||
},
|
||||
]);
|
||||
|
||||
// Status update: tool-call-update (scheduled)
|
||||
const awaitingEvent = events[4].result as TaskStatusUpdateEvent;
|
||||
expect(awaitingEvent.metadata?.['coderAgent']).toMatchObject({
|
||||
kind: 'tool-call-update',
|
||||
});
|
||||
expect(awaitingEvent.status.message?.parts).toMatchObject([
|
||||
{
|
||||
data: {
|
||||
status: 'scheduled',
|
||||
request: { callId: 'test-call-id-yolo' },
|
||||
},
|
||||
},
|
||||
]);
|
||||
|
||||
// Status update: tool-call-update (executing)
|
||||
const executingEvent = events[5].result as TaskStatusUpdateEvent;
|
||||
expect(executingEvent.metadata?.['coderAgent']).toMatchObject({
|
||||
kind: 'tool-call-update',
|
||||
});
|
||||
expect(executingEvent.status.message?.parts).toMatchObject([
|
||||
{
|
||||
data: {
|
||||
status: 'executing',
|
||||
request: { callId: 'test-call-id-yolo' },
|
||||
},
|
||||
},
|
||||
]);
|
||||
|
||||
// Status update: tool-call-update (success)
|
||||
const successEvent = events[6].result as TaskStatusUpdateEvent;
|
||||
expect(successEvent.metadata?.['coderAgent']).toMatchObject({
|
||||
kind: 'tool-call-update',
|
||||
});
|
||||
expect(successEvent.status.message?.parts).toMatchObject([
|
||||
{
|
||||
data: {
|
||||
status: 'success',
|
||||
request: { callId: 'test-call-id-yolo' },
|
||||
},
|
||||
},
|
||||
]);
|
||||
|
||||
// Status update: working (before sending tool result to LLM)
|
||||
const workingEvent3 = events[7].result as TaskStatusUpdateEvent;
|
||||
expect(workingEvent3.kind).toBe('status-update');
|
||||
expect(workingEvent3.status.state).toBe('working');
|
||||
|
||||
// Status update: text-content (final LLM response)
|
||||
const textContentEvent = events[8].result as TaskStatusUpdateEvent;
|
||||
expect(textContentEvent.metadata?.['coderAgent']).toMatchObject({
|
||||
kind: 'text-content',
|
||||
});
|
||||
expect(textContentEvent.status.message?.parts).toMatchObject([
|
||||
{ text: 'Tool executed successfully.' },
|
||||
]);
|
||||
|
||||
assertUniqueFinalEventIsLast(events);
|
||||
expect(events.length).toBe(10);
|
||||
});
|
||||
});
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,149 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2025 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { describe, it, expect, beforeAll, afterAll, vi } from 'vitest';
|
||||
import request from 'supertest';
|
||||
import type express from 'express';
|
||||
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 { AddressInfo } from 'node:net';
|
||||
|
||||
// Mock the logger to avoid polluting test output
|
||||
// Comment out to help debug
|
||||
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('../agent/task.js', () => {
|
||||
class MockTask {
|
||||
id: string;
|
||||
contextId: string;
|
||||
taskState = 'submitted';
|
||||
config = {
|
||||
getContentGeneratorConfig: vi
|
||||
.fn()
|
||||
.mockReturnValue({ model: 'gemini-pro' }),
|
||||
};
|
||||
geminiClient = {
|
||||
initialize: vi.fn().mockResolvedValue(undefined),
|
||||
};
|
||||
constructor(id: string, contextId: string) {
|
||||
this.id = id;
|
||||
this.contextId = contextId;
|
||||
}
|
||||
static create = vi
|
||||
.fn()
|
||||
.mockImplementation((id, contextId) =>
|
||||
Promise.resolve(new MockTask(id, contextId)),
|
||||
);
|
||||
getMetadata = vi.fn().mockImplementation(async () => ({
|
||||
id: this.id,
|
||||
contextId: this.contextId,
|
||||
taskState: this.taskState,
|
||||
model: 'gemini-pro',
|
||||
mcpServers: [],
|
||||
availableTools: [],
|
||||
}));
|
||||
}
|
||||
return { Task: MockTask };
|
||||
});
|
||||
|
||||
describe('Agent Server Endpoints', () => {
|
||||
let app: express.Express;
|
||||
let server: Server;
|
||||
let testWorkspace: string;
|
||||
|
||||
const createTask = (contextId: string) =>
|
||||
request(app)
|
||||
.post('/tasks')
|
||||
.send({
|
||||
contextId,
|
||||
agentSettings: {
|
||||
kind: 'agent-settings',
|
||||
workspacePath: testWorkspace,
|
||||
},
|
||||
})
|
||||
.set('Content-Type', 'application/json');
|
||||
|
||||
beforeAll(async () => {
|
||||
// Create a unique temporary directory for the workspace to avoid conflicts
|
||||
testWorkspace = fs.mkdtempSync(
|
||||
path.join(os.tmpdir(), 'gemini-agent-test-'),
|
||||
);
|
||||
app = await createApp();
|
||||
await new Promise<void>((resolve) => {
|
||||
server = app.listen(0, () => {
|
||||
const port = (server.address() as AddressInfo).port;
|
||||
updateCoderAgentCardUrl(port);
|
||||
resolve();
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
afterAll(async () => {
|
||||
if (server) {
|
||||
await new Promise<void>((resolve, reject) => {
|
||||
server.close((err) => {
|
||||
if (err) return reject(err);
|
||||
resolve();
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
if (testWorkspace) {
|
||||
try {
|
||||
fs.rmSync(testWorkspace, { recursive: true, force: true });
|
||||
} catch (e) {
|
||||
console.warn(`Could not remove temp dir '${testWorkspace}':`, e);
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
it('should create a new task via POST /tasks', async () => {
|
||||
const response = await createTask('test-context');
|
||||
expect(response.status).toBe(201);
|
||||
expect(response.body).toBeTypeOf('string'); // Should return the task ID
|
||||
}, 7000);
|
||||
|
||||
it('should get metadata for a specific task via GET /tasks/:taskId/metadata', async () => {
|
||||
const createResponse = await createTask('test-context-2');
|
||||
const taskId = createResponse.body;
|
||||
const response = await request(app).get(`/tasks/${taskId}/metadata`);
|
||||
expect(response.status).toBe(200);
|
||||
expect(response.body.metadata.id).toBe(taskId);
|
||||
}, 6000);
|
||||
|
||||
it('should get metadata for all tasks via GET /tasks/metadata', async () => {
|
||||
const createResponse = await createTask('test-context-3');
|
||||
const taskId = createResponse.body;
|
||||
const response = await request(app).get('/tasks/metadata');
|
||||
expect(response.status).toBe(200);
|
||||
expect(Array.isArray(response.body)).toBe(true);
|
||||
expect(response.body.length).toBeGreaterThan(0);
|
||||
const taskMetadata = response.body.find(
|
||||
(m: TaskMetadata) => m.id === taskId,
|
||||
);
|
||||
expect(taskMetadata).toBeDefined();
|
||||
});
|
||||
|
||||
it('should return 404 for a non-existent task', async () => {
|
||||
const response = await request(app).get('/tasks/fake-task/metadata');
|
||||
expect(response.status).toBe(404);
|
||||
});
|
||||
|
||||
it('should return agent metadata via GET /.well-known/agent-card.json', async () => {
|
||||
const response = await request(app).get('/.well-known/agent-card.json');
|
||||
const port = (server.address() as AddressInfo).port;
|
||||
expect(response.status).toBe(200);
|
||||
expect(response.body.name).toBe('Gemini SDLC Agent');
|
||||
expect(response.body.url).toBe(`http://localhost:${port}/`);
|
||||
});
|
||||
});
|
||||
@@ -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 }>();
|
||||
@@ -0,0 +1,33 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2025 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import * as url from 'node:url';
|
||||
import * as path from 'node:path';
|
||||
|
||||
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.
|
||||
const isMainModule =
|
||||
path.resolve(process.argv[1]) ===
|
||||
path.resolve(url.fileURLToPath(import.meta.url));
|
||||
|
||||
process.on('uncaughtException', (error) => {
|
||||
logger.error('Unhandled exception:', error);
|
||||
process.exit(1);
|
||||
});
|
||||
|
||||
if (
|
||||
import.meta.url.startsWith('file:') &&
|
||||
isMainModule &&
|
||||
process.env['NODE_ENV'] !== 'test'
|
||||
) {
|
||||
main().catch((error) => {
|
||||
logger.error('[CoreAgent] Unhandled error in main:', error);
|
||||
process.exit(1);
|
||||
});
|
||||
}
|
||||
Reference in New Issue
Block a user