make PWA chat resume idempotent

This commit is contained in:
2026-07-23 17:12:11 -07:00
parent abc1124d27
commit c5217b2710
10 changed files with 326 additions and 140 deletions

View File

@@ -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 },

View File

@@ -33,6 +33,7 @@ export type ChatMessage = {
export type MultiplexRequest = {
chatId?: string;
persist?: boolean;
clientRequestId?: string;
provider: Provider;
model: string;
messages: ChatMessage[];

View File

@@ -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);
}