mcpctl/src/mcplocal/tests/proxymodel-executor.test.ts

import { describe, it, expect, vi } from 'vitest';
import { executePipeline, type ExecuteOptions } from '../src/proxymodel/executor.js';
import type { ProxyModelDefinition } from '../src/proxymodel/schema.js';
import type { LLMProvider, CacheProvider, StageLogger } from '../src/proxymodel/types.js';

function mockLlm(available = false): LLMProvider {
  return {
    async complete(prompt) { return `Summary: ${prompt.slice(0, 30)}...`; },
    available: () => available,
  };
}

function mockCache(): CacheProvider {
  const store = new Map<string, string>();
  return {
    async getOrCompute(key, compute) {
      if (store.has(key)) return store.get(key)!;
      const val = await compute();
      store.set(key, val);
      return val;
    },
    hash(content) { return content.slice(0, 8); },
    async get(key) { return store.get(key) ?? null; },
    async set(key, value) { store.set(key, value); },
  };
}

function mockLog(): StageLogger {
  return {
    debug: vi.fn(),
    info: vi.fn(),
    warn: vi.fn(),
    error: vi.fn(),
  };
}

function makeModel(stages: ProxyModelDefinition['spec']['stages'], appliesTo = ['toolResult'] as const): ProxyModelDefinition {
  return {
    kind: 'ProxyModel',
    metadata: { name: 'test' },
    spec: {
      controller: 'gate',
      stages,
      appliesTo: [...appliesTo],
      cacheable: false,
    },
    source: 'built-in',
  };
}

function makeOpts(content: string, model: ProxyModelDefinition, overrides: Partial<ExecuteOptions> = {}): ExecuteOptions {
  return {
    content,
    contentType: 'toolResult',
    sourceName: 'test/tool',
    projectName: 'test',
    sessionId: 'sess-1',
    proxyModel: model,
    llm: mockLlm(),
    cache: mockCache(),
    ...overrides,
  };
}

describe('executePipeline', () => {
  it('passes content through passthrough stage unchanged', async () => {
    const model = makeModel([{ type: 'passthrough' }]);
    const result = await executePipeline(makeOpts('hello world', model));
    expect(result.content).toBe('hello world');
    expect(result.sections).toBeUndefined();
  });

  it('chains multiple stages', async () => {
    // passthrough → passthrough should still return same content
    const model = makeModel([
      { type: 'passthrough' },
      { type: 'passthrough' },
    ]);
    const result = await executePipeline(makeOpts('data', model));
    expect(result.content).toBe('data');
  });

  it('paginate splits large content', async () => {
    const model = makeModel([
      { type: 'paginate', config: { pageSize: 50 } },
    ]);
    const content = 'line\n'.repeat(100);
    const result = await executePipeline(makeOpts(content, model));
    expect(result.sections).toBeDefined();
    expect(result.sections!.length).toBeGreaterThan(1);
    expect(result.content).toContain('pages');
  });

  it('skips missing stages with warning', async () => {
    const log = mockLog();
    const model = makeModel([
      { type: 'nonexistent-stage' },
      { type: 'passthrough' },
    ]);
    const result = await executePipeline(makeOpts('data', model, { log }));
    expect(result.content).toBe('data');
    expect(log.warn).toHaveBeenCalledWith(expect.stringContaining('nonexistent-stage'));
  });

  it('continues pipeline on stage error', async () => {
    // We'll test this by verifying the pipeline doesn't throw even if something goes wrong internally
    const model = makeModel([
      { type: 'passthrough' },
    ]);
    const result = await executePipeline(makeOpts('data', model));
    expect(result.content).toBe('data');
  });

  it('preserves originalContent across all stages', async () => {
    // section-split + summarize-tree pipeline — originalContent should always be the initial input
    const model = makeModel([
      { type: 'section-split', config: { minSectionSize: 5 } },
    ]);
    const jsonContent = JSON.stringify([
      { id: 'a', label: 'First', data: 'x'.repeat(100) },
      { id: 'b', label: 'Second', data: 'y'.repeat(100) },
    ]);
    const result = await executePipeline(makeOpts(jsonContent, model));
    expect(result.sections).toBeDefined();
  });

  it('respects appliesTo filter', async () => {
    const model = makeModel(
      [{ type: 'passthrough' }],
      ['resource'],
    );
    // contentType is toolResult but model only applies to resource
    const result = await executePipeline(makeOpts('data', model));
    expect(result.content).toBe('data');
    expect(result.sections).toBeUndefined();
  });

  it('returns empty metadata when no stages set it', async () => {
    const model = makeModel([{ type: 'passthrough' }]);
    const result = await executePipeline(makeOpts('data', model));
    expect(result.metadata).toBeUndefined();
  });

  it('handles section-split + summarize-tree (subindex) pipeline', async () => {
    const model = makeModel([
      { type: 'section-split', config: { minSectionSize: 5 } },
      { type: 'summarize-tree' },
    ]);
    const items = Array.from({ length: 10 }, (_, i) => ({
      id: `item-${i}`,
      name: `Item ${i}`,
      data: 'x'.repeat(300),
    }));
    const json = JSON.stringify(items);
    const result = await executePipeline(makeOpts(json, model));
    // Should produce sections
    expect(result.sections).toBeDefined();
    expect(result.content).toContain('sections');
  });

  it('works with empty content', async () => {
    const model = makeModel([{ type: 'passthrough' }]);
    const result = await executePipeline(makeOpts('', model));
    expect(result.content).toBe('');
  });

  describe('audit event emission', () => {
    const mockCollector = { emit: vi.fn(), flush: vi.fn(), dispose: vi.fn() };

    it('emits stage_execution for each stage + pipeline_execution summary', async () => {
      mockCollector.emit.mockClear();
      const model = makeModel([
        { type: 'passthrough' },
        { type: 'paginate', config: { pageSize: 50 } },
      ]);
      const content = 'line\n'.repeat(100);
      await executePipeline(makeOpts(content, model, { auditCollector: mockCollector as never }));

      // 2 stages + 1 pipeline summary = 3 events
      expect(mockCollector.emit).toHaveBeenCalledTimes(3);

      const calls = mockCollector.emit.mock.calls.map((c: unknown[]) => c[0] as { eventKind: string; payload: Record<string, unknown> });
      expect(calls[0]!.eventKind).toBe('stage_execution');
      expect(calls[0]!.payload['stage']).toBe('passthrough');
      expect(calls[0]!.payload['durationMs']).toBeGreaterThanOrEqual(0);
      expect(calls[1]!.eventKind).toBe('stage_execution');
      expect(calls[1]!.payload['stage']).toBe('paginate');
      expect(calls[2]!.eventKind).toBe('pipeline_execution');
      expect(calls[2]!.payload['totalDurationMs']).toBeGreaterThanOrEqual(0);
      expect(calls[2]!.payload['stageCount']).toBe(2);
    });

    it('includes serverName and correlationId when provided', async () => {
      mockCollector.emit.mockClear();
      const model = makeModel([{ type: 'passthrough' }]);
      await executePipeline(makeOpts('hello', model, {
        auditCollector: mockCollector as never,
        serverName: 'ha',
        correlationId: 'req-1',
      }));

      const calls = mockCollector.emit.mock.calls.map((c: unknown[]) => c[0] as { serverName?: string; correlationId?: string });
      for (const call of calls) {
        expect(call.serverName).toBe('ha');
        expect(call.correlationId).toBe('req-1');
      }
    });

    it('does not emit when auditCollector is undefined', async () => {
      mockCollector.emit.mockClear();
      const model = makeModel([{ type: 'passthrough' }]);
      // No auditCollector — should not throw
      await executePipeline(makeOpts('hello', model));
      expect(mockCollector.emit).not.toHaveBeenCalled();
    });

    it('preserves correct inputSize/outputSize per stage', async () => {
      mockCollector.emit.mockClear();
      const model = makeModel([{ type: 'passthrough' }]);
      await executePipeline(makeOpts('hello', model, { auditCollector: mockCollector as never }));

      const stageEvent = mockCollector.emit.mock.calls[0]![0] as { payload: Record<string, unknown> };
      expect(stageEvent.payload['inputSize']).toBe(5);
      expect(stageEvent.payload['outputSize']).toBe(5);
    });

    it('emits pipeline_execution with input/output sizes', async () => {
      mockCollector.emit.mockClear();
      const model = makeModel([{ type: 'passthrough' }]);
      await executePipeline(makeOpts('hello', model, { auditCollector: mockCollector as never }));

      const pipelineEvent = mockCollector.emit.mock.calls[1]![0] as { payload: Record<string, unknown> };
      expect(pipelineEvent.payload['inputSize']).toBe(5);
      expect(pipelineEvent.payload['outputSize']).toBe(5);
      expect(pipelineEvent.payload['stageCount']).toBe(1);
    });
  });
});