@kddd/artinet-sdk
Version:
TypeScript SDK for the Agent2Agent (A2A) Protocol
434 lines • 15 kB
JavaScript
"use strict";
var __importDefault = (this && this.__importDefault) || function (mod) {
return (mod && mod.__esModule) ? mod : { "default": mod };
};
Object.defineProperty(exports, "__esModule", { value: true });
exports.A2AServer = void 0;
const util_1 = __importDefault(require("util"));
const index_js_1 = require("../utils/index.js");
const index_js_2 = require("../transport/index.js");
const json_middleware_js_1 = require("./lib/json-middleware.js");
const express_server_js_1 = require("./lib/express-server.js");
const state_js_1 = require("./lib/state.js");
const memory_js_1 = require("./lib/storage/memory.js");
/**
* Implements an A2A protocol compliant server using Express.
* Handles task creation, streaming, cancellation and more.
* Uses Jayson for JSON-RPC handling.
*/
class A2AServer {
taskHandler;
taskStore;
corsOptions;
basePath;
port;
rpcServer;
serverInstance;
app;
fallbackPath;
register;
activeCancellations = new Set();
activeStreams = new Map();
/** The agent card representing this server */
card;
/**
* Returns the base path for the server
*/
getBasePath() {
return this.basePath;
}
/**
* Returns the CORS options for the server
*/
getCorsOptions() {
return this.corsOptions;
}
/**
* Returns the agent card for the server
*/
getCard() {
return this.card;
}
/**
* Returns the task store
*/
getTaskStore() {
return this.taskStore;
}
/**
* Returns the task handler
*/
getTaskHandler() {
return this.taskHandler;
}
/**
* Returns the set of active cancellations
*/
getActiveCancellations() {
return this.activeCancellations;
}
/**
* Returns the map of active streams
*/
getActiveStreams() {
return this.activeStreams;
}
/**
* Returns the port number
*/
getPort() {
return this.port;
}
/**
* Returns the JSON-RPC server
*/
getRpcServer() {
return this.rpcServer;
}
/**
* Returns the server instance
*/
getServerInstance() {
return this.serverInstance;
}
/**
* Returns the Express app
*/
getExpressApp() {
return this.app;
}
/**
* Returns a task context for the specified task and messages
*/
getTaskContext(task, userMessage, history) {
return this.createTaskContext(task, userMessage, history);
}
/**
* Returns the default agent card
*/
defaultAgentCard() {
return A2AServer.defaultAgentCard();
}
/**
* Creates a new A2AServer.
* @param handler The task handler function that will process tasks
* @param options Options for configuring the server
*/
constructor(params) {
// Store the handler
this.taskHandler = params.handler;
// Set up store
this.taskStore = params.taskStore ?? new memory_js_1.InMemoryTaskStore();
// Configure CORS
this.corsOptions = params.corsOptions ?? {
origin: "*",
methods: ["GET", "POST"],
allowedHeaders: ["Content-Type"],
};
// Set port
this.port = params.port ?? 41241;
let basePath = params.basePath ?? "/";
if (basePath !== "/") {
basePath = `/${basePath.replace(/^\/|\/$/g, "")}/`;
}
this.basePath = basePath;
// Set up default agent card if not provided
this.card = params.card ?? A2AServer.defaultAgentCard();
// Initialize the Jayson server
this.rpcServer = params.createJSONRPCServer
? params.createJSONRPCServer({
taskStore: this.taskStore,
card: this.card,
taskHandler: this.taskHandler,
activeCancellations: this.activeCancellations,
createTaskContext: this.createTaskContext.bind(this),
closeStreamsForTask: this.closeStreamsForTask.bind(this),
})
: (0, json_middleware_js_1.defaultCreateJSONRPCServer)({
taskStore: this.taskStore,
card: this.card,
taskHandler: this.taskHandler,
activeCancellations: this.activeCancellations,
createTaskContext: this.createTaskContext.bind(this),
closeStreamsForTask: this.closeStreamsForTask.bind(this),
});
this.fallbackPath = params.fallbackPath ?? "/agent-card";
const { app } = (0, express_server_js_1.createExpressServer)({
card: this.card,
corsOptions: this.corsOptions,
basePath: this.basePath,
port: this.port,
rpcServer: this.rpcServer,
fallbackPath: this.fallbackPath,
errorHandler: index_js_1.errorHandler,
onTaskSendSubscribe: this.handleTaskSendSubscribe.bind(this),
onTaskResubscribe: this.handleTaskResubscribe.bind(this),
});
this.app = app;
//register your server with the A2A registry on startup
this.register = params.register ?? false;
(0, index_js_1.logDebug)("A2AServer", "Server initialized", {
basePath: this.basePath,
port: this.port,
corsEnabled: !!this.corsOptions,
});
}
/**
* Starts the Express server listening on the specified port.
* @returns The running Express application instance.
*/
start() {
if (this.serverInstance) {
throw new Error("Server already started");
}
const server = this.app.listen(this.port, () => {
(0, index_js_1.logInfo)("A2AServer", `A2A Server started and listening`, {
port: this.port,
path: this.basePath,
});
});
this.serverInstance = server;
//lazily register your server with the A2A registry on startup
//this is so that you can start the server without having to wait for registration
//you can call also call this.registerServer() later to register your server
if (this.register) {
this.registerServer();
}
return this.app;
}
/**
* Stops the server and closes all connections.
* @returns A promise that resolves when the server is stopped.
*/
async stop() {
if (!this.serverInstance) {
return;
}
// Close all active streams first
this.activeStreams.forEach((streams, taskId) => {
if (streams.length > 0) {
(0, index_js_1.logDebug)("A2AServer", "Closing streams for task during stop", {
taskId,
});
this.closeStreamsForTask(taskId);
}
});
this.activeStreams.clear();
const closeServer = util_1.default
.promisify(this.serverInstance.close)
.bind(this.serverInstance);
try {
await closeServer();
(0, index_js_1.logDebug)("A2AServer", "Server stopped successfully.");
this.serverInstance = undefined;
}
catch (err) {
(0, index_js_1.logDebug)("A2AServer", "Error stopping server:", err);
this.serverInstance = undefined;
throw err;
}
}
/**
* Registers the server with the A2A registry.
* @returns A promise that resolves to the registration ID or an empty string if registration fails.
*/
async registerServer() {
if (this.card) {
return await (0, index_js_1.register)(this.card);
}
return "";
}
/**
* Handles task cancellation
* @param data Task and history data
* @param res Response object
*/
async onCancel(data, res) {
const currentData = await (0, state_js_1.updateState)(this.taskStore, data, index_js_1.CANCEL_UPDATE);
// Send the canceled status
(0, index_js_2.sendSSEEvent)(res, {
id: currentData.task.id,
status: currentData.task.status,
final: true,
});
this.closeStreamsForTask(currentData.task.id);
}
/**
* Handles cleanup when a task stream ends
* @param taskId The task ID
* @param res Response object
*/
async onEnd(taskId, res) {
this.activeCancellations.delete(taskId);
this.removeStreamForTask(taskId, res);
}
/**
* Handles the tasks/sendSubscribe method.
* @param req The SendTaskRequest object
* @param res The Express Response object
*/
async handleTaskSendSubscribe(req, res) {
(0, index_js_1.validateTaskSendParams)(req.params);
const { id: taskId, message, sessionId, metadata } = req.params;
// Set up SSE stream with initial status
(0, index_js_2.setupSseStream)(res, taskId, {
id: taskId,
status: {
state: "submitted",
timestamp: (0, index_js_1.getCurrentTimestamp)(),
},
}, this.addStreamForTask.bind(this));
// Load or create task
let currentData = await (0, state_js_1.loadState)(this.taskStore, taskId, message, sessionId, metadata);
// Create task context
const context = this.createTaskContext(currentData.task, message, currentData.history);
currentData = await (0, state_js_1.updateState)(this.taskStore, currentData, index_js_1.WORKING_UPDATE);
// Send the working status
(0, index_js_2.sendSSEEvent)(res, {
id: taskId,
status: currentData.task.status,
final: false,
});
// Process the task using the shared method
await (0, index_js_2.processTaskStream)(this.taskStore, this.taskHandler, res, taskId, context, currentData, this.onCancel.bind(this), this.onEnd.bind(this));
}
/**
* Handles the tasks/resubscribe method.
* @param req The TaskResubscriptionRequest object
* @param res The Express Response object
*/
async handleTaskResubscribe(req, res) {
const { id: taskId } = req.params;
if (!taskId) {
console.error("Task ID is required", req);
throw (0, index_js_1.INVALID_PARAMS)("Missing task ID");
}
// Try to load the task
const data = await this.taskStore.load(taskId);
if (!data) {
throw (0, index_js_1.TASK_NOT_FOUND)("Task Id: " + taskId);
}
// Set up SSE stream with current task status
(0, index_js_2.setupSseStream)(res, taskId, {
id: taskId,
status: data.task.status,
final: false,
}, this.addStreamForTask.bind(this));
// Check if task is in final state
if (index_js_1.FINAL_STATES.includes(data.task.status.state)) {
// If the task is already complete, send all artifacts and close
if (data.task.artifacts && data.task.artifacts.length > 0) {
for (const artifact of data.task.artifacts) {
(0, index_js_2.sendSSEEvent)(res, {
id: taskId,
artifact,
final: true,
});
}
}
// Remove from tracking and close
this.removeStreamForTask(taskId, res);
res.write("event: close\ndata: {}\n\n");
res.end();
return;
}
// For non-final states, create context and continue processing
// We need to use the last user message as the current message
const lastUserMessage = data.history
.filter((msg) => msg.role === "user")
.pop();
if (!lastUserMessage) {
throw (0, index_js_1.INVALID_REQUEST)("No user message found");
}
const context = this.createTaskContext(data.task, lastUserMessage, data.history);
// Continue processing the task using the shared method
await (0, index_js_2.processTaskStream)(this.taskStore, this.taskHandler, res, taskId, context, data, this.onCancel.bind(this), this.onEnd.bind(this));
}
/**
* Adds a response stream to the tracking map for a task.
* @param taskId The task ID
* @param res The response stream
*/
addStreamForTask(taskId, res) {
if (!this.activeStreams.has(taskId)) {
this.activeStreams.set(taskId, []);
}
(0, index_js_1.logDebug)("A2AServer", "Adding stream for task", {
taskId,
activeStreams: this.activeStreams,
});
this.activeStreams.get(taskId)?.push(res);
}
/**
* Removes a response stream from the tracking map for a task.
* @param taskId The task ID
* @param res The response stream
*/
removeStreamForTask(taskId, res) {
const streams = this.activeStreams.get(taskId);
if (streams) {
const index = streams.indexOf(res);
if (index !== -1) {
streams.splice(index, 1);
if (streams.length === 0) {
(0, index_js_1.logDebug)("A2AServer", "Removing stream for task", {
taskId,
activeStreams: this.activeStreams,
});
this.activeStreams.delete(taskId);
}
}
}
}
/**
* Initializes the default agent card
*/
static defaultAgentCard() {
return {
name: "A2A Server",
description: "A general-purpose A2A protocol server",
version: "0.1.0",
url: "http://localhost",
capabilities: {
streaming: true,
pushNotifications: false,
stateTransitionHistory: true,
},
skills: [],
};
}
/**
* Creates a TaskContext object for a task handler.
* @param task The task
* @param userMessage The user message
* @param history The message history
* @returns A TaskContext object
*/
createTaskContext(task, userMessage, history) {
return {
task,
userMessage,
history,
isCancelled: () => this.activeCancellations.has(task.id),
};
}
/**
* Closes any active streams for a task.
* @param taskId The task ID
*/
closeStreamsForTask(taskId) {
const streams = this.activeStreams.get(taskId);
if (streams) {
// Send close event to all streams
for (const stream of streams) {
if (stream.writable) {
stream.write("event: close\ndata: {}\n\n");
stream.end();
}
}
this.activeStreams.delete(taskId);
}
}
}
exports.A2AServer = A2AServer;
//# sourceMappingURL=a2a-server.js.map