FastGPT/test/cases/service/support/mcp/utils.test.ts

import { describe, expect, it, vi } from 'vitest';
import {
  pluginNodes2InputSchema,
  workflow2InputSchema,
  getMcpServerTools,
  callMcpServerTool
} from '@/service/support/mcp/utils';
import { FlowNodeTypeEnum } from '@fastgpt/global/core/workflow/node/constant';
import { MongoMcpKey } from '@fastgpt/service/support/mcp/schema';
import { MongoApp } from '@fastgpt/service/core/app/schema';
import { AppTypeEnum } from '@fastgpt/global/core/app/constants';
import { CommonErrEnum } from '@fastgpt/global/common/error/code/common';
import { authAppByTmbId } from '@fastgpt/service/support/permission/app/auth';
import { getAppLatestVersion } from '@fastgpt/service/core/app/version/controller';
import {
  getUserChatInfoAndAuthTeamPoints,
  getRunningUserInfoByTmbId
} from '@fastgpt/service/support/permission/auth/team';
import { dispatchWorkFlow } from '@fastgpt/service/core/workflow/dispatch';
import { saveChat } from '@fastgpt/service/core/chat/saveChat';
import { createChatUsage } from '@fastgpt/service/support/wallet/usage/controller';

vi.mock('@fastgpt/service/support/mcp/schema', () => ({
  MongoMcpKey: {
    findOne: vi.fn().mockReturnValue({
      lean: vi.fn()
    })
  }
}));

vi.mock('@fastgpt/service/core/app/schema', () => ({
  MongoApp: {
    find: vi.fn().mockReturnValue({
      lean: vi.fn()
    })
  }
}));

vi.mock('@fastgpt/service/support/permission/app/auth', () => ({
  authAppByTmbId: vi.fn()
}));

vi.mock('@fastgpt/service/core/app/version/controller', () => ({
  getAppLatestVersion: vi.fn()
}));

vi.mock('@fastgpt/service/support/permission/auth/team', () => ({
  getUserChatInfoAndAuthTeamPoints: vi.fn(),
  getRunningUserInfoByTmbId: vi.fn()
}));

vi.mock('@fastgpt/service/core/workflow/dispatch', () => ({
  dispatchWorkFlow: vi.fn()
}));

vi.mock('@fastgpt/service/core/chat/saveChat', () => ({
  saveChat: vi.fn()
}));

vi.mock('@fastgpt/service/support/wallet/usage/controller', () => ({
  createChatUsage: vi.fn()
}));

describe('pluginNodes2InputSchema', () => {
  it('should generate input schema from plugin nodes', () => {
    const nodes = [
      {
        flowNodeType: FlowNodeTypeEnum.pluginInput,
        inputs: [
          {
            key: 'testKey',
            valueType: 'string',
            description: 'test description',
            required: true,
            enum: 'a\nb\nc'
          }
        ]
      }
    ];

    const schema = pluginNodes2InputSchema(nodes);

    expect(schema).toEqual({
      type: 'object',
      properties: {
        testKey: {
          type: 'string',
          description: 'test description',
          enum: ['a', 'b', 'c']
        }
      },
      required: ['testKey']
    });
  });

  it('should handle empty plugin input nodes', () => {
    const nodes = [
      {
        flowNodeType: FlowNodeTypeEnum.pluginInput,
        inputs: []
      }
    ];

    const schema = pluginNodes2InputSchema(nodes);

    expect(schema).toEqual({
      type: 'object',
      properties: {},
      required: []
    });
  });
});

describe('workflow2InputSchema', () => {
  it('should generate input schema with file config', () => {
    const chatConfig = {
      fileSelectConfig: {
        canSelectFile: true,
        canSelectImg: true
      },
      variables: [
        {
          key: 'var1',
          valueType: 'string',
          description: 'test var',
          required: true,
          enums: [{ value: 'a' }, { value: 'b' }]
        }
      ]
    };

    const schema = workflow2InputSchema(chatConfig);

    expect(schema).toEqual({
      type: 'object',
      properties: {
        question: {
          type: 'string',
          description: 'Question from user'
        },
        fileUrlList: {
          type: 'array',
          items: {
            type: 'string'
          },
          description: 'File linkage'
        },
        var1: {
          type: 'string',
          description: 'test var',
          enum: ['a', 'b']
        }
      },
      required: ['question', 'var1']
    });
  });
});

describe('getMcpServerTools', () => {
  it('should return tools list', async () => {
    const mockMcp = {
      tmbId: 'test-tmb',
      apps: [
        {
          appId: 'test-app',
          toolName: 'test-tool',
          description: 'test description'
        }
      ]
    };

    vi.mocked(MongoMcpKey.findOne).mockReturnValue({
      lean: () => mockMcp
    });

    vi.mocked(MongoApp.find).mockReturnValue({
      lean: () => [
        {
          _id: 'test-app',
          name: 'Test App',
          type: AppTypeEnum.plugin
        }
      ]
    });

    vi.mocked(authAppByTmbId).mockResolvedValue(undefined);

    vi.mocked(getAppLatestVersion).mockResolvedValue({
      nodes: [
        {
          flowNodeType: FlowNodeTypeEnum.pluginInput,
          inputs: []
        }
      ],
      edges: [],
      chatConfig: {}
    });

    const tools = await getMcpServerTools('test-key');

    expect(tools).toHaveLength(1);
    expect(tools[0].name).toBe('test-tool');
  });

  it('should reject if key not found', async () => {
    vi.mocked(MongoMcpKey.findOne).mockReturnValue({
      lean: () => null
    });

    await expect(getMcpServerTools('invalid-key')).rejects.toBe(CommonErrEnum.invalidResource);
  });
});

describe('callMcpServerTool', () => {
  it('should call tool and return response', async () => {
    const mockMcp = {
      apps: [
        {
          appId: 'test-app',
          toolName: 'test-tool'
        }
      ]
    };

    const mockApp = {
      _id: 'test-app',
      type: AppTypeEnum.plugin,
      teamId: 'test-team',
      tmbId: 'test-tmb',
      name: 'Test App'
    };

    vi.mocked(MongoMcpKey.findOne).mockReturnValue({
      lean: () => mockMcp
    });

    vi.mocked(MongoApp.find).mockReturnValue({
      lean: () => [mockApp]
    });

    vi.mocked(getUserChatInfoAndAuthTeamPoints).mockResolvedValue({
      timezone: 'UTC',
      externalProvider: {}
    });

    vi.mocked(getRunningUserInfoByTmbId).mockResolvedValue({});

    vi.mocked(getAppLatestVersion).mockResolvedValue({
      nodes: [
        {
          flowNodeType: FlowNodeTypeEnum.pluginInput,
          inputs: []
        }
      ],
      edges: [],
      chatConfig: {}
    });

    vi.mocked(dispatchWorkFlow).mockResolvedValue({
      flowUsages: [],
      assistantResponses: [],
      newVariables: {},
      flowResponses: [
        {
          moduleType: FlowNodeTypeEnum.pluginOutput,
          pluginOutput: { result: 'test' }
        }
      ],
      durationSeconds: 1,
      system_memories: []
    });

    const response = await callMcpServerTool({
      key: 'test-key',
      toolName: 'test-tool',
      inputs: {
        question: 'test question'
      }
    });

    expect(response).toBe('{"result":"test"}');
  });

  it('should reject if key not found', async () => {
    vi.mocked(MongoMcpKey.findOne).mockReturnValue({
      lean: () => null
    });

    await expect(
      callMcpServerTool({
        key: 'invalid-key',
        toolName: 'test-tool',
        inputs: {}
      })
    ).rejects.toBe(CommonErrEnum.invalidResource);
  });
});