make PWA chat resume idempotent
This commit is contained in:
@@ -119,7 +119,12 @@ export async function* runMultiplexStream(req: MultiplexRequest): AsyncGenerator
|
||||
if (shouldPersist && chatId && call) {
|
||||
await prisma.$transaction(async (tx) => {
|
||||
await tx.message.create({
|
||||
data: { chatId, role: "assistant" as any, content: text },
|
||||
data: {
|
||||
chatId,
|
||||
role: "assistant" as any,
|
||||
content: text,
|
||||
metadata: req.clientRequestId ? ({ clientRequestId: req.clientRequestId } as any) : undefined,
|
||||
},
|
||||
});
|
||||
await tx.llmCall.update({
|
||||
where: { id: call.id },
|
||||
|
||||
@@ -33,6 +33,7 @@ export type ChatMessage = {
|
||||
export type MultiplexRequest = {
|
||||
chatId?: string;
|
||||
persist?: boolean;
|
||||
clientRequestId?: string;
|
||||
provider: Provider;
|
||||
model: string;
|
||||
messages: ChatMessage[];
|
||||
|
||||
@@ -88,7 +88,7 @@ function withRequestUserLocation<T extends { userLocation?: string }>(body: T, r
|
||||
return body.userLocation ? body : { ...body, userLocation: inferRequestUserLocation(req) };
|
||||
}
|
||||
|
||||
async function storeNonAssistantMessages(chatId: string, messages: IncomingChatMessage[]) {
|
||||
async function storeNonAssistantMessages(chatId: string, messages: IncomingChatMessage[], clientRequestId?: string) {
|
||||
const incoming = messages.filter((m) => m.role !== "assistant");
|
||||
if (!incoming.length) return;
|
||||
|
||||
@@ -109,14 +109,21 @@ async function storeNonAssistantMessages(chatId: string, messages: IncomingChatM
|
||||
const toInsert = sharedPrefix === existingNonAssistant.length ? incoming.slice(existingNonAssistant.length) : incoming;
|
||||
if (!toInsert.length) return;
|
||||
|
||||
const finalUserMessageIndex = toInsert.map((message) => message.role).lastIndexOf("user");
|
||||
await prisma.message.createMany({
|
||||
data: toInsert.map((m) => ({
|
||||
chatId,
|
||||
role: m.role as any,
|
||||
content: m.content,
|
||||
name: m.name,
|
||||
metadata: m.attachments?.length ? ({ attachments: m.attachments } as any) : undefined,
|
||||
})),
|
||||
data: toInsert.map((m, index) => {
|
||||
const metadata = {
|
||||
...(m.attachments?.length ? { attachments: m.attachments } : {}),
|
||||
...(clientRequestId && index === finalUserMessageIndex ? { clientRequestId } : {}),
|
||||
};
|
||||
return {
|
||||
chatId,
|
||||
role: m.role as any,
|
||||
content: m.content,
|
||||
name: m.name,
|
||||
metadata: Object.keys(metadata).length ? (metadata as any) : undefined,
|
||||
};
|
||||
}),
|
||||
});
|
||||
}
|
||||
|
||||
@@ -169,6 +176,7 @@ const CompletionStreamBody = z
|
||||
.object({
|
||||
chatId: z.string().optional(),
|
||||
persist: z.boolean().optional(),
|
||||
clientRequestId: z.string().trim().min(1).max(128).optional(),
|
||||
provider: ProviderSchema,
|
||||
model: z.string().min(1),
|
||||
messages: z.array(CompletionMessageSchema),
|
||||
@@ -186,6 +194,13 @@ const CompletionStreamBody = z
|
||||
path: ["chatId"],
|
||||
});
|
||||
}
|
||||
if (value.clientRequestId && (value.persist === false || !value.chatId)) {
|
||||
ctx.addIssue({
|
||||
code: z.ZodIssueCode.custom,
|
||||
message: "clientRequestId requires a persisted stream with chatId",
|
||||
path: ["clientRequestId"],
|
||||
});
|
||||
}
|
||||
});
|
||||
|
||||
function mergeAttachmentsIntoMetadata(metadata: unknown, attachments?: ChatAttachment[]) {
|
||||
@@ -399,6 +414,7 @@ function buildSseHeaders(originHeader: string | undefined) {
|
||||
type SearchRunRequest = z.infer<typeof SearchRunBody>;
|
||||
|
||||
const activeChatStreams = new Map<string, ActiveSseStream>();
|
||||
const activeChatStreamRequestIds = new Map<string, string>();
|
||||
const activeSearchStreams = new Map<string, ActiveSseStream>();
|
||||
const STARRED_PROJECT_ID = "starred";
|
||||
|
||||
@@ -554,6 +570,7 @@ function writeSseEvent(reply: FastifyReply, event: SseStreamEvent) {
|
||||
}
|
||||
|
||||
async function streamActiveRun(req: FastifyRequest, reply: FastifyReply, stream: ActiveSseStream) {
|
||||
if (reply.raw.destroyed || reply.raw.writableEnded) return reply;
|
||||
reply.raw.writeHead(200, buildSseHeaders(typeof req.headers.origin === "string" ? req.headers.origin : undefined));
|
||||
reply.raw.flushHeaders?.();
|
||||
|
||||
@@ -588,10 +605,24 @@ function mapChatStreamEvent(ev: StreamEvent): SseStreamEvent {
|
||||
return { event: ev.type, data: ev };
|
||||
}
|
||||
|
||||
function startActiveChatStream(chatId: string, body: z.infer<typeof CompletionStreamBody>) {
|
||||
function registerActiveChatStream(chatId: string, clientRequestId?: string) {
|
||||
const stream = new ActiveSseStream();
|
||||
activeChatStreams.set(chatId, stream);
|
||||
if (clientRequestId) {
|
||||
activeChatStreamRequestIds.set(chatId, clientRequestId);
|
||||
} else {
|
||||
activeChatStreamRequestIds.delete(chatId);
|
||||
}
|
||||
return stream;
|
||||
}
|
||||
|
||||
function clearActiveChatStream(chatId: string, stream: ActiveSseStream) {
|
||||
if (activeChatStreams.get(chatId) !== stream) return;
|
||||
activeChatStreams.delete(chatId);
|
||||
activeChatStreamRequestIds.delete(chatId);
|
||||
}
|
||||
|
||||
function executeActiveChatStream(chatId: string, body: z.infer<typeof CompletionStreamBody>, stream: ActiveSseStream) {
|
||||
void (async () => {
|
||||
let sawTerminalEvent = false;
|
||||
try {
|
||||
@@ -611,13 +642,54 @@ function startActiveChatStream(chatId: string, body: z.infer<typeof CompletionSt
|
||||
} catch (err) {
|
||||
stream.complete({ event: "error", data: { message: getErrorMessage(err) } });
|
||||
} finally {
|
||||
activeChatStreams.delete(chatId);
|
||||
clearActiveChatStream(chatId, stream);
|
||||
}
|
||||
})();
|
||||
}
|
||||
|
||||
function startActiveChatStream(chatId: string, body: z.infer<typeof CompletionStreamBody>) {
|
||||
const stream = registerActiveChatStream(chatId, body.clientRequestId);
|
||||
executeActiveChatStream(chatId, body, stream);
|
||||
return stream;
|
||||
}
|
||||
|
||||
function getMetadataClientRequestId(metadata: unknown) {
|
||||
if (!metadata || typeof metadata !== "object" || Array.isArray(metadata)) return null;
|
||||
const clientRequestId = (metadata as Record<string, unknown>).clientRequestId;
|
||||
return typeof clientRequestId === "string" ? clientRequestId : null;
|
||||
}
|
||||
|
||||
async function findCompletedChatSubmission(chatId: string, clientRequestId: string) {
|
||||
const assistantMessages = await prisma.message.findMany({
|
||||
where: { chatId, role: "assistant" as any },
|
||||
orderBy: { createdAt: "desc" },
|
||||
select: { content: true, metadata: true },
|
||||
});
|
||||
return assistantMessages.find((message) => getMetadataClientRequestId(message.metadata) === clientRequestId) ?? null;
|
||||
}
|
||||
|
||||
function completeChatSubmissionStream(
|
||||
stream: ActiveSseStream,
|
||||
chatId: string,
|
||||
body: z.infer<typeof CompletionStreamBody>,
|
||||
assistantText: string
|
||||
) {
|
||||
stream.emit("meta", {
|
||||
type: "meta",
|
||||
chatId,
|
||||
callId: null,
|
||||
provider: body.provider,
|
||||
model: body.model,
|
||||
});
|
||||
stream.complete({
|
||||
event: "done",
|
||||
data: {
|
||||
type: "done",
|
||||
text: assistantText,
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
async function executeSearchRunStream(searchId: string, body: SearchRunRequest, stream: ActiveSseStream) {
|
||||
const startedAt = performance.now();
|
||||
const query = body.query?.trim();
|
||||
@@ -1353,15 +1425,39 @@ export async function registerRoutes(app: FastifyInstance) {
|
||||
if (!exists) return app.httpErrors.notFound("chat not found");
|
||||
}
|
||||
|
||||
// Store only new non-assistant messages to avoid duplicate history entries.
|
||||
if (body.persist !== false && body.chatId) {
|
||||
await storeNonAssistantMessages(body.chatId, body.messages);
|
||||
}
|
||||
|
||||
if (body.persist !== false && body.chatId) {
|
||||
if (activeChatStreams.has(body.chatId)) {
|
||||
const activeStream = activeChatStreams.get(body.chatId);
|
||||
if (activeStream) {
|
||||
if (body.clientRequestId && activeChatStreamRequestIds.get(body.chatId) === body.clientRequestId) {
|
||||
return streamActiveRun(req, reply, activeStream);
|
||||
}
|
||||
return app.httpErrors.conflict("chat completion already running");
|
||||
}
|
||||
|
||||
if (body.clientRequestId) {
|
||||
const reservedStream = registerActiveChatStream(body.chatId, body.clientRequestId);
|
||||
try {
|
||||
const completedSubmission = await findCompletedChatSubmission(body.chatId, body.clientRequestId);
|
||||
if (completedSubmission) {
|
||||
completeChatSubmissionStream(reservedStream, body.chatId, body, completedSubmission.content);
|
||||
clearActiveChatStream(body.chatId, reservedStream);
|
||||
return streamActiveRun(req, reply, reservedStream);
|
||||
}
|
||||
|
||||
// Store only new non-assistant messages to avoid duplicate history entries.
|
||||
await storeNonAssistantMessages(body.chatId, body.messages, body.clientRequestId);
|
||||
const configuredBody = await applyStoredChatSettings(body);
|
||||
executeActiveChatStream(body.chatId, configuredBody, reservedStream);
|
||||
return streamActiveRun(req, reply, reservedStream);
|
||||
} catch (err) {
|
||||
reservedStream.complete({ event: "error", data: { message: getErrorMessage(err) } });
|
||||
clearActiveChatStream(body.chatId, reservedStream);
|
||||
throw err;
|
||||
}
|
||||
}
|
||||
|
||||
// Legacy requests without an idempotency key retain the original behavior.
|
||||
await storeNonAssistantMessages(body.chatId, body.messages);
|
||||
const stream = startActiveChatStream(body.chatId, await applyStoredChatSettings(body));
|
||||
return streamActiveRun(req, reply, stream);
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user