@rudderhq/server 0.7.11 → 0.7.13

package/dist/routes/chats.js

@@ -18,6 +18,7 @@ import { buildChatTitlePromptFromMessages, chatTitleGenerationService, } from ".
 import { chatWorkManifestService } from "../services/chat-work-manifest.js";
 import { validateCron } from "../services/cron.js";
 import { accessService, agentService, automationService, chatService, goalService, heartbeatService, issueService, logActivity, operatorProfileService, organizationService, productIntelligenceService, projectService, sideChatService, } from "../services/index.js";
+import { NETWORK_WAIT_EXHAUSTED_ERROR, NETWORK_WAIT_EXHAUSTED_ERROR_CODE, NETWORK_WAIT_UNSAFE_ERROR_CODE, } from "../services/runtime-kernel/heartbeat.core.js";
 import { runtimeResultText, sanitizeGeneratedTitle, } from "../services/title-generation.js";
 import { assertBoard, assertCompanyAccess, getActorInfo, getAuthorizedOrgScope } from "./authz.js";
 import { createChatBackgroundRuntime, } from "./chat-background-runtime.js";
@@ -33,7 +34,7 @@ function chatVisibleOutputAdmissionClosed(error) {
         && error.status === 409
         && error.message === "Chat-visible output admission is closed for this generation";
 }
-export function chatRoutes(db, storage, backgroundRuntime = createChatBackgroundRuntime()) {
+export function chatRoutes(db, storage, backgroundRuntime = createChatBackgroundRuntime(), registerNetworkWaitingRunHandler) {
     const router = Router();
     const svc = chatService(db);
     const organizationsSvc = organizationService(db);
@@ -47,7 +48,14 @@ export function chatRoutes(db, storage, backgroundRuntime = createChatBackground
     const chatRunsSvc = chatAgentRunService(db);
     const workManifestSvc = chatWorkManifestService(db);
     const operatorProfiles = operatorProfileService(db);
-    const heartbeat = heartbeatService(db);
+    let recoverNetworkWaitingChatRun = null;
+    const heartbeat = heartbeatService(db, {
+        onNetworkWaitingRun: async (run) => {
+            if (!recoverNetworkWaitingChatRun)
+                return false;
+            return recoverNetworkWaitingChatRun(run);
+        },
+    });
     const productIntelligence = productIntelligenceService(db);
     const chatTitles = chatTitleGenerationService({ chats: svc, productIntelligence });
     const steerMessages = chatSteerMessageService(db);
@@ -1107,6 +1115,7 @@ export function chatRoutes(db, storage, backgroundRuntime = createChatBackground
                 userMessageId: userMessage.id,
                 chatTurnId: turnContext.chatTurnId,
                 turnVariant: turnContext.turnVariant,
+                runContext: { chatGenerationId: claim.generationId },
                 stream: false,
                 abortSignal: abortController.signal,
                 controlCoordinator: createChatRuntimeControlCoordinator(conversation.id, claim.generationId, {
@@ -1212,6 +1221,11 @@ export function chatRoutes(db, storage, backgroundRuntime = createChatBackground
                     });
                 }
             }
+            else if (streamed.outcome === "waiting_for_network") {
+                // The durable generation remains open; the background recovery path
+                // will reattach and project the eventual assistant result.
+                return;
+            }
             else {
                 let completionMessageId = null;
                 try {
@@ -1323,6 +1337,353 @@ export function chatRoutes(db, storage, backgroundRuntime = createChatBackground
             managedAbort.release();
         }
     }
+    /** Resume a Chat generation that the heartbeat coordinator woke from a
+     * durable network wait. This path deliberately shares the existing
+     * generation/run rows instead of creating a second active Chat run. */
+    async function runRecoveredNetworkWaitingChatRun(run) {
+        const context = run.contextSnapshot && typeof run.contextSnapshot === "object"
+            && !Array.isArray(run.contextSnapshot)
+            ? run.contextSnapshot
+            : {};
+        const conversationId = run.chatConversationId
+            ?? (typeof context.conversationId === "string" ? context.conversationId : null);
+        const generationId = typeof context.chatGenerationId === "string"
+            ? context.chatGenerationId
+            : null;
+        const nonStreamChatRun = context.chatMode === "non_stream";
+        const userMessageId = typeof context.userMessageId === "string"
+            ? context.userMessageId
+            : typeof context.messageId === "string"
+                ? context.messageId
+                : null;
+        const ownerToken = typeof run.executionOwnerToken === "string" ? run.executionOwnerToken : null;
+        if (!conversationId || (!generationId && !nonStreamChatRun) || !userMessageId || !ownerToken) {
+            logger.error({ runId: run.id }, "network-wait Chat run is missing recovery context");
+            return false;
+        }
+        const conversation = await svc.getById(conversationId);
+        const userMessage = conversation
+            ? await svc.getMessage(conversation.id, userMessageId)
+            : null;
+        if (!conversation || !userMessage) {
+            logger.error({ runId: run.id, conversationId, generationId }, "network-wait Chat recovery target is missing");
+            return false;
+        }
+        const abortController = new AbortController();
+        const managedAbort = backgroundRuntime.manageAbortController(abortController);
+        const releaseGeneration = claimChatGeneration(conversation.id, abortController, generationId);
+        if (!releaseGeneration) {
+            managedAbort.release();
+            return true;
+        }
+        const request = requestForQueuedActor({
+            type: "board",
+            source: "local_implicit",
+            userId: "local-board",
+            isInstanceAdmin: true,
+        }, conversation.orgId);
+        const actor = getActorInfo(request);
+        let assistantConversation = conversation;
+        let activeAttemptEpoch = Math.max(1, Number(context.attemptEpoch) || 1);
+        let activeChatRunId = run.id;
+        let partialBody = "";
+        let assistantProjectionMessageId = null;
+        let transcript = [];
+        let terminalStatus = "failed";
+        let terminalReason = null;
+        let waitingAgain = false;
+        const isTerminalGenerationStatus = (status) => [
+            "stop_requested",
+            "stopping",
+            "completed",
+            "failed",
+            "stopped",
+            "aborted",
+            "interrupted_unverified",
+            "control_lost",
+        ].includes(status);
+        const finishRecoveryWithoutProvider = async (status, reason) => {
+            terminalStatus = ["stop_requested", "stopping", "stopped"].includes(status)
+                ? "stopped"
+                : "failed";
+            terminalReason = reason ?? (terminalStatus === "stopped" ? "operator_stop" : "generation_terminal_before_recovery");
+            await chatRunsSvc.finalizeRun(run.id, {
+                status: "cancelled",
+                error: terminalReason,
+                errorCode: terminalStatus === "stopped" ? "chat_run_cancelled" : "chat_generation_already_terminal",
+            }).catch((error) => {
+                logger.warn({ err: error, runId: run.id }, "failed to finalize Chat run after recovery race");
+            });
+            return true;
+        };
+        try {
+            if (generationId) {
+                const frozen = await svc.generationProtocol.getFrozenVisibleProjection({
+                    orgId: conversation.orgId,
+                    conversationId: conversation.id,
+                    generationId,
+                });
+                partialBody = frozen.projection.body;
+                assistantProjectionMessageId = frozen.projection.assistantMessageId;
+                transcript = [...frozen.projection.transcript];
+                activeAttemptEpoch = Math.max(1, frozen.generation.attemptEpoch);
+                const recoveryFailure = run.networkRecoveryFailure
+                    ?? (run.networkRecoveryExhausted
+                        ? {
+                            errorCode: NETWORK_WAIT_EXHAUSTED_ERROR_CODE,
+                            error: NETWORK_WAIT_EXHAUSTED_ERROR,
+                        }
+                        : null);
+                if (isTerminalGenerationStatus(frozen.generation.status)) {
+                    return finishRecoveryWithoutProvider(frozen.generation.status, frozen.generation.terminalReason);
+                }
+                if (!recoveryFailure) {
+                    const resumedGeneration = await svc.generationProtocol.markNetworkResumed({
+                        orgId: conversation.orgId,
+                        conversationId: conversation.id,
+                        generationId,
+                        expectedAttemptEpoch: activeAttemptEpoch,
+                    });
+                    if (isTerminalGenerationStatus(resumedGeneration.status)) {
+                        return finishRecoveryWithoutProvider(resumedGeneration.status, resumedGeneration.terminalReason);
+                    }
+                    const latestGeneration = await svc.getLatestGeneration(conversation.id).catch(() => null);
+                    if (latestGeneration?.id === generationId && isTerminalGenerationStatus(latestGeneration.status)) {
+                        return finishRecoveryWithoutProvider(latestGeneration.status, latestGeneration.terminalReason);
+                    }
+                }
+                if (recoveryFailure) {
+                    throw new ChatAssistantStreamError(recoveryFailure.error, partialBody, [], {
+                        errorCode: recoveryFailure.errorCode,
+                        userMessage: recoveryFailure.errorCode === NETWORK_WAIT_UNSAFE_ERROR_CODE
+                            ? "Network recovery could not safely resume this reply. Check connectivity, then retry this reply."
+                            : "Network recovery retries were exhausted. Check connectivity, then retry this reply.",
+                        retryable: true,
+                        failurePhase: "model_generation",
+                        action: "retry",
+                    });
+                }
+            }
+            const turnContext = turnContextFromUserMessage(userMessage);
+            const assistantInput = await loadAssistantInput(conversation, actor);
+            assistantConversation = assistantInput.conversation;
+            const controlCoordinator = generationId
+                ? createChatRuntimeControlCoordinator(conversation.id, generationId, {
+                    onAttemptStarted: async ({ generationId: currentGenerationId, attemptEpoch, ownerToken: attemptOwnerToken, attempt }) => {
+                        activeAttemptEpoch = attemptEpoch;
+                        await svc.beginGenerationControlAttempt({
+                            orgId: conversation.orgId,
+                            conversationId: conversation.id,
+                            generationId: currentGenerationId,
+                            attemptEpoch,
+                            ownerToken: attemptOwnerToken,
+                            runtimeType: attempt.runtimeType,
+                        });
+                    },
+                    onHandleRegistered: async ({ generationId: currentGenerationId, attemptEpoch, ownerToken: attemptOwnerToken, handle }) => {
+                        await svc.markGenerationControlReady({
+                            generationId: currentGenerationId,
+                            attemptEpoch,
+                            ownerToken: attemptOwnerToken,
+                            runtimeType: handle.runtimeType,
+                            providerThreadId: handle.providerThreadId ?? null,
+                            providerTurnId: handle.providerTurnId ?? null,
+                        });
+                    },
+                    onAttemptLeaseRenewed: async ({ generationId: currentGenerationId, attemptEpoch, ownerToken: attemptOwnerToken }) => {
+                        await svc.renewGenerationControlLease({
+                            generationId: currentGenerationId,
+                            attemptEpoch,
+                            ownerToken: attemptOwnerToken,
+                        });
+                    },
+                    onAttemptCompleted: async ({ generationId: currentGenerationId, attemptEpoch, ownerToken: attemptOwnerToken }) => {
+                        await svc.markGenerationControlAttemptCompleted({
+                            generationId: currentGenerationId,
+                            attemptEpoch,
+                            ownerToken: attemptOwnerToken,
+                        });
+                    },
+                })
+                : undefined;
+            const streamed = await assistantSvc.streamChatAssistantReply({
+                ...assistantInput,
+                userMessageId: userMessage.id,
+                chatTurnId: turnContext.chatTurnId,
+                turnVariant: turnContext.turnVariant,
+                runContext: { chatGenerationId: generationId },
+                resumeRunId: run.id,
+                resumeRunOwnerToken: ownerToken,
+                stream: false,
+                abortSignal: abortController.signal,
+                controlCoordinator,
+                onRunCreated: (recoveredRunId) => {
+                    activeChatRunId = recoveredRunId;
+                },
+                onWaitingForNetwork: async (suspension) => {
+                    if (!generationId) {
+                        waitingAgain = true;
+                        return;
+                    }
+                    const marked = await svc.generationProtocol.markWaitingForNetwork({
+                        orgId: conversation.orgId,
+                        conversationId: conversation.id,
+                        generationId,
+                        expectedAttemptEpoch: activeAttemptEpoch,
+                        suspension: suspension,
+                    });
+                    if (marked.stopped) {
+                        terminalStatus = "stopped";
+                        terminalReason = "operator_stop";
+                        if (!abortController.signal.aborted)
+                            abortController.abort();
+                        return;
+                    }
+                    waitingAgain = true;
+                },
+                onAssistantDelta: async (delta) => {
+                    if (abortController.signal.aborted)
+                        return;
+                    const projectedBody = `${partialBody}${delta}`;
+                    if (!generationId) {
+                        partialBody = projectedBody;
+                        return;
+                    }
+                    const projection = await svc.generationProtocol.appendVisibleEventAndProject({
+                        orgId: conversation.orgId,
+                        conversationId: conversation.id,
+                        generationId,
+                        expectedAttemptEpoch: activeAttemptEpoch,
+                        eventKind: "assistant_delta",
+                        payload: { delta },
+                        bodyOffset: partialBody.length,
+                        bodyLength: delta.length,
+                        messageId: assistantProjectionMessageId,
+                        runId: activeChatRunId,
+                        bodyHash: hashChatGenerationBody(projectedBody),
+                        body: projectedBody,
+                        replyingAgentId: chatReplyingAgentId(assistantConversation),
+                        chatTurnId: turnContext.chatTurnId,
+                        turnVariant: turnContext.turnVariant,
+                    });
+                    assistantProjectionMessageId = projection.message.id;
+                    partialBody = projectedBody;
+                },
+                onTranscriptEntry: async (entry) => {
+                    if (abortController.signal.aborted)
+                        return;
+                    transcript.push(entry);
+                    if (!generationId)
+                        return;
+                    const projection = await svc.generationProtocol.appendVisibleEventAndProject({
+                        orgId: conversation.orgId,
+                        conversationId: conversation.id,
+                        generationId,
+                        expectedAttemptEpoch: activeAttemptEpoch,
+                        eventKind: "transcript",
+                        payload: { entry },
+                        messageId: assistantProjectionMessageId,
+                        runId: activeChatRunId,
+                        bodyHash: hashChatGenerationBody(partialBody),
+                        body: partialBody,
+                        replyingAgentId: chatReplyingAgentId(assistantConversation),
+                        chatTurnId: turnContext.chatTurnId,
+                        turnVariant: turnContext.turnVariant,
+                    });
+                    assistantProjectionMessageId = projection.message.id;
+                },
+            });
+            partialBody = streamed.partialBody || partialBody;
+            if (waitingAgain || streamed.outcome === "waiting_for_network")
+                return true;
+            if (abortController.signal.aborted || streamed.outcome === "stopped") {
+                terminalStatus = "stopped";
+                terminalReason = "operator_stop";
+                const stoppedMessage = await persistPartialAssistantMessage(assistantConversation, partialBody, "stopped", turnContext, transcript, streamed.replyingAgentId, assistantProjectionMessageId, activeChatRunId, undefined, false);
+                if (stoppedMessage) {
+                    await linkChatRunMessages(assistantConversation, activeChatRunId, [stoppedMessage]);
+                    await logChatMessagesAdded(assistantConversation, [stoppedMessage], {
+                        actorType: "system",
+                        actorId: "chat-assistant",
+                        agentId: streamed.replyingAgentId,
+                    });
+                }
+            }
+            else {
+                const completionMessageId = generationId
+                    ? (await svc.generationProtocol.appendVisibleEventAndProject({
+                        orgId: conversation.orgId,
+                        conversationId: conversation.id,
+                        generationId,
+                        expectedAttemptEpoch: activeAttemptEpoch,
+                        eventKind: "runtime_output",
+                        payload: { resultKind: streamed.reply.kind, body: streamed.reply.body },
+                        bodyOffset: 0,
+                        bodyLength: streamed.reply.body.length,
+                        messageId: assistantProjectionMessageId,
+                        runId: activeChatRunId,
+                        bodyHash: hashChatGenerationBody(streamed.reply.body),
+                        body: streamed.reply.body,
+                        replyingAgentId: streamed.replyingAgentId,
+                        chatTurnId: turnContext.chatTurnId,
+                        turnVariant: turnContext.turnVariant,
+                    })).message.id
+                    : null;
+                const createdMessages = await persistAssistantReply(request, assistantConversation, actor, streamed.reply, turnContext, transcript, streamed.replyingAgentId, completionMessageId, activeChatRunId, false);
+                await linkChatRunMessages(assistantConversation, activeChatRunId, createdMessages);
+                await logChatMessagesAdded(assistantConversation, createdMessages, {
+                    actorType: "system",
+                    actorId: "chat-assistant",
+                    agentId: streamed.replyingAgentId,
+                });
+                terminalStatus = "completed";
+            }
+        }
+        catch (error) {
+            terminalStatus = abortController.signal.aborted ? "aborted" : "failed";
+            terminalReason = error instanceof Error ? error.message : "network_wait_chat_recovery_failed";
+            const failurePayload = recoverableFailurePayload(error, activeChatRunId);
+            const failureBody = userVisiblePartialBodyFromError(error)
+                || recoverableFailureBody(failurePayload)
+                || CHAT_ASSISTANT_USER_ERROR_MESSAGE;
+            const failedMessage = await persistPartialAssistantMessage(assistantConversation, failureBody, "failed", null, transcript, chatReplyingAgentId(assistantConversation), assistantProjectionMessageId, activeChatRunId, failurePayload, false).catch(() => null);
+            if (failedMessage) {
+                await linkChatRunMessages(assistantConversation, activeChatRunId, [failedMessage]).catch(() => undefined);
+                await logChatMessagesAdded(assistantConversation, [failedMessage], {
+                    actorType: "system",
+                    actorId: "chat-assistant",
+                    agentId: chatReplyingAgentId(assistantConversation),
+                }).catch(() => undefined);
+            }
+        }
+        finally {
+            if (!waitingAgain && generationId) {
+                const latestGeneration = await svc.getLatestGeneration(conversation.id).catch(() => null);
+                if (latestGeneration?.id === generationId) {
+                    const terminalEvidence = await svc.generationProtocol.recordRuntimeTerminal({
+                        orgId: conversation.orgId,
+                        conversationId: conversation.id,
+                        generationId,
+                        expectedAttemptEpoch: latestGeneration.attemptEpoch,
+                        expectedOwnerToken: latestGeneration.controlOwnerToken,
+                        finalStatus: terminalStatus,
+                        terminalReason: terminalReason ?? terminalStatus,
+                    }).catch((error) => {
+                        logger.warn({ err: error, generationId }, "failed to record recovered Chat terminal evidence");
+                        return null;
+                    });
+                    if (terminalEvidence)
+                        wakeTerminalProjector();
+                }
+            }
+            chatRunsSvc.releaseOwnedRun(run.id, ownerToken);
+            releaseGeneration();
+            managedAbort.release();
+        }
+        return true;
+    }
+    recoverNetworkWaitingChatRun = runRecoveredNetworkWaitingChatRun;
+    registerNetworkWaitingRunHandler?.(runRecoveredNetworkWaitingChatRun);
     async function drainServerQueue() {
         await svc.recoverExpiredServerQueueClaims();
         while (backgroundRuntime.acceptingWork
@@ -1779,7 +2140,10 @@ export function chatRoutes(db, storage, backgroundRuntime = createChatBackground
             return;
         }
         const active = getActiveChatGeneration(conversation.id);
-        res.json(await svc.getQueueSnapshot(conversation.id, active?.generationId ?? null));
+        const snapshot = active
+            ? await svc.getQueueSnapshot(conversation.id, active.generationId)
+            : await svc.getQueueSnapshot(conversation.id);
+        res.json(snapshot);
     });
     router.post("/chats/:id/queue", async (req, res) => {
         const conversation = await assertConversationAccess(req, req.params.id);
@@ -2503,6 +2867,7 @@ export function chatRoutes(db, storage, backgroundRuntime = createChatBackground
             }
             const turnContext = turnContextFromUserMessage(userMessage);
             let activeChatRunId = null;
+            let networkWaiting = false;
             const persistedAssistantMessages = await (async () => {
                 const assistantInput = await loadAssistantInput(conversation, actor);
                 const transcript = [];
@@ -2514,6 +2879,7 @@ export function chatRoutes(db, storage, backgroundRuntime = createChatBackground
                         userMessageId: userMessage.id,
                         chatTurnId: turnContext.chatTurnId,
                         turnVariant: turnContext.turnVariant,
+                        runContext: { chatMode: "non_stream" },
                         stream: false,
                         onRunCreated: (runId) => {
                             activeChatRunId = runId;
@@ -2523,6 +2889,10 @@ export function chatRoutes(db, storage, backgroundRuntime = createChatBackground
                         },
                     });
                     fallbackOutput = streamed.partialBody;
+                    if (streamed.outcome === "waiting_for_network") {
+                        networkWaiting = true;
+                        return [];
+                    }
                     if (streamed.outcome !== "completed") {
                         throw new Error("Chat assistant reply was stopped before completion");
                     }
@@ -2557,7 +2927,10 @@ export function chatRoutes(db, storage, backgroundRuntime = createChatBackground
                 }
             })();
             const createdMessages = [userMessage, ...persistedAssistantMessages];
-            res.status(201).json({ messages: createdMessages });
+            res.status(networkWaiting ? 202 : 201).json({
+                messages: createdMessages,
+                ...(networkWaiting ? { waitingForNetwork: true, runId: activeChatRunId } : {}),
+            });
         }
         catch (err) {
             logger.warn({