LibreChat/api/app/clients/tools/structured/specs/imageTools-agent.spec.js

/**
 * Regression tests for image tool agent mode — verifies that invoke() returns
 * a ToolMessage with base64 in artifact.content rather than serialized into content.
 *
 * Root cause: DALLE3/FluxAPI/StableDiffusion extend LangChain's Tool but did not
 * set responseFormat = 'content_and_artifact'. LangChain's invoke() would then
 * JSON.stringify the entire [content, artifact] tuple into ToolMessage.content,
 * dumping base64 into token counting and causing context exhaustion.
 */

const axios = require('axios');
const OpenAI = require('openai');
const undici = require('undici');
const fetch = require('node-fetch');
const { ToolMessage } = require('@langchain/core/messages');
const { ContentTypes } = require('librechat-data-provider');
const StableDiffusionAPI = require('../StableDiffusion');
const FluxAPI = require('../FluxAPI');
const DALLE3 = require('../DALLE3');

jest.mock('axios');
jest.mock('openai');
jest.mock('node-fetch');
jest.mock('undici', () => ({
  ProxyAgent: jest.fn(),
  fetch: jest.fn(),
}));
jest.mock('@librechat/data-schemas', () => ({
  logger: { info: jest.fn(), warn: jest.fn(), debug: jest.fn(), error: jest.fn() },
}));
jest.mock('path', () => ({
  resolve: jest.fn(),
  join: jest.fn().mockReturnValue('/mock/path'),
  relative: jest.fn().mockReturnValue('relative/path'),
  extname: jest.fn().mockReturnValue('.png'),
}));
jest.mock('fs', () => ({
  existsSync: jest.fn().mockReturnValue(true),
  mkdirSync: jest.fn(),
  promises: { writeFile: jest.fn(), readFile: jest.fn(), unlink: jest.fn() },
}));

const FAKE_BASE64 = 'aGVsbG8=';

const makeToolCall = (name, args) => ({
  id: 'call_test_123',
  name,
  args,
  type: 'tool_call',
});

describe('image tools - agent mode ToolMessage format', () => {
  const ENV_KEYS = ['DALLE_API_KEY', 'FLUX_API_KEY', 'SD_WEBUI_URL', 'PROXY'];
  let savedEnv = {};

  beforeEach(() => {
    jest.clearAllMocks();
    for (const key of ENV_KEYS) {
      savedEnv[key] = process.env[key];
    }
    process.env.DALLE_API_KEY = 'test-dalle-key';
    process.env.FLUX_API_KEY = 'test-flux-key';
    process.env.SD_WEBUI_URL = 'http://localhost:7860';
    delete process.env.PROXY;
  });

  afterEach(() => {
    for (const key of ENV_KEYS) {
      if (savedEnv[key] === undefined) {
        delete process.env[key];
      } else {
        process.env[key] = savedEnv[key];
      }
    }
    savedEnv = {};
  });

  describe('DALLE3', () => {
    beforeEach(() => {
      OpenAI.mockImplementation(() => ({
        images: {
          generate: jest.fn().mockResolvedValue({
            data: [{ url: 'https://example.com/image.png' }],
          }),
        },
      }));
      undici.fetch.mockResolvedValue({
        arrayBuffer: () => Promise.resolve(Buffer.from(FAKE_BASE64, 'base64')),
      });
    });

    it('sets responseFormat to content_and_artifact when isAgent is true', () => {
      const dalle = new DALLE3({ isAgent: true });
      expect(dalle.responseFormat).toBe('content_and_artifact');
    });

    it('does not set responseFormat when isAgent is false', () => {
      const dalle = new DALLE3({ isAgent: false, processFileURL: jest.fn() });
      expect(dalle.responseFormat).not.toBe('content_and_artifact');
    });

    it('invoke() returns ToolMessage with base64 in artifact, not serialized in content', async () => {
      const dalle = new DALLE3({ isAgent: true });
      const result = await dalle.invoke(
        makeToolCall('dalle', {
          prompt: 'a box',
          quality: 'standard',
          size: '1024x1024',
          style: 'vivid',
        }),
      );

      expect(result).toBeInstanceOf(ToolMessage);

      const contentStr =
        typeof result.content === 'string' ? result.content : JSON.stringify(result.content);
      expect(contentStr).not.toContain(FAKE_BASE64);

      expect(result.artifact).toBeDefined();
      const artifactContent = result.artifact?.content;
      expect(Array.isArray(artifactContent)).toBe(true);
      expect(artifactContent[0].type).toBe(ContentTypes.IMAGE_URL);
      expect(artifactContent[0].image_url.url).toContain('base64');
    });

    it('invoke() returns ToolMessage with error string in content when API fails', async () => {
      OpenAI.mockImplementation(() => ({
        images: { generate: jest.fn().mockRejectedValue(new Error('API error')) },
      }));

      const dalle = new DALLE3({ isAgent: true });
      const result = await dalle.invoke(
        makeToolCall('dalle', {
          prompt: 'a box',
          quality: 'standard',
          size: '1024x1024',
          style: 'vivid',
        }),
      );

      expect(result).toBeInstanceOf(ToolMessage);
      const contentStr =
        typeof result.content === 'string' ? result.content : JSON.stringify(result.content);
      expect(contentStr).toContain('Something went wrong');
      expect(result.artifact).toBeDefined();
    });
  });

  describe('FluxAPI', () => {
    beforeEach(() => {
      jest.useFakeTimers();
      axios.post.mockResolvedValue({ data: { id: 'task-123' } });
      axios.get.mockResolvedValue({
        data: { status: 'Ready', result: { sample: 'https://example.com/image.png' } },
      });
      fetch.mockResolvedValue({
        arrayBuffer: () => Promise.resolve(Buffer.from(FAKE_BASE64, 'base64')),
      });
    });

    afterEach(() => {
      jest.useRealTimers();
    });

    it('sets responseFormat to content_and_artifact when isAgent is true', () => {
      const flux = new FluxAPI({ isAgent: true });
      expect(flux.responseFormat).toBe('content_and_artifact');
    });

    it('does not set responseFormat when isAgent is false', () => {
      const flux = new FluxAPI({ isAgent: false, processFileURL: jest.fn() });
      expect(flux.responseFormat).not.toBe('content_and_artifact');
    });

    it('invoke() returns ToolMessage with base64 in artifact, not serialized in content', async () => {
      const flux = new FluxAPI({ isAgent: true });
      const invokePromise = flux.invoke(
        makeToolCall('flux', { prompt: 'a box', endpoint: '/v1/flux-dev' }),
      );
      await jest.runAllTimersAsync();
      const result = await invokePromise;

      expect(result).toBeInstanceOf(ToolMessage);
      const contentStr =
        typeof result.content === 'string' ? result.content : JSON.stringify(result.content);
      expect(contentStr).not.toContain(FAKE_BASE64);

      expect(result.artifact).toBeDefined();
      const artifactContent = result.artifact?.content;
      expect(Array.isArray(artifactContent)).toBe(true);
      expect(artifactContent[0].type).toBe(ContentTypes.IMAGE_URL);
      expect(artifactContent[0].image_url.url).toContain('base64');
    });

    it('invoke() returns ToolMessage with base64 in artifact for generate_finetuned action', async () => {
      const flux = new FluxAPI({ isAgent: true });
      const invokePromise = flux.invoke(
        makeToolCall('flux', {
          action: 'generate_finetuned',
          prompt: 'a box',
          finetune_id: 'ft-abc123',
          endpoint: '/v1/flux-pro-finetuned',
        }),
      );
      await jest.runAllTimersAsync();
      const result = await invokePromise;

      expect(result).toBeInstanceOf(ToolMessage);
      const contentStr =
        typeof result.content === 'string' ? result.content : JSON.stringify(result.content);
      expect(contentStr).not.toContain(FAKE_BASE64);

      expect(result.artifact).toBeDefined();
      const artifactContent = result.artifact?.content;
      expect(Array.isArray(artifactContent)).toBe(true);
      expect(artifactContent[0].type).toBe(ContentTypes.IMAGE_URL);
      expect(artifactContent[0].image_url.url).toContain('base64');
    });

    it('invoke() returns ToolMessage with error string in content when task submission fails', async () => {
      axios.post.mockRejectedValue(new Error('Network error'));

      const flux = new FluxAPI({ isAgent: true });
      const invokePromise = flux.invoke(
        makeToolCall('flux', { prompt: 'a box', endpoint: '/v1/flux-dev' }),
      );
      await jest.runAllTimersAsync();
      const result = await invokePromise;

      expect(result).toBeInstanceOf(ToolMessage);
      const contentStr =
        typeof result.content === 'string' ? result.content : JSON.stringify(result.content);
      expect(contentStr).toContain('Something went wrong');
      expect(result.artifact).toBeDefined();
    });
  });

  describe('StableDiffusion', () => {
    beforeEach(() => {
      axios.post.mockResolvedValue({
        data: {
          images: [FAKE_BASE64],
          info: JSON.stringify({ height: 1024, width: 1024, seed: 42, infotexts: [] }),
        },
      });
    });

    it('sets responseFormat to content_and_artifact when isAgent is true', () => {
      const sd = new StableDiffusionAPI({ isAgent: true, override: true });
      expect(sd.responseFormat).toBe('content_and_artifact');
    });

    it('does not set responseFormat when isAgent is false', () => {
      const sd = new StableDiffusionAPI({
        isAgent: false,
        override: true,
        uploadImageBuffer: jest.fn(),
      });
      expect(sd.responseFormat).not.toBe('content_and_artifact');
    });

    it('invoke() returns ToolMessage with base64 in artifact, not serialized in content', async () => {
      const sd = new StableDiffusionAPI({ isAgent: true, override: true, userId: 'user-1' });
      const result = await sd.invoke(
        makeToolCall('stable-diffusion', { prompt: 'a box', negative_prompt: '' }),
      );

      expect(result).toBeInstanceOf(ToolMessage);
      const contentStr =
        typeof result.content === 'string' ? result.content : JSON.stringify(result.content);
      expect(contentStr).not.toContain(FAKE_BASE64);

      expect(result.artifact).toBeDefined();
      const artifactContent = result.artifact?.content;
      expect(Array.isArray(artifactContent)).toBe(true);
      expect(artifactContent[0].type).toBe(ContentTypes.IMAGE_URL);
      expect(artifactContent[0].image_url.url).toContain('base64');
    });

    it('invoke() returns ToolMessage with error string in content when API fails', async () => {
      axios.post.mockRejectedValue(new Error('Connection refused'));

      const sd = new StableDiffusionAPI({ isAgent: true, override: true, userId: 'user-1' });
      const result = await sd.invoke(
        makeToolCall('stable-diffusion', { prompt: 'a box', negative_prompt: '' }),
      );

      expect(result).toBeInstanceOf(ToolMessage);
      const contentStr =
        typeof result.content === 'string' ? result.content : JSON.stringify(result.content);
      expect(contentStr).toContain('Error making API request');
    });
  });
});