nestjs-a2a
Version:
NestJS module for creating Google Agent to Agent Server
613 lines (612 loc) • 32.7 kB
JavaScript
"use strict";
var __decorate = (this && this.__decorate) || function (decorators, target, key, desc) {
var c = arguments.length, r = c < 3 ? target : desc === null ? desc = Object.getOwnPropertyDescriptor(target, key) : desc, d;
if (typeof Reflect === "object" && typeof Reflect.decorate === "function") r = Reflect.decorate(decorators, target, key, desc);
else for (var i = decorators.length - 1; i >= 0; i--) if (d = decorators[i]) r = (c < 3 ? d(r) : c > 3 ? d(target, key, r) : d(target, key)) || r;
return c > 3 && r && Object.defineProperty(target, key, r), r;
};
var __metadata = (this && this.__metadata) || function (k, v) {
if (typeof Reflect === "object" && typeof Reflect.metadata === "function") return Reflect.metadata(k, v);
};
var __param = (this && this.__param) || function (paramIndex, decorator) {
return function (target, key) { decorator(target, key, paramIndex); }
};
var __awaiter = (this && this.__awaiter) || function (thisArg, _arguments, P, generator) {
function adopt(value) { return value instanceof P ? value : new P(function (resolve) { resolve(value); }); }
return new (P || (P = Promise))(function (resolve, reject) {
function fulfilled(value) { try { step(generator.next(value)); } catch (e) { reject(e); } }
function rejected(value) { try { step(generator["throw"](value)); } catch (e) { reject(e); } }
function step(result) { result.done ? resolve(result.value) : adopt(result.value).then(fulfilled, rejected); }
step((generator = generator.apply(thisArg, _arguments || [])).next());
});
};
var __asyncValues = (this && this.__asyncValues) || function (o) {
if (!Symbol.asyncIterator) throw new TypeError("Symbol.asyncIterator is not defined.");
var m = o[Symbol.asyncIterator], i;
return m ? m.call(o) : (o = typeof __values === "function" ? __values(o) : o[Symbol.iterator](), i = {}, verb("next"), verb("throw"), verb("return"), i[Symbol.asyncIterator] = function () { return this; }, i);
function verb(n) { i[n] = o[n] && function (v) { return new Promise(function (resolve, reject) { v = o[n](v), settle(resolve, reject, v.done, v.value); }); }; }
function settle(resolve, reject, d, v) { Promise.resolve(v).then(function(v) { resolve({ value: v, done: d }); }, reject); }
};
var A2AExecutor_1;
Object.defineProperty(exports, "__esModule", { value: true });
exports.A2AExecutor = void 0;
const common_1 = require("@nestjs/common");
const core_1 = require("@nestjs/core");
const constant_1 = require("../constant");
const a2a_exception_1 = require("../interfaces/a2a.exception");
const a2a_store_1 = require("../interfaces/a2a.store");
const a2a_types_1 = require("../interfaces/a2a.types");
const a2a_registry_1 = require("./a2a.registry");
/**
* Implements the A2A executor for handling A2A protocol requests in NestJS.
*/
let A2AExecutor = A2AExecutor_1 = class A2AExecutor {
constructor(options, request, registry, moduleRef) {
var _a;
this.options = options;
this.request = request;
this.registry = registry;
this.moduleRef = moduleRef;
this.logger = new common_1.Logger(A2AExecutor_1.name);
this.activeCancellations = new Set();
// Initialize the task store from options or use in-memory store as fallback
this.taskStore = (_a = this.options.taskStore) !== null && _a !== void 0 ? _a : new a2a_store_1.InMemoryTaskStore();
}
/**
* Handles a task send request (non-streaming).
*/
handleTaskSend(body) {
return __awaiter(this, void 0, void 0, function* () {
var _a, e_1, _b, _c;
try {
this.validateTaskSendParams(body.params);
const { id: taskId, message, sessionId, metadata } = body.params;
// Load or create task AND history
let currentData = yield this.loadOrCreateTaskAndHistory(taskId, message, sessionId, metadata);
// Find the appropriate handler for this task
const handler = yield this.findHandlerForTask(currentData.task, message);
if (!handler) {
throw a2a_exception_1.A2AException.unsupportedOperation(`No handler found for task ${taskId}`);
}
// Create context and run handler
const context = this.createTaskContext(currentData.task, message, currentData.history);
const generator = handler(context);
// Process generator yields
try {
try {
for (var _d = true, generator_1 = __asyncValues(generator), generator_1_1; generator_1_1 = yield generator_1.next(), _a = generator_1_1.done, !_a; _d = true) {
_c = generator_1_1.value;
_d = false;
const yieldValue = _c;
// Apply update immutably
currentData = this.applyUpdateToTaskAndHistory(currentData, yieldValue);
// Save the updated state
yield this.taskStore.save(currentData);
// Update context snapshot for next iteration
context.task = currentData.task;
}
}
catch (e_1_1) { e_1 = { error: e_1_1 }; }
finally {
try {
if (!_d && !_a && (_b = generator_1.return)) yield _b.call(generator_1);
}
finally { if (e_1) throw e_1.error; }
}
}
catch (handlerError) {
// If handler throws, apply 'failed' status, save, and rethrow
const failureStatusUpdate = {
state: a2a_types_1.TaskState.FAILED,
message: {
role: a2a_types_1.MessageRole.AGENT,
parts: [
{
text: `Handler failed: ${handlerError instanceof Error
? handlerError.message
: String(handlerError)}`,
},
],
},
};
currentData = this.applyUpdateToTaskAndHistory(currentData, failureStatusUpdate);
try {
yield this.taskStore.save(currentData);
}
catch (saveError) {
this.logger.debug(`Failed to save task ${taskId} after handler error:`, saveError);
throw a2a_exception_1.A2AException.internalError(handlerError.message, {
stack: handlerError.stack,
});
}
throw a2a_exception_1.A2AException.internalError(handlerError.message, {
stack: handlerError.stack,
});
}
// The loop finished, return the final task state
return a2a_types_1.JSONRPCResponse.success(currentData.task);
}
catch (error) {
return this.handleError(error, body.id);
}
});
}
/**
* Handles a task send request with streaming response.
*/
handleTaskSendSubscribe(body, response) {
return __awaiter(this, void 0, void 0, function* () {
var _a, e_2, _b, _c;
try {
this.validateTaskSendParams(body.params);
const { id: taskId, message, sessionId, metadata } = body.params;
// Load or create task AND history
let currentData = yield this.loadOrCreateTaskAndHistory(taskId, message, sessionId, metadata);
// Find the appropriate handler for this task
const handler = yield this.findHandlerForTask(currentData.task, message);
if (!handler) {
throw a2a_exception_1.A2AException.unsupportedOperation(`No handler found for task ${taskId}`);
}
// Create context and run handler
const context = this.createTaskContext(currentData.task, message, currentData.history);
const generator = handler(context);
// --- Setup SSE ---
response.setHeader('Content-Type', 'text/event-stream');
response.setHeader('Cache-Control', 'no-cache');
response.setHeader('Connection', 'keep-alive');
// Function to send SSE data
const sendEvent = (eventData) => {
response.write(`data: ${JSON.stringify(eventData)}\n\n`);
};
let lastEventWasFinal = false; // Track if the last sent event was marked final
try {
try {
// Process generator yields
for (var _d = true, generator_2 = __asyncValues(generator), generator_2_1; generator_2_1 = yield generator_2.next(), _a = generator_2_1.done, !_a; _d = true) {
_c = generator_2_1.value;
_d = false;
const yieldValue = _c;
// Apply update immutably
currentData = this.applyUpdateToTaskAndHistory(currentData, yieldValue);
// Save the updated state
yield this.taskStore.save(currentData);
// Update context snapshot
context.task = currentData.task;
let event;
let isFinal = false;
// Determine event type and check for final state based on the *updated* task
if (this.isTaskStatusUpdate(yieldValue)) {
const terminalStates = [
a2a_types_1.TaskState.COMPLETED,
a2a_types_1.TaskState.FAILED,
a2a_types_1.TaskState.CANCELED,
a2a_types_1.TaskState.INPUT_REQUIRED, // Treat input-required as potentially final for streaming?
];
isFinal = terminalStates.includes(currentData.task.status.state);
event = this.createTaskStatusEvent(taskId, currentData.task.status, isFinal);
if (isFinal) {
this.logger.debug(`[SSE ${taskId}] Yielded terminal state ${currentData.task.status.state}, marking event as final.`);
}
}
else {
// It's an artifact update
// Find the updated artifact in the new task object
// Commented out to avoid unused variable warning
// const updatedArtifact =
// currentData.task.artifacts?.find(
// (a) =>
// (a.index !== undefined && a.index === (yieldValue as Artifact).index) ||
// (a.name && a.name === (yieldValue as Artifact).name),
// ) ?? (yieldValue as Artifact); // Fallback
// For now, we'll just send status updates, not artifact updates
// This could be enhanced to support artifact streaming in the future
event = this.createTaskStatusEvent(taskId, currentData.task.status, false);
}
sendEvent(a2a_types_1.JSONRPCResponse.success(event));
lastEventWasFinal = isFinal;
// If the status update resulted in a final state, stop processing
if (isFinal)
break;
}
}
catch (e_2_1) { e_2 = { error: e_2_1 }; }
finally {
try {
if (!_d && !_a && (_b = generator_2.return)) yield _b.call(generator_2);
}
finally { if (e_2) throw e_2.error; }
}
// Loop finished. Check if a final event was already sent.
if (!lastEventWasFinal) {
this.logger.debug(`[SSE ${taskId}] Handler finished without yielding terminal state. Sending final state: ${currentData.task.status.state}`);
// Ensure the task is actually in a recognized final state before sending.
const finalStates = [
a2a_types_1.TaskState.COMPLETED,
a2a_types_1.TaskState.FAILED,
a2a_types_1.TaskState.CANCELED,
a2a_types_1.TaskState.INPUT_REQUIRED, // Consider input-required final for SSE end?
];
if (!finalStates.includes(currentData.task.status.state)) {
this.logger.warn(`[SSE ${taskId}] Task ended non-terminally (${currentData.task.status.state}). Forcing 'completed'.`);
// Apply 'completed' state update
currentData = this.applyUpdateToTaskAndHistory(currentData, {
state: a2a_types_1.TaskState.COMPLETED,
});
// Save the forced final state
yield this.taskStore.save(currentData);
}
// Send the final status event
const finalEvent = this.createTaskStatusEvent(taskId, currentData.task.status, true);
sendEvent(a2a_types_1.JSONRPCResponse.success(finalEvent));
}
}
catch (handlerError) {
// Handler threw an error
this.logger.debug(`[SSE ${taskId}] Handler error during streaming:`, handlerError);
// Apply 'failed' status update
const failureUpdate = {
state: a2a_types_1.TaskState.FAILED,
message: {
role: a2a_types_1.MessageRole.AGENT,
parts: [
{
text: `Handler failed: ${handlerError instanceof Error
? handlerError.message
: String(handlerError)}`,
},
],
},
};
currentData = this.applyUpdateToTaskAndHistory(currentData, failureUpdate);
try {
// Save the failed state
yield this.taskStore.save(currentData);
}
catch (saveError) {
this.logger.debug(`[SSE ${taskId}] Failed to save task after handler error:`, saveError);
}
// Send final error status event via SSE
const errorEvent = this.createTaskStatusEvent(taskId, currentData.task.status, // Use the updated status
true);
sendEvent(a2a_types_1.JSONRPCResponse.success(errorEvent));
}
finally {
// End the SSE stream if it hasn't already been closed by sending a final event
if (!response.writableEnded) {
response.end();
}
}
}
catch (error) {
// Handle initial setup errors (before streaming starts)
const errorResponse = this.handleError(error, body.id);
response.status(200).json(errorResponse);
}
});
}
/**
* Handles a task get request.
*/
handleTaskGet(body) {
return __awaiter(this, void 0, void 0, function* () {
try {
const { id: taskId } = body.params;
if (!taskId)
throw a2a_exception_1.A2AException.invalidParams('Missing task ID.');
// Load both task and history
const data = yield this.taskStore.load(taskId);
if (!data) {
throw a2a_exception_1.A2AException.taskNotFound(taskId);
}
// Return only the task object as per spec
return a2a_types_1.JSONRPCResponse.success(data.task);
}
catch (error) {
return this.handleError(error, body.id);
}
});
}
/**
* Handles a task cancel request.
*/
handleTaskCancel(body) {
return __awaiter(this, void 0, void 0, function* () {
try {
const { id: taskId } = body.params;
if (!taskId)
throw a2a_exception_1.A2AException.invalidParams('Missing task ID.');
// Load task and history
let data = yield this.taskStore.load(taskId);
if (!data) {
throw a2a_exception_1.A2AException.taskNotFound(taskId);
}
// Check if cancelable (not already in a final state)
const finalStates = [
a2a_types_1.TaskState.COMPLETED,
a2a_types_1.TaskState.FAILED,
a2a_types_1.TaskState.CANCELED,
];
if (finalStates.includes(data.task.status.state)) {
this.logger.debug(`Task ${taskId} already in final state ${data.task.status.state}, cannot cancel.`);
return a2a_types_1.JSONRPCResponse.success(data.task); // Return current state
}
// Signal cancellation
this.activeCancellations.add(taskId);
// Apply 'canceled' state update
const cancelUpdate = {
state: a2a_types_1.TaskState.CANCELED,
message: {
role: a2a_types_1.MessageRole.AGENT,
parts: [{ text: 'Task cancelled by request.' }],
},
};
data = this.applyUpdateToTaskAndHistory(data, cancelUpdate);
// Save the updated state
yield this.taskStore.save(data);
// Remove from active cancellations *after* saving
this.activeCancellations.delete(taskId);
// Return the updated task object
return a2a_types_1.JSONRPCResponse.success(data.task);
}
catch (error) {
return this.handleError(error, body.id);
}
});
}
/**
* Handles a task resubscribe request.
* Not fully implemented yet - returns unsupported operation error.
*/
handleTaskResubscribe(_body) {
return __awaiter(this, void 0, void 0, function* () {
return a2a_types_1.JSONRPCResponse.error(a2a_exception_1.A2AException.unsupportedOperation('Task resubscription not supported').toJSONRPCError());
});
}
/**
* Handles a task get push notification request.
* Not fully implemented yet - returns unsupported operation error.
*/
handleTaskGetPushNotification(_body) {
return __awaiter(this, void 0, void 0, function* () {
return a2a_types_1.JSONRPCResponse.error(a2a_exception_1.A2AException.pushNotificationNotSupported().toJSONRPCError());
});
}
// --- Helper Methods ---
/**
* Loads an existing task or creates a new one with the given message.
*/
loadOrCreateTaskAndHistory(taskId, initialMessage, sessionId, // Allow null
metadata) {
return __awaiter(this, void 0, void 0, function* () {
let data = yield this.taskStore.load(taskId);
let needsSave = false;
if (!data) {
// Create new task and history
const initialTask = {
id: taskId,
sessionId: sessionId !== null && sessionId !== void 0 ? sessionId : undefined, // Store undefined if null
status: {
state: a2a_types_1.TaskState.SUBMITTED, // Start as submitted
timestamp: this.getCurrentTimestamp(),
message: null, // Initial user message goes only to history for now
},
artifacts: [],
metadata: metadata !== null && metadata !== void 0 ? metadata : undefined, // Store undefined if null
};
const initialHistory = [initialMessage]; // History starts with user message
data = { task: initialTask, history: initialHistory };
needsSave = true; // Mark for saving
this.logger.debug(`[Task ${taskId}] Created new task and history.`);
}
else {
this.logger.debug(`[Task ${taskId}] Loaded existing task and history.`);
// Add current user message to history
// Make a copy before potentially modifying
data = { task: data.task, history: [...data.history, initialMessage] };
needsSave = true; // History updated, mark for saving
// Handle state transitions for existing tasks
const finalStates = [
a2a_types_1.TaskState.COMPLETED,
a2a_types_1.TaskState.FAILED,
a2a_types_1.TaskState.CANCELED,
];
if (finalStates.includes(data.task.status.state)) {
this.logger.warn(`[Task ${taskId}] Received message for task already in final state ${data.task.status.state}. Handling as new submission (keeping history).`);
// Option 1: Reset state to 'submitted' (keeps history, effectively restarts)
const resetUpdate = {
state: a2a_types_1.TaskState.SUBMITTED,
message: null, // Clear old agent message
};
data = this.applyUpdateToTaskAndHistory(data, resetUpdate);
// needsSave is already true
}
else if (data.task.status.state === a2a_types_1.TaskState.INPUT_REQUIRED) {
this.logger.debug(`[Task ${taskId}] Received message while 'input-required', changing state to 'working'.`);
// If it was waiting for input, update state to 'working'
const workingUpdate = {
state: a2a_types_1.TaskState.WORKING,
};
data = this.applyUpdateToTaskAndHistory(data, workingUpdate);
// needsSave is already true
}
else if (data.task.status.state === a2a_types_1.TaskState.WORKING) {
// If already working, maybe warn but allow? Or force back to submitted?
this.logger.warn(`[Task ${taskId}] Received message while already 'working'. Proceeding.`);
// No state change needed, but history was updated, so needsSave is true.
}
// If 'submitted', receiving another message might be odd, but proceed.
}
// Save if created or modified before returning
if (needsSave) {
yield this.taskStore.save(data);
}
// Return copies to prevent mutation by caller before handler runs
return { task: Object.assign({}, data.task), history: [...data.history] };
});
}
/**
* Creates a task context object for the handler.
*/
createTaskContext(task, userMessage, history) {
return {
task: Object.assign({}, task), // Pass a copy
userMessage: userMessage,
history: [...history], // Pass a copy of the history
isCancelled: () => this.activeCancellations.has(task.id),
};
}
/**
* Finds an appropriate handler for the given task and message.
* Currently just returns a default handler, but could be enhanced to select
* based on task type, message content, etc.
*/
findHandlerForTask(task, _message) {
return __awaiter(this, void 0, void 0, function* () {
var _a, _b, _c, _d, _e, _f;
try {
const preferredSkill = (_e = (_b = (_a = task.metadata) === null || _a === void 0 ? void 0 : _a.preferredSkill) !== null && _b !== void 0 ? _b : (yield ((_d = (_c = this.options).selectSkill) === null || _d === void 0 ? void 0 : _d.call(_c, task)))) !== null && _e !== void 0 ? _e : '';
/**
* Get the skill from the registry.
* If no preferred skill is set, use the fallback skill.
* If no fallback skill is set, use the first skill in the registry.
*/
const skill = (_f = this.registry.getSkill(preferredSkill)) !== null && _f !== void 0 ? _f : this.registry.getFallbackSkill();
if (!skill) {
this.logger.warn(`Skill ${preferredSkill} not found in the A2A registry`);
throw a2a_exception_1.A2AException.unsupportedOperation(`Skill ${preferredSkill} not found in the A2A registry`);
}
const contextId = core_1.ContextIdFactory.getByRequest(this.request);
this.moduleRef.registerRequestByContextId(this.request, contextId);
const instance = yield this.moduleRef.resolve(skill.providerClass, contextId, { strict: false });
return instance[skill.methodName];
}
catch (error) {
this.logger.debug('Error finding handler for task:', error);
throw a2a_exception_1.A2AException.unsupportedOperation(`Error finding handler for task: ${error}`);
}
});
}
/**
* Validates the parameters for a task send request.
*/
validateTaskSendParams(params) {
if (!params || typeof params !== 'object') {
throw a2a_exception_1.A2AException.invalidParams('Missing or invalid params object.');
}
if (typeof params.id !== 'string' || params.id === '') {
throw a2a_exception_1.A2AException.invalidParams('Invalid or missing task ID (params.id).');
}
if (!params.message ||
typeof params.message !== 'object' ||
!Array.isArray(params.message.parts)) {
throw a2a_exception_1.A2AException.invalidParams('Invalid or missing message object (params.message).');
}
// Add more checks for message structure, sessionID, metadata, etc. if needed
}
/**
* Helper to apply updates (status or artifact) immutably to a task and history.
*/
applyUpdateToTaskAndHistory(current, update) {
var _a, _b;
const newTask = Object.assign({}, current.task); // Shallow copy task
const newHistory = [...current.history]; // Shallow copy history
if (this.isTaskStatusUpdate(update)) {
// Merge status update
newTask.status = Object.assign(Object.assign(Object.assign({}, newTask.status), update), { timestamp: this.getCurrentTimestamp() });
// If the update includes an agent message, add it to history
if (((_a = update.message) === null || _a === void 0 ? void 0 : _a.role) === a2a_types_1.MessageRole.AGENT) {
newHistory.push(update.message);
}
}
else {
// Handle artifact update
const artifact = update;
if (!newTask.artifacts) {
newTask.artifacts = [];
}
else {
// Ensure we're working with a copy of the artifacts array
newTask.artifacts = [...newTask.artifacts];
}
const existingIndex = (_b = artifact.index) !== null && _b !== void 0 ? _b : -1; // Use index if provided
let replaced = false;
if (existingIndex >= 0 && existingIndex < newTask.artifacts.length) {
const existingArtifact = newTask.artifacts[existingIndex];
if (artifact.append) {
// Create a deep copy for modification to avoid mutating original
const appendedArtifact = JSON.parse(JSON.stringify(existingArtifact));
appendedArtifact.parts.push(...artifact.parts);
if (artifact.metadata) {
appendedArtifact.metadata = Object.assign(Object.assign({}, (appendedArtifact.metadata || {})), artifact.metadata);
}
if (artifact.lastChunk !== undefined)
appendedArtifact.lastChunk = artifact.lastChunk;
if (artifact.description)
appendedArtifact.description = artifact.description;
newTask.artifacts[existingIndex] = appendedArtifact; // Replace with appended version
replaced = true;
}
else {
// Overwrite artifact at index (with a copy of the update)
newTask.artifacts[existingIndex] = Object.assign({}, artifact);
replaced = true;
}
}
else if (artifact.name) {
const namedIndex = newTask.artifacts.findIndex((a) => a.name === artifact.name);
if (namedIndex >= 0) {
newTask.artifacts[namedIndex] = Object.assign({}, artifact); // Replace by name (with copy)
replaced = true;
}
}
if (!replaced) {
newTask.artifacts.push(Object.assign({}, artifact)); // Add as a new artifact (copy)
// Sort if indices are present
if (newTask.artifacts.some((a) => a.index !== undefined)) {
newTask.artifacts.sort((a, b) => { var _a, _b; return ((_a = a.index) !== null && _a !== void 0 ? _a : 0) - ((_b = b.index) !== null && _b !== void 0 ? _b : 0); });
}
}
}
return { task: newTask, history: newHistory };
}
/**
* Type guard to check if an update is a task status update.
*/
isTaskStatusUpdate(update) {
return 'state' in update;
}
/**
* Creates a TaskStatusUpdateEvent object.
*/
createTaskStatusEvent(taskId, status, final) {
return {
id: taskId,
status: status, // Assumes status already has timestamp from applyUpdate
final: final,
};
}
/**
* Gets the current timestamp in ISO format.
*/
getCurrentTimestamp() {
return new Date().toISOString();
}
/**
* Handles errors by converting them to appropriate JSON-RPC responses.
*/
handleError(error, _reqId) {
var _a;
if (!(error instanceof a2a_exception_1.A2AException)) {
error = a2a_exception_1.A2AException.internalError((_a = error.message) !== null && _a !== void 0 ? _a : 'Unknown error');
}
return a2a_types_1.JSONRPCResponse.error(error.toJSONRPCError());
}
};
exports.A2AExecutor = A2AExecutor;
exports.A2AExecutor = A2AExecutor = A2AExecutor_1 = __decorate([
(0, common_1.Injectable)({ scope: common_1.Scope.REQUEST }),
__param(0, (0, common_1.Inject)(constant_1.A2A_OPTIONS_TOKEN)),
__param(1, (0, common_1.Inject)(core_1.REQUEST)),
__metadata("design:paramtypes", [Object, Object, a2a_registry_1.A2ARegistry,
core_1.ModuleRef])
], A2AExecutor);