@dexwox-labs/a2a-server
Version:
TypeScript server implementation for Google's Agent-to-Agent (A2A) protocol - includes Express/WebSocket handlers, request validation and queue management
355 lines (337 loc) • 11.5 kB
text/typescript
/**
* @module TaskManager
* @description Task lifecycle and state management for the A2A server
*
* This module provides a task manager implementation that handles creating,
* updating, and managing tasks throughout their lifecycle. It also provides
* support for task artifacts and push notifications.
*/
import { Task, TaskState, Artifact, TaskNotFoundError, InvalidTaskStateError } from '@dexwox-labs/a2a-core';
import { TaskStore } from './task-store';
import { PushNotificationService } from '../push-notifications/push-service';
import { ResultAggregator } from '../agent-execution/result-aggregator';
import logger from '../utils/logger';
/**
* Manages task lifecycle and state
*
* The TaskManager is responsible for creating, updating, and managing tasks
* throughout their lifecycle. It provides methods for task creation, status
* updates, cancellation, and artifact management.
*
* @example
* ```typescript
* // Create a task manager with in-memory storage
* const taskStore = new InMemoryTaskStore();
* const taskManager = new TaskManager(taskStore);
*
* // Create a new task
* const task = await taskManager.createTask({
* name: 'ProcessData',
* agentId: 'data-processor',
* parts: [{ type: 'text', content: 'Process this data', format: 'plain' }],
* expectedParts: 1,
* createdAt: new Date().toISOString(),
* updatedAt: new Date().toISOString()
* });
*
* console.log('Created task:', task.id);
* ```
*/
export class TaskManager {
/** Map of task aggregators for collecting and processing task results */
private readonly aggregators = new Map<string, ResultAggregator>();
/**
* Creates a new TaskManager
*
* @param taskStore - Storage backend for tasks
* @param pushService - Optional service for push notifications
*/
constructor(
private readonly taskStore: TaskStore,
private readonly pushService?: PushNotificationService
) {
logger.info('TaskManager initialized', { hasPushService: !!pushService });
}
/**
* Creates a new task
*
* Creates a task with the provided parameters and initializes a result
* aggregator if expectedParts is specified.
*
* @param task - Task parameters (without id and status)
* @returns Promise resolving to the created task
*
* @example
* ```typescript
* const task = await taskManager.createTask({
* name: 'ProcessImage',
* agentId: 'image-processor',
* parts: [{
* type: 'file',
* content: 'base64-encoded-image-data',
* mimeType: 'image/jpeg',
* name: 'image.jpg'
* }],
* expectedParts: 2,
* createdAt: new Date().toISOString(),
* updatedAt: new Date().toISOString()
* });
* ```
*/
async createTask(task: Omit<Task, 'id'|'status'>): Promise<Task> {
logger.debug('Creating new task', { taskParams: task });
const createdTask = await this.taskStore.createTask({
...task,
status: 'submitted' as const,
parts: task.parts ?? [] // Ensure parts array exists
});
logger.info('Successfully created task', { taskId: createdTask.id });
// Initialize result aggregation if expectedParts is specified
if (task.expectedParts) {
logger.debug('Initializing result aggregator', { taskId: createdTask.id });
this.aggregators.set(createdTask.id, new ResultAggregator(createdTask));
}
return createdTask;
}
/**
* Gets the result aggregator for a task
*
* Result aggregators collect and process task results, including
* handling streaming responses and combining multiple parts.
*
* @param taskId - ID of the task
* @returns The result aggregator for the task, or undefined if none exists
*/
getAggregator(taskId: string): ResultAggregator | undefined {
return this.aggregators.get(taskId);
}
/**
* Gets a task by ID
*
* Retrieves a task from the task store by its ID.
*
* @param taskId - ID of the task to retrieve
* @returns Promise resolving to the task
* @throws {TaskNotFoundError} If the task is not found
*
* @example
* ```typescript
* try {
* const task = await taskManager.getTask('task-123');
* console.log('Task status:', task.status);
* } catch (error) {
* if (error instanceof TaskNotFoundError) {
* console.error('Task not found');
* } else {
* console.error('Error retrieving task:', error);
* }
* }
* ```
*/
async getTask(taskId: string): Promise<Task> {
logger.debug('Getting task', { taskId });
const task = await this.taskStore.getTask(taskId);
if (!task) {
logger.error('Task not found', { taskId });
throw new TaskNotFoundError(taskId);
}
logger.debug('Successfully retrieved task', { taskId });
return task;
}
/**
* Updates a task with new data
*
* Updates a task with the provided partial task data and sets
* the updatedAt timestamp to the current time.
*
* @param taskId - ID of the task to update
* @param updates - Partial task data to update
* @returns Promise resolving to the updated task
* @throws {TaskNotFoundError} If the task is not found
*
* @example
* ```typescript
* const updatedTask = await taskManager.updateTask('task-123', {
* name: 'Updated Task Name',
* metadata: { priority: 'high' }
* });
* ```
*/
async updateTask(taskId: string, updates: Partial<Omit<Task, 'id'>>): Promise<Task> {
const currentTask = await this.getTask(taskId);
return this.taskStore.updateTask(taskId, {
...currentTask,
...updates,
updatedAt: new Date().toISOString()
});
}
/**
* Updates a task's status
*
* Updates the status of a task and sets the updatedAt timestamp
* to the current time.
*
* @param taskId - ID of the task to update
* @param status - New status for the task
* @returns Promise resolving to the updated task
* @throws {TaskNotFoundError} If the task is not found
*
* @example
* ```typescript
* // Update task status to 'completed'
* const task = await taskManager.updateTaskStatus('task-123', 'completed');
* console.log('Task status updated:', task.status);
* ```
*/
async updateTaskStatus(taskId: string, status: TaskState): Promise<Task> {
logger.info('Updating task status', { taskId, newStatus: status });
const updatedTask = await this.updateTask(taskId, { status });
logger.info('Successfully updated task status', { taskId });
return updatedTask;
}
/**
* Cancels a task
*
* Cancels a task that is not already in a terminal state (completed, failed, or canceled).
* If a push service is configured, it will send a push notification about the cancellation.
*
* @param taskId - ID of the task to cancel
* @returns Promise resolving when the task is canceled
* @throws {TaskNotFoundError} If the task is not found
* @throws {InvalidTaskStateError} If the task is already in a terminal state
*
* @example
* ```typescript
* try {
* await taskManager.cancelTask('task-123');
* console.log('Task canceled successfully');
* } catch (error) {
* if (error instanceof InvalidTaskStateError) {
* console.error('Cannot cancel task in terminal state');
* } else {
* console.error('Error canceling task:', error);
* }
* }
* ```
*/
async cancelTask(taskId: string): Promise<void> {
logger.info('Attempting to cancel task', { taskId });
const task = await this.getTask(taskId);
if (['completed', 'failed', 'canceled'].includes(task.status)) {
logger.error('Cannot cancel task in terminal state', { taskId, currentStatus: task.status });
throw new InvalidTaskStateError(
`Task ${taskId} is already in terminal state: ${task.status}`
);
}
await this.taskStore.cancelTask(taskId);
logger.info('Successfully canceled task', { taskId });
if (this.pushService) {
logger.debug('Sending push notification for canceled task', { taskId });
await this.pushService.notifyStatusChange(taskId, 'canceled');
}
}
/**
* Lists all active tasks
*
* Retrieves all tasks that are not in a terminal state (completed, failed, or canceled).
* Active tasks include those with status 'submitted', 'working', or 'input_required'.
*
* @returns Promise resolving to an array of active tasks
*
* @example
* ```typescript
* const activeTasks = await taskManager.listActiveTasks();
* console.log(`Found ${activeTasks.length} active tasks`);
*
* // Process each active task
* for (const task of activeTasks) {
* console.log(`Task ${task.id} is ${task.status}`);
* }
* ```
*/
async listActiveTasks(): Promise<Task[]> {
logger.debug('Listing active tasks');
const tasks = await this.taskStore.getTasksByStatus([
'submitted', 'working', 'input_required'
]);
logger.debug('Found active tasks', { count: tasks.length });
return tasks;
}
/**
* Adds an artifact to a task
*
* Artifacts represent files, data, or other content associated with a task.
* This method adds an artifact to a specific task.
*
* @param taskId - ID of the task
* @param artifact - Artifact to add
* @returns Promise resolving when the artifact is added
* @throws {TaskNotFoundError} If the task is not found
*
* @example
* ```typescript
* await taskManager.addArtifact('task-123', {
* id: 'artifact-456',
* type: 'file',
* content: 'base64-encoded-file-data',
* createdAt: new Date().toISOString(),
* updatedAt: new Date().toISOString(),
* metadata: {
* filename: 'report.pdf',
* mimeType: 'application/pdf'
* }
* });
* ```
*/
async addArtifact(taskId: string, artifact: Artifact): Promise<void> {
await this.taskStore.addArtifact(taskId, artifact);
}
/**
* Gets all artifacts for a task
*
* Retrieves all artifacts associated with a specific task.
*
* @param taskId - ID of the task
* @returns Promise resolving to an array of artifacts
* @throws {TaskNotFoundError} If the task is not found
*
* @example
* ```typescript
* const artifacts = await taskManager.getArtifacts('task-123');
* console.log(`Task has ${artifacts.length} artifacts`);
*
* // Process each artifact
* for (const artifact of artifacts) {
* console.log(`Artifact ${artifact.id} of type ${artifact.type}`);
* }
* ```
*/
async getArtifacts(taskId: string): Promise<Artifact[]> {
return this.taskStore.getArtifacts(taskId);
}
/**
* Gets a specific artifact for a task
*
* Retrieves a specific artifact by ID from a specific task.
*
* @param taskId - ID of the task
* @param artifactId - ID of the artifact
* @returns Promise resolving to the artifact, or null if not found
* @throws {TaskNotFoundError} If the task is not found
*
* @example
* ```typescript
* const artifact = await taskManager.getArtifact('task-123', 'artifact-456');
* if (artifact) {
* console.log('Found artifact:', artifact.id);
* console.log('Type:', artifact.type);
* console.log('Created at:', artifact.createdAt);
* } else {
* console.log('Artifact not found');
* }
* ```
*/
async getArtifact(taskId: string, artifactId: string): Promise<Artifact | null> {
return this.taskStore.getArtifact(taskId, artifactId);
}
}