flynn/src/gateway/session-bridge.test.ts

import { describe, it, expect, vi, beforeEach } from 'vitest';
import { SessionBridge } from './session-bridge.js';
import type { SessionBridgeConfig } from './session-bridge.js';

// Minimal mocks
const mockSession = {
  id: 'test',
  addMessage: vi.fn(),
  getHistory: vi.fn(() => []),
  clear: vi.fn(),
  replaceHistory: vi.fn(),
  getConfig: vi.fn((_key: string) => undefined as string | undefined),
  setConfig: vi.fn(),
  deleteConfig: vi.fn(),
};

const mockSessionManager = {
  getSession: vi.fn(() => mockSession),
  listSessions: vi.fn(() => []),
  transferSession: vi.fn(),
  closeSession: vi.fn(),
};

const mockModelClient = {
  chat: vi.fn(async () => ({
    content: 'test',
    stopReason: 'end_turn',
    usage: { inputTokens: 0, outputTokens: 0 },
  })),
};

const mockToolRegistry = {
  register: vi.fn(),
  get: vi.fn(),
  list: vi.fn(() => []),
  toAnthropicFormat: vi.fn(() => []),
  toOpenAIFormat: vi.fn(() => []),
};

const mockToolExecutor = {
  execute: vi.fn(),
};

function createBridge(): SessionBridge {
  return new SessionBridge({
    sessionManager: mockSessionManager as unknown as SessionBridgeConfig['sessionManager'],
    modelClient: mockModelClient,
    systemPrompt: 'test prompt',
    toolRegistry: mockToolRegistry as unknown as SessionBridgeConfig['toolRegistry'],
    toolExecutor: mockToolExecutor as unknown as SessionBridgeConfig['toolExecutor'],
  });
}

function createBridgeWithConfig(config: SessionBridgeConfig['config']): SessionBridge {
  return new SessionBridge({
    sessionManager: mockSessionManager as unknown as SessionBridgeConfig['sessionManager'],
    modelClient: mockModelClient,
    systemPrompt: 'test prompt',
    toolRegistry: mockToolRegistry as unknown as SessionBridgeConfig['toolRegistry'],
    toolExecutor: mockToolExecutor as unknown as SessionBridgeConfig['toolExecutor'],
    config,
  });
}

describe('SessionBridge', () => {
  beforeEach(() => {
    vi.clearAllMocks();
    mockSession.getConfig.mockImplementation((_key: string) => undefined);
  });

  it('connect assigns a connection ID', () => {
    const bridge = createBridge();
    const id = bridge.connect('conn-1');
    expect(id).toBe('conn-1');
    expect(bridge.connectionCount).toBe(1);
  });

  it('connect auto-generates ID when not provided', () => {
    const bridge = createBridge();
    const id = bridge.connect();
    expect(typeof id).toBe('string');
    expect(id.length).toBeGreaterThan(0);
  });

  it('getAgent returns agent for connected client', () => {
    const bridge = createBridge();
    bridge.connect('conn-1');
    const agent = bridge.getAgent('conn-1');
    expect(agent).toBeDefined();
  });

  it('getSessionId returns session ID for connected client', () => {
    const bridge = createBridge();
    bridge.connect('conn-1');
    expect(bridge.getSessionId('conn-1')).toBe('ws:conn-1');
  });

  it('disconnect removes client but preserves session', () => {
    const bridge = createBridge();
    bridge.connect('conn-1');
    bridge.disconnect('conn-1');
    expect(bridge.connectionCount).toBe(0);
    expect(bridge.getAgent('conn-1')).toBeUndefined();
  });

  it('tracks busy state', () => {
    const bridge = createBridge();
    bridge.connect('conn-1');
    expect(bridge.isBusy('conn-1')).toBe(false);
    bridge.setBusy('conn-1', true);
    expect(bridge.isBusy('conn-1')).toBe(true);
    bridge.setBusy('conn-1', false);
    expect(bridge.isBusy('conn-1')).toBe(false);
  });

  it('cancel returns false when no active operation exists', () => {
    const bridge = createBridge();
    bridge.connect('conn-1');

    expect(bridge.cancel('conn-1')).toBe(false);
  });

  it('cancel requests cancellation when connection is busy', () => {
    const bridge = createBridge();
    bridge.connect('conn-1');
    const agent = bridge.getAgent('conn-1');
    const cancelSpy = vi.spyOn(agent!, 'cancel');

    bridge.setBusy('conn-1', true);
    expect(bridge.cancel('conn-1')).toBe(true);
    expect(cancelSpy).toHaveBeenCalledTimes(1);
  });

  it('switchSession changes session for a connection', () => {
    const bridge = createBridge();
    bridge.connect('conn-1');
    expect(bridge.getSessionId('conn-1')).toBe('ws:conn-1');

    bridge.switchSession('conn-1', 'custom-session');
    expect(bridge.getSessionId('conn-1')).toBe('custom-session');
  });

  it('switchSession throws when busy', () => {
    const bridge = createBridge();
    bridge.connect('conn-1');
    bridge.setBusy('conn-1', true);

    expect(() => bridge.switchSession('conn-1', 'other')).toThrow('Cannot switch session while agent is busy');
  });

  it('switchSession throws for unknown connection', () => {
    const bridge = createBridge();
    expect(() => bridge.switchSession('unknown', 'other')).toThrow('Unknown connection');
  });

  it('listSessions groups connections by session', () => {
    const bridge = createBridge();
    bridge.connect('conn-1');
    bridge.connect('conn-2');
    // Switch conn-2 to share conn-1's session
    bridge.switchSession('conn-2', 'ws:conn-1');

    const sessions = bridge.listSessions();
    expect(sessions).toEqual([{ sessionId: 'ws:conn-1', connections: 2 }]);
  });

  it('shared sessions keep agent alive when one client disconnects', () => {
    const bridge = createBridge();
    bridge.connect('conn-1');
    bridge.connect('conn-2');
    bridge.switchSession('conn-2', 'ws:conn-1');

    bridge.disconnect('conn-1');
    // conn-2 still has the session
    expect(bridge.getAgent('conn-2')).toBeDefined();
    expect(bridge.connectionCount).toBe(1);
  });

  it('loads model tier from per-session config when creating a session agent', () => {
    mockSession.getConfig.mockImplementation((key: string) => (key === 'modelTier' ? 'local' : undefined));

    const bridge = createBridgeWithConfig({
      agents: {
        primary_tier: 'default',
        delegation: {
          compaction: 'fast',
          memory_extraction: 'fast',
          classification: 'fast',
          tool_summarisation: 'fast',
          complex_reasoning: 'complex',
        },
        max_delegation_depth: 3,
      },
      compaction: { enabled: false },
      models: { default: { provider: 'anthropic', model: 'claude-3-haiku' } },
    } as any);

    bridge.connect('conn-tier');
    const agent = bridge.getAgent('conn-tier');
    expect(agent?.getModelTier()).toBe('local');
  });

  it('keeps different sessions isolated by persisted model tier', () => {
    const sessionById: Record<string, any> = {};
    const localSessionManager = {
      ...mockSessionManager,
      getSession: vi.fn((frontend: string, sessionId: string) => {
        const fullId = `${frontend}:${sessionId}`;
        if (!sessionById[fullId]) {
          const tier = fullId === 'ws:ws:conn-a' ? 'fast' : 'complex';
          sessionById[fullId] = {
            ...mockSession,
            id: fullId,
            getConfig: vi.fn((key: string) => (key === 'modelTier' ? tier : undefined)),
          };
        }
        return sessionById[fullId];
      }),
    };

    const bridge = new SessionBridge({
      sessionManager: localSessionManager as unknown as SessionBridgeConfig['sessionManager'],
      modelClient: mockModelClient,
      systemPrompt: 'test prompt',
      toolRegistry: mockToolRegistry as unknown as SessionBridgeConfig['toolRegistry'],
      toolExecutor: mockToolExecutor as unknown as SessionBridgeConfig['toolExecutor'],
      config: {
        agents: {
          primary_tier: 'default',
          delegation: {
            compaction: 'fast',
            memory_extraction: 'fast',
            classification: 'fast',
            tool_summarisation: 'fast',
            complex_reasoning: 'complex',
          },
          max_delegation_depth: 3,
        },
        compaction: { enabled: false },
        models: { default: { provider: 'anthropic', model: 'claude-3-haiku' } },
      } as any,
    });

    bridge.connect('conn-a');
    bridge.connect('conn-b');
    const agentA = bridge.getAgent('conn-a');
    const agentB = bridge.getAgent('conn-b');

    expect(agentA?.getModelTier()).toBe('fast');
    expect(agentB?.getModelTier()).toBe('complex');
  });
});