@adonis-agora/agent/dist/src/redis-token-stream-sink.js

Governed, durable-backed AI agent (chat + tool-calling + governance) for AdonisJS — part of the Agora ecosystem.

/** Marker stored at the state key once a run's stream has finished. */
const ENDED = 'ended';
/** Default TTL (seconds) for a run's Redis keys: 1h — a generous reconnect/replay window past finish. */
const DEFAULT_TTL_SECONDS = 3600;
/**
 * A multi-replica {@link TokenStreamSink} over Redis — the shared-transport counterpart to the default
 * {@link import('./in-process-sink.js').InProcessTokenStreamSink}. The replica running the model turn
 * `open`s the writer and appends each token chunk to a Redis LIST (`${keyPrefix}:${runId}:chunks`) so
 * ANY replica can replay the whole stream so far (the same late-subscriber/reconnect guarantee the
 * in-process sink gives, but shared across pods), and PUBLISHes to a per-run channel
 * (`${keyPrefix}:${runId}`) to wake live subscribers. The run's terminal state lives in a marker key
 * (`${keyPrefix}:${runId}:state`), so an SSE handler on ANOTHER replica that subscribes after the run
 * finished still sees the ending and closes the stream.
 *
 * The yielded frames are identical (frame-for-frame) to the in-process sink's, so the agent-protocol
 * SSE the provider pipes is unchanged — swapping this in is transparent to clients.
 *
 * The host supplies a {@link RedisStreamClient} adapter over its own Redis driver; `subscribe` needs a
 * connection in subscriber mode (see the client docs). Wire it with `defineConfig({ sink:
 * tokenSinks.redis({...}) })`, which builds the adapter over `@adonisjs/redis` lazily.
 */
export class RedisTokenStreamSink {
    client;
    keyPrefix;
    ttlSeconds;
    constructor(client, options = {}) {
        this.client = client;
        this.keyPrefix = options.keyPrefix ?? 'agent:stream';
        this.ttlSeconds = options.ttlSeconds ?? DEFAULT_TTL_SECONDS;
    }
    /** Refresh the sliding TTL on `key` (no-op when TTL is disabled or the client can't `expire`). */
    async touchTtl(key) {
        if (this.ttlSeconds > 0) {
            await this.client.expire?.(key, this.ttlSeconds);
        }
    }
    chunksKey(runId) {
        return `${this.keyPrefix}:${runId}:chunks`;
    }
    stateKey(runId) {
        return `${this.keyPrefix}:${runId}:state`;
    }
    channel(runId) {
        return `${this.keyPrefix}:${runId}`;
    }
    open(runId) {
        return {
            write: async (frame) => {
                await this.client.rpush(this.chunksKey(runId), JSON.stringify(frame));
                await this.touchTtl(this.chunksKey(runId));
                await this.client.publish(this.channel(runId), 'chunk');
            },
            end: async () => {
                await this.client.set(this.stateKey(runId), ENDED);
                await this.touchTtl(this.stateKey(runId));
                await this.touchTtl(this.chunksKey(runId));
                await this.client.publish(this.channel(runId), 'end');
            },
        };
    }
    async *subscribe(runId) {
        // Notify bridge: the pub/sub handler resolves the current `wake` promise so the drain loop
        // re-reads the list. A promise is armed BEFORE each drain, so a publish arriving mid-drain is
        // never lost.
        let resolveWake = null;
        const onMessage = () => {
            const resolve = resolveWake;
            resolveWake = null;
            resolve?.();
        };
        const unsubscribe = await this.client.subscribe(this.channel(runId), onMessage);
        try {
            let index = 0;
            while (true) {
                const wake = new Promise((resolve) => {
                    resolveWake = resolve;
                });
                const chunks = await this.client.lrange(this.chunksKey(runId), index, -1);
                for (const encoded of chunks) {
                    index += 1;
                    yield JSON.parse(encoded);
                }
                if (await this.hasEnded(runId)) {
                    return;
                }
                await wake;
            }
        }
        finally {
            await unsubscribe();
        }
    }
    async has(runId) {
        if ((await this.client.get(this.stateKey(runId))) !== null)
            return true;
        return (await this.client.lrange(this.chunksKey(runId), 0, 0)).length > 0;
    }
    async close(runId) {
        await this.client.del(this.chunksKey(runId), this.stateKey(runId));
    }
    async hasEnded(runId) {
        const raw = await this.client.get(this.stateKey(runId));
        return raw === ENDED;
    }
}
//# sourceMappingURL=redis-token-stream-sink.js.map