170 lines
9.1 KiB
JavaScript
170 lines
9.1 KiB
JavaScript
import test from "node:test";
|
|
import assert from "node:assert/strict";
|
|
import {createApp} from "../src/app.js";
|
|
|
|
const deferred = () => {
|
|
let resolve;
|
|
const promise = new Promise(done => { resolve = done; });
|
|
return {promise, resolve};
|
|
};
|
|
|
|
async function fixture(t, config = {}) {
|
|
const calls = [];
|
|
const gates = [];
|
|
let nextGate = null;
|
|
const modelService = {
|
|
async generateReply(model, messages, options) {
|
|
calls.push({model, messages, options});
|
|
const gate = nextGate;
|
|
nextGate = null;
|
|
if (gate) {
|
|
gate.started.resolve();
|
|
const abort = () => { gate.aborted.resolve(); };
|
|
options.signal.addEventListener("abort", abort, {once: true});
|
|
// Intentionally allow a late model result, exercising the commit-time abort check.
|
|
try { await gate.release.promise; }
|
|
finally { options.signal.removeEventListener("abort", abort); }
|
|
}
|
|
if (options.reasoning) options.onThinking?.("Reasoning", "Reasoning");
|
|
return "The guide waits beside the station.";
|
|
},
|
|
async streamReply(model, messages, onChunk, options) {
|
|
const reply = await this.generateReply(model, messages, options);
|
|
onChunk(reply, reply);
|
|
return reply;
|
|
},
|
|
mapError(error) { return {status: 500, body: {error: error.message}}; }
|
|
};
|
|
const app = createApp({modelService, config: {dbPath: ":memory:", testMode: true, sessionSecret: "generation-tests",
|
|
chatSummaryUpdateEveryMessages: 0, eventLoopLagMonitorEnabled: false, structuredRequestLoggingEnabled: false,
|
|
slowRequestLoggingEnabled: false, ...config}});
|
|
const server = app.listen(0, "127.0.0.1");
|
|
await new Promise(resolve => server.once("listening", resolve));
|
|
t.after(async () => {
|
|
for (const gate of gates) gate.release.resolve();
|
|
await new Promise(resolve => { server.closeAllConnections(); server.close(resolve); });
|
|
await app.locals.close();
|
|
});
|
|
const base = `http://127.0.0.1:${server.address().port}`;
|
|
let cookie = "";
|
|
const api = async (path, body, signal) => {
|
|
const response = await fetch(base + path, {method: body ? "POST" : "GET", signal,
|
|
headers: {"Content-Type": "application/json", Cookie: cookie}, body: body ? JSON.stringify(body) : undefined});
|
|
if (response.headers.get("set-cookie")) cookie = response.headers.get("set-cookie").split(";")[0];
|
|
return response;
|
|
};
|
|
assert.equal((await api("/register", {username: "generator", password: "test-password"})).status, 200);
|
|
const persona = (await (await api("/personas", {name: "Guide", personaType: "assistant"})).json()).persona;
|
|
const published = await (await api(`/personas/${persona.id}/publish`, {})).json();
|
|
const marketId = published.persona.id;
|
|
const chat = (await (await api("/chats", {title: "Regular chat"})).json()).chat;
|
|
const repo = app.locals.repositories;
|
|
const promptId = Number(repo.insertBranchChatMessage(chat.id, "user", "Original question").lastInsertRowid);
|
|
const answerId = Number(repo.insertBranchChatMessage(chat.id, "bot", "Original answer", null, null, "complete", null, promptId).lastInsertRowid);
|
|
return {api, repo, calls, persona, chat, promptId, answerId, marketId,
|
|
roleplay: {path: "/roleplays/start", body: {assistantPersonaId: persona.id, scenarioPrompt: "First scene"}},
|
|
market: {path: `/personas/market/${marketId}/chat`, body: {scenarioPrompt: "First scene"}},
|
|
edit: {path: `/chats/${chat.id}/messages/${promptId}/edit-resend`, body: {chatId: chat.id, message: "Edited question", confirmed: true, reasoning: true, model: "qwen3:14b"}},
|
|
hold() {
|
|
const gate = {started: deferred(), release: deferred(), aborted: deferred()};
|
|
gates.push(gate); nextGate = gate;
|
|
return gate;
|
|
}};
|
|
}
|
|
|
|
for (const limitedBy of ["chatRateLimitMaxPerAccount", "chatRateLimitMaxPerIp"]) {
|
|
test(`roleplay, marketplace and edit/resend share normal chat's ${limitedBy} bucket`, {timeout: 10000}, async t => {
|
|
const s = await fixture(t, {[limitedBy]: 1});
|
|
assert.equal((await s.api("/chat", {chatId: s.chat.id, message: "Use the budget"})).status, 200);
|
|
const before = s.repo.exportWorkspace("generator");
|
|
const attempts = [s.roleplay, s.roleplay, s.market, s.edit];
|
|
const responses = await Promise.all(attempts.map(({path, body}) => s.api(path, body)));
|
|
assert.deepEqual(responses.map(response => response.status), [429, 429, 429, 429]);
|
|
assert.equal(s.calls.length, 1);
|
|
assert.deepEqual(s.repo.exportWorkspace("generator").chats, before.chats);
|
|
assert.deepEqual(s.repo.exportWorkspace("generator").personas, before.personas);
|
|
});
|
|
}
|
|
|
|
for (const kind of ["roleplay", "market", "edit"]) {
|
|
test(`${kind} consumes the shared generation budget`, {timeout: 10000}, async t => {
|
|
const s = await fixture(t, {chatRateLimitMaxPerAccount: 1});
|
|
const operation = s[kind];
|
|
assert.equal((await s.api(operation.path, operation.body)).status, 200);
|
|
assert.equal((await s.api("/chat", {chatId: s.chat.id, message: "Over budget"})).status, 429);
|
|
assert.equal(s.calls.length, 1);
|
|
});
|
|
}
|
|
|
|
for (const kind of ["roleplay", "market"]) {
|
|
test(`concurrent ${kind} starts save exactly one opener and preserve the first scene`, {timeout: 10000}, async t => {
|
|
const s = await fixture(t);
|
|
const gate = s.hold();
|
|
const operation = s[kind];
|
|
const first = s.api(operation.path, operation.body);
|
|
await gate.started.promise;
|
|
assert.equal((await s.api(operation.path, {...operation.body, scenarioPrompt: "Overlapping scene"})).status, 409);
|
|
// The other opener route must use the very same lock.
|
|
const openingChat = s.repo.listChats("generator").find(chat => chat.id !== s.chat.id);
|
|
assert.equal((await s.api("/roleplays/start", {assistantPersonaId: openingChat.assistant_persona_id})).status, 409);
|
|
gate.release.resolve();
|
|
const response = await first;
|
|
assert.equal(response.status, 200);
|
|
const result = await response.json();
|
|
assert.equal(result.generatedInitialMessage, true);
|
|
assert.equal(result.chat.scenario_prompt, "First scene");
|
|
assert.equal(s.repo.listChatMessages(result.chat.id).length, 1);
|
|
assert.equal((await (await s.api(operation.path, operation.body)).json()).generatedInitialMessage, false);
|
|
assert.equal(s.calls.length, 1);
|
|
});
|
|
}
|
|
|
|
for (const kind of ["roleplay", "market", "edit"]) {
|
|
for (const cancellation of ["stop", "disconnect"]) {
|
|
test(`${kind} ${cancellation} aborts the model, prevents late saves and releases the lock`, {timeout: 10000}, async t => {
|
|
const s = await fixture(t);
|
|
const gate = s.hold();
|
|
const operation = s[kind];
|
|
const controller = new AbortController();
|
|
const pending = s.api(operation.path, operation.body, controller.signal);
|
|
await gate.started.promise;
|
|
const generatedChat = s.repo.listChats("generator").find(chat => chat.id !== s.chat.id);
|
|
const stopChatId = kind === "edit" ? s.chat.id : generatedChat.id;
|
|
if (cancellation === "stop") {
|
|
assert.equal((await (await s.api(`/chats/${stopChatId}/stop`, {})).json()).stopped, true);
|
|
} else {
|
|
const rejected = assert.rejects(pending, {name: "AbortError"});
|
|
controller.abort();
|
|
await rejected;
|
|
}
|
|
await gate.aborted.promise;
|
|
gate.release.resolve();
|
|
if (cancellation === "stop") await (await pending).json();
|
|
// Wait for the disconnected handler to finish before checking its database effects.
|
|
for (let attempt = 0; attempt < 100; attempt++) {
|
|
if (!(await (await s.api(`/chats/${stopChatId}/stop`, {})).json()).stopped) break;
|
|
await new Promise(resolve => setTimeout(resolve, 5));
|
|
}
|
|
assert.equal(s.repo.listChatMessages(generatedChat.id).filter(message => message.role === "bot").length, 0);
|
|
assert.equal(s.repo.listChatMessages(s.chat.id).filter(message => message.role === "bot").length, 1);
|
|
assert.equal((await s.api(operation.path, operation.body)).status, 200);
|
|
assert.equal(s.calls.length, 2);
|
|
});
|
|
}
|
|
}
|
|
|
|
test("edit/resend locks its source before cloning and shares exclusion with regeneration", {timeout: 10000}, async t => {
|
|
const s = await fixture(t);
|
|
const gate = s.hold();
|
|
const first = s.api(s.edit.path, s.edit.body);
|
|
await gate.started.promise;
|
|
assert.equal((await s.api(s.edit.path, s.edit.body)).status, 409);
|
|
assert.equal((await s.api(`/chats/${s.chat.id}/messages/${s.answerId}/retry`, {})).status, 409);
|
|
assert.equal(s.repo.listChats("generator").length, 2);
|
|
gate.release.resolve();
|
|
const result = await (await first).json();
|
|
assert.equal(result.messages.at(-1).thinking, "Reasoning");
|
|
assert.equal(s.calls[0].options.reasoning, true);
|
|
assert.equal(s.repo.listChatMessages(s.chat.id).length, 2);
|
|
});
|