@artinet/sdk/dist/services/a2a/repository.js

TypeScript SDK for Agentic Communication

import { CANCEL_UPDATE } from "../../utils/index.js";
import { processUpdate } from "../../server/lib/state.js";
import { sendSSEEvent } from "../../transport/streaming/stream.js";
import { logDebug } from "../../utils/logging/log.js";
import { InMemoryTaskStore } from "../../server/lib/storage/memory.js";
import { v4 as uuidv4 } from "uuid";
/**
 * @description The A2ARepository class.
 */
export class A2ARepository {
    taskStore;
    _activeCancellations;
    _activeStreams;
    card;
    constructor(params) {
        this.taskStore = params.taskStore ?? new InMemoryTaskStore();
        this.card = params.card;
        this._activeCancellations = new Set();
        this._activeStreams = new Map();
    }
    /**
     * @description Gets the agent card.
     * @returns {AgentCard} The agent card.
     */
    getCard() {
        return this.card;
    }
    /**
     * @description Gets the active streams.
     * @returns {Map<string, Response[]>} The active streams.
     */
    get activeStreams() {
        return this._activeStreams;
    }
    /**
     * @description Gets the active cancellations.
     * @returns {Set<string>} The active cancellations.
     */
    get activeCancellations() {
        return this._activeCancellations;
    }
    /**
     * @description Gets the task store.
     * @returns {TaskStore} The task store.
     */
    getTaskStore() {
        return this.taskStore;
    }
    /**
     * 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, configuration) {
        return {
            contextId: task.contextId ?? userMessage.contextId ?? uuidv4(),
            task,
            userMessage,
            history,
            configuration,
            isCancelled: () => this._activeCancellations.has(task.id),
        };
    }
    /**
     * Handles task cancellation
     * @param data Task and history data
     * @param res Response object
     */
    async onCancel(context, data, res) {
        const cancelUpdate = CANCEL_UPDATE(data.task.id, context.contextId);
        const currentData = await processUpdate(this.taskStore, {
            context: context,
            current: data,
            update: cancelUpdate,
        });
        // Send the canceled status
        sendSSEEvent(res, currentData.task.id, cancelUpdate);
        this.closeStreamsForTask(currentData.task.id);
    }
    /**
     * 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, []);
        }
        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) {
                    logDebug("A2AServer", "Removing stream for task", {
                        taskId,
                        activeStreams: this._activeStreams,
                    });
                    this._activeStreams.delete(taskId);
                }
            }
        }
    }
    /**
     * 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);
        }
    }
    /**
     * 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);
    }
    /**
     * Destroys the repository.
     */
    async destroy() {
        this._activeStreams.forEach((streams, taskId) => {
            if (streams.length > 0) {
                logDebug("A2ARepository", "Closing streams for task during stop", {
                    taskId,
                });
                this.closeStreamsForTask(taskId);
            }
        });
        this._activeStreams.clear();
        this._activeCancellations.clear();
    }
}
//# sourceMappingURL=repository.js.map