mcpctl/src/mcplocal/tests/session-gate.test.ts

import { describe, it, expect } from 'vitest';
import { SessionGate } from '../src/gate/session-gate.js';
import type { TagMatchResult, PromptIndexEntry } from '../src/gate/tag-matcher.js';

function makeMatchResult(names: string[]): TagMatchResult {
  return {
    fullContent: names.map((name) => ({
      name,
      priority: 5,
      summary: null,
      chapters: null,
      content: `Content of ${name}`,
    })),
    indexOnly: [],
    remaining: [],
  };
}

describe('SessionGate', () => {
  it('creates a gated session when project is gated', () => {
    const gate = new SessionGate();
    gate.createSession('s1', true);
    expect(gate.isGated('s1')).toBe(true);
  });

  it('creates an ungated session when project is not gated', () => {
    const gate = new SessionGate();
    gate.createSession('s1', false);
    expect(gate.isGated('s1')).toBe(false);
  });

  it('unknown sessions are treated as ungated', () => {
    const gate = new SessionGate();
    expect(gate.isGated('nonexistent')).toBe(false);
  });

  it('getSession returns null for unknown sessions', () => {
    const gate = new SessionGate();
    expect(gate.getSession('nonexistent')).toBeNull();
  });

  it('getSession returns session state', () => {
    const gate = new SessionGate();
    gate.createSession('s1', true);
    const state = gate.getSession('s1');
    expect(state).not.toBeNull();
    expect(state!.gated).toBe(true);
    expect(state!.tags).toEqual([]);
    expect(state!.retrievedPrompts.size).toBe(0);
    expect(state!.briefing).toBeNull();
  });

  it('ungate marks session as ungated and records tags', () => {
    const gate = new SessionGate();
    gate.createSession('s1', true);

    gate.ungate('s1', ['zigbee', 'mqtt'], makeMatchResult(['prompt-a', 'prompt-b']));

    expect(gate.isGated('s1')).toBe(false);
    const state = gate.getSession('s1');
    expect(state!.tags).toEqual(['zigbee', 'mqtt']);
    expect(state!.retrievedPrompts.has('prompt-a')).toBe(true);
    expect(state!.retrievedPrompts.has('prompt-b')).toBe(true);
  });

  it('ungate appends tags on repeated calls', () => {
    const gate = new SessionGate();
    gate.createSession('s1', true);

    gate.ungate('s1', ['zigbee'], makeMatchResult(['p1']));
    gate.ungate('s1', ['mqtt'], makeMatchResult(['p2']));

    const state = gate.getSession('s1');
    expect(state!.tags).toEqual(['zigbee', 'mqtt']);
    expect(state!.retrievedPrompts.has('p1')).toBe(true);
    expect(state!.retrievedPrompts.has('p2')).toBe(true);
  });

  it('ungate is no-op for unknown sessions', () => {
    const gate = new SessionGate();
    // Should not throw
    gate.ungate('nonexistent', ['tag'], makeMatchResult(['p']));
  });

  it('addRetrievedPrompts records additional prompts', () => {
    const gate = new SessionGate();
    gate.createSession('s1', true);
    gate.ungate('s1', ['zigbee'], makeMatchResult(['p1']));

    gate.addRetrievedPrompts('s1', ['mqtt', 'lights'], ['p2', 'p3']);

    const state = gate.getSession('s1');
    expect(state!.tags).toEqual(['zigbee', 'mqtt', 'lights']);
    expect(state!.retrievedPrompts.has('p2')).toBe(true);
    expect(state!.retrievedPrompts.has('p3')).toBe(true);
  });

  it('addRetrievedPrompts is no-op for unknown sessions', () => {
    const gate = new SessionGate();
    gate.addRetrievedPrompts('nonexistent', ['tag'], ['p']);
  });

  it('filterAlreadySent removes already-sent prompts', () => {
    const gate = new SessionGate();
    gate.createSession('s1', true);
    gate.ungate('s1', ['zigbee'], makeMatchResult(['p1']));

    const prompts: PromptIndexEntry[] = [
      { name: 'p1', priority: 5, summary: 'already sent', chapters: null, content: 'x' },
      { name: 'p2', priority: 5, summary: 'new', chapters: null, content: 'y' },
    ];

    const filtered = gate.filterAlreadySent('s1', prompts);
    expect(filtered).toHaveLength(1);
    expect(filtered[0]!.name).toBe('p2');
  });

  it('filterAlreadySent returns all prompts for unknown sessions', () => {
    const gate = new SessionGate();
    const prompts: PromptIndexEntry[] = [
      { name: 'p1', priority: 5, summary: null, chapters: null, content: 'x' },
    ];

    const filtered = gate.filterAlreadySent('nonexistent', prompts);
    expect(filtered).toHaveLength(1);
  });

  it('removeSession cleans up state', () => {
    const gate = new SessionGate();
    gate.createSession('s1', true);
    expect(gate.getSession('s1')).not.toBeNull();

    gate.removeSession('s1');
    expect(gate.getSession('s1')).toBeNull();
    expect(gate.isGated('s1')).toBe(false);
  });

  it('removeSession is safe for unknown sessions', () => {
    const gate = new SessionGate();
    gate.removeSession('nonexistent'); // Should not throw
  });

  it('manages multiple sessions independently', () => {
    const gate = new SessionGate();
    gate.createSession('s1', true);
    gate.createSession('s2', false);

    expect(gate.isGated('s1')).toBe(true);
    expect(gate.isGated('s2')).toBe(false);

    gate.ungate('s1', ['zigbee'], makeMatchResult(['p1']));
    expect(gate.isGated('s1')).toBe(false);
    expect(gate.getSession('s2')!.tags).toEqual([]); // s2 untouched
  });

  describe('per-McpToken ungate cache', () => {
    it('new session from an already-ungated token starts ungated, with prior tags + prompts', () => {
      const gate = new SessionGate();
      gate.createSession('session-1', true, 'tokA');
      expect(gate.isGated('session-1')).toBe(true);

      gate.ungate('session-1', ['ops'], makeMatchResult(['runbook']), 'tokA');
      expect(gate.isTokenUngated('tokA')).toBe(true);

      // LiteLLM semantics: same token, brand-new session-id.
      gate.createSession('session-2', true, 'tokA');
      expect(gate.isGated('session-2')).toBe(false);
      const s2 = gate.getSession('session-2')!;
      expect(s2.tags).toContain('ops');
      expect(s2.retrievedPrompts.has('runbook')).toBe(true);
    });

    it('does not persist across tokens', () => {
      const gate = new SessionGate();
      gate.createSession('s1', true, 'tokA');
      gate.ungate('s1', ['ops'], makeMatchResult(['p']), 'tokA');

      // Different token → fresh gated session.
      gate.createSession('s2', true, 'tokB');
      expect(gate.isGated('s2')).toBe(true);
      expect(gate.isTokenUngated('tokB')).toBe(false);
    });

    it('is not triggered when no tokenSha is supplied (STDIO path)', () => {
      const gate = new SessionGate();
      gate.createSession('s1', true);
      gate.ungate('s1', ['ops'], makeMatchResult(['p']));

      // A second session with no token starts gated — STDIO semantics preserved.
      gate.createSession('s2', true);
      expect(gate.isGated('s2')).toBe(true);
    });

    it('honors the TTL window and expires', () => {
      const gate = new SessionGate(50); // 50ms TTL for the test
      gate.createSession('s1', true, 'tokA');
      gate.ungate('s1', ['ops'], makeMatchResult(['p']), 'tokA');
      expect(gate.isTokenUngated('tokA')).toBe(true);

      return new Promise<void>((resolve) => setTimeout(() => {
        expect(gate.isTokenUngated('tokA')).toBe(false);
        gate.createSession('s2', true, 'tokA');
        expect(gate.isGated('s2')).toBe(true);
        resolve();
      }, 70));
    });

    it('revokeToken clears the ungate entry immediately', () => {
      const gate = new SessionGate();
      gate.createSession('s1', true, 'tokA');
      gate.ungate('s1', ['ops'], makeMatchResult(['p']), 'tokA');
      expect(gate.isTokenUngated('tokA')).toBe(true);

      gate.revokeToken('tokA');
      expect(gate.isTokenUngated('tokA')).toBe(false);
      gate.createSession('s2', true, 'tokA');
      expect(gate.isGated('s2')).toBe(true);
    });

    it('empty-string tokenSha does not register an ungate entry', () => {
      const gate = new SessionGate();
      gate.createSession('s1', true, '');
      gate.ungate('s1', ['ops'], makeMatchResult(['p']), '');
      expect(gate.isTokenUngated('')).toBe(false);
    });
  });
});