add generation cancellation middleware and refactor chat/message workflows to support HTTP client disconnections
This commit is contained in:
@@ -0,0 +1,22 @@
|
||||
import test from "node:test";
|
||||
import assert from "node:assert/strict";
|
||||
|
||||
test("conversation navigation", async t => {
|
||||
const original = Object.fromEntries(["document", "window", "history", "location"].map(key => [key, globalThis[key]]));
|
||||
const elements = new Map();
|
||||
globalThis.document = {
|
||||
getElementById(id) {
|
||||
if (!elements.has(id)) elements.set(id, {value: "", textContent: "", querySelectorAll: () => []});
|
||||
return elements.get(id);
|
||||
},
|
||||
querySelectorAll: () => [], querySelector: () => null
|
||||
};
|
||||
globalThis.location = {pathname: "/app/chats/1", search: ""};
|
||||
globalThis.window = {location: globalThis.location};
|
||||
globalThis.history = {replaceState(_data, _unused, path) { globalThis.location.pathname = path; }};
|
||||
t.after(() => { for (const [key, value] of Object.entries(original)) globalThis[key] = value; });
|
||||
const {registerConversationTests} = await import("../browser-tests/conversationCases.js");
|
||||
const cases = [];
|
||||
registerConversationTests((name, run) => cases.push({name, run}), assert);
|
||||
for (const {name, run} of cases) await t.test(name, run);
|
||||
});
|
||||
@@ -0,0 +1,169 @@
|
||||
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);
|
||||
});
|
||||
@@ -0,0 +1,163 @@
|
||||
import test from "node:test";
|
||||
import assert from "node:assert/strict";
|
||||
import {createDatabase} from "../src/db/index.js";
|
||||
import {createRepositories} from "../src/db/repositories.js";
|
||||
import {createConfig} from "../src/config.js";
|
||||
import {validateWorkspaceImport} from "../src/services/workspaceImportService.js";
|
||||
import {createWorkspaceArchive, parseWorkspaceArchive} from "../src/services/workspaceArchive.js";
|
||||
|
||||
function fixture(t) {
|
||||
const config = createConfig({dbPath: ":memory:", testMode: true});
|
||||
const db = createDatabase(config);
|
||||
t.after(() => db.close());
|
||||
const repo = createRepositories(db, config);
|
||||
for (const user of ["source", "destination", "second-copy"]) repo.insertUser(user, "unused");
|
||||
return {db, repo};
|
||||
}
|
||||
|
||||
function restore(repo, workspace, user = "destination") {
|
||||
const validation = validateWorkspaceImport(workspace);
|
||||
assert.deepEqual(validation.errors, []);
|
||||
repo.importWorkspace(user, validation.workspace);
|
||||
return repo.exportWorkspace(user);
|
||||
}
|
||||
|
||||
for (const format of ["JSON", "archive"]) {
|
||||
test(`${format} export/import round trips remap personas, branches, prompts, state, memories and versions`, t => {
|
||||
const {db, repo} = fixture(t);
|
||||
// Identical display names must not be used as relationship keys.
|
||||
const assistant = repo.createPersona("source", {name: "Alex", details: "Guide"}, "assistant");
|
||||
const userPersona = repo.createPersona("source", {name: "Alex", details: "Explorer"}, "user");
|
||||
const original = repo.createChat("source", "Original", assistant.id, userPersona.id, "At the station", "Meeting");
|
||||
const prompt = Number(repo.insertBranchChatMessage(original.id, "user", "Question").lastInsertRowid);
|
||||
const answer = Number(repo.insertBranchChatMessage(original.id, "bot", "Answer", {
|
||||
retryVariants: ["Answer"], retryPromptMessageId: prompt
|
||||
}, "qwen3:14b", "complete", null, prompt).lastInsertRowid);
|
||||
repo.updateChatMessageThinking(original.id, answer, "Reasoning");
|
||||
repo.updateChatMessageSources(original.id, answer, [{title: "Source", url: "https://example.org/", snippet: "Evidence"}]);
|
||||
repo.saveMessageVersion(original.id, answer, {content: "Revised", thinking: "Revised reasoning", sources: [], modelName: "qwen3:14b"});
|
||||
repo.updateChatContextSummary(original.id, "source", "Summary", answer);
|
||||
repo.updatePersonaChatState("source", original, {tone: "Curious", current_location: "Station"});
|
||||
repo.createChatMemory("source", original, "Remember the station");
|
||||
repo.createPromptTemplate("source", {name: "Guide prompt", personaId: assistant.id, promptText: "Begin here"});
|
||||
const branch = repo.createChat("source", "Branch", assistant.id, userPersona.id);
|
||||
repo.updateChatBranchSource(branch.id, "source", original.id, answer);
|
||||
repo.updateChatGenerationSettings(branch.id, "source", {temperature: 0, contextLength: 4096, responseLength: 256});
|
||||
repo.insertBranchChatMessage(branch.id, "bot", "Another root");
|
||||
|
||||
let workspace = repo.exportWorkspace("source");
|
||||
if (format === "archive") workspace = parseWorkspaceArchive(createWorkspaceArchive(workspace, [])).workspace;
|
||||
else workspace = JSON.parse(JSON.stringify(workspace));
|
||||
// A branch can be exported before its source chat.
|
||||
workspace.chats.sort((a, b) => b.id - a.id);
|
||||
for (const destination of ["destination", "second-copy"]) {
|
||||
const restored = restore(repo, workspace, destination);
|
||||
const copied = restored.chats.find(chat => chat.stored_title === "Original");
|
||||
const copiedBranch = restored.chats.find(chat => chat.stored_title === "Branch");
|
||||
const copiedAssistant = restored.personas.find(persona => persona.persona_type === "assistant");
|
||||
const copiedUser = restored.personas.find(persona => persona.persona_type === "user");
|
||||
assert.notEqual(copied.id, original.id);
|
||||
assert.notEqual(copiedAssistant.id, assistant.id);
|
||||
assert.notEqual(copiedUser.id, userPersona.id);
|
||||
for (const record of [copied, copiedBranch, copied.memories[0], copied.persona_state]) {
|
||||
assert.equal(record.assistant_persona_id, copiedAssistant.id);
|
||||
assert.equal(record.user_persona_id, copiedUser.id);
|
||||
}
|
||||
const question = copied.messages.find(message => message.role === "user");
|
||||
const reply = copied.messages.find(message => message.role === "bot");
|
||||
assert.equal(reply.parent_message_id, question.id);
|
||||
assert.equal(reply.retry_prompt_message_id, question.id);
|
||||
assert.equal(copied.active_leaf_message_id, reply.id);
|
||||
assert.equal(copied.context_summary_message_id, reply.id);
|
||||
assert.equal(copied.context_summary, "Summary");
|
||||
assert.equal(copiedBranch.source_chat_id, copied.id);
|
||||
assert.equal(copiedBranch.branched_from_message_id, reply.id);
|
||||
assert.deepEqual(repo.listActiveBranchMessages(copied.id).map(message => message.content), ["Question", "Revised"]);
|
||||
assert.deepEqual(reply.responseVersions, workspace.chats.find(chat => chat.stored_title === "Original").messages.find(message => message.role === "bot").responseVersions);
|
||||
assert.equal(reply.thinking, "Revised reasoning");
|
||||
assert.equal(copied.temperature, null);
|
||||
assert.equal(copied.context_length, null);
|
||||
assert.equal(copied.response_length, null);
|
||||
assert.equal(copiedBranch.temperature, 0);
|
||||
assert.equal(copiedBranch.context_length, 4096);
|
||||
assert.equal(copiedBranch.response_length, 256);
|
||||
assert.equal(copied.persona_state.tone, "Curious");
|
||||
assert.equal(copied.memories[0].fact, "Remember the station");
|
||||
assert.equal(restored.templates.find(template => template.name === "Guide prompt").persona_id, copiedAssistant.id);
|
||||
assert.deepEqual(db.pragma("foreign_key_check"), []);
|
||||
workspace = restored;
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
test("legacy two-message histories with and without IDs restore a complete active chain", t => {
|
||||
const {repo} = fixture(t);
|
||||
for (const withIds of [false, true]) {
|
||||
const title = `Legacy ${withIds}`;
|
||||
const restored = restore(repo, {chats: [{title, temperature: null, messages: [
|
||||
{role: "user", content: "Question", ...(withIds ? {id: 41} : {})},
|
||||
{role: "assistant", content: "Answer", ...(withIds ? {id: 42} : {})}
|
||||
]}]});
|
||||
const chat = restored.chats.find(item => item.stored_title === title);
|
||||
assert.deepEqual(repo.listActiveBranchMessages(chat.id).map(message => message.content), ["Question", "Answer"]);
|
||||
assert.equal(chat.messages[1].parent_message_id, chat.messages[0].id);
|
||||
assert.equal(chat.active_leaf_message_id, chat.messages[1].id);
|
||||
assert.equal(chat.temperature, null);
|
||||
}
|
||||
});
|
||||
|
||||
test("explicit roots stay separate and unmapped references never bind existing records", t => {
|
||||
const {repo} = fixture(t);
|
||||
const foreignPersona = repo.createPersona("source", {name: "Other account"}, "assistant");
|
||||
const foreignChat = repo.createChat("source", "Other chat");
|
||||
const foreignMessage = Number(repo.insertChatMessage(foreignChat.id, "user", "Private").lastInsertRowid);
|
||||
const restored = restore(repo, {templates: [{name: "Detached", prompt_text: "Hello", persona_id: foreignPersona.id}], chats: [{
|
||||
title: "Roots", assistant_persona_id: foreignPersona.id, source_chat_id: foreignChat.id, branched_from_message_id: foreignMessage,
|
||||
active_leaf_message_id: 20, messages: [
|
||||
{id: 10, role: "user", content: "Root one", parent_message_id: null},
|
||||
{id: 20, role: "bot", content: "Root two", parent_message_id: null, retry_prompt_message_id: foreignMessage}
|
||||
]
|
||||
}]});
|
||||
const chat = restored.chats[0];
|
||||
assert.equal(chat.assistant_persona_id, null);
|
||||
assert.equal(chat.source_chat_id, null);
|
||||
assert.equal(chat.branched_from_message_id, null);
|
||||
assert.equal(chat.messages[1].retry_prompt_message_id, null);
|
||||
assert.ok(chat.messages.every(message => message.parent_message_id === null));
|
||||
assert.deepEqual(repo.listActiveBranchMessages(chat.id).map(message => message.content), ["Root two"]);
|
||||
assert.equal(restored.templates.find(template => template.name === "Detached").persona_id, null);
|
||||
});
|
||||
|
||||
test("message references are remapped even when the parent is imported later", t => {
|
||||
const {repo} = fixture(t);
|
||||
const restored = restore(repo, {chats: [{title: "Forward reference", active_leaf_message_id: 20, messages: [
|
||||
{id: 20, role: "bot", content: "Answer", parentMessageId: 10, retryPromptMessageId: 10},
|
||||
{id: 10, role: "user", content: "Question", parentMessageId: null}
|
||||
]}]});
|
||||
const [answer, prompt] = restored.chats[0].messages;
|
||||
assert.equal(answer.parent_message_id, prompt.id);
|
||||
assert.equal(answer.retry_prompt_message_id, prompt.id);
|
||||
assert.equal(restored.chats[0].active_leaf_message_id, answer.id);
|
||||
});
|
||||
|
||||
test("optional generation settings preserve unset values and numeric zero", () => {
|
||||
for (const value of [null, undefined, "", " ", 0, "0", 0.7]) {
|
||||
const {workspace} = validateWorkspaceImport({chats: [{title: "Settings", temperature: value, context_length: value, response_length: value}]});
|
||||
const chat = workspace.chats[0];
|
||||
const expected = value === 0 || value === "0" ? 0 : value === 0.7 ? 0.7 : null;
|
||||
assert.equal(chat.temperature, expected);
|
||||
assert.equal(chat.context_length, expected === 0.7 ? null : expected);
|
||||
assert.equal(chat.response_length, expected === 0.7 ? null : expected);
|
||||
}
|
||||
});
|
||||
|
||||
test("ambiguous IDs and cyclic message chains fail validation", () => {
|
||||
for (const messages of [
|
||||
[{id: 1}, {id: 1}],
|
||||
[{id: 1, parent_message_id: 2}, {id: 2, parent_message_id: 1}]
|
||||
]) {
|
||||
const validation = validateWorkspaceImport({chats: [{messages: messages.map(message => ({role: "user", content: "Text", ...message}))}]});
|
||||
assert.equal(validation.ok, false);
|
||||
assert.match(validation.errors.join(" "), /duplicate message ID|cyclic message chain/);
|
||||
}
|
||||
});
|
||||
Reference in New Issue
Block a user