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

A TypeScript SDK for building collaborative AI agents.

import { sleep } from "../../utils/utils.js";
import { logger } from "../../config/index.js";
const STREAM_INTERVAL = 5;
export class Stream {
    _contextId;
    _context;
    _updates;
    _completed;
    running = null;
    constructor(_contextId, _context, _updates = [], _completed = false) {
        this._contextId = _contextId;
        this._context = _context;
        this._updates = _updates;
        this._completed = _completed;
    }
    get contextId() {
        return this._contextId;
    }
    get isAlive() {
        return !this._completed;
    }
    async kill() {
        this._completed = true;
    }
    get context() {
        return this._context;
    }
    set context(context) {
        this._context = context;
        this._contextId = context.contextId;
    }
    get updates() {
        return this._updates;
    }
    async *run({ service, }) {
        this.context.publisher.on("complete", async () => {
            logger.debug(`Stream[run:${this.contextId}]: complete`);
            await this.kill();
            this.running = null;
        });
        this.context.publisher.on("error", async (err) => {
            logger.error(`Stream[run:${this.contextId}]: error`, err);
            await this.context.publisher.onComplete();
        });
        if (!this.running) {
            this.running = this._run({ service });
            yield* this.running;
        }
        yield* this.subscribe();
    }
    async *subscribe() {
        const subscription = this.updates;
        this.context.publisher.on("update", async (_, update) => {
            subscription.push(update);
        });
        let updatesRead = 0;
        while (!this._completed || updatesRead < subscription.length) {
            if (updatesRead < subscription.length) {
                yield subscription[updatesRead++];
            }
            await sleep(STREAM_INTERVAL);
        }
    }
    async *_run({ service, }) {
        let executionError = null;
        const context = this.context;
        // NOTE: isCancelled() returns a Promise, so must be awaited.
        // Without async/await here, updates would never be pushed due to
        // Promise being truthy (making !Promise always false).
        context.publisher.on("update", async (_, update) => {
            if (!(await context.isCancelled())) {
                this.updates.push(update);
            }
        });
        const executePromise = service
            .execute({ engine: service.engine, context })
            .catch((error) => {
            executionError = error;
        })
            .finally(() => {
            this._completed = true;
        });
        while (!this._completed || this.updates.length > 0) {
            if (executionError) {
                throw executionError;
            }
            if (this.updates.length > 0) {
                yield this.updates.shift();
            }
            await sleep(STREAM_INTERVAL);
        }
        await executePromise;
        this._completed = true;
        if (executionError) {
            throw executionError;
        }
    }
}