KrisHD_AI/test/modelStreaming.test.js

import test from "node:test";
import assert from "node:assert/strict";
import {createServer} from "node:http";
import {createModelService} from "../src/services/modelService.js";
import {createApp} from "../src/app.js";
import {stream} from "../public/app/api.js";

const config = {
    modelApiBaseUrl: "http://fixture", modelRequestTimeoutMs: 20,
    modelConnectTimeoutMs: 1000, modelInactivityTimeoutMs: 1000,
    modelMaxGenerationMs: 3000, modelUnloadAfterMs: 60000
};
const line = data => new TextEncoder().encode(JSON.stringify(data) + "\n");

test("active tokens outlive the legacy total timeout and arrive before completion", async () => {
    let finish;
    let first;
    const firstChunk = new Promise(resolve => { first = resolve; });
    let interval;
    const service = createModelService(config, {fetch: async () => new Response(new ReadableStream({
        start(controller) {
            interval = setInterval(() => controller.enqueue(line({message: {content: "x"}})), 15);
            finish = () => { clearInterval(interval); controller.enqueue(line({done: true})); controller.close(); };
        },
        cancel() { clearInterval(interval); }
    }))});
    let completed = false;
    const pending = service.streamReply("qwen3:14b", [], () => first()).then(reply => { completed = true; return reply; });
    try {
        await firstChunk;
        assert.equal(completed, false);
        await new Promise(resolve => setTimeout(resolve, 80));
        assert.equal(completed, false);
        finish();
        assert.ok((await pending).length >= 2);
    } finally { clearInterval(interval); }
});

test("connection, inactivity and overall deadlines abort generation", async (t) => {
    for (const kind of ["connection", "inactivity", "overall"]) await t.test(kind, async () => {
        let cancelled = false;
        let interval;
        const service = createModelService({
            ...config,
            modelConnectTimeoutMs: kind === "connection" ? 25 : 1000,
            modelInactivityTimeoutMs: kind === "inactivity" ? 25 : 1000,
            modelMaxGenerationMs: kind === "overall" ? 70 : 1000
        }, {fetch: async (_url, {signal}) => {
            if (kind === "connection") return new Promise((_resolve, reject) => {
                signal.addEventListener("abort", () => { cancelled = true; reject(new DOMException("Aborted", "AbortError")); });
            });
            return new Response(new ReadableStream({
                start(controller) {
                    if (kind === "overall") interval = setInterval(() => controller.enqueue(line({message: {content: "x"}})), 10);
                },
                cancel() { cancelled = true; clearInterval(interval); }
            }));
        }});
        await assert.rejects(service.streamReply("qwen3:14b", [], () => {}), {code: "MODEL_TIMEOUT"});
        assert.equal(cancelled, true);
    });
});

test("stream validation rejects truncated output, upstream errors and invalid fields", async () => {
    for (const body of [
        '{"message":{"content":"partial"}}\n',
        '{"error":"internal upstream details"}\n',
        '{"message":{"content":42},"done":true}\n',
        '{broken}\n'
    ]) {
        const service = createModelService(config, {fetch: async () => new Response(body)});
        await assert.rejects(service.generateReply("qwen3:14b", []));
    }
    const service = createModelService(config, {fetch: async () => { throw new Error("Connection refused"); }});
    await assert.rejects(service.generateReply("qwen3:14b", []), {code: "MODEL_UNAVAILABLE"});
});

test("SSE EOF is a failure rather than successful completion", async (t) => {
    const original = globalThis.fetch;
    t.after(() => { globalThis.fetch = original; });
    globalThis.fetch = async () => new Response('event: chunk\ndata: {"fullReply":"partial"}\n\n');
    const chunks = [];
    await assert.rejects(stream("/fixture", {}, {onChunk: data => chunks.push(data.fullReply)}), /interrupted/);
    assert.deepEqual(chunks, ["partial"]);
});

test("browser disconnect aborts Ollama and incomplete output is not persisted", async (t) => {
    let signal;
    let releaseAbort;
    const aborted = new Promise(resolve => { releaseAbort = resolve; });
    const app = createApp({config: {
        dbPath: ":memory:", testMode: true, sessionSecret: "test-only",
        chatSummaryUpdateEveryMessages: 0, structuredRequestLoggingEnabled: false,
        eventLoopLagMonitorEnabled: false, slowRequestLoggingEnabled: false
    }, modelOverrides: {fetch: async (_url, options) => {
        signal = options.signal;
        signal.addEventListener("abort", releaseAbort, {once: true});
        return new Response(new ReadableStream({start(controller) {
            controller.enqueue(line({message: {content: "partial"}}));
        }}));
    }}});
    const server = createServer(app);
    await new Promise(resolve => server.listen(0, "127.0.0.1", resolve));
    t.after(async () => {
        await new Promise(resolve => { server.closeAllConnections(); server.close(resolve); });
        app.locals.close();
    });
    const base = `http://127.0.0.1:${server.address().port}`;
    const registered = await fetch(base + "/register", {method: "POST", headers: {"Content-Type": "application/json"}, body: JSON.stringify({username: "disconnect", password: "test-password"})});
    const cookie = registered.headers.get("set-cookie").split(";")[0];
    const headers = {"Content-Type": "application/json", Cookie: cookie};
    const chat = await (await fetch(base + "/chats", {method: "POST", headers, body: '{"title":"Disconnect"}'})).json();
    const controller = new AbortController();
    const response = await fetch(base + "/chat/stream", {
        method: "POST", headers, signal: controller.signal,
        body: JSON.stringify({chatId: chat.chat.id, message: "Hello"})
    });
    const reader = response.body.getReader();
    let received = "";
    while (!received.includes("event: chunk")) received += new TextDecoder().decode((await reader.read()).value);
    controller.abort();
    await Promise.race([aborted, new Promise((_resolve, reject) => {
        const timer = setTimeout(() => reject(new Error("Upstream was not aborted")), 1000);
        timer.unref();
    })]);
    assert.equal(signal.aborted, true);
    assert.equal(app.locals.db.prepare("SELECT COUNT(*) AS count FROM chat_messages WHERE role = 'bot'").get().count, 0);
});