metona-ai-desktop/electron/ipc/__tests__/agent.test.ts

/**
 * IPC Agent Handlers — sendMessage 编排链路测试（v0.4.1 测试补齐）
 *
 * 覆盖 sendMessage 的主编排逻辑：
 * 1. 参数校验（无效 sessionId / userMessage → ERROR+DONE 流事件，防止前端 isStreaming 卡死）
 * 2. Adapter 加载失败中止
 * 3. Prompt 注入阻断（riskScore >= 7）
 * 4. 成功路径（消息持久化 / 审计 / 记忆固化 / 摘要评估 / Token 统计）
 * 5. 引擎异常路径（审计错误 + ERROR 流事件）
 */

import { describe, it, expect, vi, beforeEach, type Mock } from 'vitest';
import { EventEmitter } from 'events';

// ===== Mock electron（ipcMain） =====
const ipcMainHandleMock = vi.fn();
const ipcMainOnMock = vi.fn();
vi.mock('electron', () => ({
  ipcMain: {
    handle: (...args: unknown[]) => ipcMainHandleMock(...args),
    on: (...args: unknown[]) => ipcMainOnMock(...args),
  },
}));

// ===== Mock broadcast（ipc/context） =====
const broadcastMock = vi.fn();
vi.mock('../context', () => ({
  broadcast: (...args: unknown[]) => broadcastMock(...args),
}));

import { registerAgentHandlers } from '../agent';
import type { IPCContext } from '../context';
import type { MetonaMessage } from '../../harness/types';

// ===== Mock 依赖工厂 =====

function makeEngineMock(overrides: Record<string, unknown> = {}) {
  return {
    runStream: vi.fn().mockResolvedValue({
      finalAnswer: '这是最终回答',
      terminationReason: 'completed',
      iterations: [
        {
          iteration: 1,
          state: 'OBSERVING',
          startedAt: 1,
          completedAt: 2,
          thought: {
            id: 'thought-1',
            content: '本轮思考文本',
            reasoningContent: '推理过程',
            timestamp: 1,
            iteration: 1,
          },
          toolCalls: [],
          toolResults: [],
        },
      ],
      totalTokenUsage: { promptTokens: 100, completionTokens: 50, totalTokens: 150 },
      durationMs: 1234,
      metadata: {},
    }),
    ...overrides,
  };
}

function makeCtx(overrides: Record<string, unknown> = {}) {
  const engine = makeEngineMock();
  const engineManager = new EventEmitter() as EventEmitter & {
    getEngine: Mock;
    abort: Mock;
    waitForAbort: Mock;
  };
  (engineManager as unknown as { getEngine: Mock }).getEngine = vi.fn(() => engine);
  (engineManager as unknown as { abort: Mock }).abort = vi.fn();
  (engineManager as unknown as { waitForAbort: Mock }).waitForAbort = vi
    .fn()
    .mockResolvedValue(true);

  const ctx = {
    agentEngineManager: engineManager,
    sessionRecorder: {
      startRecording: vi.fn(),
      stopRecording: vi.fn(),
      recordContextBuilt: vi.fn(),
      recordToolCall: vi.fn(),
      recordToolResult: vi.fn(),
      recordLLMResponse: vi.fn(),
      recordIterationStart: vi.fn(),
      recordIterationEnd: vi.fn(),
      recordLLMRequest: vi.fn(),
    },
    configService: { get: vi.fn(() => '') },
    sessionService: {
      saveMessage: vi.fn(),
      getMessages: vi.fn(() => []),
      updateTokenUsage: vi.fn(),
    },
    workspaceService: {
      getFiles: vi.fn(() => ({ soul: '# Metona', memory: '# Memory' })),
      getPath: vi.fn(() => '/workspace'),
      updateMemoryTimestamp: vi.fn(),
    },
    contextBuilder: {
      buildSystemPrompt: vi.fn(() => ({
        roleDefinition: 'role',
        outputConstraints: 'constraints',
        safetyGuidelines: 'safety',
        dynamicReminders: 'reminders',
      })),
      isUsingFallbackRole: vi.fn(() => false),
    },
    auditService: {
      logSessionStart: vi.fn(),
      logSessionEnd: vi.fn(),
      log: vi.fn(),
    },
    memoryManager: { search: vi.fn(() => []) },
    promptInjectionDefender: {
      detect: vi.fn(() => ({
        isInjection: false,
        riskScore: 0,
        findings: [],
        recommendation: 'PASS: ok',
      })),
    },
    outputValidator: {
      validate: vi.fn().mockResolvedValue({ valid: true, issues: [], score: 1 }),
    },
    memoryConsolidator: {
      consolidate: vi.fn().mockResolvedValue({ appended: 0, entries: [], skipped: 0 }),
      isRunning: vi.fn(() => false),
    },
    sessionSummaryService: {
      buildHistoryMessages: vi.fn(() => []),
      maybeSummarize: vi.fn().mockResolvedValue(undefined),
    },
    // v0.5.0: orchestrator 增加事件监听（SubAgent 可观测性广播/录制），
    // mock 改用 EventEmitter 提供真实 on/emit 行为
    orchestrator: Object.assign(new EventEmitter(), { abortByParent: vi.fn() }),
    confirmationHook: { clearPending: vi.fn() },
    reloadAdapter: vi.fn(() => true),
    ...overrides,
  };
  return { ctx: ctx as unknown as IPCContext, engine, engineManager, ctxRaw: ctx };
}

function getHandler(channel: string): (...args: unknown[]) => Promise<unknown> {
  const call = ipcMainHandleMock.mock.calls.find(([ch]) => ch === channel);
  if (!call) throw new Error(`IPC handler not registered: ${channel}`);
  return call[1] as (...args: unknown[]) => Promise<unknown>;
}

const VALID_MESSAGE: MetonaMessage = {
  role: 'user',
  content: '你好，请帮我分析这个项目',
  timestamp: Date.now(),
};

beforeEach(() => {
  ipcMainHandleMock.mockClear();
  ipcMainOnMock.mockClear();
  broadcastMock.mockClear();
});

describe('agent:sendMessage — 参数校验', () => {
  it('无效 sessionId 拒绝并发送 ERROR + DONE 流事件（防止前端 isStreaming 卡死）', async () => {
    const { ctx } = makeCtx();
    registerAgentHandlers(ctx);
    const handler = getHandler('agent:sendMessage');

    const result = await handler(null, VALID_MESSAGE, '');
    expect(result).toEqual({ success: false, error: 'Invalid sessionId' });

    // ERROR + DONE 两个流事件都应广播
    const eventTypes = broadcastMock.mock.calls.map(([, ev]) => (ev as { type: string }).type);
    expect(eventTypes).toContain('error');
    expect(eventTypes).toContain('done');
  });

  it('无效 userMessage（非对象 / content 非字符串）拒绝', async () => {
    const { ctx } = makeCtx();
    registerAgentHandlers(ctx);
    const handler = getHandler('agent:sendMessage');

    const result = await handler(null, { content: 123 }, 'sess_1');
    expect(result).toEqual({ success: false, error: 'Invalid message format' });
    expect(broadcastMock).toHaveBeenCalled();
  });
});

describe('agent:sendMessage — 前置检查', () => {
  it('Adapter 加载失败时中止并停止录制', async () => {
    const { ctx, ctxRaw } = makeCtx({ reloadAdapter: vi.fn(() => false) });
    registerAgentHandlers(ctx);
    const handler = getHandler('agent:sendMessage');

    const result = await handler(null, VALID_MESSAGE, 'sess_1');
    expect((result as { success: boolean }).success).toBe(false);
    expect(ctxRaw.sessionRecorder.stopRecording).toHaveBeenCalled();
    // 不应调用引擎
    expect(ctxRaw.agentEngineManager.getEngine).not.toHaveBeenCalled();
  });

  it('注入风险 riskScore >= 7 时阻断消息', async () => {
    const { ctx, ctxRaw } = makeCtx({
      promptInjectionDefender: {
        detect: vi.fn(() => ({
          isInjection: true,
          riskScore: 8,
          findings: [{ pattern: 'x', matched: 'ignore previous instructions', severity: 'high' }],
          recommendation: 'BLOCK: High-risk injection detected',
        })),
      },
    });
    registerAgentHandlers(ctx);
    const handler = getHandler('agent:sendMessage');

    const result = await handler(null, VALID_MESSAGE, 'sess_1');
    expect((result as { success: boolean }).success).toBe(false);
    expect((result as { error: string }).error).toContain('blocked by prompt injection defense');
    // 用户消息不保存（在注入检测前已保存？—— 现实现:先保存再检测，验证已保存）
    expect(ctxRaw.sessionService.saveMessage).toHaveBeenCalledWith(
      expect.objectContaining({ role: 'user', sessionId: 'sess_1' }),
    );
    // 引擎不启动
    expect(ctxRaw.agentEngineManager.getEngine).not.toHaveBeenCalled();
  });

  it('SOUL.md 缺失降级时发送 toast 提示', async () => {
    const { ctx } = makeCtx({
      contextBuilder: {
        buildSystemPrompt: vi.fn(() => ({
          roleDefinition: 'fallback',
          outputConstraints: 'c',
          safetyGuidelines: 's',
        })),
        isUsingFallbackRole: vi.fn(() => true),
      },
    });
    registerAgentHandlers(ctx);
    const handler = getHandler('agent:sendMessage');

    await handler(null, VALID_MESSAGE, 'sess_1');
    const toastCall = broadcastMock.mock.calls.find(([ch]) => ch === 'toast:show');
    expect(toastCall).toBeDefined();
    expect((toastCall![1] as { message: string }).message).toContain('SOUL.md');
  });
});

describe('agent:sendMessage — 成功路径', () => {
  it('完整编排：保存消息 → 运行引擎 → 持久化 assistant 消息 → 审计 → 异步固化', async () => {
    const { ctx, ctxRaw, engine } = makeCtx();
    registerAgentHandlers(ctx);
    const handler = getHandler('agent:sendMessage');

    const result = await handler(null, VALID_MESSAGE, 'sess_1');
    expect(result).toEqual({ success: true });

    // 1. 用户消息保存
    expect(ctxRaw.sessionService.saveMessage).toHaveBeenCalledWith(
      expect.objectContaining({
        role: 'user',
        content: VALID_MESSAGE.content,
        sessionId: 'sess_1',
      }),
    );
    // 2. 引擎启动（每会话引擎）
    expect(ctxRaw.agentEngineManager.getEngine).toHaveBeenCalledWith('sess_1');
    expect(engine.runStream).toHaveBeenCalledWith(
      VALID_MESSAGE,
      'sess_1',
      [],
      expect.objectContaining({ roleDefinition: 'role' }),
    );
    // 3. assistant 消息保存（含思考内容）
    expect(ctxRaw.sessionService.saveMessage).toHaveBeenCalledWith(
      expect.objectContaining({
        role: 'assistant',
        content: '本轮思考文本',
        reasoningContent: '推理过程',
      }),
    );
    // 4. Token 统计更新
    expect(ctxRaw.sessionService.updateTokenUsage).toHaveBeenCalledWith('sess_1', 150);
    // 5. MEMORY.md 时间戳更新
    expect(ctxRaw.workspaceService.updateMemoryTimestamp).toHaveBeenCalled();
    // 6. 审计 + 录制结束
    expect(ctxRaw.auditService.logSessionEnd).toHaveBeenCalledWith(
      expect.objectContaining({ sessionId: 'sess_1', terminationReason: 'completed' }),
    );
    expect(ctxRaw.sessionRecorder.stopRecording).toHaveBeenCalled();
    // 7. 输出验证执行
    expect(ctxRaw.outputValidator.validate).toHaveBeenCalledWith('这是最终回答', expect.anything());
    // 8. 摘要评估（异步触发）
    await vi.waitFor(() =>
      expect(ctxRaw.sessionSummaryService.maybeSummarize).toHaveBeenCalledWith('sess_1'),
    );
    // 9. 记忆固化（异步触发）
    await vi.waitFor(() => expect(ctxRaw.memoryConsolidator.consolidate).toHaveBeenCalled());
  });

  it('注入相关记忆到 System Prompt 动态区', async () => {
    const { ctx, ctxRaw } = makeCtx({
      memoryManager: {
        search: vi.fn(() => [
          {
            id: 'm1',
            type: 'semantic',
            content: '用户偏好深色主题',
            importance: 0.9,
            createdAt: Date.now(),
            score: 0.8,
          },
        ]),
      },
    });
    registerAgentHandlers(ctx);
    const handler = getHandler('agent:sendMessage');

    await handler(null, VALID_MESSAGE, 'sess_1');
    expect(ctxRaw.memoryManager.search).toHaveBeenCalled();
    // 引擎收到的 systemPrompt 应包含记忆块
    const prompt = engine_runStreamPrompt(ctxRaw);
    expect(prompt.dynamicReminders).toContain('用户偏好深色主题');
  });

  it('验证发现 warning 级问题时广播 VALIDATION 流事件', async () => {
    const { ctx, ctxRaw } = makeCtx({
      outputValidator: {
        validate: vi.fn().mockResolvedValue({
          valid: false,
          score: 0.7,
          issues: [
            { severity: 'warning', type: 'hallucination', message: 'Path not found in context' },
            { severity: 'info', type: 'format', message: 'noise' },
          ],
        }),
      },
    });
    registerAgentHandlers(ctx);
    const handler = getHandler('agent:sendMessage');

    await handler(null, VALID_MESSAGE, 'sess_1');
    const validationCall = broadcastMock.mock.calls.find(
      ([ch, ev]) => ch === 'agent:streamEvent' && (ev as { type: string }).type === 'validation',
    );
    expect(validationCall).toBeDefined();
    const payload = (validationCall![1] as { validation: { issues: unknown[] } }).validation;
    // info 级噪声不推送
    expect(payload.issues).toHaveLength(1);
    expect(ctxRaw.outputValidator.validate).toHaveBeenCalled();
  });
});

describe('agent:sendMessage — 异常路径', () => {
  it('引擎抛错时返回失败并记录审计错误', async () => {
    const engine = makeEngineMock({
      runStream: vi.fn().mockRejectedValue(new Error('LLM connection failed')),
    });
    const engineManager = new EventEmitter() as EventEmitter & {
      getEngine: Mock;
      abort: Mock;
      waitForAbort: Mock;
    };
    (engineManager as unknown as { getEngine: Mock }).getEngine = vi.fn(() => engine);
    (engineManager as unknown as { abort: Mock }).abort = vi.fn();
    (engineManager as unknown as { waitForAbort: Mock }).waitForAbort = vi
      .fn()
      .mockResolvedValue(true);
    const { ctx, ctxRaw } = makeCtx({
      agentEngineManager: engineManager,
    });
    registerAgentHandlers(ctx);
    const handler = getHandler('agent:sendMessage');

    const result = await handler(null, VALID_MESSAGE, 'sess_1');
    expect((result as { success: boolean }).success).toBe(false);
    expect((result as { error: string }).error).toBe('LLM connection failed');
    // 审计记录错误
    expect(ctxRaw.auditService.log).toHaveBeenCalledWith(
      expect.objectContaining({ sessionId: 'sess_1', eventType: 'error', outcome: 'error' }),
    );
    // ERROR 流事件广播
    const errorCall = broadcastMock.mock.calls.find(
      ([ch, ev]) => ch === 'agent:streamEvent' && (ev as { type: string }).type === 'error',
    );
    expect(errorCall).toBeDefined();
  });
});

describe('agent:abortSession — 中断编排', () => {
  it('联动 SubAgent 中断 + 引擎中断 + 清理确认', async () => {
    const { ctx, ctxRaw } = makeCtx();
    registerAgentHandlers(ctx);
    const handler = getHandler('agent:abortSession');

    const result = await handler(null, 'sess_1');
    expect(result).toEqual({ success: true });
    expect(ctxRaw.orchestrator.abortByParent).toHaveBeenCalledWith('sess_1');
    expect(ctxRaw.confirmationHook.clearPending).toHaveBeenCalled();
    expect(ctxRaw.auditService.log).toHaveBeenCalledWith(
      expect.objectContaining({ eventType: 'session_end', outcome: 'denied' }),
    );
  });
});

/** 从 runStream 调用参数中提取 systemPrompt */
function engine_runStreamPrompt(ctxRaw: Record<string, unknown>): {
  roleDefinition: string;
  dynamicReminders?: string;
} {
  const engine = (ctxRaw.agentEngineManager as unknown as { getEngine: Mock }).getEngine() as {
    runStream: Mock;
  };
  return engine.runStream.mock.calls[0][3];
}