@rudderhq/server 0.5.2-canary.2 → 0.6.0

package/dist/routes/chats.stream-routes.js

@@ -1,65 +1,21 @@
-import { addChatMessageSchema, chatClientCheckpointSchema, convertChatToIssueSchema, createChatAttachmentMetadataSchema, createChatContextLinkSchema, createChatFirstTurnSchema, resolveChatOperationProposalSchema, setChatProjectContextSchema, stopChatGenerationSchema, updateChatConversationUserStateSchema } from "@rudderhq/shared";
+import { addChatMessageSchema, chatClientCheckpointSchema, convertChatToIssueSchema, createChatAttachmentMetadataSchema, createChatContextLinkSchema, createChatFirstTurnSchema, resolveChatOperationProposalSchema, setChatProjectContextSchema, stopChatGenerationSchema, updateChatConversationUserStateSchema, } from "@rudderhq/shared";
+import { coalesceChatTranscriptTextEntries, withChatTranscriptGenerationProvenance, } from "@rudderhq/shared/chat-transcript-provenance";
 import multer from "multer";
 import { randomUUID } from "node:crypto";
 import { isAllowedContentType, MAX_ATTACHMENT_BYTES } from "../attachment-types.js";
 import { conflict } from "../errors.js";
 import { logger } from "../middleware/logger.js";
 import { validate } from "../middleware/validate.js";
-import { chatAssistantErrorForLog, createAssistantTextAccumulator, } from "../services/chat-assistant.helpers.js";
+import { chatAssistantErrorForLog } from "../services/chat-assistant.helpers.js";
 import { CHAT_ASSISTANT_USER_ERROR_MESSAGE, ChatAssistantStreamError, userVisiblePartialBodyFromError } from "../services/chat-assistant.js";
 import { cancelActiveChatGeneration, claimChatGeneration, createChatRuntimeControlCoordinator, getActiveChatGeneration, setActiveChatGenerationId, } from "../services/chat-generation-locks.js";
 import { hashChatGenerationBody } from "../services/chat-generation-protocol.js";
 import { logActivity } from "../services/index.js";
 import { assertCompanyAccess, getActorInfo } from "./authz.js";
 import { wakeIssueAssigneeAfterChatConversion } from "./chat-issue-assignment-wakeup.js";
-const startingChatGenerationGates = new Map();
-function createStartingChatGenerationGate() {
-    let resolveGeneration;
-    let resolveStopApplied;
-    return {
-        generationReady: new Promise((resolve) => {
-            resolveGeneration = resolve;
-        }),
-        stopApplied: new Promise((resolve) => {
-            resolveStopApplied = resolve;
-        }),
-        stopRequested: false,
-        resolveGeneration,
-        resolveStopApplied,
-    };
-}
-function outputAdmissionClosed(error) {
-    return Boolean(error
-        && typeof error === "object"
-        && "status" in error
-        && Number(error.status) === 409
-        && "message" in error
-        && error.message === "Chat-visible output admission is closed for this generation");
-}
-function normalizeMultipartFirstTurnBody(body) {
-    if (!body)
-        return {};
-    const normalized = { ...body };
-    if (typeof normalized.planMode === "string") {
-        if (normalized.planMode === "true")
-            normalized.planMode = true;
-        else if (normalized.planMode === "false")
-            normalized.planMode = false;
-    }
-    if (typeof normalized.contextLinks === "string") {
-        try {
-            normalized.contextLinks = JSON.parse(normalized.contextLinks);
-        }
-        catch {
-            // Leave invalid JSON in place so the shared schema returns a normal 400.
-        }
-    }
-    return normalized;
-}
+import { CHAT_ASSISTANT_RECOVERABLE_FAILURE_FALLBACK_MESSAGE, CHAT_ASSISTANT_STOPPED_FALLBACK_MESSAGE, createChatStreamFileStaging, createStartingChatGenerationGate, normalizeMultipartFirstTurnBody, normalizeMultipartMessageBody, outputAdmissionClosed, resolveStoppedAssistantState, startingChatGenerationGates, withMergedChatMessageAttachments, } from "./chats.stream-support.js";
 export function registerChatStreamRoutes(ctx) {
-    const CHAT_ASSISTANT_RECOVERABLE_FAILURE_FALLBACK_MESSAGE = "The assistant reply could not be completed. Rudder saved this attempt for diagnostics; retry when ready.";
-    const CHAT_ASSISTANT_STOPPED_FALLBACK_MESSAGE = "Chat run stopped before a final reply. Continue the conversation to resume from the preserved context.";
-    const { router, db, storage, svc, assistantSvc, agentsSvc, issuesSvc, projectsSvc, goalsSvc, access, operatorProfiles, heartbeat, assertConversationAccess, assertChatLocalMutationAllowed, assertSideChatMutationAllowed, touchSideChat, boardUserId, assertCanAssignTasks, runSingleFileUpload, runMessageFileUpload, isMultipartRequest, uploadedMessageFiles, validateUploadedMessageFiles, preflightChatDraft, logChatMessagesAdded, assertContextLinksBelongToCompany, turnContextFromUserMessage, addUserMessage, startChatTitleGeneration, attachFilesToUserMessage, loadAssistantInput, chatReplyingAgentId, assertCanConvertIssueProposal, persistAssistantReply, linkChatRunMessages, attachGeneratedFilesToPartialMessage, persistPartialAssistantMessage, recoverableFailurePayload, recoverableFailureBody, writeStreamEvent, queueRequestActor, wakeServerQueue, wakeTerminalProjector, } = ctx;
+    const { router, db, storage, svc, assistantSvc, agentsSvc, issuesSvc, projectsSvc, goalsSvc, access, operatorProfiles, heartbeat, assertConversationAccess, assertChatLocalMutationAllowed, assertSideChatMutationAllowed, touchSideChat, boardUserId, assertCanAssignTasks, runSingleFileUpload, runMessageFileUpload, isMultipartRequest, uploadedMessageFiles, validateUploadedMessageFiles, preflightChatDraft, logChatMessagesAdded, assertContextLinksBelongToCompany, turnContextFromUserMessage, addUserMessage, inlineAnnotations, storeUserMessageFiles, cleanupStoredUserMessageFiles, storeQueuedAnnotationFiles, cleanupUncommittedQueuedAnnotationFiles, startChatTitleGeneration, attachFilesToUserMessage, loadAssistantInput, chatReplyingAgentId, assertCanConvertIssueProposal, persistAssistantReply, linkChatRunMessages, attachGeneratedFilesToPartialMessage, persistPartialAssistantMessage, recoverableFailurePayload, recoverableFailureBody, writeStreamEvent, queueRequestActor, wakeServerQueue, wakeTerminalProjector, } = ctx;
     const handleChatMessageStream = async (req, res) => {
         const atomicFirstTurn = req.atomicFirstTurn;
         if (isMultipartRequest(req) && !atomicFirstTurn?.uploadPrepared) {
@@ -78,7 +34,11 @@ export function registerChatStreamRoutes(ctx) {
                 throw err;
             }
         }
-        const parsedBody = addChatMessageSchema.safeParse(req.body ?? {});
+        const normalizedBody = isMultipartRequest(req)
+            ? normalizeMultipartMessageBody(req.body)
+            : (req.body ?? {});
+        const inlineAnnotationsProvided = Object.hasOwn(normalizedBody, "inlineAnnotations");
+        const parsedBody = addChatMessageSchema.safeParse(normalizedBody);
         if (!parsedBody.success) {
             res.status(400).json({ error: "Invalid chat message", details: parsedBody.error.issues });
             return;
@@ -107,6 +67,21 @@ export function registerChatStreamRoutes(ctx) {
             return;
         }
         const queuedMessageId = parsedBody.data.queuedMessageId ?? null;
+        if (queuedMessageId) {
+            res.status(409).json({
+                error: "Queued messages are delivered only by Rudder's server-owned Queue worker",
+            });
+            return;
+        }
+        const preparedAnnotations = !atomicFirstTurn && inlineAnnotationsProvided
+            ? await inlineAnnotations.prepare({
+                orgId: conversation.orgId,
+                conversationId: conversation.id,
+                annotations: parsedBody.data.inlineAnnotations ?? [],
+                uploadedFileCount: messageFiles.length,
+                editUserMessageId: parsedBody.data.editUserMessageId ?? null,
+            })
+            : null;
         if (!atomicFirstTurn) {
             const assistantAvailability = await assistantSvc.getChatAssistantAvailability(conversation);
             if (!assistantAvailability.available) {
@@ -125,27 +100,61 @@ export function registerChatStreamRoutes(ctx) {
                 res.status(409).json({ error: "A chat reply is already being generated for this conversation" });
                 return;
             }
-            if (messageFiles.length > 0) {
-                res.status(422).json({ error: "Queue does not support new files yet" });
-                return;
+            const clientMutationId = `stream:${randomUUID()}`;
+            const storedQueueFiles = await storeQueuedAnnotationFiles(conversation, messageFiles);
+            let queuedResult;
+            try {
+                queuedResult = await svc.createQueuedMessageWithStagedAttachments({
+                    orgId: conversation.orgId,
+                    conversationId: conversation.id,
+                    clientMutationId,
+                    expectedGenerationId: getActiveChatGeneration(conversation.id)?.generationId ?? null,
+                    requestActor: queueRequestActor(req),
+                    payload: {
+                        body: parsedBody.data.body,
+                        attachmentIds: [],
+                        ...(inlineAnnotationsProvided
+                            ? { inlineAnnotations: preparedAnnotations?.annotations ?? [] }
+                            : {}),
+                        skillRefs: [],
+                        projectId: null,
+                        accessMode: null,
+                        model: null,
+                        effort: null,
+                        metadata: {
+                            source: "stream_endpoint_during_active_generation",
+                        },
+                    },
+                    stagedAttachments: storedQueueFiles.map((attachment) => ({
+                        ...attachment,
+                        createdByAgentId: null,
+                        createdByUserId: actor.actorId,
+                    })),
+                    attachmentFileIndexesByAnnotationId: preparedAnnotations?.attachmentFileIndexesByAnnotationId ?? new Map(),
+                });
             }
-            const item = await svc.createQueuedMessage({
+            catch (error) {
+                await cleanupUncommittedQueuedAnnotationFiles(conversation.orgId, clientMutationId, storedQueueFiles);
+                throw error;
+            }
+            await cleanupUncommittedQueuedAnnotationFiles(conversation.orgId, clientMutationId, queuedResult.cleanupAttachments);
+            const item = queuedResult.item;
+            await logActivity(db, {
                 orgId: conversation.orgId,
-                conversationId: conversation.id,
-                clientMutationId: `stream:${randomUUID()}`,
-                expectedGenerationId: getActiveChatGeneration(conversation.id)?.generationId ?? null,
-                requestActor: queueRequestActor(req),
-                payload: {
-                    body: parsedBody.data.body,
-                    attachmentIds: [],
-                    skillRefs: [],
-                    projectId: null,
-                    accessMode: null,
-                    model: null,
-                    effort: null,
-                    metadata: {
-                        source: "stream_endpoint_during_active_generation",
-                    },
+                actorType: actor.actorType,
+                actorId: actor.actorId,
+                agentId: actor.agentId,
+                runId: actor.runId,
+                action: "chat.queue.created",
+                entityType: "chat",
+                entityId: conversation.id,
+                details: {
+                    queuedMessageId: item.id,
+                    position: item.position,
+                    annotationCount: item.annotationCount ?? 0,
+                    annotationSourceMessageIds: [
+                        ...new Set((item.payload.inlineAnnotations ?? []).map((annotation) => annotation.sourceMessageId)),
+                    ],
                 },
             });
             res.status(202);
@@ -156,6 +165,20 @@ export function registerChatStreamRoutes(ctx) {
             res.end();
             return;
         }
+        const stagedMessageFiles = createChatStreamFileStaging({
+            conversation: conversation,
+            shouldStage: !atomicFirstTurn,
+            files: messageFiles,
+            store: storeUserMessageFiles,
+            cleanup: cleanupStoredUserMessageFiles,
+        });
+        try {
+            await stagedMessageFiles.stage();
+        }
+        catch (error) {
+            releaseGeneration();
+            throw error;
+        }
         const startupGate = createStartingChatGenerationGate();
         startingChatGenerationGates.set(conversation.id, startupGate);
         if (queuedMessageId) {
@@ -170,6 +193,7 @@ export function registerChatStreamRoutes(ctx) {
                 startupGate.resolveGeneration(null);
                 startingChatGenerationGates.delete(conversation.id);
                 releaseGeneration();
+                await stagedMessageFiles.cleanup();
                 throw error;
             }
         }
@@ -184,6 +208,7 @@ export function registerChatStreamRoutes(ctx) {
             startupGate.resolveGeneration(null);
             startingChatGenerationGates.delete(conversation.id);
             releaseGeneration();
+            await stagedMessageFiles.cleanup();
             if (atomicFirstTurn) {
                 const failurePayload = recoverableFailurePayload(error, null);
                 const failureBody = recoverableFailureBody(failurePayload) || CHAT_ASSISTANT_RECOVERABLE_FAILURE_FALLBACK_MESSAGE;
@@ -229,6 +254,8 @@ export function registerChatStreamRoutes(ctx) {
         const transcript = [];
         let assistantProgressMessageId = null;
         let activeChatRunId = null;
+        let userMessagePersisted = Boolean(atomicFirstTurn);
+        let committedUserMessageId = atomicFirstTurn?.userMessage.id ?? null;
         let generationTerminalStatus = "failed";
         let admittedAssistantBody = "";
         let stopCutoff = null;
@@ -247,24 +274,12 @@ export function registerChatStreamRoutes(ctx) {
             };
         };
         abortController.signal.addEventListener("abort", freezeStopCutoff, { once: true });
-        const stoppedState = (fallbackBody) => {
-            if (!stopCutoff) {
-                return {
-                    body: admittedAssistantBody.trim() ? admittedAssistantBody : fallbackBody,
-                    transcript: [...transcript],
-                };
-            }
-            const accumulator = createAssistantTextAccumulator();
-            for (const entry of stopCutoff.transcript) {
-                if (entry.kind === "assistant") {
-                    accumulator.push(entry.text, entry.delta === true);
-                }
-            }
-            return {
-                body: stopCutoff.body.trim() ? stopCutoff.body : accumulator.fullText,
-                transcript: stopCutoff.transcript,
-            };
-        };
+        const stoppedState = (fallbackBody) => resolveStoppedAssistantState({
+            stopCutoff,
+            admittedAssistantBody,
+            transcript,
+            fallbackBody,
+        });
         let stoppedPersistencePromise = null;
         const persistStoppedAssistant = (stoppedConversation, replyingAgentId, fallbackBody) => {
             if (stoppedPersistencePromise)
@@ -293,7 +308,7 @@ export function registerChatStreamRoutes(ctx) {
                 let lastProjectionError;
                 for (let attempt = 1; attempt <= 3; attempt += 1) {
                     try {
-                        stoppedMessage = await persistPartialAssistantMessage(stoppedConversation, stoppedBody, "stopped", turnContextForPartial, frozen.transcript, replyingAgentId, assistantProgressMessageId, activeChatRunId);
+                        stoppedMessage = await persistPartialAssistantMessage(stoppedConversation, stoppedBody, "stopped", turnContextForPartial, frozen.transcript, replyingAgentId, assistantProgressMessageId, activeChatRunId, undefined, false);
                         lastProjectionError = undefined;
                         break;
                     }
@@ -338,7 +353,19 @@ export function registerChatStreamRoutes(ctx) {
         res.setHeader("X-Accel-Buffering", "no");
         res.flushHeaders();
         try {
-            const userMessage = atomicFirstTurn?.userMessage ?? await addUserMessage(conversation, parsedBody.data.body, actor, parsedBody.data.editUserMessageId ?? null);
+            const userMessage = atomicFirstTurn?.userMessage ?? await addUserMessage(conversation, parsedBody.data.body, actor, parsedBody.data.editUserMessageId ?? null, {
+                provided: inlineAnnotationsProvided,
+                prepared: preparedAnnotations,
+                storedAttachments: stagedMessageFiles.files,
+                onPersisted: (messageId) => {
+                    userMessagePersisted = true;
+                    committedUserMessageId = messageId;
+                    stagedMessageFiles.markCommitted();
+                },
+            });
+            userMessagePersisted = true;
+            committedUserMessageId = userMessage.id;
+            stagedMessageFiles.markCommitted();
             await touchSideChat(req, conversation);
             if (queuedMessageId) {
                 await svc.markQueuedMessageRunning({
@@ -350,19 +377,10 @@ export function registerChatStreamRoutes(ctx) {
             if (!parsedBody.data.editUserMessageId) {
                 startChatTitleGeneration(conversation, userMessage);
             }
-            const userAttachments = await attachFilesToUserMessage(conversation, userMessage.id, messageFiles, actor);
-            const mergedUserAttachmentsById = new Map();
-            for (const attachment of userMessage.attachments ?? []) {
-                mergedUserAttachmentsById.set(attachment.id, attachment);
-            }
-            for (const attachment of userAttachments) {
-                mergedUserAttachmentsById.set(attachment.id, attachment);
-            }
-            const mergedUserAttachments = [...mergedUserAttachmentsById.values()];
-            const hydratedUserMessage = {
-                ...userMessage,
-                attachments: mergedUserAttachments,
-            };
+            const userAttachments = atomicFirstTurn
+                ? await attachFilesToUserMessage(conversation, userMessage.id, messageFiles, actor)
+                : [];
+            const hydratedUserMessage = withMergedChatMessageAttachments(userMessage, userAttachments);
             turnContextForPartial = turnContextFromUserMessage(userMessage);
             writeStreamEvent(res, {
                 type: "ack",
@@ -448,7 +466,6 @@ export function registerChatStreamRoutes(ctx) {
                                     runId: activeChatRunId,
                                     bodyHash: hashChatGenerationBody(projectedBody),
                                     body: projectedBody,
-                                    transcript: [...transcript],
                                     replyingAgentId: chatReplyingAgentId(assistantInput.conversation),
                                     chatTurnId: turnContextForPartial.chatTurnId,
                                     turnVariant: turnContextForPartial.turnVariant,
@@ -487,7 +504,6 @@ export function registerChatStreamRoutes(ctx) {
                                 return;
                             const activeControl = getActiveChatGeneration(conversation.id);
                             const attemptEpoch = Math.max(1, activeControl?.attemptEpoch ?? generation?.attemptEpoch ?? 1);
-                            const projectedTranscript = [...transcript, entry];
                             let committed;
                             try {
                                 committed = await svc.generationProtocol.appendVisibleEventAndProject({
@@ -501,7 +517,6 @@ export function registerChatStreamRoutes(ctx) {
                                     runId: activeChatRunId,
                                     bodyHash: hashChatGenerationBody(admittedAssistantBody),
                                     body: admittedAssistantBody,
-                                    transcript: projectedTranscript,
                                     replyingAgentId: chatReplyingAgentId(assistantInput.conversation),
                                     chatTurnId: turnContextForPartial.chatTurnId,
                                     turnVariant: turnContextForPartial.turnVariant,
@@ -513,12 +528,20 @@ export function registerChatStreamRoutes(ctx) {
                                 throw error;
                             }
                             assistantProgressMessageId = committed.message.id;
-                            transcript.push(entry);
+                            const durableEntry = withChatTranscriptGenerationProvenance(entry, {
+                                generationId: generation.id,
+                                generationSeq: committed.event.generationSeq,
+                            });
+                            const normalizedTranscript = coalesceChatTranscriptTextEntries([
+                                ...transcript,
+                                durableEntry,
+                            ]);
+                            transcript.splice(0, transcript.length, ...normalizedTranscript);
                             if (abortController.signal.aborted || clientClosed)
                                 return;
                             writeStreamEvent(res, {
                                 type: "transcript_entry",
-                                entry,
+                                entry: durableEntry,
                                 generationId: generation.id,
                                 attemptEpoch,
                                 generationSeq: committed.event.generationSeq,
@@ -552,7 +575,6 @@ export function registerChatStreamRoutes(ctx) {
                                 runId: activeChatRunId,
                                 bodyHash: hashChatGenerationBody(streamed.reply.body),
                                 body: streamed.reply.body,
-                                transcript: [...transcript],
                                 replyingAgentId: streamed.replyingAgentId,
                                 chatTurnId: turnContextForPartial.chatTurnId,
                                 turnVariant: turnContextForPartial.turnVariant,
@@ -572,7 +594,7 @@ export function registerChatStreamRoutes(ctx) {
                         await persistStoppedAssistant(assistantInput.conversation, streamed.replyingAgentId, streamed.partialBody);
                         return;
                     }
-                    const createdMessages = await persistAssistantReply(req, assistantInput.conversation, actor, streamed.reply, turnContextForPartial, transcript, streamed.replyingAgentId, assistantProgressMessageId, activeChatRunId);
+                    const createdMessages = await persistAssistantReply(req, assistantInput.conversation, actor, streamed.reply, turnContextForPartial, transcript, streamed.replyingAgentId, assistantProgressMessageId, activeChatRunId, false);
                     await linkChatRunMessages(assistantInput.conversation, activeChatRunId, createdMessages);
                     generationTerminalStatus = "completed";
                     await logChatMessagesAdded(assistantInput.conversation, createdMessages, {
@@ -599,6 +621,41 @@ export function registerChatStreamRoutes(ctx) {
             }
         }
         catch (err) {
+            if (!userMessagePersisted) {
+                logger.warn({
+                    err: chatAssistantErrorForLog(err),
+                    conversationId: conversation.id,
+                }, "chat user message persistence failed");
+                if (!clientClosed) {
+                    writeStreamEvent(res, {
+                        type: "error",
+                        error: CHAT_ASSISTANT_USER_ERROR_MESSAGE,
+                        errorCode: "chat_runtime_exception",
+                        runId: null,
+                        messageId: null,
+                    });
+                    res.end();
+                }
+                return;
+            }
+            if (!turnContextForPartial && committedUserMessageId) {
+                logger.warn({
+                    err: chatAssistantErrorForLog(err),
+                    conversationId: conversation.id,
+                    userMessageId: committedUserMessageId,
+                }, "chat user message hydration failed after commit");
+                if (!clientClosed) {
+                    writeStreamEvent(res, {
+                        type: "error",
+                        error: CHAT_ASSISTANT_USER_ERROR_MESSAGE,
+                        errorCode: "chat_runtime_exception",
+                        runId: null,
+                        messageId: committedUserMessageId,
+                    });
+                    res.end();
+                }
+                return;
+            }
             if (abortController.signal.aborted) {
                 try {
                     if (stoppedPersistencePromise) {
@@ -621,7 +678,7 @@ export function registerChatStreamRoutes(ctx) {
                 || CHAT_ASSISTANT_USER_ERROR_MESSAGE;
             const generatedAttachments = err instanceof ChatAssistantStreamError ? err.generatedAttachments : [];
             const failedReplyingAgentId = chatReplyingAgentId(assistantConversationForPartial);
-            let failedMessage = await persistPartialAssistantMessage(assistantConversationForPartial ?? conversation, partialBody, "failed", turnContextForPartial, transcript, failedReplyingAgentId, assistantProgressMessageId, activeChatRunId, failurePayload).catch(() => null);
+            let failedMessage = await persistPartialAssistantMessage(assistantConversationForPartial ?? conversation, partialBody, "failed", turnContextForPartial, transcript, failedReplyingAgentId, assistantProgressMessageId, activeChatRunId, failurePayload, false).catch(() => null);
             await linkChatRunMessages(assistantConversationForPartial ?? conversation, activeChatRunId, failedMessage ? [failedMessage] : []).catch(() => { });
             failedMessage = await attachGeneratedFilesToPartialMessage(assistantConversationForPartial ?? conversation, failedMessage, generatedAttachments, failedReplyingAgentId).catch(() => failedMessage);
             if (failedMessage && assistantConversationForPartial) {
@@ -648,6 +705,7 @@ export function registerChatStreamRoutes(ctx) {
             }
         }
         finally {
+            await stagedMessageFiles.cleanup();
             abortController.signal.removeEventListener("abort", freezeStopCutoff);
             req.off("aborted", handleClosed);
             res.off("close", handleClosed);