zn-ai/tests/chat-provider-tool-loop.test.ts

// @vitest-environment node

import { beforeEach, describe, expect, it, vi } from 'vitest';

const mocks = vi.hoisted(() => {
  const sessionMessages: any[] = [];
  let activeRun:
    | {
        runId: string;
        abortController: AbortController;
      }
    | undefined;

  return {
    sessionMessages,
    getActiveRun: vi.fn(() => activeRun),
    providerChat: vi.fn(),
    openUrlInBrowser: vi.fn(),
    appendTranscriptLine: vi.fn(),
    logger: {
      error: vi.fn(),
      warn: vi.fn(),
    },
    sessionStore: {
      appendMessage: vi.fn((_: string, message: unknown) => {
        sessionMessages.push(message);
      }),
      getOrCreate: vi.fn(() => ({
        key: 'agent:test:main',
        messages: [...sessionMessages],
        updatedAt: Date.now(),
      })),
      setActiveRun: vi.fn((_: string, runId: string, abortController: AbortController) => {
        activeRun = { runId, abortController };
      }),
      clearActiveRun: vi.fn(() => {
        activeRun = undefined;
      }),
      getActiveRun: vi.fn(() => activeRun),
    },
  };
});

vi.mock('@electron/providers', () => ({
  createProvider: vi.fn(() => ({
    getCapabilities: () => ({
      structuredMessages: true,
      toolCalls: true,
      toolResults: true,
      thinking: false,
    }),
    chat: mocks.providerChat,
  })),
}));

vi.mock('@electron/service/provider-api-service', () => ({
  providerApiService: {
    getDefault: () => ({ accountId: 'provider-1' }),
    getAccounts: () => [
      {
        id: 'provider-1',
        model: 'gpt-4o-mini',
        vendorId: 'openai',
        label: 'OpenAI',
      },
    ],
  },
}));

vi.mock('../electron/gateway/session-store', () => ({
  sessionStore: mocks.sessionStore,
}));

vi.mock('@electron/service/browser-open-service', () => ({
  openUrlInBrowser: mocks.openUrlInBrowser,
}));

vi.mock('@electron/utils/token-usage-writer', () => ({
  appendTranscriptLine: mocks.appendTranscriptLine,
}));

vi.mock('../electron/gateway/skill-capability-registry', () => ({
  getEnabledSkillCapabilities: () => [],
}));

vi.mock('@electron/service/logger', () => ({
  default: mocks.logger,
}));

function createStream(chunks: Array<Record<string, unknown>>) {
  return {
    async *[Symbol.asyncIterator]() {
      for (const chunk of chunks) {
        yield chunk;
      }
    },
  };
}

function flushAsyncTasks(): Promise<void> {
  return new Promise((resolve) => setTimeout(resolve, 0));
}

describe('chat provider tool loop', () => {
  beforeEach(() => {
    vi.clearAllMocks();
    mocks.sessionMessages.length = 0;
    mocks.openUrlInBrowser.mockResolvedValue({
      url: 'https://example.com/',
      pageUrl: 'https://example.com/',
      title: 'Example Domain',
    });

    mocks.providerChat
      .mockResolvedValueOnce(createStream([
        {
          toolCalls: [
            {
              index: 0,
              id: 'call_browser_1',
              name: 'browser.open_url',
              argumentsDelta: '{"url":"https://example.com/"}',
            },
          ],
          finishReason: 'tool_calls',
        },
      ]))
      .mockResolvedValueOnce(createStream([
        {
          result: '已打开页面并获取标题：Example Domain',
          finishReason: 'stop',
        },
      ]));
  });

  it('executes provider-requested tools and resumes the model with tool_result context', async () => {
    const { handleChatSend } = await import('../electron/gateway/handlers/chat');
    const broadcast = vi.fn();

    const result = handleChatSend(
      {
        sessionKey: 'agent:test:main',
        message: {
          role: 'user',
          content: '请帮我查看一下官网首页信息',
        },
      },
      broadcast,
    );

    expect(result.runId).toBeTypeOf('string');

    await flushAsyncTasks();
    await flushAsyncTasks();

    expect(mocks.providerChat).toHaveBeenCalledTimes(2);
    expect(mocks.providerChat.mock.calls[0]?.[2]).toEqual(expect.objectContaining({
      tools: expect.arrayContaining([
        expect.objectContaining({
          name: 'browser.open_url',
        }),
      ]),
      toolChoice: 'auto',
    }));

    const secondCallMessages = mocks.providerChat.mock.calls[1]?.[0] as Array<Record<string, unknown>>;
    expect(secondCallMessages).toEqual(expect.arrayContaining([
      expect.objectContaining({
        role: 'assistant',
        content: expect.arrayContaining([
          expect.objectContaining({
            type: 'tool_use',
            name: 'browser.open_url',
          }),
        ]),
      }),
      expect.objectContaining({
        role: 'tool_result',
      }),
    ]));

    expect(mocks.openUrlInBrowser).toHaveBeenCalledWith(
      'https://example.com/',
      expect.objectContaining({
        signal: expect.any(AbortSignal),
      }),
    );
    expect(broadcast).toHaveBeenCalledWith(expect.objectContaining({
      type: 'tool:status',
      toolName: 'browser.open_url',
      status: 'running',
    }));
    expect(broadcast).toHaveBeenCalledWith(expect.objectContaining({
      type: 'tool:status',
      toolName: 'browser.open_url',
      status: 'completed',
    }));
    expect(broadcast).toHaveBeenCalledWith(expect.objectContaining({
      type: 'chat:final',
      message: expect.objectContaining({
        role: 'assistant',
        content: '已打开页面并获取标题：Example Domain',
      }),
    }));
  });
});