LibreChat/api/app/clients/specs/BaseClient.test.js

const { Constants } = require('librechat-data-provider');
const { FakeClient, initializeFakeClient } = require('./FakeClient');

function deferred() {
  let resolve;
  const promise = new Promise((resolvePromise) => {
    resolve = resolvePromise;
  });
  return { promise, resolve };
}

jest.mock('~/db/connect');
jest.mock('~/server/services/Config', () => ({
  getAppConfig: jest.fn().mockResolvedValue({
    // Default app config for tests
    paths: { uploads: '/tmp' },
    fileStrategy: 'local',
    memory: { disabled: false },
  }),
}));
jest.mock('~/models', () => ({
  User: jest.fn(),
  Key: jest.fn(),
  Session: jest.fn(),
  Balance: jest.fn(),
  Transaction: jest.fn(),
  getMessages: jest.fn().mockResolvedValue([]),
  saveMessage: jest.fn(),
  updateMessage: jest.fn(),
  deleteMessagesSince: jest.fn(),
  deleteMessages: jest.fn(),
  getConvoTitle: jest.fn(),
  getConvo: jest.fn(),
  saveConvo: jest.fn(),
  deleteConvos: jest.fn(),
  getPreset: jest.fn(),
  getPresets: jest.fn(),
  savePreset: jest.fn(),
  deletePresets: jest.fn(),
  findFileById: jest.fn(),
  createFile: jest.fn(),
  updateFile: jest.fn(),
  deleteFile: jest.fn(),
  deleteFiles: jest.fn(),
  getFiles: jest.fn(),
  updateFileUsage: jest.fn(),
}));

const { getConvo, getFiles, getMessages, saveConvo, saveMessage } = require('~/models');

jest.mock('@librechat/agents', () => {
  const actual = jest.requireActual('@librechat/agents');
  return {
    ...actual,
    ChatOpenAI: jest.fn().mockImplementation(() => {
      return {};
    }),
  };
});

let parentMessageId;
let conversationId;
const fakeMessages = [];
const userMessage = 'Hello, ChatGPT!';
const apiKey = 'fake-api-key';

const messageHistory = [
  { role: 'user', isCreatedByUser: true, text: 'Hello', messageId: '1' },
  { role: 'assistant', isCreatedByUser: false, text: 'Hi', messageId: '2', parentMessageId: '1' },
  {
    role: 'user',
    isCreatedByUser: true,
    text: "What's up",
    messageId: '3',
    parentMessageId: '2',
  },
];

describe('BaseClient', () => {
  let TestClient;
  const options = {
    // debug: true,
    modelOptions: {
      model: 'gpt-4o-mini',
      temperature: 0,
    },
  };

  beforeEach(() => {
    TestClient = initializeFakeClient(apiKey, options, fakeMessages);
    TestClient.summarizeMessages = jest.fn().mockResolvedValue({
      summaryMessage: {
        role: 'system',
        content: 'Refined answer',
      },
      summaryTokenCount: 5,
    });
  });

  test('returns the input messages without instructions when addInstructions() is called with empty instructions', () => {
    const messages = [{ content: 'Hello' }, { content: 'How are you?' }, { content: 'Goodbye' }];
    const instructions = '';
    const result = TestClient.addInstructions(messages, instructions);
    expect(result).toEqual(messages);
  });

  test('returns the input messages with instructions properly added when addInstructions() is called with non-empty instructions', () => {
    const messages = [{ content: 'Hello' }, { content: 'How are you?' }, { content: 'Goodbye' }];
    const instructions = { content: 'Please respond to the question.' };
    const result = TestClient.addInstructions(messages, instructions);
    const expected = [
      { content: 'Please respond to the question.' },
      { content: 'Hello' },
      { content: 'How are you?' },
      { content: 'Goodbye' },
    ];
    expect(result).toEqual(expected);
  });

  test('returns the input messages with instructions properly added when addInstructions() with legacy flag', () => {
    const messages = [{ content: 'Hello' }, { content: 'How are you?' }, { content: 'Goodbye' }];
    const instructions = { content: 'Please respond to the question.' };
    const result = TestClient.addInstructions(messages, instructions, true);
    const expected = [
      { content: 'Hello' },
      { content: 'How are you?' },
      { content: 'Please respond to the question.' },
      { content: 'Goodbye' },
    ];
    expect(result).toEqual(expected);
  });

  test('concats messages correctly in concatenateMessages()', () => {
    const messages = [
      { name: 'User', content: 'Hello' },
      { name: 'Assistant', content: 'How can I help you?' },
      { name: 'User', content: 'I have a question.' },
    ];
    const result = TestClient.concatenateMessages(messages);
    const expected =
      'User:\nHello\n\nAssistant:\nHow can I help you?\n\nUser:\nI have a question.\n\n';
    expect(result).toBe(expected);
  });

  test('refines messages correctly in summarizeMessages()', async () => {
    const messagesToRefine = [
      { role: 'user', content: 'Hello', tokenCount: 10 },
      { role: 'assistant', content: 'How can I help you?', tokenCount: 20 },
    ];
    const remainingContextTokens = 100;
    const expectedRefinedMessage = {
      role: 'system',
      content: 'Refined answer',
    };

    const result = await TestClient.summarizeMessages({ messagesToRefine, remainingContextTokens });
    expect(result.summaryMessage).toEqual(expectedRefinedMessage);
  });

  test('gets messages within token limit (under limit) correctly in getMessagesWithinTokenLimit()', async () => {
    TestClient.maxContextTokens = 100;
    TestClient.shouldSummarize = true;

    const messages = [
      { role: 'user', content: 'Hello', tokenCount: 5 },
      { role: 'assistant', content: 'How can I help you?', tokenCount: 19 },
      { role: 'user', content: 'I have a question.', tokenCount: 18 },
    ];
    const expectedContext = [
      { role: 'user', content: 'Hello', tokenCount: 5 }, // 'Hello'.length
      { role: 'assistant', content: 'How can I help you?', tokenCount: 19 },
      { role: 'user', content: 'I have a question.', tokenCount: 18 },
    ];
    // Subtract 3 tokens for Assistant Label priming after all messages have been counted.
    const expectedRemainingContextTokens = 58 - 3; // (100 - 5 - 19 - 18) - 3
    const expectedMessagesToRefine = [];

    const lastExpectedMessage =
      expectedMessagesToRefine?.[expectedMessagesToRefine.length - 1] ?? {};
    const expectedIndex = messages.findIndex((msg) => msg.content === lastExpectedMessage?.content);

    const result = await TestClient.getMessagesWithinTokenLimit({ messages });

    expect(result.context).toEqual(expectedContext);
    expect(result.messagesToRefine.length - 1).toEqual(expectedIndex);
    expect(result.remainingContextTokens).toBe(expectedRemainingContextTokens);
    expect(result.messagesToRefine).toEqual(expectedMessagesToRefine);
  });

  test('gets result over token limit correctly in getMessagesWithinTokenLimit()', async () => {
    TestClient.maxContextTokens = 50; // Set a lower limit
    TestClient.shouldSummarize = true;

    const messages = [
      { role: 'user', content: 'Hello', tokenCount: 30 },
      { role: 'assistant', content: 'How can I help you?', tokenCount: 30 },
      { role: 'user', content: 'I have a question.', tokenCount: 5 },
      { role: 'user', content: 'I need a coffee, stat!', tokenCount: 19 },
      { role: 'assistant', content: 'Sure, I can help with that.', tokenCount: 18 },
    ];

    // Subtract 3 tokens for Assistant Label priming after all messages have been counted.
    const expectedRemainingContextTokens = 5; // (50 - 18 - 19 - 5) - 3
    const expectedMessagesToRefine = [
      { role: 'user', content: 'Hello', tokenCount: 30 },
      { role: 'assistant', content: 'How can I help you?', tokenCount: 30 },
    ];
    const expectedContext = [
      { role: 'user', content: 'I have a question.', tokenCount: 5 },
      { role: 'user', content: 'I need a coffee, stat!', tokenCount: 19 },
      { role: 'assistant', content: 'Sure, I can help with that.', tokenCount: 18 },
    ];

    const lastExpectedMessage =
      expectedMessagesToRefine?.[expectedMessagesToRefine.length - 1] ?? {};
    const expectedIndex = messages.findIndex((msg) => msg.content === lastExpectedMessage?.content);

    const result = await TestClient.getMessagesWithinTokenLimit({ messages });

    expect(result.context).toEqual(expectedContext);
    expect(result.messagesToRefine.length - 1).toEqual(expectedIndex);
    expect(result.remainingContextTokens).toBe(expectedRemainingContextTokens);
    expect(result.messagesToRefine).toEqual(expectedMessagesToRefine);
  });

  describe('getMessagesForConversation', () => {
    it('should return an empty array if the parentMessageId does not exist', () => {
      const result = TestClient.constructor.getMessagesForConversation({
        messages: unorderedMessages,
        parentMessageId: '999',
      });
      expect(result).toEqual([]);
    });

    it('should handle messages with messageId property', () => {
      const messagesWithMessageId = [
        { messageId: '1', parentMessageId: null, text: 'Message 1' },
        { messageId: '2', parentMessageId: '1', text: 'Message 2' },
      ];
      const result = TestClient.constructor.getMessagesForConversation({
        messages: messagesWithMessageId,
        parentMessageId: '2',
      });
      expect(result).toEqual([
        { messageId: '1', parentMessageId: null, text: 'Message 1' },
        { messageId: '2', parentMessageId: '1', text: 'Message 2' },
      ]);
    });

    const messagesWithNullParent = [
      { id: '1', parentMessageId: null, text: 'Message 1' },
      { id: '2', parentMessageId: null, text: 'Message 2' },
    ];

    it('should handle messages with null parentMessageId that are not root', () => {
      const result = TestClient.constructor.getMessagesForConversation({
        messages: messagesWithNullParent,
        parentMessageId: '2',
      });
      expect(result).toEqual([{ id: '2', parentMessageId: null, text: 'Message 2' }]);
    });

    const cyclicMessages = [
      { id: '3', parentMessageId: '2', text: 'Message 3' },
      { id: '1', parentMessageId: '3', text: 'Message 1' },
      { id: '2', parentMessageId: '1', text: 'Message 2' },
    ];

    it('should handle cyclic references without going into an infinite loop', () => {
      const result = TestClient.constructor.getMessagesForConversation({
        messages: cyclicMessages,
        parentMessageId: '3',
      });
      expect(result).toEqual([
        { id: '1', parentMessageId: '3', text: 'Message 1' },
        { id: '2', parentMessageId: '1', text: 'Message 2' },
        { id: '3', parentMessageId: '2', text: 'Message 3' },
      ]);
    });

    const unorderedMessages = [
      { id: '3', parentMessageId: '2', text: 'Message 3' },
      { id: '2', parentMessageId: '1', text: 'Message 2' },
      { id: '1', parentMessageId: Constants.NO_PARENT, text: 'Message 1' },
    ];

    it('should return ordered messages based on parentMessageId', () => {
      const result = TestClient.constructor.getMessagesForConversation({
        messages: unorderedMessages,
        parentMessageId: '3',
      });
      expect(result).toEqual([
        { id: '1', parentMessageId: Constants.NO_PARENT, text: 'Message 1' },
        { id: '2', parentMessageId: '1', text: 'Message 2' },
        { id: '3', parentMessageId: '2', text: 'Message 3' },
      ]);
    });

    const unorderedBranchedMessages = [
      { id: '4', parentMessageId: '2', text: 'Message 4', summary: 'Summary for Message 4' },
      { id: '10', parentMessageId: '7', text: 'Message 10' },
      { id: '1', parentMessageId: null, text: 'Message 1' },
      { id: '6', parentMessageId: '5', text: 'Message 7' },
      { id: '7', parentMessageId: '5', text: 'Message 7' },
      { id: '2', parentMessageId: '1', text: 'Message 2' },
      { id: '8', parentMessageId: '6', text: 'Message 8' },
      { id: '5', parentMessageId: '3', text: 'Message 5' },
      { id: '3', parentMessageId: '1', text: 'Message 3' },
      { id: '6', parentMessageId: '4', text: 'Message 6' },
      { id: '8', parentMessageId: '7', text: 'Message 9' },
      { id: '9', parentMessageId: '7', text: 'Message 9' },
      { id: '11', parentMessageId: '2', text: 'Message 11', summary: 'Summary for Message 11' },
    ];

    it('should return ordered messages from a branched array based on parentMessageId', () => {
      const result = TestClient.constructor.getMessagesForConversation({
        messages: unorderedBranchedMessages,
        parentMessageId: '10',
        summary: true,
      });
      expect(result).toEqual([
        { id: '1', parentMessageId: null, text: 'Message 1' },
        { id: '3', parentMessageId: '1', text: 'Message 3' },
        { id: '5', parentMessageId: '3', text: 'Message 5' },
        { id: '7', parentMessageId: '5', text: 'Message 7' },
        { id: '10', parentMessageId: '7', text: 'Message 10' },
      ]);
    });

    it('should return an empty array if no messages are provided', () => {
      const result = TestClient.constructor.getMessagesForConversation({
        messages: [],
        parentMessageId: '3',
      });
      expect(result).toEqual([]);
    });

    it('should map over the ordered messages if mapMethod is provided', () => {
      const mapMethod = (msg) => msg.text;
      const result = TestClient.constructor.getMessagesForConversation({
        messages: unorderedMessages,
        parentMessageId: '3',
        mapMethod,
      });
      expect(result).toEqual(['Message 1', 'Message 2', 'Message 3']);
    });

    let unorderedMessagesWithSummary = [
      { id: '4', parentMessageId: '3', text: 'Message 4' },
      { id: '2', parentMessageId: '1', text: 'Message 2', summary: 'Summary for Message 2' },
      { id: '3', parentMessageId: '2', text: 'Message 3', summary: 'Summary for Message 3' },
      { id: '1', parentMessageId: null, text: 'Message 1' },
    ];

    it('should start with the message that has a summary property and continue until the specified parentMessageId', () => {
      const result = TestClient.constructor.getMessagesForConversation({
        messages: unorderedMessagesWithSummary,
        parentMessageId: '4',
        summary: true,
      });
      expect(result).toEqual([
        {
          id: '3',
          parentMessageId: '2',
          role: 'system',
          text: 'Message 3',
          content: [{ type: 'text', text: 'Summary for Message 3' }],
          summary: 'Summary for Message 3',
        },
        { id: '4', parentMessageId: '3', text: 'Message 4' },
      ]);
    });

    it('should handle multiple summaries and return the branch from the latest to the parentMessageId', () => {
      unorderedMessagesWithSummary = [
        { id: '5', parentMessageId: '4', text: 'Message 5' },
        { id: '2', parentMessageId: '1', text: 'Message 2', summary: 'Summary for Message 2' },
        { id: '3', parentMessageId: '2', text: 'Message 3', summary: 'Summary for Message 3' },
        { id: '4', parentMessageId: '3', text: 'Message 4', summary: 'Summary for Message 4' },
        { id: '1', parentMessageId: null, text: 'Message 1' },
      ];
      const result = TestClient.constructor.getMessagesForConversation({
        messages: unorderedMessagesWithSummary,
        parentMessageId: '5',
        summary: true,
      });
      expect(result).toEqual([
        {
          id: '4',
          parentMessageId: '3',
          role: 'system',
          text: 'Message 4',
          content: [{ type: 'text', text: 'Summary for Message 4' }],
          summary: 'Summary for Message 4',
        },
        { id: '5', parentMessageId: '4', text: 'Message 5' },
      ]);
    });

    it('should handle summary at root edge case and continue until the parentMessageId', () => {
      unorderedMessagesWithSummary = [
        { id: '5', parentMessageId: '4', text: 'Message 5' },
        { id: '1', parentMessageId: null, text: 'Message 1', summary: 'Summary for Message 1' },
        { id: '4', parentMessageId: '3', text: 'Message 4', summary: 'Summary for Message 4' },
        { id: '2', parentMessageId: '1', text: 'Message 2', summary: 'Summary for Message 2' },
        { id: '3', parentMessageId: '2', text: 'Message 3', summary: 'Summary for Message 3' },
      ];
      const result = TestClient.constructor.getMessagesForConversation({
        messages: unorderedMessagesWithSummary,
        parentMessageId: '5',
        summary: true,
      });
      expect(result).toEqual([
        {
          id: '4',
          parentMessageId: '3',
          role: 'system',
          text: 'Message 4',
          content: [{ type: 'text', text: 'Summary for Message 4' }],
          summary: 'Summary for Message 4',
        },
        { id: '5', parentMessageId: '4', text: 'Message 5' },
      ]);
    });

    it('should detect summary content block and use it over legacy fields (summary mode)', () => {
      const messagesWithContentBlock = [
        { id: '3', parentMessageId: '2', text: 'Message 3' },
        {
          id: '2',
          parentMessageId: '1',
          text: 'Message 2',
          content: [
            { type: 'text', text: 'Original text' },
            { type: 'summary', text: 'Content block summary', tokenCount: 42 },
          ],
        },
        { id: '1', parentMessageId: null, text: 'Message 1' },
      ];
      const result = TestClient.constructor.getMessagesForConversation({
        messages: messagesWithContentBlock,
        parentMessageId: '3',
        summary: true,
      });
      expect(result).toHaveLength(2);
      expect(result[0].role).toBe('system');
      expect(result[0].content).toEqual([{ type: 'text', text: 'Content block summary' }]);
      expect(result[0].tokenCount).toBe(42);
    });

    it('should prefer content block summary over legacy summary field', () => {
      const messagesWithBoth = [
        { id: '2', parentMessageId: '1', text: 'Message 2' },
        {
          id: '1',
          parentMessageId: null,
          text: 'Message 1',
          summary: 'Legacy summary',
          summaryTokenCount: 10,
          content: [{ type: 'summary', text: 'Content block summary', tokenCount: 20 }],
        },
      ];
      const result = TestClient.constructor.getMessagesForConversation({
        messages: messagesWithBoth,
        parentMessageId: '2',
        summary: true,
      });
      expect(result).toHaveLength(2);
      expect(result[0].content).toEqual([{ type: 'text', text: 'Content block summary' }]);
      expect(result[0].tokenCount).toBe(20);
    });

    it('should fallback to legacy summary when no content block exists', () => {
      const messagesWithLegacy = [
        { id: '2', parentMessageId: '1', text: 'Message 2' },
        {
          id: '1',
          parentMessageId: null,
          text: 'Message 1',
          summary: 'Legacy summary only',
          summaryTokenCount: 15,
        },
      ];
      const result = TestClient.constructor.getMessagesForConversation({
        messages: messagesWithLegacy,
        parentMessageId: '2',
        summary: true,
      });
      expect(result).toHaveLength(2);
      expect(result[0].content).toEqual([{ type: 'text', text: 'Legacy summary only' }]);
      expect(result[0].tokenCount).toBe(15);
    });
  });

  describe('findSummaryContentBlock', () => {
    it('should find a summary block in the content array', () => {
      const message = {
        content: [
          { type: 'text', text: 'some text' },
          { type: 'summary', text: 'Summary of conversation', tokenCount: 50 },
        ],
      };
      const result = TestClient.constructor.findSummaryContentBlock(message);
      expect(result).toBeTruthy();
      expect(result.text).toBe('Summary of conversation');
      expect(result.tokenCount).toBe(50);
    });

    it('should return null when no summary block exists', () => {
      const message = {
        content: [
          { type: 'text', text: 'some text' },
          { type: 'tool_call', tool_call: {} },
        ],
      };
      expect(TestClient.constructor.findSummaryContentBlock(message)).toBeNull();
    });

    it('should return null for string content', () => {
      const message = { content: 'just a string' };
      expect(TestClient.constructor.findSummaryContentBlock(message)).toBeNull();
    });

    it('should return null for missing content', () => {
      expect(TestClient.constructor.findSummaryContentBlock({})).toBeNull();
      expect(TestClient.constructor.findSummaryContentBlock(null)).toBeNull();
    });

    it('should skip summary blocks with no text', () => {
      const message = {
        content: [{ type: 'summary', tokenCount: 10 }],
      };
      expect(TestClient.constructor.findSummaryContentBlock(message)).toBeNull();
    });
  });

  describe('sendMessage', () => {
    test('sendMessage should return a response message', async () => {
      const expectedResult = expect.objectContaining({
        sender: TestClient.sender,
        text: expect.any(String),
        isCreatedByUser: false,
        messageId: expect.any(String),
        parentMessageId: expect.any(String),
        conversationId: expect.any(String),
      });

      const response = await TestClient.sendMessage(userMessage);
      parentMessageId = response.messageId;
      conversationId = response.conversationId;
      expect(response).toEqual(expectedResult);
    });

    test('should replace responseMessageId with new UUID when isRegenerate is true and messageId ends with underscore', async () => {
      const mockCrypto = require('crypto');
      const newUUID = 'new-uuid-1234';
      jest.spyOn(mockCrypto, 'randomUUID').mockReturnValue(newUUID);

      const opts = {
        isRegenerate: true,
        responseMessageId: 'existing-message-id_',
      };

      await TestClient.setMessageOptions(opts);

      expect(TestClient.responseMessageId).toBe(newUUID);
      expect(TestClient.responseMessageId).not.toBe('existing-message-id_');

      mockCrypto.randomUUID.mockRestore();
    });

    test('should not replace responseMessageId when isRegenerate is false', async () => {
      const opts = {
        isRegenerate: false,
        responseMessageId: 'existing-message-id_',
      };

      await TestClient.setMessageOptions(opts);

      expect(TestClient.responseMessageId).toBe('existing-message-id_');
    });

    test('should not replace responseMessageId when it does not end with underscore', async () => {
      const opts = {
        isRegenerate: true,
        responseMessageId: 'existing-message-id',
      };

      await TestClient.setMessageOptions(opts);

      expect(TestClient.responseMessageId).toBe('existing-message-id');
    });

    test('sendMessage should work with provided conversationId and parentMessageId', async () => {
      const userMessage = 'Second message in the conversation';
      const opts = {
        conversationId,
        parentMessageId,
        getReqData: jest.fn(),
        onStart: jest.fn(),
      };

      const expectedResult = expect.objectContaining({
        sender: TestClient.sender,
        text: expect.any(String),
        isCreatedByUser: false,
        messageId: expect.any(String),
        parentMessageId: expect.any(String),
        conversationId: opts.conversationId,
      });

      const response = await TestClient.sendMessage(userMessage, opts);
      parentMessageId = response.messageId;
      expect(response.conversationId).toEqual(conversationId);
      expect(response).toEqual(expectedResult);
      expect(opts.getReqData).toHaveBeenCalled();
      expect(opts.onStart).toHaveBeenCalled();
      expect(TestClient.getBuildMessagesOptions).toHaveBeenCalled();
      expect(TestClient.getSaveOptions).toHaveBeenCalled();
    });

    test('should return chat history', async () => {
      TestClient = initializeFakeClient(apiKey, options, messageHistory);
      const chatMessages = await TestClient.loadHistory(conversationId, '2');
      expect(TestClient.currentMessages).toHaveLength(2);
      expect(chatMessages[0].text).toEqual('Hello');

      const chatMessages2 = await TestClient.loadHistory(conversationId, '3');
      expect(TestClient.currentMessages).toHaveLength(3);
      expect(chatMessages2[chatMessages2.length - 1].text).toEqual("What's up");
    });

    test('loadHistory should scope database reads to the current user', async () => {
      const user = 'user-123';
      TestClient = new FakeClient(apiKey, options);
      TestClient.user = user;
      getMessages.mockResolvedValueOnce([
        {
          role: 'user',
          isCreatedByUser: true,
          text: 'Hello',
          messageId: '1',
          conversationId,
        },
      ]);

      const chatMessages = await TestClient.loadHistory(conversationId, '1');

      expect(getMessages).toHaveBeenCalledWith({ conversationId, user });
      expect(chatMessages).toHaveLength(1);
      expect(chatMessages[0].text).toBe('Hello');
    });

    /* Most of the new sendMessage logic revolving around edited/continued AI messages
     *  can be summarized by the following test. The condition will load the entire history up to
     *  the message that is being edited, which will trigger the AI API to 'continue' the response.
     *  The 'userMessage' is only passed by convention and is not necessary for the generation.
     */
    it('should not push userMessage to currentMessages when isEdited is true and vice versa', async () => {
      const overrideParentMessageId = 'user-message-id';
      const responseMessageId = 'response-message-id';
      const newHistory = messageHistory.slice();
      newHistory.push({
        role: 'assistant',
        isCreatedByUser: false,
        text: 'test message',
        messageId: responseMessageId,
        parentMessageId: '3',
      });

      TestClient = initializeFakeClient(apiKey, options, newHistory);
      const sendMessageOptions = {
        isEdited: true,
        overrideParentMessageId,
        parentMessageId: '3',
        responseMessageId,
      };

      await TestClient.sendMessage('test message', sendMessageOptions);
      const currentMessages = TestClient.currentMessages;
      expect(currentMessages[currentMessages.length - 1].messageId).not.toEqual(
        overrideParentMessageId,
      );

      // Test the opposite case
      sendMessageOptions.isEdited = false;
      await TestClient.sendMessage('test message', sendMessageOptions);
      const currentMessages2 = TestClient.currentMessages;
      expect(currentMessages2[currentMessages2.length - 1].messageId).toEqual(
        overrideParentMessageId,
      );
    });

    test('setOptions is called with the correct arguments only when replaceOptions is set to true', async () => {
      TestClient.setOptions = jest.fn();
      const opts = { conversationId: '123', parentMessageId: '456', replaceOptions: true };
      await TestClient.sendMessage('Hello, world!', opts);
      expect(TestClient.setOptions).toHaveBeenCalledWith(opts);
      TestClient.setOptions.mockClear();
    });

    test('loadHistory is called with the correct arguments', async () => {
      const opts = { conversationId: '123', parentMessageId: '456' };
      await TestClient.sendMessage('Hello, world!', opts);
      expect(TestClient.loadHistory).toHaveBeenCalledWith(
        opts.conversationId,
        opts.parentMessageId,
      );
    });

    test('getReqData is called with the correct arguments', async () => {
      const getReqData = jest.fn();
      const opts = { getReqData };
      const response = await TestClient.sendMessage('Hello, world!', opts);
      expect(getReqData).toHaveBeenCalledWith(
        expect.objectContaining({
          userMessage: expect.objectContaining({ text: 'Hello, world!' }),
          conversationId: response.conversationId,
          responseMessageId: response.messageId,
        }),
      );
    });

    test('onStart is called with the correct arguments', async () => {
      const onStart = jest.fn();
      const opts = { onStart };
      await TestClient.sendMessage('Hello, world!', opts);

      expect(onStart).toHaveBeenCalledWith(
        expect.objectContaining({ text: 'Hello, world!' }),
        expect.any(String),
        /** `isNewConvo` */
        true,
      );
    });

    test('saveMessageToDatabase is called with the correct arguments', async () => {
      const saveOptions = TestClient.getSaveOptions();
      const user = {};
      const opts = { user };
      const saveSpy = jest.spyOn(TestClient, 'saveMessageToDatabase');
      await TestClient.sendMessage('Hello, world!', opts);
      expect(saveSpy).toHaveBeenCalledWith(
        expect.objectContaining({
          sender: expect.any(String),
          text: expect.any(String),
          isCreatedByUser: expect.any(Boolean),
          messageId: expect.any(String),
          parentMessageId: expect.any(String),
          conversationId: expect.any(String),
        }),
        saveOptions,
        user,
      );
    });

    test('does not start the completed response write when terminal ownership is denied', async () => {
      const hookStarted = deferred();
      const terminalDecision = deferred();
      const beforeResponsePersistence = jest.fn(() => {
        hookStarted.resolve();
        return terminalDecision.promise;
      });
      const saveSpy = jest.spyOn(TestClient, 'saveMessageToDatabase');

      const responsePromise = TestClient.sendMessage('Race Stop against completion.', {
        user: {},
        beforeResponsePersistence,
      });
      await hookStarted.promise;

      expect(beforeResponsePersistence).toHaveBeenCalledTimes(1);
      expect(
        saveSpy.mock.calls.filter(([message]) => message?.isCreatedByUser === false),
      ).toHaveLength(0);

      terminalDecision.resolve(false);
      const response = await responsePromise;

      expect(beforeResponsePersistence).toHaveBeenCalledWith(response);
      expect(
        saveSpy.mock.calls.filter(([message]) => message?.isCreatedByUser === false),
      ).toHaveLength(0);
      expect(TestClient.savedMessageIds.has(response.messageId)).toBe(false);
      await expect(response.databasePromise).resolves.toEqual({ persistenceSkipped: true });
    });

    test('starts the completed response write only after terminal ownership is granted', async () => {
      const hookStarted = deferred();
      const terminalDecision = deferred();
      const beforeResponsePersistence = jest.fn(() => {
        hookStarted.resolve();
        return terminalDecision.promise;
      });
      const saveSpy = jest.spyOn(TestClient, 'saveMessageToDatabase');

      const responsePromise = TestClient.sendMessage('Complete after winning ownership.', {
        user: {},
        beforeResponsePersistence,
      });
      await hookStarted.promise;
      expect(
        saveSpy.mock.calls.filter(([message]) => message?.isCreatedByUser === false),
      ).toHaveLength(0);

      terminalDecision.resolve(true);
      const response = await responsePromise;

      expect(
        saveSpy.mock.calls.filter(([message]) => message?.isCreatedByUser === false),
      ).toHaveLength(1);
      await expect(response.databasePromise).resolves.toEqual(expect.any(Object));
    });

    test('persists the generation-time Langfuse sampling decision for agent responses', async () => {
      const previousSampleRate = process.env.LANGFUSE_SAMPLE_RATE;
      process.env.LANGFUSE_SAMPLE_RATE = '0';
      TestClient.options.endpoint = 'agents';
      const saveSpy = jest.spyOn(TestClient, 'saveMessageToDatabase');

      try {
        const response = await TestClient.sendMessage('Hello, world!', { user: {} });

        expect(response.langfuseSampled).toBe(false);
        expect(response.langfuseDestinationIds).toEqual([]);
        expect(saveSpy).toHaveBeenCalledWith(
          expect.objectContaining({
            langfuseSampled: false,
            langfuseDestinationIds: [],
          }),
          expect.any(Object),
          expect.any(Object),
        );
      } finally {
        if (previousSampleRate == null) {
          delete process.env.LANGFUSE_SAMPLE_RATE;
        } else {
          process.env.LANGFUSE_SAMPLE_RATE = previousSampleRate;
        }
      }
    });

    test('persists no Langfuse destination when a sampled trace has no configured export', async () => {
      const envKeys = [
        'LANGFUSE_PUBLIC_KEY',
        'LANGFUSE_SECRET_KEY',
        'LANGFUSE_FANOUT_ENABLED',
        'LANGFUSE_FANOUT_COLLECTOR_URL',
        'TENANT_ISOLATION_STRICT',
      ];
      const previousEnv = Object.fromEntries(envKeys.map((key) => [key, process.env[key]]));
      const previousSampleRate = process.env.LANGFUSE_SAMPLE_RATE;
      envKeys.forEach((key) => delete process.env[key]);
      process.env.LANGFUSE_SAMPLE_RATE = '1';
      TestClient.options.endpoint = 'agents';
      const saveSpy = jest.spyOn(TestClient, 'saveMessageToDatabase');

      try {
        const response = await TestClient.sendMessage('Hello, world!', { user: {} });

        expect(response.langfuseSampled).toBe(true);
        expect(response.langfuseDestinationIds).toEqual([]);
        expect(saveSpy).toHaveBeenCalledWith(
          expect.objectContaining({
            langfuseSampled: true,
            langfuseDestinationIds: [],
          }),
          expect.any(Object),
          expect.any(Object),
        );
      } finally {
        for (const [key, value] of Object.entries(previousEnv)) {
          if (value == null) {
            delete process.env[key];
          } else {
            process.env[key] = value;
          }
        }
        if (previousSampleRate == null) {
          delete process.env.LANGFUSE_SAMPLE_RATE;
        } else {
          process.env.LANGFUSE_SAMPLE_RATE = previousSampleRate;
        }
      }
    });

    test('should handle existing conversation when getConvo retrieves one', async () => {
      const existingConvo = {
        conversationId: 'existing-convo-id',
        endpoint: 'openai',
        endpointType: 'openai',
        model: 'gpt-3.5-turbo',
        messages: [
          { role: 'user', content: 'Existing message 1' },
          { role: 'assistant', content: 'Existing response 1' },
        ],
        temperature: 1,
      };

      const { temperature: _temp, ...newConvo } = existingConvo;

      const user = {
        id: 'user-id',
      };

      getConvo.mockResolvedValue(existingConvo);
      saveConvo.mockResolvedValue(newConvo);

      TestClient = initializeFakeClient(
        apiKey,
        {
          ...options,
          req: {
            user,
          },
        },
        [],
      );

      const saveSpy = jest.spyOn(TestClient, 'saveMessageToDatabase');

      const newMessage = 'New message in existing conversation';
      const response = await TestClient.sendMessage(newMessage, {
        user,
        conversationId: existingConvo.conversationId,
      });

      expect(getConvo).toHaveBeenCalledWith(user.id, existingConvo.conversationId);
      expect(TestClient.conversationId).toBe(existingConvo.conversationId);
      expect(response.conversationId).toBe(existingConvo.conversationId);
      expect(TestClient.fetchedConvo).toBe(true);

      expect(saveSpy).toHaveBeenCalledWith(
        expect.objectContaining({
          conversationId: existingConvo.conversationId,
          text: newMessage,
        }),
        expect.any(Object),
        expect.any(Object),
      );

      expect(saveConvo).toHaveBeenCalledTimes(2);
      expect(saveConvo).toHaveBeenCalledWith(
        expect.any(Object),
        expect.objectContaining({
          conversationId: existingConvo.conversationId,
        }),
        expect.objectContaining({
          context: 'api/app/clients/BaseClient.js - saveMessageToDatabase #saveConvo',
          unsetFields: {
            temperature: 1,
          },
        }),
      );

      await TestClient.sendMessage('Another message', {
        conversationId: existingConvo.conversationId,
      });
      expect(getConvo).toHaveBeenCalledTimes(1);
    });

    test('should correctly handle existing conversation and unset fields appropriately', async () => {
      const existingConvo = {
        conversationId: 'existing-convo-id',
        endpoint: 'openai',
        endpointType: 'openai',
        model: 'gpt-3.5-turbo',
        messages: [
          { role: 'user', content: 'Existing message 1' },
          { role: 'assistant', content: 'Existing response 1' },
        ],
        title: 'Existing Conversation',
        someExistingField: 'existingValue',
        anotherExistingField: 'anotherValue',
        temperature: 0.7,
        modelLabel: 'GPT-3.5',
      };

      getConvo.mockResolvedValue(existingConvo);
      saveConvo.mockResolvedValue(existingConvo);

      TestClient = initializeFakeClient(
        apiKey,
        {
          ...options,
          modelOptions: {
            model: 'gpt-4',
            temperature: 0.5,
          },
        },
        [],
      );

      const newMessage = 'New message in existing conversation';
      await TestClient.sendMessage(newMessage, {
        conversationId: existingConvo.conversationId,
      });

      expect(saveConvo).toHaveBeenCalledTimes(2);

      const saveConvoCall = saveConvo.mock.calls[0];
      const [, savedFields, saveOptions] = saveConvoCall;

      // Instead of checking all excludedKeys, we'll just check specific fields
      // that we know should be excluded
      expect(savedFields).not.toHaveProperty('messages');
      expect(savedFields).not.toHaveProperty('title');

      // Only check that someExistingField is in unsetFields
      expect(saveOptions.unsetFields).toHaveProperty('someExistingField', 1);

      // Mock saveConvo to return the expected fields
      saveConvo.mockImplementation((req, fields) => {
        return Promise.resolve({
          ...fields,
          endpoint: 'openai',
          endpointType: 'openai',
          model: 'gpt-4',
          temperature: 0.5,
        });
      });

      // Only check the conversationId since that's the only field we can be sure about
      expect(savedFields).toHaveProperty('conversationId', 'existing-convo-id');

      expect(TestClient.fetchedConvo).toBe(true);

      await TestClient.sendMessage('Another message', {
        conversationId: existingConvo.conversationId,
      });

      expect(getConvo).toHaveBeenCalledTimes(1);

      const secondSaveConvoCall = saveConvo.mock.calls[1];
      expect(secondSaveConvoCall[2]).toHaveProperty('unsetFields', {});
    });

    test('sendCompletion is called with the correct arguments', async () => {
      const payload = {}; // Mock payload
      TestClient.buildMessages.mockReturnValue({ prompt: payload, tokenCountMap: null });
      const opts = {};
      await TestClient.sendMessage('Hello, world!', opts);
      expect(TestClient.sendCompletion).toHaveBeenCalledWith(payload, opts);
    });

    test('records history and message-build startup milestones', async () => {
      const startupTelemetry = { mark: jest.fn() };
      TestClient.options.startupTelemetry = startupTelemetry;

      await TestClient.sendMessage('Hello, world!', {});

      expect(startupTelemetry.mark.mock.calls.map(([milestone]) => milestone)).toEqual([
        'history_loaded',
        'messages_built',
      ]);
    });

    test('getTokenCount for response is called with the correct arguments', async () => {
      const tokenCountMap = {}; // Mock tokenCountMap
      TestClient.buildMessages.mockReturnValue({ prompt: [], tokenCountMap });
      TestClient.getTokenCountForResponse = jest.fn();
      const response = await TestClient.sendMessage('Hello, world!', {});
      expect(TestClient.getTokenCountForResponse).toHaveBeenCalledWith(response);
    });

    test('returns an object with the correct shape', async () => {
      const response = await TestClient.sendMessage('Hello, world!', {});
      expect(response).toEqual(
        expect.objectContaining({
          sender: expect.any(String),
          text: expect.any(String),
          isCreatedByUser: expect.any(Boolean),
          messageId: expect.any(String),
          parentMessageId: expect.any(String),
          conversationId: expect.any(String),
        }),
      );
    });

    test('saveMessageToDatabase returns early when this.options is null (client disposed)', async () => {
      const savedOptions = TestClient.options;
      TestClient.options = null;
      saveMessage.mockClear();

      const result = await TestClient.saveMessageToDatabase(
        { messageId: 'msg-1', conversationId: 'conv-1', isCreatedByUser: true, text: 'hi' },
        {},
        null,
      );

      expect(result).toEqual({});
      expect(saveMessage).not.toHaveBeenCalled();

      TestClient.options = savedOptions;
    });

    test('saveMessageToDatabase uses snapshot of options, immune to mid-await disposal', async () => {
      const savedOptions = TestClient.options;
      saveMessage.mockClear();
      saveConvo.mockClear();

      // Make db.saveMessage yield, simulating I/O suspension during which disposal occurs
      saveMessage.mockImplementation(async (_reqCtx, msgData) => {
        // Simulate disposeClient nullifying client.options while awaiting
        TestClient.options = null;
        return msgData;
      });
      saveConvo.mockResolvedValue({ conversationId: 'conv-1' });

      const result = await TestClient.saveMessageToDatabase(
        { messageId: 'msg-1', conversationId: 'conv-1', isCreatedByUser: true, text: 'hi' },
        { endpoint: 'openAI' },
        null,
      );

      // Should complete without TypeError, using the snapshotted options
      expect(result).toHaveProperty('message');
      expect(result).toHaveProperty('conversation');
      expect(saveMessage).toHaveBeenCalled();

      TestClient.options = savedOptions;
      saveMessage.mockReset();
      saveConvo.mockReset();
    });

    test('saveMessageToDatabase reuses conversation resolved on the request', async () => {
      const existingConvo = {
        conversationId: 'cached-convo-id',
        endpoint: 'openai',
        endpointType: 'openai',
        temperature: 0.7,
      };
      const user = { id: 'user-id' };
      const req = { user, resolvedConversation: existingConvo };

      getConvo.mockClear();
      saveMessage.mockResolvedValue({ messageId: 'msg-1' });
      saveConvo.mockResolvedValue(existingConvo);

      TestClient = initializeFakeClient(apiKey, { ...options, endpoint: 'openai', req }, []);

      await TestClient.saveMessageToDatabase(
        {
          messageId: 'msg-1',
          conversationId: existingConvo.conversationId,
          isCreatedByUser: true,
          text: 'hi',
        },
        { endpoint: 'openai' },
        user,
      );

      expect(getConvo).not.toHaveBeenCalled();
      expect(req).not.toHaveProperty('resolvedConversation');
      expect(TestClient.fetchedConvo).toBe(true);
      expect(saveConvo).toHaveBeenCalledWith(
        expect.any(Object),
        expect.objectContaining({ conversationId: existingConvo.conversationId }),
        expect.objectContaining({
          unsetFields: expect.objectContaining({ temperature: 1 }),
        }),
      );
    });

    test('userMessagePromise is awaited before saving response message', async () => {
      // Mock the saveMessageToDatabase method
      TestClient.saveMessageToDatabase = jest.fn().mockImplementation(() => {
        return new Promise((resolve) => setTimeout(resolve, 100)); // Simulate a delay
      });

      // Send a message
      const messagePromise = TestClient.sendMessage('Hello, world!');

      // Wait a short time to ensure the user message save has started
      await new Promise((resolve) => setTimeout(resolve, 50));

      // Check that saveMessageToDatabase has been called once (for the user message)
      expect(TestClient.saveMessageToDatabase).toHaveBeenCalledTimes(1);

      // Wait for the message to be fully processed
      await messagePromise;

      // Check that saveMessageToDatabase has been called twice (once for user message, once for response)
      expect(TestClient.saveMessageToDatabase).toHaveBeenCalledTimes(2);

      // Check the order of calls
      const calls = TestClient.saveMessageToDatabase.mock.calls;
      expect(calls[0][0].isCreatedByUser).toBe(true); // First call should be for user message
      expect(calls[1][0].isCreatedByUser).toBe(false); // Second call should be for response message
    });
  });

  describe('recordTokenUsage model assignment', () => {
    test('should pass this.model to recordTokenUsage, not the agent ID from responseMessage.model', async () => {
      const actualModel = 'claude-opus-4-5';
      const agentId = 'agent_p5Z_IU6EIxBoqn1BoqLBp';

      TestClient.model = actualModel;
      TestClient.options.endpoint = 'agents';
      TestClient.options.agent = { id: agentId };

      TestClient.getTokenCountForResponse = jest.fn().mockReturnValue(50);
      TestClient.recordTokenUsage = jest.fn().mockResolvedValue(undefined);
      TestClient.buildMessages.mockReturnValue({
        prompt: [],
        tokenCountMap: { res: 50 },
      });

      await TestClient.sendMessage('Hello', {});

      expect(TestClient.recordTokenUsage).toHaveBeenCalledWith(
        expect.objectContaining({
          model: actualModel,
        }),
      );

      const callArgs = TestClient.recordTokenUsage.mock.calls[0][0];
      expect(callArgs.model).not.toBe(agentId);
    });

    test('should pass this.model even when this.model differs from modelOptions.model', async () => {
      const instanceModel = 'gpt-4o';
      TestClient.model = instanceModel;
      TestClient.modelOptions = { model: 'gpt-4o-mini' };

      TestClient.getTokenCountForResponse = jest.fn().mockReturnValue(50);
      TestClient.recordTokenUsage = jest.fn().mockResolvedValue(undefined);
      TestClient.buildMessages.mockReturnValue({
        prompt: [],
        tokenCountMap: { res: 50 },
      });

      await TestClient.sendMessage('Hello', {});

      expect(TestClient.recordTokenUsage).toHaveBeenCalledWith(
        expect.objectContaining({
          model: instanceModel,
        }),
      );
    });
  });

  describe('getMessagesWithinTokenLimit with instructions', () => {
    test('should always include instructions when present', async () => {
      TestClient.maxContextTokens = 50;
      const instructions = {
        role: 'system',
        content: 'System instructions',
        tokenCount: 20,
      };

      const messages = [
        instructions,
        { role: 'user', content: 'Hello', tokenCount: 10 },
        { role: 'assistant', content: 'Hi there', tokenCount: 15 },
      ];

      const result = await TestClient.getMessagesWithinTokenLimit({
        messages,
        instructions,
      });

      expect(result.context[0]).toBe(instructions);
      expect(result.remainingContextTokens).toBe(2);
    });

    test('should handle case when messages exceed limit but instructions must be preserved', async () => {
      TestClient.maxContextTokens = 30;
      const instructions = {
        role: 'system',
        content: 'System instructions',
        tokenCount: 20,
      };

      const messages = [
        instructions,
        { role: 'user', content: 'Hello', tokenCount: 10 },
        { role: 'assistant', content: 'Hi there', tokenCount: 15 },
      ];

      const result = await TestClient.getMessagesWithinTokenLimit({
        messages,
        instructions,
      });

      // Should only include instructions and the last message that fits
      expect(result.context).toHaveLength(1);
      expect(result.context[0].content).toBe(instructions.content);
      expect(result.messagesToRefine).toHaveLength(2);
      expect(result.remainingContextTokens).toBe(7); // 30 - 20 - 3 (assistant label)
    });

    test('should work correctly without instructions (1/2)', async () => {
      TestClient.maxContextTokens = 50;
      const messages = [
        { role: 'user', content: 'Hello', tokenCount: 10 },
        { role: 'assistant', content: 'Hi there', tokenCount: 15 },
      ];

      const result = await TestClient.getMessagesWithinTokenLimit({
        messages,
      });

      expect(result.context).toHaveLength(2);
      expect(result.remainingContextTokens).toBe(22); // 50 - 10 - 15 - 3(assistant label)
      expect(result.messagesToRefine).toHaveLength(0);
    });

    test('should work correctly without instructions (2/2)', async () => {
      TestClient.maxContextTokens = 30;
      const messages = [
        { role: 'user', content: 'Hello', tokenCount: 10 },
        { role: 'assistant', content: 'Hi there', tokenCount: 20 },
      ];

      const result = await TestClient.getMessagesWithinTokenLimit({
        messages,
      });

      expect(result.context).toHaveLength(1);
      expect(result.remainingContextTokens).toBe(7);
      expect(result.messagesToRefine).toHaveLength(1);
    });

    test('should handle case when only instructions fit within limit', async () => {
      TestClient.maxContextTokens = 25;
      const instructions = {
        role: 'system',
        content: 'System instructions',
        tokenCount: 20,
      };

      const messages = [
        instructions,
        { role: 'user', content: 'Hello', tokenCount: 10 },
        { role: 'assistant', content: 'Hi there', tokenCount: 15 },
      ];

      const result = await TestClient.getMessagesWithinTokenLimit({
        messages,
        instructions,
      });

      expect(result.context).toHaveLength(1);
      expect(result.context[0]).toBe(instructions);
      expect(result.messagesToRefine).toHaveLength(2);
      expect(result.remainingContextTokens).toBe(2); // 25 - 20 - 3(assistant label)
    });
  });

  describe('sendMessage file population', () => {
    const attachment = {
      file_id: 'file-abc',
      filename: 'image.png',
      filepath: '/uploads/image.png',
      type: 'image/png',
      bytes: 1024,
      object: 'file',
      user: 'user-1',
      embedded: false,
      usage: 0,
      text: 'large ocr blob that should be stripped',
      _id: 'mongo-id-1',
    };

    beforeEach(() => {
      TestClient.options.req = { body: { files: [{ file_id: 'file-abc' }] } };
      TestClient.options.attachments = [attachment];
    });

    test('populates userMessage.files before saveMessageToDatabase is called', async () => {
      TestClient.saveMessageToDatabase = jest.fn().mockImplementation((msg) => {
        return Promise.resolve({ message: msg });
      });

      await TestClient.sendMessage('Hello');

      const userSave = TestClient.saveMessageToDatabase.mock.calls.find(
        ([msg]) => msg.isCreatedByUser,
      );
      expect(userSave).toBeDefined();
      expect(userSave[0].files).toBeDefined();
      expect(userSave[0].files).toHaveLength(1);
      expect(userSave[0].files[0].file_id).toBe('file-abc');
    });

    test('strips text and _id from files before saving', async () => {
      TestClient.saveMessageToDatabase = jest.fn().mockResolvedValue({ message: {} });

      await TestClient.sendMessage('Hello');

      const userSave = TestClient.saveMessageToDatabase.mock.calls.find(
        ([msg]) => msg.isCreatedByUser,
      );
      expect(userSave[0].files[0].text).toBeUndefined();
      expect(userSave[0].files[0]._id).toBeUndefined();
      expect(userSave[0].files[0].filename).toBe('image.png');
    });

    test('deletes image_urls from userMessage when files are present', async () => {
      TestClient.saveMessageToDatabase = jest.fn().mockResolvedValue({ message: {} });
      TestClient.options.attachments = [
        { ...attachment, image_urls: ['data:image/png;base64,...'] },
      ];

      await TestClient.sendMessage('Hello');

      const userSave = TestClient.saveMessageToDatabase.mock.calls.find(
        ([msg]) => msg.isCreatedByUser,
      );
      expect(userSave[0].image_urls).toBeUndefined();
    });

    test('does not set files when no attachments match request file IDs', async () => {
      TestClient.options.req = { body: { files: [{ file_id: 'file-nomatch' }] } };
      TestClient.saveMessageToDatabase = jest.fn().mockResolvedValue({ message: {} });

      await TestClient.sendMessage('Hello');

      const userSave = TestClient.saveMessageToDatabase.mock.calls.find(
        ([msg]) => msg.isCreatedByUser,
      );
      expect(userSave[0].files).toBeUndefined();
    });

    test('skips file population when attachments is not an array (Promise case)', async () => {
      TestClient.options.attachments = Promise.resolve([attachment]);
      TestClient.saveMessageToDatabase = jest.fn().mockResolvedValue({ message: {} });

      await TestClient.sendMessage('Hello');

      const userSave = TestClient.saveMessageToDatabase.mock.calls.find(
        ([msg]) => msg.isCreatedByUser,
      );
      expect(userSave[0].files).toBeUndefined();
    });

    test('skips file population when skipSaveUserMessage is true', async () => {
      TestClient.skipSaveUserMessage = true;
      TestClient.saveMessageToDatabase = jest.fn().mockResolvedValue({ message: {} });

      await TestClient.sendMessage('Hello');

      const userSave = TestClient.saveMessageToDatabase.mock.calls.find(
        ([msg]) => msg?.isCreatedByUser,
      );
      expect(userSave).toBeUndefined();
    });

    test('ignores file_id: undefined entries in req.body.files (no set poisoning)', async () => {
      TestClient.options.req = {
        body: { files: [{ file_id: undefined }, { file_id: 'file-abc' }] },
      };
      TestClient.options.attachments = [
        { ...attachment, file_id: undefined },
        { ...attachment, file_id: 'file-abc' },
      ];
      TestClient.saveMessageToDatabase = jest.fn().mockResolvedValue({ message: {} });

      await TestClient.sendMessage('Hello');

      const userSave = TestClient.saveMessageToDatabase.mock.calls.find(
        ([msg]) => msg.isCreatedByUser,
      );
      expect(userSave[0].files).toHaveLength(1);
      expect(userSave[0].files[0].file_id).toBe('file-abc');
    });
  });

  describe('addPreviousAttachments authorization', () => {
    const ownerFile = {
      file_id: 'owner-file',
      filename: 'owner.txt',
      filepath: '/uploads/owner.txt',
      source: 'local',
      type: 'text/plain',
      bytes: 100,
      object: 'file',
      user: 'user-1',
      embedded: false,
      usage: 0,
      text: 'authorized owner text',
      _id: 'owner-mongo-id',
      metadata: {
        codeEnvRef: {
          kind: 'user',
          id: 'user-1',
          storage_session_id: 'owner-session',
          file_id: 'owner-code-file',
        },
      },
    };

    beforeEach(() => {
      getFiles.mockReset();
      TestClient.options.resendFiles = true;
      TestClient.options.attachments = undefined;
      TestClient.options.req = {
        user: {
          id: 'user-1',
          tenantId: 'tenant-a',
        },
      };
      TestClient.addFileContextToMessage = jest.fn(async (message, files) => {
        const text = files
          .map((file) => file.text)
          .filter(Boolean)
          .join('\n');
        if (text) {
          message.fileContext = text;
        }
      });
      TestClient.processAttachments = jest.fn(async (_message, files) => files);
      TestClient.checkVisionRequest = jest.fn();
    });

    test('rehydrates historical file refs from owner-scoped DB rows only', async () => {
      getFiles.mockResolvedValueOnce([ownerFile]);

      const [message] = await TestClient.addPreviousAttachments([
        {
          messageId: 'msg-1',
          text: 'Use the attachment',
          files: [
            {
              file_id: 'owner-file',
              filename: 'attacker-controlled-owner-name.txt',
              filepath: '/forged/owner.txt',
              text: 'forged owner text',
            },
            {
              file_id: 'victim-file',
              filename: 'victim.txt',
              filepath: '/victim/private.txt',
              text: 'victim private text',
            },
          ],
          attachments: [
            {
              file_id: 'victim-file',
              filename: 'victim-output.csv',
              text: 'victim output text',
            },
          ],
          fileContext: 'stale victim private text',
        },
      ]);

      expect(getFiles).toHaveBeenCalledWith(
        {
          file_id: { $in: ['owner-file', 'victim-file'] },
          user: 'user-1',
          tenantId: 'tenant-a',
        },
        {},
        {},
      );
      expect(TestClient.addFileContextToMessage).toHaveBeenCalledWith(message, [ownerFile]);
      expect(TestClient.processAttachments).toHaveBeenCalledWith(message, [ownerFile]);
      expect(message.fileContext).toBe('authorized owner text');
      expect(message.files).toEqual([
        expect.objectContaining({
          file_id: 'owner-file',
          filename: 'owner.txt',
          filepath: '/uploads/owner.txt',
          source: 'local',
          metadata: ownerFile.metadata,
        }),
      ]);
      expect(message.files[0].text).toBeUndefined();
      expect(message.files[0]._id).toBeUndefined();
      expect(message.attachments).toBeUndefined();
      expect(JSON.stringify(message)).not.toContain('victim');
      expect(JSON.stringify(message)).not.toContain('forged owner text');
    });

    test('strips historical file context when no authenticated owner scope is available', async () => {
      TestClient.options.req = {};

      const [message] = await TestClient.addPreviousAttachments([
        {
          messageId: 'msg-2',
          files: [{ file_id: 'victim-file', filename: 'victim.txt' }],
          fileContext: 'stale victim private text',
        },
      ]);

      expect(getFiles).not.toHaveBeenCalled();
      expect(message.files).toBeUndefined();
      expect(message.fileContext).toBeUndefined();
    });

    test('preserves repeated owner-authorized historical file refs after the first context use', async () => {
      getFiles.mockResolvedValueOnce([ownerFile]);

      const [firstMessage, secondMessage] = await TestClient.addPreviousAttachments([
        {
          messageId: 'msg-repeat-1',
          files: [{ file_id: 'owner-file', filename: 'first-forged.txt' }],
        },
        {
          messageId: 'msg-repeat-2',
          files: [{ file_id: 'owner-file', filename: 'second-forged.txt' }],
        },
      ]);

      expect(getFiles).toHaveBeenCalledTimes(1);
      expect(getFiles).toHaveBeenCalledWith(
        {
          file_id: { $in: ['owner-file'] },
          user: 'user-1',
          tenantId: 'tenant-a',
        },
        {},
        {},
      );
      expect(TestClient.addFileContextToMessage).toHaveBeenCalledTimes(1);
      expect(TestClient.addFileContextToMessage).toHaveBeenCalledWith(firstMessage, [ownerFile]);
      expect(secondMessage.fileContext).toBeUndefined();
      expect(firstMessage.files).toEqual([
        expect.objectContaining({ file_id: 'owner-file', filename: 'owner.txt' }),
      ]);
      expect(secondMessage.files).toEqual([
        expect.objectContaining({ file_id: 'owner-file', filename: 'owner.txt' }),
      ]);
      expect(JSON.stringify(secondMessage)).not.toContain('second-forged');
    });

    test('extracts historical file context while encoding provider attachments', async () => {
      getFiles.mockResolvedValueOnce([ownerFile]);
      const fileContext = deferred();
      const providerAttachments = deferred();
      let completed = false;

      TestClient.addFileContextToMessage.mockImplementation(async (message) => {
        await fileContext.promise;
        message.fileContext = 'authorized owner text';
      });
      TestClient.processAttachments.mockImplementation(() => providerAttachments.promise);

      const messagesPromise = TestClient.addPreviousAttachments([
        {
          messageId: 'msg-concurrent-file-work',
          files: [{ file_id: 'owner-file', filename: 'owner.txt' }],
        },
      ]).then((messages) => {
        completed = true;
        return messages;
      });

      await Promise.resolve();
      await Promise.resolve();

      expect(TestClient.addFileContextToMessage).toHaveBeenCalledTimes(1);
      expect(TestClient.processAttachments).toHaveBeenCalledTimes(1);

      providerAttachments.resolve([ownerFile]);
      await Promise.resolve();
      expect(completed).toBe(false);

      fileContext.resolve();
      const [message] = await messagesPromise;

      expect(message.fileContext).toBe('authorized owner text');
      expect(TestClient.message_file_map['msg-concurrent-file-work']).toEqual([ownerFile]);
    });

    test('preserves download-only historical attachments without trusting file fields', async () => {
      const [message] = await TestClient.addPreviousAttachments([
        {
          messageId: 'msg-download-only',
          attachments: [
            {
              filename: 'report.csv',
              filepath: '/api/files/code/download/session/file',
              expiresAt: 123456,
              conversationId: 'conversation-1',
              messageId: 'assistant-message',
              toolCallId: 'tool-call-1',
              text: 'untrusted text should not survive',
              source: 'forged-source',
              metadata: { codeEnvRef: { id: 'victim' } },
            },
          ],
          fileContext: 'stale context',
        },
      ]);

      expect(getFiles).not.toHaveBeenCalled();
      expect(message.fileContext).toBeUndefined();
      expect(message.attachments).toEqual([
        {
          filename: 'report.csv',
          filepath: '/api/files/code/download/session/file',
          expiresAt: 123456,
          conversationId: 'conversation-1',
          messageId: 'assistant-message',
          toolCallId: 'tool-call-1',
        },
      ]);
      expect(JSON.stringify(message)).not.toContain('untrusted text');
      expect(JSON.stringify(message)).not.toContain('forged-source');
      expect(JSON.stringify(message)).not.toContain('victim');
    });

    test('merges safe per-message metadata onto authorized DB-backed attachments', async () => {
      getFiles.mockResolvedValueOnce([ownerFile]);

      const [message] = await TestClient.addPreviousAttachments([
        {
          messageId: 'msg-artifact',
          attachments: [
            {
              file_id: 'owner-file',
              filename: 'forged-artifact.csv',
              filepath: '/forged/artifact.csv',
              source: 'forged-source',
              metadata: { codeEnvRef: { id: 'victim' } },
              text: 'forged artifact text',
              messageId: 'assistant-message',
              toolCallId: 'tool-call-2',
            },
          ],
        },
      ]);

      expect(message.attachments).toEqual([
        expect.objectContaining({
          file_id: 'owner-file',
          filename: 'owner.txt',
          filepath: '/uploads/owner.txt',
          source: 'local',
          metadata: ownerFile.metadata,
          messageId: 'assistant-message',
          toolCallId: 'tool-call-2',
        }),
      ]);
      expect(message.attachments[0].text).toBeUndefined();
      expect(message.attachments[0]._id).toBeUndefined();
      expect(JSON.stringify(message)).not.toContain('forged-artifact');
      expect(JSON.stringify(message)).not.toContain('forged artifact text');
    });
  });

  describe('sendMessage quote references', () => {
    // The blockquote merge itself lives in AgentClient.buildMessages / prependQuotes
    // (covered by packages/api specs). BaseClient's job is to attach the normalized
    // quotes onto the user message early and keep the stored text clean.
    test('attaches normalized quotes before getReqData fires and keeps stored text clean', async () => {
      TestClient.options.req = { body: { quotes: ['  the selected text  ', '', 42] } };
      TestClient.saveMessageToDatabase = jest.fn().mockResolvedValue({ message: {} });
      let captured;
      await TestClient.sendMessage('What does this mean?', {
        getReqData: (data) => {
          if (data.userMessage) {
            captured = { text: data.userMessage.text, quotes: data.userMessage.quotes };
          }
        },
      });

      // Quotes are present (trimmed, non-strings dropped) at getReqData time, and
      // the user text is never mutated by the merge.
      expect(captured).toBeDefined();
      expect(captured.quotes).toEqual(['the selected text']);
      expect(captured.text).toBe('What does this mean?');

      const userSave = TestClient.saveMessageToDatabase.mock.calls.find(
        ([msg]) => msg.isCreatedByUser,
      );
      expect(userSave[0].text).toBe('What does this mean?');
      expect(userSave[0].quotes).toEqual(['the selected text']);
    });

    test('persists multiple quotes in order on the saved message', async () => {
      TestClient.options.req = { body: { quotes: ['first excerpt', 'second excerpt'] } };
      TestClient.saveMessageToDatabase = jest.fn().mockResolvedValue({ message: {} });

      await TestClient.sendMessage('Compare these');

      const userSave = TestClient.saveMessageToDatabase.mock.calls.find(
        ([msg]) => msg.isCreatedByUser,
      );
      expect(userSave[0].text).toBe('Compare these');
      expect(userSave[0].quotes).toEqual(['first excerpt', 'second excerpt']);
    });

    test('leaves the message untouched when no quotes are provided', async () => {
      TestClient.options.req = { body: {} };
      TestClient.saveMessageToDatabase = jest.fn().mockResolvedValue({ message: {} });

      await TestClient.sendMessage('Just a question');

      const userSave = TestClient.saveMessageToDatabase.mock.calls.find(
        ([msg]) => msg.isCreatedByUser,
      );
      expect(userSave[0].text).toBe('Just a question');
      expect(userSave[0].quotes).toBeUndefined();
    });
  });
});