From 1c0f06716ea9c02c382ed5d2d9ae6ef9d91597d8 Mon Sep 17 00:00:00 2001 From: fretuks Date: Mon, 5 Oct 2026 23:21:42 +0200 Subject: [PATCH] add generation cancellation middleware and refactor chat/message workflows to support HTTP client disconnections --- README.md | 13 +- package.json | 3 +- public/js/conversationController.js | 90 ++++++------ public/js/market.js | 2 +- public/js/personaController.js | 2 +- public/js/script.js | 12 +- scripts/browser-tests.js | 10 ++ src/app.js | 2 +- src/db/repositories.js | 45 ++++-- src/middleware/generationCancellation.js | 13 ++ src/routes/chats.js | 28 ++-- src/routes/personas.js | 58 ++------ src/services/chatService.js | 91 +++++++----- src/services/workspaceImportService.js | 69 +++++++-- test/conversation.test.js | 22 +++ test/generationProtection.test.js | 169 +++++++++++++++++++++++ test/workspaceImport.test.js | 163 ++++++++++++++++++++++ 17 files changed, 620 insertions(+), 172 deletions(-) create mode 100644 scripts/browser-tests.js create mode 100644 src/middleware/generationCancellation.js create mode 100644 test/conversation.test.js create mode 100644 test/generationProtection.test.js create mode 100644 test/workspaceImport.test.js diff --git a/README.md b/README.md index 70e0f99..43b54e9 100644 --- a/README.md +++ b/README.md @@ -29,6 +29,7 @@ on Node 22 and 24. - `npm run dev`: run the server with Node watch mode. - `npm run lint`: syntax-check the backend and browser modules. - `npm test`: run API tests with Node's built-in test runner. +- `npm run test:browser`: serve the navigation regression harness locally; open the printed URL in a browser. The same cases run with a lightweight DOM fixture in `npm test`. - `npm run user:role -- `: assign or revoke a role from the server console. ## Administrative access @@ -75,7 +76,7 @@ Duplicate registrations, including concurrent inserts, return HTTP 409 with - `MODEL_UNLOAD_AFTER_MS`: delay before idle local request bookkeeping is cleared. Model unloading is managed by Ollama through `OLLAMA_KEEP_ALIVE`. - `MODELS_CACHE_TTL_MS`: model list cache TTL. - `AUTH_RATE_LIMIT_*`: rate-limit window and thresholds for registration/login. -- `CHAT_RATE_LIMIT_*`: rate-limit window and thresholds for chat and retry endpoints. +- `CHAT_RATE_LIMIT_*`: shared rate-limit window and thresholds for chat, retry, roleplay openers, marketplace openers, and edit/resend endpoints. - `RATE_LIMIT_CLEANUP_INTERVAL_MS`: interval for pruning stale in-memory rate-limit buckets. - `DROP_LEGACY_CHATS`: when `true`, migrate any legacy `chats` table into `chat_sessions` / `chat_messages` and remove the old table. @@ -107,6 +108,11 @@ accepts older JSON exports, which contain attachment references only and reconne images only when those IDs already belong to the importing account. Archives are limited to 250 MB uncompressed, and each image must meet the destination's upload limit. +Imports remap persona, chat and message IDs, including template, memory, continuity, +branch, retry and summary references. References absent from the import are left +unassigned. Legacy histories without parent fields are restored as ordered message +chains; explicit branch roots remain separate. Unset generation settings remain unset. + `CHAT_LIMIT` defaults to 10 active chats per account. Archived chats remain available without using an active slot; restoring one requires a free slot. @@ -215,6 +221,11 @@ request still uses `stream: true` and the existing capability-checked reasoning Stop/disconnect cancels the current vision or final generation request. Logging includes IDs, sizes, types and cache events only, never images/Base64/private analysis. +Roleplay and marketplace initialization use the same chat generation lock and +cancellation tracking as other replies. Concurrent initializations return HTTP 409 +and save at most one opener. Edit/resend holds the source chat's lock while generating +its branch, and Stop can cancel it through the source chat or the new branch. + Dependencies added: `multer` (bounded multipart handling), `sharp` (content detection, dimensions and full image decoding). Install through the normal `npm install`/`npm ci` workflow, including Sharp's platform-specific optional dependencies in Docker. The diff --git a/package.json b/package.json index 264195c..868094a 100644 --- a/package.json +++ b/package.json @@ -11,7 +11,8 @@ "dev": "node --env-file-if-exists=.env --watch server.js", "user:role": "node --env-file-if-exists=.env scripts/set-role.js", "lint": "eslint .", - "test": "node --test" + "test": "node --test", + "test:browser": "node scripts/browser-tests.js" }, "dependencies": { "bcrypt": "^6.0.0", diff --git a/public/js/conversationController.js b/public/js/conversationController.js index 5ad98b1..e23bad1 100644 --- a/public/js/conversationController.js +++ b/public/js/conversationController.js @@ -40,54 +40,60 @@ export function createConversationController({ async function setActiveChat(id, {preserveComposer = false, matchMessageId = null} = {}) { const openRequest = ++state.chatOpenRequest; - if (!preserveComposer && id !== state.activeChatId) { - if (state.activeChatId != null) { - state.chatDrafts.set(state.activeChatId, msgInput.value); - attachmentComposer.saveDraft(state.activeChatId); - msgInput.value = state.chatDrafts.get(id) || ""; - attachmentComposer.restoreDraft(id); - resizeComposerInput(); - } else if (!msgInput.value && !attachmentComposer.pending && !attachmentComposer.attachments.length) { - msgInput.value = state.chatDrafts.get(id) || ""; - attachmentComposer.restoreDraft(id); - resizeComposerInput(); - } - } if (isMobileLayout()) closeMobileSidebar(); setChatLoading(true, "Opening chat..."); setChatActivity(true, { - eyebrow: "Opening chat", - title: "Loading conversation", + eyebrow: "Opening chat", title: "Loading conversation", detail: "Fetching messages and restoring the current scene." }); - state.activeChatId = id; - history.replaceState(null, "", `/app/chats/${id}`); - state.messageRetryState = new Map(); - updateWorkspaceCopy(); renderChatList(); updateChatActionState(); - let res = await get(`/chats/${id}/messages`); - if (openRequest !== state.chatOpenRequest) { setChatLoading(false); setChatActivity(false); return; } - if (res.error) { - setChatLoading(false); - setChatActivity(false); - return setNotice(res.error, "error"); + try { + let res = await get(`/chats/${id}/messages`); + if (openRequest !== state.chatOpenRequest) return; + if (res.error) return setNotice(res.error, "error"); + if (matchMessageId && !res.messages?.some(message => message.id === matchMessageId)) { + const activated = await post(`/chats/${id}/branches/${matchMessageId}/activate`, {includeDescendants: true}); + if (openRequest !== state.chatOpenRequest) return; + if (activated.error) return setNotice(activated.error, "error"); + res = await get(`/chats/${id}/messages`); + if (openRequest !== state.chatOpenRequest) return; + if (res.error) return setNotice(res.error, "error"); + } + // Commit identity, content and drafts together, only for the latest successful load. + if (!preserveComposer && id !== state.activeChatId) { + if (state.activeChatId != null) { + state.chatDrafts.set(state.activeChatId, msgInput.value); + attachmentComposer.saveDraft(state.activeChatId); + msgInput.value = state.chatDrafts.get(id) || ""; + attachmentComposer.restoreDraft(id); + resizeComposerInput(); + } else if (!msgInput.value && !attachmentComposer.pending && !attachmentComposer.attachments.length) { + msgInput.value = state.chatDrafts.get(id) || ""; + attachmentComposer.restoreDraft(id); + resizeComposerInput(); + } + } + const existing = getChatById(id); + if (res.chat && existing) Object.assign(existing, res.chat); + state.activeChatId = id; + state.currentMessages = res.messages || []; + history.replaceState(null, "", `/app/chats/${id}`); + state.messageRetryState = new Map(); + updateWorkspaceCopy(); renderChatList(); updateChatActionState(); + populateGenerationSettings(res.chat || existing); + renderMessages(); + if (matchMessageId) requestAnimationFrame(() => { + if (openRequest !== state.chatOpenRequest) return; + const target = [...messagesDiv.querySelectorAll("[data-message-id]")] + .find(node => Number(node.dataset.messageId) === Number(matchMessageId)); + if (target) { target.scrollIntoView({block: "center"}); updateJumpToLatest(); } + }); + updateSendState(); + } finally { + if (openRequest === state.chatOpenRequest) { + setChatLoading(false); + setChatActivity(false); + } } - if (matchMessageId && !res.messages?.some(message => message.id === matchMessageId)) { - const activated = await post(`/chats/${id}/branches/${matchMessageId}/activate`, {includeDescendants: true}); - if (!activated.error) res = await get(`/chats/${id}/messages`); - if (openRequest !== state.chatOpenRequest) { setChatLoading(false); setChatActivity(false); return; } - } - const existing = getChatById(id); - if (res.chat && existing) Object.assign(existing, res.chat); - populateGenerationSettings(res.chat || existing); - state.currentMessages = res.messages || []; renderMessages(); - if (matchMessageId) requestAnimationFrame(() => { - const target = [...messagesDiv.querySelectorAll("[data-message-id]")] - .find(node => Number(node.dataset.messageId) === Number(matchMessageId)); - if (target) { target.scrollIntoView({block: "center"}); updateJumpToLatest(); } - }); - updateSendState(); - setChatLoading(false); - setChatActivity(false); } async function loadChatSessions() { diff --git a/public/js/market.js b/public/js/market.js index 314412d..a34708b 100644 --- a/public/js/market.js +++ b/public/js/market.js @@ -554,7 +554,7 @@ async function startMarketPersonaChat(marketId, userPersonaSelection, scenarioPr const res = await post(`/personas/market/${marketId}/chat`, { userPersonaId, scenarioPrompt: scenarioPrompt.trim() - }); + }, {timeoutMs: 660000}); if (res.error || !res.chat) { setMarketActivity(false); setMarketStatus(res.error || "Unable to start chat.", "error"); diff --git a/public/js/personaController.js b/public/js/personaController.js index a3074cf..1bd6c2f 100644 --- a/public/js/personaController.js +++ b/public/js/personaController.js @@ -385,7 +385,7 @@ export function createPersonaController({ userPersonaId, scenarioPrompt, model: modelSelect.value - }); + }, {timeoutMs: 660000}); roleplayStarterConfirm.disabled = false; if (res.error || !res.chat) { setChatLoading(false); diff --git a/public/js/script.js b/public/js/script.js index 54de54f..f2fcd6f 100644 --- a/public/js/script.js +++ b/public/js/script.js @@ -496,6 +496,7 @@ async function editChatMessage(messageId, fallbackIndex = -1) { const target = await resolveMessageTarget(messageId, fallbackIndex); if (!target) return setNotice("Message is not ready yet. Try again.", "error"); const { index, message } = target; + const sourceChatId = conversation.activeChatId; const nextContent = await promptPopup({ eyebrow: message.role === "user" ? "Your message" : "Assistant message", title: "Edit message", @@ -513,8 +514,8 @@ async function editChatMessage(messageId, fallbackIndex = -1) { description: "This creates a new chat from the messages before this one and sends your edited message there. The original chat stays unchanged.", confirmLabel: "Create branch" }); - if (!confirmed) return; - const sourceChatId = conversation.activeChatId; + if (!confirmed || conversation.isProcessing) return; + conversation.regenerationChatId = sourceChatId; setLoadingState(true, { eyebrow: "Branching", title: "Editing and resending", @@ -526,11 +527,12 @@ async function editChatMessage(messageId, fallbackIndex = -1) { message: content, model: modelSelect.value, confirmed: true - }); + }, {timeoutMs: 660000}); if (res.error) return setNotice(res.error, "error"); - await openCreatedChat(res.chat); - return setNotice("Created and opened the edited chat branch.", "success"); + if (conversation.activeChatId === sourceChatId) await openCreatedChat(res.chat); + return setNotice("Created the edited chat branch.", "success"); } finally { + conversation.regenerationChatId = null; setLoadingState(false); } } diff --git a/scripts/browser-tests.js b/scripts/browser-tests.js new file mode 100644 index 0000000..cc10302 --- /dev/null +++ b/scripts/browser-tests.js @@ -0,0 +1,10 @@ +/** Local-only browser regression harness; serves no account data or workspace files. */ +import express from "express"; +import {fileURLToPath} from "node:url"; + +const app = express(); +app.use("/public/js", express.static(fileURLToPath(new URL("../public/js", import.meta.url)))); +app.use("/browser-tests", express.static(fileURLToPath(new URL("../browser-tests", import.meta.url)))); +const server = app.listen(0, "127.0.0.1", () => { + console.log(`Navigation tests: http://127.0.0.1:${server.address().port}/browser-tests/conversation.html`); +}); diff --git a/src/app.js b/src/app.js index 4027425..8cef7f5 100644 --- a/src/app.js +++ b/src/app.js @@ -321,7 +321,7 @@ export function createApp(options = {}) { if (config.visionDebug) app.use(createVisionDebugRouter({visionDebugService, repositories})); app.use(createSettingsRouter({repositories})); app.use(createChatsRouter({repositories, chatService, modelService, config, chatRateLimiters, attachmentService})); - app.use(createPersonasRouter({repositories, chatService, modelService, config})); + app.use(createPersonasRouter({repositories, chatService, modelService, config, chatRateLimiters})); app.use(createSystemRouter({modelService, config})); app.use(notFoundHandler); app.use(errorHandler); diff --git a/src/db/repositories.js b/src/db/repositories.js index b7dbf97..2c77260 100644 --- a/src/db/repositories.js +++ b/src/db/repositories.js @@ -201,6 +201,7 @@ export function createRepositories(db, config) { updateChatMessageDelivery: db.prepare("UPDATE chat_messages SET delivery_status = ?, error_message = ? WHERE id = ? AND chat_id = ?"), updateChatMessageSources: db.prepare("UPDATE chat_messages SET sources = ? WHERE id = ? AND chat_id = ?"), updateChatMessageThinking: db.prepare("UPDATE chat_messages SET thinking = ? WHERE id = ? AND chat_id = ?"), + updateChatMessageLinks: db.prepare("UPDATE chat_messages SET parent_message_id = ?, retry_prompt_message_id = ? WHERE id = ? AND chat_id = ?"), getChatMessages: db.prepare(` SELECT thinking, sources, id, role, content, model_name, retry_variants, retry_active_index, retry_retries_used, retry_prompt_message_id, delivery_status, error_message, parent_message_id, created_at FROM chat_messages @@ -821,6 +822,8 @@ export function createRepositories(db, config) { const importWorkspaceTransaction = db.transaction((username, workspace) => { const imported = {personas: 0, chats: 0, templates: 0}; + const personaIds = new Map(); + const importedChats = []; for (const persona of workspace.personas || []) { const personaId = statements.insertPersona.run( username, @@ -832,14 +835,15 @@ export function createRepositories(db, config) { persona.example_dialogues ?? null, persona.persona_type || "assistant" ).lastInsertRowid; - statements.insertPersonaVersion.run(personaId, username, 1, JSON.stringify(persona)); + if (persona.sourceId != null) personaIds.set(persona.sourceId, Number(personaId)); + statements.insertPersonaVersion.run(personaId, username, 1, buildPersonaSnapshot(statements.getPersona.get(personaId, username))); imported.personas += 1; } for (const template of workspace.templates || []) { statements.insertPromptTemplate.run( username, - null, + personaIds.get(template.persona_id) ?? null, template.name, template.description ?? null, template.category ?? null, @@ -851,11 +855,13 @@ export function createRepositories(db, config) { } for (const chat of workspace.chats || []) { + const assistantPersonaId = personaIds.get(chat.assistant_persona_id) ?? null; + const userPersonaId = personaIds.get(chat.user_persona_id) ?? null; const chatId = statements.insertChatSession.run( username, chat.stored_title || chat.title || "Imported chat", - null, - null, + assistantPersonaId, + userPersonaId, chat.folder_name ?? null, chat.is_pinned ? 1 : 0, chat.archived_at ?? null, @@ -863,9 +869,8 @@ export function createRepositories(db, config) { chat.scenario_summary ?? null ).lastInsertRowid; const importedMessageIds = new Map(); + importedChats.push({chat, chatId, messageIds: importedMessageIds}); for (const message of chat.messages || []) { - const mappedRetryPromptId = importedMessageIds.get(message.retryPromptMessageId) ?? null; - const mappedParentId = importedMessageIds.get(message.parentMessageId) ?? null; const result = statements.insertChatMessage.run( chatId, message.role, @@ -874,10 +879,10 @@ export function createRepositories(db, config) { JSON.stringify(message.retryVariants || [message.content || ""]), Number(message.retryActiveIndex || 0), Number(message.retryRetriesUsed || 0), - mappedRetryPromptId, + null, message.deliveryStatus || "complete", message.errorMessage || null, - mappedParentId + null ); if (message.sourceId != null) importedMessageIds.set(message.sourceId, Number(result.lastInsertRowid)); if (message.sources) statements.updateChatMessageSources.run(JSON.stringify(message.sources), result.lastInsertRowid, chatId); @@ -893,6 +898,15 @@ export function createRepositories(db, config) { db.prepare("UPDATE attachments SET status='attached' WHERE id=? AND username=?").run(attachment.id, username); } } + // Resolve links after every message has an ID, including forward references. + for (const message of chat.messages || []) { + statements.updateChatMessageLinks.run( + importedMessageIds.get(message.parentMessageId) ?? null, + importedMessageIds.get(message.retryPromptMessageId) ?? null, + importedMessageIds.get(message.sourceId), + chatId + ); + } statements.updateChatGenerationSettings.run( chat.preferred_model ?? null, chat.temperature ?? null, @@ -918,15 +932,17 @@ export function createRepositories(db, config) { for (const memory of chat.memories || []) { const fact = String(memory?.fact || "").trim(); if (!fact) continue; - statements.insertChatMemory.run(username, chatId, null, null, fact); + statements.insertChatMemory.run(username, chatId, + memory.assistant_persona_id === undefined ? assistantPersonaId : personaIds.get(memory.assistant_persona_id) ?? null, + memory.user_persona_id === undefined ? userPersonaId : personaIds.get(memory.user_persona_id) ?? null, fact); } const personaState = chat.persona_state || chat.personaChatState; if (personaState) { statements.upsertPersonaChatState.run( username, chatId, - null, - null, + personaIds.get(personaState.assistant_persona_id) ?? null, + personaIds.get(personaState.user_persona_id) ?? null, personaState.relationship_notes ?? personaState.relationshipNotes ?? null, personaState.tone ?? null, personaState.current_location ?? personaState.currentLocation ?? null, @@ -938,6 +954,13 @@ export function createRepositories(db, config) { imported.chats += 1; } + const chatsBySourceId = new Map(importedChats.filter(({chat}) => chat.sourceId != null).map(entry => [entry.chat.sourceId, entry])); + for (const {chat, chatId} of importedChats) { + const sourceChat = chatsBySourceId.get(chat.source_chat_id); + statements.updateChatBranchSource.run(sourceChat?.chatId ?? null, + sourceChat?.messageIds.get(chat.branched_from_message_id) ?? null, chatId, username); + } + return imported; }); diff --git a/src/middleware/generationCancellation.js b/src/middleware/generationCancellation.js new file mode 100644 index 0000000..1b08155 --- /dev/null +++ b/src/middleware/generationCancellation.js @@ -0,0 +1,13 @@ +/** Cancel generation when its HTTP client disconnects, for JSON and streamed replies. */ +export function generationCancellation(req, res, next) { + const controller = new AbortController(); + const onClose = () => { + if (!res.writableEnded) controller.abort(); + res.removeListener("finish", onFinish); + }; + const onFinish = () => res.removeListener("close", onClose); + res.once("close", onClose); + res.once("finish", onFinish); + req.generationSignal = controller.signal; + next(); +} diff --git a/src/routes/chats.js b/src/routes/chats.js index 339b71a..f806092 100644 --- a/src/routes/chats.js +++ b/src/routes/chats.js @@ -4,6 +4,7 @@ import express from "express"; import {createTextStreamEvents} from "../services/streamEvents.js"; import {asyncHandler} from "../middleware/asyncHandler.js"; +import {generationCancellation} from "../middleware/generationCancellation.js"; import multer from "multer"; import {requireLogin} from "../middleware/auth.js"; import {validateBody} from "../middleware/validate.js"; @@ -264,7 +265,7 @@ export function createChatsRouter({repositories, chatService, modelService, conf return res.json(result); }); - router.post("/chats/:id/messages/:messageId/edit-resend", validateBody(sendMessageValidator), asyncHandler(async (req, res) => { + router.post("/chats/:id/messages/:messageId/edit-resend", ...chatRateLimiters, generationCancellation, validateBody(sendMessageValidator), asyncHandler(async (req, res) => { if (req.body?.confirmed !== true) return res.status(400).json({error: "Confirmation is required before creating an edited branch"}); const chatId = Number(req.params.id); const targetMessage = repositories.getChatMessage(chatId, Number(req.params.messageId)); @@ -272,6 +273,7 @@ export function createChatsRouter({repositories, chatService, modelService, conf if (!targetMessage) return res.status(404).json({error: "Message not found"}); try { const result = await chatService.editAndResendMessage({ + signal: req.generationSignal, user: req.session.user, requestId: req.id, chatId, @@ -448,9 +450,10 @@ export function createChatsRouter({repositories, chatService, modelService, conf return res.json({state}); }); - router.post("/chat", ...chatRateLimiters, validateBody(sendMessageValidator), asyncHandler(async (req, res) => { + router.post("/chat", ...chatRateLimiters, generationCancellation, validateBody(sendMessageValidator), asyncHandler(async (req, res) => { try { const result = await chatService.sendChatMessage({ + signal: req.generationSignal, user: req.session.user, requestId: req.id, chatId: req.validatedBody.chatId, @@ -468,12 +471,7 @@ export function createChatsRouter({repositories, chatService, modelService, conf } })); - router.post("/chat/stream", ...chatRateLimiters, validateBody(sendMessageValidator), asyncHandler(async (req, res) => { - const generationController = new AbortController(); - let generationFinished = false; - res.once("close", () => { - if (!generationFinished) generationController.abort(); - }); + router.post("/chat/stream", ...chatRateLimiters, generationCancellation, validateBody(sendMessageValidator), asyncHandler(async (req, res) => { res.writeHead(200, { "Content-Type": "text/event-stream", "Cache-Control": "no-cache, no-transform", @@ -499,7 +497,7 @@ export function createChatsRouter({repositories, chatService, modelService, conf webSearch: req.validatedBody.webSearch, reasoning: req.validatedBody.reasoning, onStatus: (status) => sendEvent("status", status), - signal: generationController.signal + signal: req.generationSignal }); if (result.error) { sendEvent("error", {error: result.error, status: result.status, requestId: req.id, ...(result.searchState ? {searchState: result.searchState} : {})}); @@ -510,7 +508,6 @@ export function createChatsRouter({repositories, chatService, modelService, conf const mapped = modelService.mapError(error); sendEvent("error", {...mapped.body, status: mapped.status, requestId: req.id}); } finally { - generationFinished = true; res.end(); } })); @@ -529,7 +526,7 @@ export function createChatsRouter({repositories, chatService, modelService, conf res.json({message: chatService.formatChatMessage(repositories.getChatMessage(chatId, messageId))}); }); - router.post(["/chats/:id/messages/:messageId/retry", "/chats/:id/messages/:messageId/retry/stream"], ...chatRateLimiters, asyncHandler(async (req, res) => { + router.post(["/chats/:id/messages/:messageId/retry", "/chats/:id/messages/:messageId/retry/stream"], ...chatRateLimiters, generationCancellation, asyncHandler(async (req, res) => { const chatId = Number(req.params.id); const messageId = Number(req.params.messageId); const chat = repositories.getChat(chatId, req.session.user); @@ -539,13 +536,10 @@ export function createChatsRouter({repositories, chatService, modelService, conf if (targetMessage.role !== "bot") return res.status(400).json({error: "Only assistant messages can be regenerated"}); const streaming = req.path.endsWith("/stream"); - const controller = new AbortController(); - let finished = false; const sendEvent = (event, data) => { if (!res.destroyed) res.write(`event: ${event}\ndata: ${JSON.stringify(data)}\n\n`); }; if (streaming) { - res.once("close", () => { if (!finished) controller.abort(); }); res.writeHead(200, {"Content-Type": "text/event-stream", "Cache-Control": "no-cache, no-transform", Connection: "keep-alive", "X-Accel-Buffering": "no"}); } try { @@ -558,7 +552,7 @@ export function createChatsRouter({repositories, chatService, modelService, conf retryStyle: String(req.body?.style || "").trim() || null, webSearch: req.body?.webSearch === true, reasoning: req.body?.reasoning === true, - signal: controller.signal, + signal: req.generationSignal, ...(streaming ? { ...createTextStreamEvents(sendEvent), onStatus: status => sendEvent("status", status) @@ -575,12 +569,11 @@ export function createChatsRouter({repositories, chatService, modelService, conf if (streaming) { sendEvent("error", {...mapped.body, status: mapped.status, requestId: req.id}); return; } return res.status(mapped.status).json(mapped.body); } finally { - finished = true; if (streaming) res.end(); } })); - router.post("/chats/:id/messages/by-index/:index/retry", ...chatRateLimiters, asyncHandler(async (req, res) => { + router.post("/chats/:id/messages/by-index/:index/retry", ...chatRateLimiters, generationCancellation, asyncHandler(async (req, res) => { const chatId = Number(req.params.id); const chat = repositories.getChat(chatId, req.session.user); if (!chat) return res.status(404).json({error: "Chat not found"}); @@ -590,6 +583,7 @@ export function createChatsRouter({repositories, chatService, modelService, conf try { const result = await chatService.retryMessage({ + signal: req.generationSignal, user: req.session.user, requestId: req.id, chatId, diff --git a/src/routes/personas.js b/src/routes/personas.js index 923f157..419c47b 100644 --- a/src/routes/personas.js +++ b/src/routes/personas.js @@ -3,6 +3,7 @@ */ import express from "express"; import {asyncHandler} from "../middleware/asyncHandler.js"; +import {generationCancellation} from "../middleware/generationCancellation.js"; import {requireLogin} from "../middleware/auth.js"; import {validateBody} from "../middleware/validate.js"; import {normalizeOptionalText, validatePersonaPayload} from "../services/personaService.js"; @@ -47,7 +48,7 @@ const templateValidator = (body) => { }; }; -export function createPersonasRouter({repositories, chatService, modelService, config}) { +export function createPersonasRouter({repositories, chatService, modelService, config, chatRateLimiters}) { const router = express.Router(); router.use(requireLogin); const isAdmin = (req) => req.account?.role === "admin"; @@ -215,7 +216,7 @@ export function createPersonasRouter({repositories, chatService, modelService, c return res.json({message: "Report submitted"}); }); - router.post("/personas/market/:id/chat", asyncHandler(async (req, res) => { + router.post("/personas/market/:id/chat", ...chatRateLimiters, generationCancellation, asyncHandler(async (req, res) => { const user = req.session.user; const marketPersona = repositories.getMarketPersona(Number(req.params.id)); const userPersonaId = req.body?.userPersonaId ? Number(req.body.userPersonaId) : null; @@ -251,40 +252,23 @@ export function createPersonasRouter({repositories, chatService, modelService, c throw error; } - if (repositories.listChatMessages(chat.id).length > 0) { - return res.json({persona, chat, existing: true, generatedInitialMessage: false}); - } - - repositories.updateChatScene(chat.id, user, scenarioPrompt, sceneSummary); - chat = repositories.getChat(chat.id, user); try { - let thinking = ""; - const opener = await chatService.generateRoleplayOpener({ - onThinking: (_chunk, fullThinking) => { thinking = fullThinking; }, + const result = await chatService.initializeRoleplay({ + user, chatId: chat.id, signal: req.generationSignal, selectedModel: "qwen3:14b", assistantPersona: persona, userPersona, scenarioPrompt, sceneSummary }); - if (!opener || chatService.containsUserVoiceInRoleplayOpener(opener, userPersona)) { - return res.json({persona, chat, existing: false, generatedInitialMessage: false}); - } - const inserted = chatService.insertChatMessage(chat.id, "bot", opener, { - retryVariants: [opener], - retryActiveIndex: 0, - retryRetriesUsed: 0, - retryPromptMessageId: null - }); - if (thinking) repositories.updateChatMessageThinking(chat.id, Number(inserted.lastInsertRowid), thinking); - repositories.touchChat(chat.id, user); - return res.json({persona, chat: repositories.getChat(chat.id, user), existing: false, generatedInitialMessage: true, opener}); + if (result.error) return res.status(result.status).json({error: result.error}); + return res.json({persona, ...result}); } catch { return res.json({persona, chat, existing: false, generatedInitialMessage: false, degraded: true}); } })); - router.post("/roleplays/start", validateBody(roleplayStartValidator), asyncHandler(async (req, res) => { + router.post("/roleplays/start", ...chatRateLimiters, generationCancellation, validateBody(roleplayStartValidator), asyncHandler(async (req, res) => { const user = req.session.user; const {assistantPersonaId, userPersonaId, scenarioPrompt, model} = req.validatedBody; const assistantPersona = repositories.getPersona(assistantPersonaId, user); @@ -306,35 +290,17 @@ export function createPersonasRouter({repositories, chatService, modelService, c throw error; } - if (repositories.listChatMessages(chat.id).length > 0) { - return res.json({chat, existing: true, generatedInitialMessage: false}); - } - - repositories.updateChatScene(chat.id, user, scenarioPrompt, sceneSummary); - chat = repositories.getChat(chat.id, user); - try { - let thinking = ""; - const opener = await chatService.generateRoleplayOpener({ - onThinking: (_chunk, fullThinking) => { thinking = fullThinking; }, + const result = await chatService.initializeRoleplay({ + user, chatId: chat.id, signal: req.generationSignal, selectedModel: model, assistantPersona, userPersona, scenarioPrompt, sceneSummary }); - if (!opener || chatService.containsUserVoiceInRoleplayOpener(opener, userPersona)) { - return res.json({chat, existing: false, generatedInitialMessage: false, degraded: true}); - } - const inserted = chatService.insertChatMessage(chat.id, "bot", opener, { - retryVariants: [opener], - retryActiveIndex: 0, - retryRetriesUsed: 0, - retryPromptMessageId: null - }); - if (thinking) repositories.updateChatMessageThinking(chat.id, Number(inserted.lastInsertRowid), thinking); - repositories.touchChat(chat.id, user); - return res.json({chat: repositories.getChat(chat.id, user), existing: false, generatedInitialMessage: true, opener}); + if (result.error) return res.status(result.status).json({error: result.error}); + return res.json(result); } catch (error) { const mapped = modelService.mapError(error); return res.json({ diff --git a/src/services/chatService.js b/src/services/chatService.js index 73dbbc1..98e5df3 100644 --- a/src/services/chatService.js +++ b/src/services/chatService.js @@ -115,7 +115,10 @@ export function createChatService(repositories, modelService, config, webSearchP else externalSignal?.addEventListener("abort", abortFromCaller, { once: true }); activeGenerations.set(chatId, controller); try { - return await operation(controller.signal); + controller.signal.throwIfAborted(); + const result = await operation(controller.signal); + controller.signal.throwIfAborted(); + return result; } finally { if (activeGenerations.get(chatId) === controller) activeGenerations.delete(chatId); externalSignal?.removeEventListener("abort", abortFromCaller); @@ -238,6 +241,7 @@ export function createChatService(repositories, modelService, config, webSearchP ? modelService.streamReply(selectedModel, messages, onChunk, { onThinking, signal, generation, reasoning }) : modelService.generateReply(selectedModel, messages, { onThinking, signal, generation, reasoning }); let opener = await generate(baseMessages); + signal?.throwIfAborted(); if (!containsUserVoiceInRoleplayOpener(opener, userPersona)) { return opener; } @@ -658,10 +662,31 @@ export function createChatService(repositories, modelService, config, webSearchP generateRoleplayOpener, insertChatMessage, containsUserVoiceInRoleplayOpener, + async initializeRoleplay({user, chatId, assistantPersona, userPersona, scenarioPrompt, sceneSummary, selectedModel, signal = null}) { + if (chatGenerationQueues.has(chatId)) return {error: "A response is already being generated", status: 409}; + return withChatGenerationLock(chatId, () => withActiveGeneration(chatId, signal, async activeSignal => { + let chat = repositories.getChat(chatId, user); + if (!chat) return {error: "Chat not found", status: 404}; + // Check inside the same lock that protects generation and persistence. + if (repositories.listChatMessages(chatId).length) return {chat, existing: true, generatedInitialMessage: false}; + repositories.updateChatScene(chatId, user, scenarioPrompt, sceneSummary); + chat = repositories.getChat(chatId, user); + const {model, generation} = getGenerationSettings(chat, selectedModel); + const opener = await generateRoleplayOpener({selectedModel: model, generation, + assistantPersona, userPersona, scenarioPrompt, sceneSummary, signal: activeSignal}); + activeSignal.throwIfAborted(); + if (!opener || containsUserVoiceInRoleplayOpener(opener, userPersona)) { + return {chat, existing: false, generatedInitialMessage: false, degraded: true}; + } + insertChatMessage(chatId, "bot", opener, null, model); + repositories.touchChat(chatId, user); + return {chat: repositories.getChat(chatId, user), existing: false, generatedInitialMessage: true, opener}; + })); + }, persistChatMessageRetryState, buildMessagesPayload: buildConversationContext, - async sendChatMessage({ user, chatId, message, model, attachments = [], webSearch = false, reasoning = false, requestId = null }) { - return generateChatTurn({ user, chatId, message, model, attachments, webSearch, reasoning, requestId }); + async sendChatMessage({ user, chatId, message, model, attachments = [], webSearch = false, reasoning = false, signal = null, requestId = null }) { + return generateChatTurn({ user, chatId, message, model, attachments, webSearch, reasoning, signal, requestId }); }, async streamChatMessage({ user, chatId, message, model, attachments = [], webSearch = false, reasoning = false, onStatus, onChunk, onThinking, signal = null, requestId = null }) { return generateChatTurn({ user, chatId, message, model, attachments, webSearch, reasoning, onStatus, onChunk, onThinking, signal, requestId }); @@ -706,37 +731,39 @@ export function createChatService(repositories, modelService, config, webSearchP throw error; }).finally(() => regeneratingMessages.delete(chatId)); }, - async editAndResendMessage({ user, chatId, targetMessage, content, model, webSearch = false, onChunk = null, signal = null, requestId = null }) { + async editAndResendMessage({ user, chatId, targetMessage, content, model, webSearch = false, reasoning = false, onChunk = null, signal = null, requestId = null }) { if (!targetMessage || targetMessage.role !== "user") return { error: "Only user messages can be edited and resent", status: 400 }; - const sourceChat = repositories.getChat(chatId, user); - if (!sourceChat) return { error: "Chat not found", status: 404 }; - const cloned = cloneChatThroughMessage({ - user, - sourceChat, - messageId: targetMessage.parent_message_id ?? null - }); - if (cloned.error) return cloned; + if (chatGenerationQueues.has(chatId)) return { error: "A response is already being generated", status: 409 }; + return withChatGenerationLock(chatId, () => withActiveGeneration(chatId, signal, async activeSignal => { + const sourceChat = repositories.getChat(chatId, user); + if (!sourceChat) return { error: "Chat not found", status: 404 }; + const freshTarget = repositories.getChatMessage(chatId, targetMessage.id); + if (!freshTarget) return { error: "Message not found", status: 404 }; + const cloned = cloneChatThroughMessage({user, sourceChat, messageId: freshTarget.parent_message_id ?? null}); + if (cloned.error) return cloned; - const generated = await generateChatTurn({ - user, - chatId: cloned.chat.id, - attachments: attachmentService?.list(targetMessage.id).map(item => item.id) || [], - reuseChatId: chatId, - message: content, - model, - webSearch, - onChunk, - signal, - requestId - }); - if (generated.error) return { ...generated, chat: cloned.chat }; - return { - ...generated, - chat: repositories.getChat(cloned.chat.id, user), - messages: listActiveBranchMessages(cloned.chat.id), - sourceChatId: chatId, - branchedFromMessageId: targetMessage.id - }; + const generated = await generateChatTurn({ + user, + chatId: cloned.chat.id, + attachments: attachmentService?.list(freshTarget.id).map(item => item.id) || [], + reuseChatId: chatId, + message: content, + model, + webSearch, + reasoning, + onChunk, + signal: activeSignal, + requestId + }); + if (generated.error) return { ...generated, chat: cloned.chat }; + return { + ...generated, + chat: repositories.getChat(cloned.chat.id, user), + messages: listActiveBranchMessages(cloned.chat.id), + sourceChatId: chatId, + branchedFromMessageId: freshTarget.id + }; + })); }, createBranchedChat({ user, chatId, messageId }) { const sourceChat = repositories.getChat(chatId, user); diff --git a/src/services/workspaceImportService.js b/src/services/workspaceImportService.js index 7ad9171..9c18bf3 100644 --- a/src/services/workspaceImportService.js +++ b/src/services/workspaceImportService.js @@ -13,6 +13,14 @@ const optionalText = (value) => { return normalized || null; }; const normalizeKey = (value) => normalizeText(value).toLowerCase(); +const optionalNumber = (value) => value == null || normalizeText(value) === "" || typeof value === "boolean" + ? null : Number.isFinite(Number(value)) ? Number(value) : null; +const optionalInteger = (value) => Number.isSafeInteger(optionalNumber(value)) ? Number(value) : null; +const sourceId = (value) => optionalInteger(value) > 0 ? Number(value) : null; +const personaReferences = (record) => ({ + assistant_persona_id: sourceId(record.assistant_persona_id ?? record.assistantPersonaId), + user_persona_id: sourceId(record.user_persona_id ?? record.userPersonaId) +}); function countBy(items, getKey) { const counts = new Map(); @@ -40,6 +48,7 @@ function validatePersona(persona, index, errors) { if (!name) errors.push(`Persona ${index + 1} is missing a name.`); if (!VALID_PERSONA_TYPES.has(personaType)) errors.push(`Persona ${index + 1} has invalid persona_type.`); return { + sourceId: sourceId(persona.id), name, pronouns: optionalText(persona.pronouns), appearance: optionalText(persona.appearance), @@ -60,6 +69,7 @@ function validateTemplate(template, index, errors) { if (!name) errors.push(`Template ${index + 1} is missing a name.`); if (!promptText) errors.push(`Template ${index + 1} is missing prompt_text.`); return { + persona_id: sourceId(template.persona_id ?? template.personaId), name, description: optionalText(template.description), category: optionalText(template.category), @@ -84,7 +94,8 @@ function validateMessage(message, chatIndex, messageIndex, errors) { ? message.retry_variants.map((item) => String(item ?? "")).filter(Boolean) : [content].filter(Boolean); return { - sourceId: Number.isInteger(Number(message.id)) ? Number(message.id) : null, + // Synthetic IDs keep ID-less legacy messages addressable within this chat. + sourceId: sourceId(message.id) ?? -(messageIndex + 1), role: role === "assistant" ? "bot" : role, content, model_name: optionalText(message.model_name || message.modelName), @@ -102,12 +113,8 @@ function validateMessage(message, chatIndex, messageIndex, errors) { retryRetriesUsed: Number.isInteger(Number(message.retryRetriesUsed ?? message.retry_retries_used)) ? Number(message.retryRetriesUsed ?? message.retry_retries_used) : 0, - retryPromptMessageId: Number.isInteger(Number(message.retryPromptMessageId ?? message.retry_prompt_message_id)) - ? Number(message.retryPromptMessageId ?? message.retry_prompt_message_id) - : null, - parentMessageId: Number.isInteger(Number(message.parentMessageId ?? message.parent_message_id)) - ? Number(message.parentMessageId ?? message.parent_message_id) - : null, + retryPromptMessageId: sourceId(message.retryPromptMessageId ?? message.retry_prompt_message_id), + parentMessageId: sourceId(message.parentMessageId ?? message.parent_message_id), deliveryStatus: ["pending", "complete", "failed"].includes(message.deliveryStatus ?? message.delivery_status) ? (message.deliveryStatus ?? message.delivery_status) : "complete", @@ -127,7 +134,7 @@ function validateMemory(memory, chatIndex, memoryIndex, errors) { } const fact = normalizeText(memory.fact); if (!fact) errors.push(`Chat ${chatIndex + 1}, memory ${memoryIndex + 1} is missing fact.`); - return {fact}; + return {fact, ...personaReferences(memory)}; } function validatePersonaState(state, chatIndex, errors) { @@ -137,6 +144,7 @@ function validatePersonaState(state, chatIndex, errors) { return null; } return { + ...personaReferences(state), relationship_notes: optionalText(state.relationship_notes || state.relationshipNotes), tone: optionalText(state.tone), current_location: optionalText(state.current_location || state.currentLocation), @@ -153,9 +161,39 @@ function validateChat(chat, index, errors) { } const title = normalizeText(chat.stored_title || chat.title || "Imported chat") || "Imported chat"; const messages = asArray(chat.messages).map((message, messageIndex) => validateMessage(message, index, messageIndex, errors)).filter(Boolean); + const legacyChain = asArray(chat.messages).every(message => message && + !Object.hasOwn(message, "parent_message_id") && !Object.hasOwn(message, "parentMessageId")); + if (legacyChain) { + for (const [messageIndex, message] of messages.entries()) { + message.parentMessageId = messages[messageIndex - 1]?.sourceId ?? null; + } + } + for (const id of duplicateNames(messages, message => message.sourceId)) { + errors.push(`Chat ${index + 1} has duplicate message ID: ${id}`); + } + // Two-pass persistence accepts forward references, but recursive history must stay acyclic. + const byId = new Map(messages.map(message => [message.sourceId, message])); + const visited = new Set(); + for (const message of messages) { + const chain = new Set(); + let current = message; + while (current && !visited.has(current.sourceId)) { + if (chain.has(current.sourceId)) { + errors.push(`Chat ${index + 1} has a cyclic message chain.`); + break; + } + chain.add(current.sourceId); + current = byId.get(current.parentMessageId); + } + for (const id of chain) visited.add(id); + } const memories = asArray(chat.memories).map((memory, memoryIndex) => validateMemory(memory, index, memoryIndex, errors)).filter(Boolean); const personaState = validatePersonaState(chat.persona_state || chat.personaChatState, index, errors); return { + sourceId: sourceId(chat.id), + ...personaReferences(chat), + source_chat_id: sourceId(chat.source_chat_id ?? chat.sourceChatId), + branched_from_message_id: sourceId(chat.branched_from_message_id ?? chat.branchedFromMessageId), title, stored_title: title, folder_name: optionalText(chat.folder_name || chat.folderName), @@ -168,13 +206,11 @@ function validateChat(chat, index, errors) { ? Number(chat.context_summary_message_id || chat.contextSummaryMessageId) : 0, preferred_model: optionalText(chat.preferred_model || chat.preferredModel), - temperature: Number.isFinite(Number(chat.temperature)) ? Number(chat.temperature) : null, - context_length: Number.isInteger(Number(chat.context_length ?? chat.contextLength)) ? Number(chat.context_length ?? chat.contextLength) : null, - response_length: Number.isInteger(Number(chat.response_length ?? chat.responseLength)) ? Number(chat.response_length ?? chat.responseLength) : null, + temperature: optionalNumber(chat.temperature), + context_length: optionalInteger(chat.context_length ?? chat.contextLength), + response_length: optionalInteger(chat.response_length ?? chat.responseLength), system_instruction: optionalText(chat.system_instruction || chat.systemInstruction), - active_leaf_message_id: Number.isInteger(Number(chat.active_leaf_message_id ?? chat.activeLeafMessageId)) - ? Number(chat.active_leaf_message_id ?? chat.activeLeafMessageId) - : null, + active_leaf_message_id: sourceId(chat.active_leaf_message_id ?? chat.activeLeafMessageId), messages, memories, persona_state: personaState @@ -206,6 +242,11 @@ export function validateWorkspaceImport(workspace, existingWorkspace = null) { chats: asArray(workspace.chats).map((chat, index) => validateChat(chat, index, errors)).filter(Boolean), templates: asArray(workspace.templates).map((template, index) => validateTemplate(template, index, errors)).filter(Boolean) }; + for (const field of ["personas", "chats"]) { + for (const id of duplicateNames(sanitized[field], record => record.sourceId)) { + errors.push(`Duplicate ${field} ID: ${id}`); + } + } const duplicates = [ ...duplicateNames(sanitized.personas, (persona) => normalizeKey(`${persona.persona_type}:${persona.name}`)).map((name) => `Duplicate persona in import: ${name}`), diff --git a/test/conversation.test.js b/test/conversation.test.js new file mode 100644 index 0000000..822fe50 --- /dev/null +++ b/test/conversation.test.js @@ -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); +}); diff --git a/test/generationProtection.test.js b/test/generationProtection.test.js new file mode 100644 index 0000000..627bf3b --- /dev/null +++ b/test/generationProtection.test.js @@ -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); +}); diff --git a/test/workspaceImport.test.js b/test/workspaceImport.test.js new file mode 100644 index 0000000..b0e8529 --- /dev/null +++ b/test/workspaceImport.test.js @@ -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/); + } +});