add generation cancellation middleware and refactor chat/message workflows to support HTTP client disconnections
This commit is contained in:
+11
-17
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user