118 lines
4.4 KiB
TypeScript
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");
|
|
});
|