🧵 fix: Prevent Message Loading Race During Streaming (#13295)

api/server/middleware/__tests__/validateMessageReq.spec.js

@@ -2,8 +2,22 @@ jest.mock('~/models', () => ({
   getConvo: jest.fn(),
 }));
 
+jest.mock('@librechat/api', () => ({
+  GenerationJobManager: {
+    getJob: jest.fn(),
+  },
+}));
+
+jest.mock('@librechat/data-schemas', () => ({
+  logger: {
+    warn: jest.fn(),
+  },
+}));
+
 const validateMessageReq = require('../validateMessageReq');
 const { getConvo } = require('~/models');
+const { GenerationJobManager } = require('@librechat/api');
+const { logger } = require('@librechat/data-schemas');
 
 function createResponse() {
   const res = {
@@ -71,4 +85,153 @@ describe('validateMessageReq', () => {
     expect(getConvo).toHaveBeenCalledWith(userId, 'convo-owned');
     expect(next).toHaveBeenCalledTimes(1);
   });
+
+  it('should allow message reads for an owned active generation job before the conversation is saved', async () => {
+    const req = {
+      method: 'GET',
+      params: { conversationId: 'active-convo' },
+      body: {},
+      user: { id: userId, tenantId: 'tenant-a' },
+    };
+    const res = createResponse();
+    const next = jest.fn();
+    getConvo.mockResolvedValue(null);
+    GenerationJobManager.getJob.mockResolvedValue({
+      status: 'running',
+      metadata: { userId, tenantId: 'tenant-a' },
+    });
+
+    await validateMessageReq(req, res, next);
+
+    expect(GenerationJobManager.getJob).toHaveBeenCalledWith('active-convo');
+    expect(next).toHaveBeenCalledTimes(1);
+    expect(res.status).not.toHaveBeenCalled();
+  });
+
+  it('should allow message reads for an owned active generation job without tenant metadata', async () => {
+    const req = {
+      method: 'GET',
+      params: { conversationId: 'active-convo' },
+      body: {},
+      user: { id: userId },
+    };
+    const res = createResponse();
+    const next = jest.fn();
+    getConvo.mockResolvedValue(null);
+    GenerationJobManager.getJob.mockResolvedValue({
+      status: 'running',
+      metadata: { userId },
+    });
+
+    await validateMessageReq(req, res, next);
+
+    expect(next).toHaveBeenCalledTimes(1);
+    expect(res.status).not.toHaveBeenCalled();
+  });
+
+  it('should reject active job message reads owned by another user', async () => {
+    const req = {
+      method: 'GET',
+      params: { conversationId: 'active-convo' },
+      body: {},
+      user: { id: userId },
+    };
+    const res = createResponse();
+    const next = jest.fn();
+    getConvo.mockResolvedValue(null);
+    GenerationJobManager.getJob.mockResolvedValue({
+      status: 'running',
+      metadata: { userId: 'another-user' },
+    });
+
+    await validateMessageReq(req, res, next);
+
+    expect(res.status).toHaveBeenCalledWith(404);
+    expect(res.json).toHaveBeenCalledWith({ error: 'Conversation not found' });
+    expect(next).not.toHaveBeenCalled();
+  });
+
+  it('should reject active job message reads from another tenant', async () => {
+    const req = {
+      method: 'GET',
+      params: { conversationId: 'active-convo' },
+      body: {},
+      user: { id: userId, tenantId: 'tenant-a' },
+    };
+    const res = createResponse();
+    const next = jest.fn();
+    getConvo.mockResolvedValue(null);
+    GenerationJobManager.getJob.mockResolvedValue({
+      status: 'running',
+      metadata: { userId, tenantId: 'tenant-b' },
+    });
+
+    await validateMessageReq(req, res, next);
+
+    expect(res.status).toHaveBeenCalledWith(404);
+    expect(res.json).toHaveBeenCalledWith({ error: 'Conversation not found' });
+    expect(next).not.toHaveBeenCalled();
+  });
+
+  it('should reject message-by-id reads before the conversation is saved', async () => {
+    const req = {
+      method: 'GET',
+      params: { conversationId: 'active-convo', messageId: 'message-id' },
+      body: {},
+      user: { id: userId },
+    };
+    const res = createResponse();
+    const next = jest.fn();
+    getConvo.mockResolvedValue(null);
+
+    await validateMessageReq(req, res, next);
+
+    expect(GenerationJobManager.getJob).not.toHaveBeenCalled();
+    expect(res.status).toHaveBeenCalledWith(404);
+    expect(res.json).toHaveBeenCalledWith({ error: 'Conversation not found' });
+    expect(next).not.toHaveBeenCalled();
+  });
+
+  it('should return not found when active job lookup fails', async () => {
+    const req = {
+      method: 'GET',
+      params: { conversationId: 'active-convo' },
+      body: {},
+      user: { id: userId },
+    };
+    const res = createResponse();
+    const next = jest.fn();
+    const error = new Error('job store unavailable');
+    getConvo.mockResolvedValue(null);
+    GenerationJobManager.getJob.mockRejectedValue(error);
+
+    await validateMessageReq(req, res, next);
+
+    expect(GenerationJobManager.getJob).toHaveBeenCalledWith('active-convo');
+    expect(logger.warn).toHaveBeenCalledWith(
+      '[validateMessageReq] Active job lookup failed for active-convo:',
+      error,
+    );
+    expect(res.status).toHaveBeenCalledWith(404);
+    expect(res.json).toHaveBeenCalledWith({ error: 'Conversation not found' });
+    expect(next).not.toHaveBeenCalled();
+  });
+
+  it('should not allow unsaved conversation writes through active job ownership', async () => {
+    const req = {
+      method: 'POST',
+      params: { conversationId: 'active-convo' },
+      body: {},
+      user: { id: userId },
+    };
+    const res = createResponse();
+    const next = jest.fn();
+    getConvo.mockResolvedValue(null);
+
+    await validateMessageReq(req, res, next);
+
+    expect(GenerationJobManager.getJob).not.toHaveBeenCalled();
+    expect(res.status).toHaveBeenCalledWith(404);
+    expect(next).not.toHaveBeenCalled();
+  });
 });

api/server/middleware/validateMessageReq.js

@@ -1,5 +1,32 @@
+const { GenerationJobManager } = require('@librechat/api');
+const { logger } = require('@librechat/data-schemas');
 const { getConvo } = require('~/models');
 
+function hasTenantMismatch(job, user) {
+  // Untenanted jobs remain readable by their owner for pre-multi-tenancy deployments.
+  return job.metadata?.tenantId != null && job.metadata.tenantId !== user.tenantId;
+}
+
+async function canReadActiveJobConversation(req, conversationId) {
+  if (req.method !== 'GET' || req.params?.messageId) {
+    return false;
+  }
+
+  let job;
+  try {
+    job = await GenerationJobManager.getJob(conversationId);
+  } catch (error) {
+    logger.warn(`[validateMessageReq] Active job lookup failed for ${conversationId}:`, error);
+    return false;
+  }
+
+  if (!job || job.status !== 'running') {
+    return false;
+  }
+
+  return job.metadata?.userId === req.user.id && !hasTenantMismatch(job, req.user);
+}
+
 // Middleware to validate conversationId and user relationship
 const validateMessageReq = async (req, res, next) => {
   const body = req.body ?? {};
@@ -25,6 +52,10 @@ const validateMessageReq = async (req, res, next) => {
   const conversation = await getConvo(req.user.id, conversationId);
 
   if (!conversation) {
+    if (await canReadActiveJobConversation(req, conversationId)) {
+      return next();
+    }
+
     return res.status(404).json({ error: 'Conversation not found' });
   }
 

client/src/data-provider/Messages/__tests__/queries.test.ts

@@ -1,5 +1,38 @@
+import { createElement } from 'react';
+import { act, renderHook, waitFor } from '@testing-library/react';
+import { MemoryRouter } from 'react-router-dom';
+import { QueryClient, QueryClientProvider } from '@tanstack/react-query';
+import { QueryKeys, dataService } from 'librechat-data-provider';
+import type { ReactNode } from 'react';
 import type { TMessage } from 'librechat-data-provider';
-import { getStableMessages } from '../queries';
+import { logger } from '~/utils';
+import {
+  getStableMessages,
+  shouldPreserveMessagesOnNotFound,
+  useGetMessagesByConvoId,
+} from '../queries';
+
+jest.mock('librechat-data-provider', () => {
+  const actual = jest.requireActual('librechat-data-provider');
+  return {
+    ...actual,
+    dataService: {
+      ...actual.dataService,
+      getMessagesByConvoId: jest.fn(),
+    },
+  };
+});
+
+jest.mock('~/utils', () => {
+  const actual = jest.requireActual('~/utils');
+  return {
+    ...actual,
+    logger: {
+      ...actual.logger,
+      warn: jest.fn(),
+    },
+  };
+});
 
 const message = (overrides: Partial<TMessage>): TMessage =>
   ({
@@ -14,6 +47,20 @@ const message = (overrides: Partial<TMessage>): TMessage =>
     ...overrides,
   }) as TMessage;
 
+function createWrapper(queryClient: QueryClient, initialEntry: string) {
+  return function Wrapper({ children }: { children: ReactNode }) {
+    return createElement(
+      QueryClientProvider,
+      { client: queryClient },
+      createElement(MemoryRouter, { initialEntries: [initialEntry] }, children),
+    );
+  };
+}
+
+afterEach(() => {
+  jest.clearAllMocks();
+});
+
 describe('getStableMessages', () => {
   it('keeps cache when an empty result races with unhydrated stream messages', () => {
     const currentMessages = [
@@ -179,3 +226,176 @@ describe('getStableMessages', () => {
     expect(result).toEqual([]);
   });
 });
+
+describe('shouldPreserveMessagesOnNotFound', () => {
+  it('keeps cache when a transient 404 races with a pending assistant tail', () => {
+    const currentMessages = [
+      message({ messageId: 'persisted-1' }),
+      message({ messageId: 'user-2' }),
+      message({
+        messageId: 'user-2_',
+        parentMessageId: 'user-2',
+        isCreatedByUser: false,
+        createdAt: undefined,
+        updatedAt: undefined,
+      }),
+    ];
+
+    expect(
+      shouldPreserveMessagesOnNotFound({
+        pathname: '/c/convo-id',
+        currentMessages,
+        isStreaming: true,
+      }),
+    ).toBe(true);
+  });
+
+  it('does not preserve cache when no stream or active job is live', () => {
+    const currentMessages = [
+      message({ messageId: 'persisted-1' }),
+      message({ messageId: 'user-2' }),
+      message({
+        messageId: 'user-2_',
+        parentMessageId: 'user-2',
+        isCreatedByUser: false,
+        createdAt: undefined,
+        updatedAt: undefined,
+      }),
+    ];
+
+    expect(
+      shouldPreserveMessagesOnNotFound({
+        pathname: '/c/convo-id',
+        currentMessages,
+        isStreaming: false,
+      }),
+    ).toBe(false);
+  });
+
+  it('does not preserve cache on the new conversation route', () => {
+    const currentMessages = [
+      message({ messageId: 'user-1' }),
+      message({
+        messageId: 'user-1_',
+        parentMessageId: 'user-1',
+        isCreatedByUser: false,
+        createdAt: undefined,
+        updatedAt: undefined,
+      }),
+    ];
+
+    expect(
+      shouldPreserveMessagesOnNotFound({
+        pathname: '/c/new',
+        currentMessages,
+        isStreaming: true,
+      }),
+    ).toBe(false);
+  });
+
+  it('does not preserve cache when there is no pending assistant tail', () => {
+    const currentMessages = [message({ messageId: 'persisted-1' })];
+
+    expect(
+      shouldPreserveMessagesOnNotFound({
+        pathname: '/c/convo-id',
+        currentMessages,
+        isStreaming: true,
+      }),
+    ).toBe(false);
+  });
+});
+
+describe('useGetMessagesByConvoId', () => {
+  it('keeps cache during submitting cleanup when active job cache still marks the stream active', async () => {
+    const conversationId = 'convo-id';
+    const currentMessages = [
+      message({ messageId: 'persisted-1' }),
+      message({ messageId: 'user-2', createdAt: undefined, updatedAt: undefined }),
+      message({
+        messageId: 'user-2_',
+        parentMessageId: 'user-2',
+        isCreatedByUser: false,
+        createdAt: undefined,
+        updatedAt: undefined,
+      }),
+    ];
+    const serverMessages = [currentMessages[0]];
+    const queryClient = new QueryClient({
+      defaultOptions: { queries: { retry: false } },
+    });
+    queryClient.setQueryData([QueryKeys.messages, conversationId], currentMessages);
+    queryClient.setQueryData([QueryKeys.activeJobs], { activeJobIds: [conversationId] });
+
+    const mockGetMessagesByConvoId = dataService.getMessagesByConvoId as jest.MockedFunction<
+      typeof dataService.getMessagesByConvoId
+    >;
+    mockGetMessagesByConvoId.mockResolvedValue(serverMessages);
+
+    const { result, unmount } = renderHook(
+      () => useGetMessagesByConvoId(conversationId, undefined, { isStreaming: false }),
+      { wrapper: createWrapper(queryClient, `/c/${conversationId}`) },
+    );
+
+    await act(async () => {
+      await result.current.refetch();
+    });
+
+    await waitFor(() => {
+      expect(result.current.data).toBe(currentMessages);
+    });
+    expect(dataService.getMessagesByConvoId).toHaveBeenCalledWith(conversationId);
+    expect(queryClient.getQueryData([QueryKeys.messages, conversationId])).toBe(currentMessages);
+
+    unmount();
+  });
+
+  it('keeps cache when a transient 404 races with a regenerated pending assistant tail', async () => {
+    const conversationId = 'convo-id';
+    const currentMessages = [
+      message({ messageId: 'persisted-1' }),
+      message({
+        messageId: 'user-2',
+        parentMessageId: 'persisted-1',
+      }),
+      message({
+        messageId: 'user-2_',
+        parentMessageId: 'user-2',
+        isCreatedByUser: false,
+        createdAt: undefined,
+        updatedAt: undefined,
+      }),
+    ];
+    const queryClient = new QueryClient({
+      defaultOptions: { queries: { retry: false } },
+    });
+    queryClient.setQueryData([QueryKeys.messages, conversationId], currentMessages);
+
+    const mockGetMessagesByConvoId = dataService.getMessagesByConvoId as jest.MockedFunction<
+      typeof dataService.getMessagesByConvoId
+    >;
+    mockGetMessagesByConvoId.mockRejectedValueOnce({ status: 404 });
+
+    const { result, unmount } = renderHook(
+      () => useGetMessagesByConvoId(conversationId, { enabled: false }, { isStreaming: true }),
+      { wrapper: createWrapper(queryClient, `/c/${conversationId}`) },
+    );
+
+    await act(async () => {
+      await result.current.refetch();
+    });
+
+    await waitFor(() => {
+      expect(result.current.data).toBe(currentMessages);
+    });
+    expect(dataService.getMessagesByConvoId).toHaveBeenCalledWith(conversationId);
+    expect(queryClient.getQueryData([QueryKeys.messages, conversationId])).toBe(currentMessages);
+    expect(logger.warn).toHaveBeenCalledWith(
+      'messages',
+      expect.stringContaining('returned 404 while cache has a pending assistant tail'),
+      currentMessages,
+    );
+
+    unmount();
+  });
+});

client/src/data-provider/Messages/queries.ts

@@ -4,7 +4,7 @@ import { useQuery, useQueryClient } from '@tanstack/react-query';
 import type { UseQueryOptions, QueryObserverResult, QueryClient } from '@tanstack/react-query';
 import { Constants, QueryKeys, dataService } from 'librechat-data-provider';
 import type * as t from 'librechat-data-provider';
-import { logger } from '~/utils';
+import { isNotFoundError, logger } from '~/utils';
 
 type StableMessagesParams = {
   pathname: string;
@@ -62,6 +62,18 @@ export function getStableMessages({
   return result;
 }
 
+export function shouldPreserveMessagesOnNotFound({
+  pathname,
+  isStreaming = false,
+  currentMessages,
+}: Pick<StableMessagesParams, 'pathname' | 'isStreaming' | 'currentMessages'>): boolean {
+  if (!isStreaming || pathname.includes('/c/new') || !currentMessages?.length) {
+    return false;
+  }
+
+  return hasPendingAssistantTail(currentMessages);
+}
+
 function hasActiveJob(queryClient: QueryClient, id: string) {
   if (!id) {
     return false;
@@ -87,7 +99,32 @@ export const useGetMessagesByConvoId = <TData = t.TMessage[]>(
   return useQuery<t.TMessage[], unknown, TData>(
     [QueryKeys.messages, id],
     async () => {
-      const result = await dataService.getMessagesByConvoId(id);
+      let result: t.TMessage[];
+      try {
+        result = await dataService.getMessagesByConvoId(id);
+      } catch (error) {
+        const currentMessages = queryClient.getQueryData<t.TMessage[]>([QueryKeys.messages, id]);
+        const hasLiveStream = isStreamingRef.current || hasActiveJob(queryClient, id);
+        if (
+          currentMessages &&
+          isNotFoundError(error) &&
+          shouldPreserveMessagesOnNotFound({
+            pathname: location.pathname,
+            currentMessages,
+            isStreaming: hasLiveStream,
+          })
+        ) {
+          logger.warn(
+            'messages',
+            `Messages query for convo ${id} returned 404 while cache has a pending assistant tail; path: "${location.pathname}"`,
+            currentMessages,
+          );
+          return currentMessages;
+        }
+
+        throw error;
+      }
+
       const currentMessages = queryClient.getQueryData<t.TMessage[]>([QueryKeys.messages, id]);
       const stableMessages = getStableMessages({
         pathname: location.pathname,

client/src/hooks/Chat/__tests__/cache.spec.ts

@@ -0,0 +1,41 @@
+import { Constants } from 'librechat-data-provider';
+import type { TMessage } from 'librechat-data-provider';
+import { getMessageCacheIds, getMessagesConversationId } from '../cache';
+
+const message = (conversationId?: string | null): TMessage =>
+  ({
+    messageId: 'message-id',
+    conversationId,
+  }) as TMessage;
+
+describe('chat message cache helpers', () => {
+  it('uses the latest concrete conversation id from streamed messages', () => {
+    expect(
+      getMessagesConversationId([
+        message(Constants.NEW_CONVO),
+        message(null),
+        message('generated-convo-id'),
+      ]),
+    ).toBe('generated-convo-id');
+  });
+
+  it('mirrors new-chat messages into the generated conversation cache', () => {
+    expect(
+      getMessageCacheIds({
+        queryParam: Constants.NEW_CONVO,
+        conversationId: Constants.NEW_CONVO,
+        messages: [message('generated-convo-id')],
+      }),
+    ).toEqual([Constants.NEW_CONVO, 'generated-convo-id']);
+  });
+
+  it('keeps the current conversation cache id while avoiding duplicate ids', () => {
+    expect(
+      getMessageCacheIds({
+        queryParam: 'generated-convo-id',
+        conversationId: 'generated-convo-id',
+        messages: [message('generated-convo-id')],
+      }),
+    ).toEqual(['generated-convo-id']);
+  });
+});

client/src/hooks/Chat/cache.ts

@@ -0,0 +1,44 @@
+import { Constants } from 'librechat-data-provider';
+import type { TMessage } from 'librechat-data-provider';
+
+type MessageCacheIdsParams = {
+  queryParam: string;
+  conversationId?: string | null;
+  messages: TMessage[];
+};
+
+function isConcreteConversationId(conversationId?: string | null) {
+  return (
+    !!conversationId &&
+    conversationId !== Constants.NEW_CONVO &&
+    conversationId !== Constants.PENDING_CONVO
+  );
+}
+
+export function getMessagesConversationId(messages: TMessage[]): string | undefined {
+  for (let i = messages.length - 1; i >= 0; i--) {
+    const conversationId = messages[i]?.conversationId;
+    if (isConcreteConversationId(conversationId)) {
+      return conversationId;
+    }
+  }
+}
+
+export function getMessageCacheIds({
+  queryParam,
+  conversationId,
+  messages,
+}: MessageCacheIdsParams): string[] {
+  const ids = [queryParam];
+  const messageConversationId = getMessagesConversationId(messages);
+
+  if (queryParam === Constants.NEW_CONVO && isConcreteConversationId(conversationId)) {
+    ids.push(conversationId);
+  }
+
+  if (isConcreteConversationId(messageConversationId) && !ids.includes(messageConversationId)) {
+    ids.push(messageConversationId);
+  }
+
+  return ids;
+}

client/src/hooks/Chat/useChatHelpers.ts

@@ -8,6 +8,7 @@ import useChatFunctions from '~/hooks/Chat/useChatFunctions';
 import { useAbortStreamMutation } from '~/data-provider';
 import useNewConvo from '~/hooks/useNewConvo';
 import { useLatestMessage, useLatestMessageId } from '~/hooks/Messages/useLatestMessage';
+import { getMessageCacheIds } from './cache';
 import store from '~/store';
 
 // this to be set somewhere else
@@ -42,9 +43,9 @@ export default function useChatHelpers(index = 0, paramId?: string) {
 
   const setMessages = useCallback(
     (messages: TMessage[]) => {
-      queryClient.setQueryData<TMessage[]>([QueryKeys.messages, queryParam], messages);
-      if (queryParam === 'new' && conversationId && conversationId !== 'new') {
-        queryClient.setQueryData<TMessage[]>([QueryKeys.messages, conversationId], messages);
+      const messageCacheIds = getMessageCacheIds({ queryParam, conversationId, messages });
+      for (const messageCacheId of messageCacheIds) {
+        queryClient.setQueryData<TMessage[]>([QueryKeys.messages, messageCacheId], messages);
       }
     },
     [queryParam, queryClient, conversationId],

packages/api/src/stream/__tests__/RedisEventTransport.spec.ts

@@ -0,0 +1,56 @@
+import type { Redis } from 'ioredis';
+import { logger } from '@librechat/data-schemas';
+import { RedisEventTransport } from '~/stream/implementations/RedisEventTransport';
+import { createMockPublisher } from './helpers/publisher';
+
+logger.silent = true;
+
+function createMockSubscriber() {
+  return {
+    on: jest.fn(),
+    subscribe: jest.fn().mockResolvedValue(undefined),
+    unsubscribe: jest.fn().mockResolvedValue(undefined),
+  };
+}
+
+function getMessageHandler(mockSubscriber: ReturnType<typeof createMockSubscriber>) {
+  return mockSubscriber.on.mock.calls.find((call) => call[0] === 'message')?.[1] as (
+    channel: string,
+    message: string,
+  ) => void;
+}
+
+describe('RedisEventTransport', () => {
+  it('resets stale abort-listener reorder state before the next real subscriber', async () => {
+    const mockPublisher = createMockPublisher();
+    const mockSubscriber = createMockSubscriber();
+    const transport = new RedisEventTransport(
+      mockPublisher as unknown as Redis,
+      mockSubscriber as unknown as Redis,
+    );
+
+    const streamId = 'reorder-abort-listener-reuse-test';
+    transport.onAbort(streamId, () => {});
+
+    const messageHandler = getMessageHandler(mockSubscriber);
+    const channel = `stream:{${streamId}}:events`;
+
+    for (let i = 0; i < 5; i++) {
+      await transport.emitChunk(streamId, { index: i });
+      messageHandler(channel, JSON.stringify({ type: 'chunk', seq: i, data: { index: i } }));
+    }
+
+    await mockPublisher.del(`stream:{${streamId}}:seq`);
+
+    const secondRunChunks: unknown[] = [];
+    transport.subscribe(streamId, {
+      onChunk: (event) => secondRunChunks.push(event),
+    });
+
+    messageHandler(channel, JSON.stringify({ type: 'chunk', seq: 0, data: { index: 0 } }));
+
+    expect(secondRunChunks.map((chunk) => (chunk as { index: number }).index)).toEqual([0]);
+
+    transport.destroy();
+  });
+});

packages/api/src/stream/implementations/RedisEventTransport.ts

@@ -439,6 +439,12 @@ export class RedisEventTransport implements IEventTransport {
     }
 
     const streamState = this.streams.get(streamId)!;
+    // Internal listeners (for example cross-replica abort) can leave ordering
+    // state behind with no real SSE subscribers. A new subscriber is a fresh
+    // attachment and must not inherit that prior generation's expected seq.
+    if (streamState.count === 0) {
+      this.resetReorderBuffer(streamId);
+    }
     streamState.count++;
     streamState.handlers.set(subscriberId, handlers);