mcpctl/src/mcplocal/tests/prompt-section-drilldown.test.ts

/**
 * Tests for prompt section drill-down.
 *
 * Verifies that large prompt responses are split into sections via the
 * content-pipeline plugin, and that clients can drill into individual
 * sections using _resultId + _section parameters on subsequent prompts/get calls.
 */
import { describe, it, expect, vi } from 'vitest';
import { createContentPipelinePlugin } from '../src/proxymodel/plugins/content-pipeline.js';
import type { PluginSessionContext } from '../src/proxymodel/plugin.js';
import type { JsonRpcRequest, JsonRpcResponse } from '../src/types.js';
import type { Section, LLMProvider, CacheProvider, StageLogger, ContentType } from '../src/proxymodel/types.js';

// ── Helpers ──

/** Generate a large markdown prompt with multiple sections. */
function generateLargePrompt(sectionCount: number): string {
  const sections: string[] = [];
  for (let i = 1; i <= sectionCount; i++) {
    const body = `This is the content of section ${i}. `.repeat(80); // ~3200 chars each
    sections.push(`## Section ${i}\n\n${body}`);
  }
  return sections.join('\n\n');
}

/** Generate a large JSON prompt (array of objects). */
function generateLargeJsonPrompt(itemCount: number): string {
  const items = Array.from({ length: itemCount }, (_, i) => ({
    name: `item-${i + 1}`,
    description: `Description for item ${i + 1}. `.repeat(60),
    config: { enabled: true, priority: i + 1 },
  }));
  return JSON.stringify(items, null, 2);
}

function createMockCtx(): PluginSessionContext {
  const state = new Map<string, unknown>();

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

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

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

  return {
    sessionId: 'test-session',
    projectName: 'test-project',
    state,
    llm: mockLlm,
    cache: mockCache,
    log: mockLog,
    discoverTools: async () => [],
    registerVirtualTool: vi.fn(),
    registerVirtualServer: vi.fn(),
    routeToUpstream: async (req: JsonRpcRequest) => ({ jsonrpc: '2.0' as const, id: req.id, result: {} }),
    fetchPromptIndex: async () => [],
    getSystemPrompt: async (_name: string, fallback: string) => fallback,
    processContent: async (toolName: string, content: string, contentType: ContentType) => {
      // Use real section-split stage for realistic testing
      const sectionSplit = (await import('../src/proxymodel/stages/section-split.js')).default;
      const ctx = {
        contentType,
        sourceName: toolName,
        projectName: 'test',
        sessionId: 'test-session',
        originalContent: content,
        llm: mockLlm,
        cache: mockCache,
        log: mockLog,
        getSystemPrompt: async (_n: string, fb: string) => fb,
        config: { minSectionSize: 500 },
      };
      return sectionSplit(content, ctx);
    },
    queueNotification: vi.fn(),
    postToMcpd: async () => ({}),
  };
}

function makePromptGetRequest(name: string, args?: Record<string, unknown>): JsonRpcRequest {
  return {
    jsonrpc: '2.0',
    id: 1,
    method: 'prompts/get',
    params: { name, ...(args ? { arguments: args } : {}) },
  };
}

function makePromptResponse(id: number | string, text: string): JsonRpcResponse {
  return {
    jsonrpc: '2.0',
    id,
    result: {
      prompt: { name: 'test-prompt', description: 'A test prompt' },
      messages: [
        { role: 'user', content: { type: 'text', text } },
      ],
    },
  };
}

function extractResponseText(response: JsonRpcResponse): string {
  const result = response.result as Record<string, unknown>;
  // Prompt format: { messages: [{ content: { text } }] }
  if (result['messages']) {
    const messages = result['messages'] as Array<{ role: string; content: unknown }>;
    const c = messages[0].content;
    if (typeof c === 'string') return c;
    return (c as { text: string }).text;
  }
  // Tool result format: { content: [{ type: 'text', text }] }
  if (Array.isArray(result['content'])) {
    const parts = result['content'] as Array<{ type: string; text: string }>;
    return parts.map((p) => p.text).join('\n');
  }
  throw new Error(`Unexpected response format: ${JSON.stringify(result)}`);
}

// ── Tests ──

describe('Prompt section drill-down', () => {
  describe('onPromptGet hook', () => {
    it('returns null when no _resultId/_section params', async () => {
      const plugin = createContentPipelinePlugin();
      const ctx = createMockCtx();
      const request = makePromptGetRequest('my-prompt');

      const result = await plugin.onPromptGet!('my-prompt', request, ctx);
      expect(result).toBeNull();
    });

    it('intercepts _resultId + _section and returns cached section', async () => {
      const plugin = createContentPipelinePlugin();
      const ctx = createMockCtx();

      // Manually store sections in the shared store (simulating prior pipeline run)
      const sections: Section[] = [
        { id: 'overview', title: 'Overview', content: 'This is the overview section with lots of detail.' },
        { id: 'setup', title: 'Setup', content: 'Step-by-step setup instructions here.' },
        { id: 'api', title: 'API Reference', content: 'Full API documentation content.' },
      ];
      const store = new Map<string, { sections: Section[]; createdAt: number }>();
      store.set('pm-test123', { sections, createdAt: Date.now() });
      ctx.state.set('_contentPipeline_sections', store);

      // Drill into 'setup' section
      const request = makePromptGetRequest('my-prompt', { _resultId: 'pm-test123', _section: 'setup' });
      const result = await plugin.onPromptGet!('my-prompt', request, ctx);

      expect(result).not.toBeNull();
      const text = extractResponseText(result!);
      expect(text).toBe('Step-by-step setup instructions here.');
    });

    it('returns error for expired/invalid _resultId', async () => {
      const plugin = createContentPipelinePlugin();
      const ctx = createMockCtx();

      const request = makePromptGetRequest('my-prompt', { _resultId: 'pm-nonexistent', _section: 'setup' });
      const result = await plugin.onPromptGet!('my-prompt', request, ctx);

      expect(result).not.toBeNull();
      const text = extractResponseText(result!);
      expect(text).toContain('Cached result not found');
    });

    it('returns error for unknown section ID', async () => {
      const plugin = createContentPipelinePlugin();
      const ctx = createMockCtx();

      const sections: Section[] = [
        { id: 'intro', title: 'Intro', content: 'Introduction content.' },
      ];
      const store = new Map<string, { sections: Section[]; createdAt: number }>();
      store.set('pm-abc', { sections, createdAt: Date.now() });
      ctx.state.set('_contentPipeline_sections', store);

      const request = makePromptGetRequest('my-prompt', { _resultId: 'pm-abc', _section: 'nonexistent' });
      const result = await plugin.onPromptGet!('my-prompt', request, ctx);

      expect(result).not.toBeNull();
      const text = extractResponseText(result!);
      expect(text).toContain("Section 'nonexistent' not found");
      expect(text).toContain('intro');
    });
  });

  describe('shared section store with tools', () => {
    it('prompt drill-down can access sections stored by tool pipeline', async () => {
      const plugin = createContentPipelinePlugin();
      const ctx = createMockCtx();

      // Simulate a tool call that stored sections (onToolCallAfter would do this)
      const toolSections: Section[] = [
        { id: 'nodes', title: 'Nodes', content: 'List of all flow nodes.' },
        { id: 'connections', title: 'Connections', content: 'Node connection map.' },
      ];
      const store = new Map<string, { sections: Section[]; createdAt: number }>();
      store.set('pm-fromtool', { sections: toolSections, createdAt: Date.now() });
      ctx.state.set('_contentPipeline_sections', store);

      // Prompt drill-down using a resultId that came from a tool result
      const request = makePromptGetRequest('any-prompt', { _resultId: 'pm-fromtool', _section: 'connections' });
      const result = await plugin.onPromptGet!('any-prompt', request, ctx);

      expect(result).not.toBeNull();
      const text = extractResponseText(result!);
      expect(text).toBe('Node connection map.');
    });
  });

  describe('large prompt → section-split → drill-down (full cycle)', () => {
    it('markdown prompt is split into sections with TOC', async () => {
      const largePrompt = generateLargePrompt(5);
      expect(largePrompt.length).toBeGreaterThan(10000);

      // Run through section-split stage directly
      const sectionSplit = (await import('../src/proxymodel/stages/section-split.js')).default;
      const mockLog: StageLogger = { debug: vi.fn(), info: vi.fn(), warn: vi.fn(), error: vi.fn() };
      const mockLlm: LLMProvider = { complete: async () => '', available: () => false };
      const mockCache: CacheProvider = {
        getOrCompute: async (_, c) => c(),
        hash: (s) => s.slice(0, 8),
        get: async () => null,
        set: async () => {},
      };

      const result = await sectionSplit(largePrompt, {
        contentType: 'prompt',
        sourceName: 'test-prompt',
        projectName: 'test',
        sessionId: 'sess-1',
        originalContent: largePrompt,
        llm: mockLlm,
        cache: mockCache,
        log: mockLog,
        getSystemPrompt: async (_n: string, fb: string) => fb,
        config: { minSectionSize: 500 },
      });

      // Should have produced sections
      expect(result.sections).toBeDefined();
      expect(result.sections!.length).toBeGreaterThanOrEqual(3);

      // TOC should list sections
      expect(result.content).toContain('sections');
      expect(result.content).toContain('Use section parameter');

      // Original was ~16K, TOC should be much shorter
      expect(result.content.length).toBeLessThan(largePrompt.length);

      // Log the transformation for visibility
      process.stderr.write(`\n--- Original prompt size: ${largePrompt.length} chars ---\n`);
      process.stderr.write(`--- Transformed TOC size: ${result.content.length} chars ---\n`);
      process.stderr.write(`--- Sections: ${result.sections!.map((s) => `${s.id} (${s.content.length})`).join(', ')} ---\n`);
    });

    it('JSON prompt is split into per-item sections', async () => {
      const largeJson = generateLargeJsonPrompt(8);
      expect(largeJson.length).toBeGreaterThan(5000);

      const sectionSplit = (await import('../src/proxymodel/stages/section-split.js')).default;
      const mockLog: StageLogger = { debug: vi.fn(), info: vi.fn(), warn: vi.fn(), error: vi.fn() };
      const mockLlm: LLMProvider = { complete: async () => '', available: () => false };
      const mockCache: CacheProvider = {
        getOrCompute: async (_, c) => c(),
        hash: (s) => s.slice(0, 8),
        get: async () => null,
        set: async () => {},
      };

      const result = await sectionSplit(largeJson, {
        contentType: 'prompt',
        sourceName: 'json-prompt',
        projectName: 'test',
        sessionId: 'sess-1',
        originalContent: largeJson,
        llm: mockLlm,
        cache: mockCache,
        log: mockLog,
        getSystemPrompt: async (_n: string, fb: string) => fb,
        config: { minSectionSize: 200 },
      });

      expect(result.sections).toBeDefined();
      expect(result.sections!.length).toBeGreaterThanOrEqual(4);
      // JSON array items should use name as section id
      expect(result.sections!.some((s) => s.id.startsWith('item-'))).toBe(true);

      process.stderr.write(`\n--- JSON prompt: ${largeJson.length} chars → ${result.sections!.length} sections ---\n`);
    });

    it('full cycle: pipeline → store → drill-down → content', async () => {
      const plugin = createContentPipelinePlugin();
      const ctx = createMockCtx();

      // Step 1: Generate a large prompt response
      const largePrompt = generateLargePrompt(4);
      const originalResponse = makePromptResponse(1, largePrompt);

      process.stderr.write(`\n=== Full cycle test ===\n`);
      process.stderr.write(`Original prompt size: ${largePrompt.length} chars\n`);

      // Step 2: Simulate the router post-processing (processContent + store)
      const pipelineResult = await ctx.processContent('test-prompt', largePrompt, 'prompt');

      expect(pipelineResult.sections).toBeDefined();
      expect(pipelineResult.sections!.length).toBeGreaterThan(1);

      // Store sections like the router would
      const resultId = `pm-${Date.now().toString(36)}`;
      const store = new Map<string, { sections: Section[]; createdAt: number }>();
      store.set(resultId, { sections: pipelineResult.sections!, createdAt: Date.now() });
      ctx.state.set('_contentPipeline_sections', store);

      process.stderr.write(`Pipeline output (TOC):\n${pipelineResult.content}\n`);
      process.stderr.write(`Stored ${pipelineResult.sections!.length} sections under ${resultId}\n`);

      // Step 3: Drill into first section
      const firstSection = pipelineResult.sections![0];
      const drillRequest = makePromptGetRequest('test-prompt', {
        _resultId: resultId,
        _section: firstSection.id,
      });
      const drillResult = await plugin.onPromptGet!('test-prompt', drillRequest, ctx);

      expect(drillResult).not.toBeNull();
      const drillText = extractResponseText(drillResult!);
      expect(drillText).toBe(firstSection.content);

      process.stderr.write(`Drill-down into "${firstSection.id}": ${drillText.length} chars\n`);

      // Step 4: Drill into last section
      const lastSection = pipelineResult.sections![pipelineResult.sections!.length - 1];
      const drillRequest2 = makePromptGetRequest('test-prompt', {
        _resultId: resultId,
        _section: lastSection.id,
      });
      const drillResult2 = await plugin.onPromptGet!('test-prompt', drillRequest2, ctx);

      expect(drillResult2).not.toBeNull();
      const drillText2 = extractResponseText(drillResult2!);
      expect(drillText2).toBe(lastSection.content);

      process.stderr.write(`Drill-down into "${lastSection.id}": ${drillText2.length} chars\n`);
      process.stderr.write(`=== Full cycle complete ===\n`);
    });
  });
});