Files
diegosouzapw--omniroute/tests/unit/stream-continuation-wiring.test.ts
2026-07-13 13:39:12 +08:00

118 lines
4.4 KiB
TypeScript

import { test } from "node:test";
import assert from "node:assert/strict";
import {
createRecoverableStream,
scanOpenAiSseText,
} from "../../open-sse/services/streamRecovery.ts";
const enc = new TextEncoder();
/** A clock that jumps +1000ms per call so the holdback commits on the very first chunk. */
function steppingClock() {
let t = 0;
return () => {
t += 1000;
return t;
};
}
/** Build a ReadableStream from string chunks, optionally erroring (truncating) at the end. */
function streamFrom(chunks: string[], opts: { truncateError?: Error } = {}) {
let i = 0;
return new ReadableStream<Uint8Array>({
pull(controller) {
if (i < chunks.length) {
controller.enqueue(enc.encode(chunks[i++]));
return;
}
if (opts.truncateError) controller.error(opts.truncateError);
else controller.close();
},
});
}
async function collectText(stream: ReadableStream<Uint8Array>): Promise<string> {
const reader = stream.getReader();
const dec = new TextDecoder();
let out = "";
for (;;) {
const { done, value } = await reader.read();
if (done) break;
if (value) out += dec.decode(value, { stream: true });
}
return out;
}
const ROLE = 'data: {"choices":[{"delta":{"role":"assistant"}}]}\n\n';
const content = (s: string) => `data: {"choices":[{"delta":{"content":${JSON.stringify(s)}}}]}\n\n`;
test("mid-stream continuation: stitches the suffix after a silent post-commit truncation", async () => {
// Commits on chunk 1, emits "Hello wor", then ends WITHOUT a terminal marker (silent cut).
const initial = streamFrom([ROLE, content("Hello wor")]);
let finalizeCount = 0;
let continueArg = "";
const stream = createRecoverableStream(initial, async () => null, {
finalize: () => {
finalizeCount += 1;
},
now: steppingClock(),
continueStream: async (soFar: string) => {
continueArg = soFar;
// The model re-emits a small overlap ("wor") which must be trimmed away.
return streamFrom([ROLE, content("world!"), "data: [DONE]\n\n"]);
},
});
const out = await collectText(stream);
const scan = scanOpenAiSseText(out);
assert.equal(continueArg, "Hello wor", "continuation is prefilled with the text already sent");
assert.equal(scan.text, "Hello world!", "client sees the full answer, overlap trimmed, exactly once");
assert.equal(scan.terminal, true, "the recovered stream ends with a terminal marker");
assert.equal(finalizeCount, 1, "finalize runs exactly once");
});
test("mid-stream continuation: recovers a post-commit transport error too", async () => {
const initial = streamFrom([ROLE, content("Partial ")], {
truncateError: Object.assign(new Error("terminated"), { code: "UND_ERR_SOCKET" }),
});
const stream = createRecoverableStream(initial, async () => null, {
finalize: () => {},
now: steppingClock(),
continueStream: async () => streamFrom([content("answer done."), "data: [DONE]\n\n"]),
});
const scan = scanOpenAiSseText(await collectText(stream));
assert.equal(scan.text, "Partial answer done.");
assert.equal(scan.terminal, true);
});
test("no continuation configured: a silent post-commit truncation closes as before (no recovery)", async () => {
const initial = streamFrom([ROLE, content("Hello wor")]);
const stream = createRecoverableStream(initial, async () => null, {
finalize: () => {},
now: steppingClock(),
// continueStream omitted → behavior unchanged from #4131
});
const scan = scanOpenAiSseText(await collectText(stream));
assert.equal(scan.text, "Hello wor", "only the committed text is delivered");
assert.equal(scan.terminal, false, "no synthetic terminal is injected when continuation is off");
});
test("tool-call in flight is never continued (would corrupt tool JSON)", async () => {
const toolDelta =
'data: {"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"name":"f","arguments":"{\\"a\\":"}}]}}]}\n\n';
const initial = streamFrom([ROLE, toolDelta]); // commits, emits a partial tool call, then truncates
let continued = false;
const stream = createRecoverableStream(initial, async () => null, {
finalize: () => {},
now: steppingClock(),
continueStream: async () => {
continued = true;
return streamFrom([content("nope"), "data: [DONE]\n\n"]);
},
});
await collectText(stream);
assert.equal(continued, false, "continuation must NOT fire once a tool call has started streaming");
});