zn-ai/tests/chat-runtime-context.test.ts

// @vitest-environment node

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

const mocks = vi.hoisted(() => {
  const sessionMessages: any[] = [];

  return {
    sessionMessages,
    providerChat: vi.fn(),
    appendMessage: vi.fn((_: string, message: unknown) => {
      sessionMessages.push(message);
    }),
    getOrCreate: vi.fn(() => ({
      key: 'agent:test:main',
      messages: [...sessionMessages],
      updatedAt: Date.now(),
    })),
    setActiveRun: vi.fn(),
    clearActiveRun: vi.fn(),
    appendTranscriptLine: vi.fn(),
    maybeHandleBrowserOpenMessage: vi.fn(() => false),
    maybeHandleSkillInstallMessage: vi.fn(() => false),
    logger: {
      error: vi.fn(),
    },
  };
});

vi.mock('@electron/providers', () => ({
  createProvider: vi.fn(() => ({
    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: {
    appendMessage: mocks.appendMessage,
    getOrCreate: mocks.getOrCreate,
    setActiveRun: mocks.setActiveRun,
    clearActiveRun: mocks.clearActiveRun,
  },
}));

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

vi.mock('../electron/gateway/browser-shortcut', () => ({
  maybeHandleBrowserOpenMessage: mocks.maybeHandleBrowserOpenMessage,
}));

vi.mock('../electron/gateway/skill-install-shortcut', () => ({
  maybeHandleSkillInstallMessage: mocks.maybeHandleSkillInstallMessage,
}));

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

function createStream(chunks: Array<{ result?: string; usage?: 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 runtime context', () => {
  beforeEach(() => {
    vi.clearAllMocks();
    mocks.sessionMessages.length = 0;
    mocks.maybeHandleBrowserOpenMessage.mockReturnValue(false);
    mocks.maybeHandleSkillInstallMessage.mockReturnValue(false);
    mocks.providerChat.mockResolvedValue(createStream([{ result: 'done' }]));
  });

  it('prepends the zn-ai runtime context before provider chat runs', async () => {
    const { handleChatSend } = await import('../electron/gateway/handlers/chat');

    const result = handleChatSend(
      {
        sessionKey: 'agent:test:main',
        message: {
          role: 'user',
          content: '帮我看一下这个网页',
        },
      },
      vi.fn(),
    );

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

    await flushAsyncTasks();

    expect(mocks.providerChat).toHaveBeenCalledTimes(1);
    const [messages, model] = mocks.providerChat.mock.calls[0] ?? [];

    expect(model).toBe('gpt-4o-mini');
    expect(messages).toEqual(
      expect.arrayContaining([
        expect.objectContaining({
          role: 'system',
          content: expect.stringContaining('browser.open_url'),
        }),
        expect.objectContaining({
          role: 'user',
          content: '帮我看一下这个网页',
        }),
      ]),
    );
    expect(messages[0]).toMatchObject({
      role: 'system',
      content: expect.stringContaining('skills.install'),
    });
  });
});