SnowLuma/packages/onebot/tests/streaming.test.ts

// Phase 1 — Stream API plumbing (#163). Verifies the dispatch seam:
// defineStreamAction registers a streaming handler, processStreamRequest
// emits N intermediate frames + a terminal frame (all echo-tagged, marked
// `stream:'stream-action'`), normal actions stay single-frame, and failures
// terminate with an error frame. Wire format mirrors NapCat's Stream API.
import { describe, it, expect } from 'vitest';
import type { ApiHandler, ApiActionContext } from '../src/api-handler';
import { defineAction, defineStreamAction, f, type RegisteredActionSpec } from '../src/action-kit';
import { okResponse, failedResponse } from '../src/types';
import { wrapStreamFrame, wrapStreamTerminal, STREAM_MARK } from '../src/streaming';
import { createCompiledTestHandler } from './helpers/compiled-action-handler';

function newHandler(...specs: RegisteredActionSpec[]): ApiHandler {
    // ActionSpecs bind during the constructor but call ctx only when dispatched.
    return createCompiledTestHandler({} as ApiActionContext, specs);
}

const parse = (frames: string[]) => frames.map((f) => JSON.parse(f) as Record<string, unknown>);

describe('Stream API plumbing', () => {
    it('streams N intermediate frames + a terminal frame, all echo-tagged', async () => {
        const action = defineStreamAction({
            name: 'demo_stream',
            params: { n: f.uint() },
            run: async (p, _ctx, sink) => {
                for (let i = 0; i < p.n; i++) {
                    await sink.send({ type: 'stream', index: i });
                }

                return okResponse({ type: 'response', total: p.n });
            },
        });
        const h = newHandler(action);

        const frames: string[] = [];
        await h.processStreamRequest(
            JSON.stringify({ action: 'demo_stream', params: { n: 2 }, echo: 'e1' }),
            (j) => frames.push(j),
        );

        const got = parse(frames);
        expect(got).toHaveLength(3);
        expect(got[0]).toMatchObject({
            status: 'ok',
            retcode: 0,
            stream: STREAM_MARK,
            echo: 'e1',
            data: { type: 'stream', index: 0 },
        });
        expect(got[1]).toMatchObject({
            stream: STREAM_MARK,
            echo: 'e1',
            data: { type: 'stream', index: 1 },
        });
        expect(got[2]).toMatchObject({
            status: 'ok',
            stream: STREAM_MARK,
            echo: 'e1',
            data: { type: 'response', total: 2 },
        });
    });

    it('marks the action as a stream action; normal actions are not', () => {
        const streamAction = defineStreamAction({
            name: 'demo_stream',
            params: {},
            run: () => okResponse({ type: 'response' }),
        });
        const normalAction = defineAction({
            name: 'demo_normal',
            params: {},
            run: () => okResponse({ ok: true }),
        });
        const h = newHandler(streamAction, normalAction);
        expect(h.isStreamAction('demo_stream')).toBe(true);
        expect(h.isStreamAction('demo_normal')).toBe(false);
        expect(h.isStreamAction('nope')).toBe(false);
    });

    it('a normal action through processStreamRequest emits exactly one un-marked frame', async () => {
        const action = defineAction({
            name: 'demo_normal',
            params: {},
            run: () => okResponse({ ok: true }),
        });
        const h = newHandler(action);

        const frames: string[] = [];
        await h.processStreamRequest(
            JSON.stringify({ action: 'demo_normal', params: {}, echo: 'e2' }),
            (j) => frames.push(j),
        );

        expect(frames).toHaveLength(1);
        const r = parse(frames)[0];
        expect(r).toMatchObject({ status: 'ok', echo: 'e2', data: { ok: true } });
        expect(r.stream).toBeUndefined();
    });

    it('a throwing stream action terminates with a single error frame', async () => {
        const action = defineStreamAction({
            name: 'demo_boom',
            params: {},
            run: async () => {
                throw new Error('boom');
            },
        });
        const h = newHandler(action);

        const frames: string[] = [];
        await h.processStreamRequest(
            JSON.stringify({ action: 'demo_boom', params: {}, echo: 'e3' }),
            (j) => frames.push(j),
        );

        expect(frames).toHaveLength(1);
        expect(parse(frames)[0]).toMatchObject({
            status: 'failed',
            stream: STREAM_MARK,
            echo: 'e3',
            wording: 'boom',
            data: { type: 'error', data_type: 'error' },
        });
    });

    it('emits already-sent frames before the terminal error when a stream throws mid-way', async () => {
        const action = defineStreamAction({
            name: 'demo_partial',
            params: {},
            run: async (_p, _ctx, sink) => {
                await sink.send({ type: 'stream', index: 0 });
                throw new Error('mid');
            },
        });
        const h = newHandler(action);

        const frames: string[] = [];
        await h.processStreamRequest(JSON.stringify({ action: 'demo_partial', params: {} }), (j) =>
            frames.push(j),
        );

        const got = parse(frames);
        expect(got).toHaveLength(2);
        expect(got[0]).toMatchObject({ stream: STREAM_MARK, data: { type: 'stream', index: 0 } });
        expect(got[1]).toMatchObject({ status: 'failed', data: { type: 'error' } });
    });

    it('stream param validation fails terminate with an error frame', async () => {
        const action = defineStreamAction({
            name: 'demo_needs',
            params: { n: f.uint() },
            run: () => okResponse({ type: 'response' }),
        });
        const h = newHandler(action);

        const frames: string[] = [];
        await h.processStreamRequest(JSON.stringify({ action: 'demo_needs', params: {} }), (j) =>
            frames.push(j),
        );
        expect(frames).toHaveLength(1);
        expect(parse(frames)[0]).toMatchObject({
            status: 'failed',
            stream: STREAM_MARK,
            data: { type: 'error' },
        });
    });

    it('aborts a stream when isAlive turns false (client disconnect)', async () => {
        let sent = 0;
        const action = defineStreamAction({
            name: 'demo_abort',
            params: {},
            run: async (_p, _ctx, sink) => {
                for (let i = 0; i < 5; i++) {
                    await sink.send({ type: 'stream', index: i });
                    sent++;
                }

                return okResponse({ type: 'response' });
            },
        });
        const h = newHandler(action);

        const frames: string[] = [];
        let alive = true;
        await h.processStreamRequest(
            JSON.stringify({ action: 'demo_abort', params: {} }),
            (j) => {
                frames.push(j);

                if (frames.length >= 2) {
                    alive = false;
                }
            },
            () => alive,
        );

        expect(sent).toBeLessThan(5); // loop aborted once the client "left"
        const got = parse(frames);
        expect(got[got.length - 1]).toMatchObject({ status: 'failed', data: { type: 'error' } });
    });

    it('wrapStreamFrame / wrapStreamTerminal shape (NapCat message+wording mirror)', () => {
        expect(wrapStreamFrame({ type: 'stream', x: 1 }, 'e')).toEqual({
            status: 'ok',
            retcode: 0,
            data: { type: 'stream', x: 1 },
            message: '',
            wording: '',
            stream: STREAM_MARK,
            echo: 'e',
        });
        // no echo → no echo key
        expect(wrapStreamFrame({ type: 'stream' }, undefined)).toEqual({
            status: 'ok',
            retcode: 0,
            data: { type: 'stream' },
            message: '',
            wording: '',
            stream: STREAM_MARK,
        });
        // failed terminal → data normalised to an error packet; reason in BOTH message + wording
        expect(wrapStreamTerminal(failedResponse(1404, 'nope'), 'e')).toMatchObject({
            status: 'failed',
            retcode: 1404,
            message: 'nope',
            wording: 'nope',
            stream: STREAM_MARK,
            echo: 'e',
            data: { type: 'error', data_type: 'error' },
        });
        // ok terminal → data passes through, message/wording empty
        expect(wrapStreamTerminal(okResponse({ type: 'response', a: 1 }), undefined)).toMatchObject(
            {
                status: 'ok',
                stream: STREAM_MARK,
                message: '',
                data: { type: 'response', a: 1 },
            },
        );
    });

    it('malformed stream requests emit one bad-request frame', async () => {
        const h = newHandler();
        const frames: string[] = [];
        await h.processStreamRequest('not json', (j) => frames.push(j));
        await h.processStreamRequest('', (j) => frames.push(j));
        await h.processStreamRequest(JSON.stringify({ params: {} }), (j) => frames.push(j)); // no action
        expect(frames).toHaveLength(3);

        for (const f of parse(frames)) {
            expect(f).toMatchObject({ status: 'failed' });
        }
    });

    it('waits for an async bad-request frame to flush', async () => {
        const h = newHandler();
        let release!: () => void;
        const flushed = new Promise<void>((resolve) => {
            release = resolve;
        });
        let completed = false;

        const request = h.processStreamRequest('', async () => flushed);
        void request.then(() => {
            completed = true;
        });
        await Promise.resolve();
        expect(completed).toBe(false);

        release();
        await request;
        expect(completed).toBe(true);
    });
});