flynn/src/gateway/handlers/agent.test.ts

import { describe, it, expect, vi, beforeEach } from 'vitest';
import type { GatewayEvent, GatewayRequest, OutboundMessage } from '../protocol.js';
import { LaneQueue, LaneQueueRejectedError } from '../lane-queue.js';
import { createAgentHandlers } from './agent.js';
import type { AgentHandlerDeps } from './agent.js';
import { CommandRegistry, registerBuiltinCommands } from '../../commands/index.js';
import { initAuditLogger } from '../../audit/index.js';

describe('createAgentHandlers command fast-path', () => {
  const mockAgent = {
    process: vi.fn(async () => 'agent response'),
    consumeContextAlert: vi.fn(() => undefined as unknown),
    getContextBudget: vi.fn(() => ({
      estimatedTokens: 100,
      contextWindow: 200000,
      remainingTokens: 199900,
      usagePct: 0.05,
      thresholdPct: 80,
      thresholdTokens: 160000,
      shouldCompact: false,
    })),
    getUsage: vi.fn(() => ({
      primary: { inputTokens: 10, outputTokens: 5, calls: 1 },
      delegation: {},
      total: { inputTokens: 10, outputTokens: 5, calls: 1, estimatedCost: 0 },
    })),
    getModelTier: vi.fn(() => 'default'),
    setModelTier: vi.fn(),
    compact: vi.fn(async () => null),
    reset: vi.fn(),
  };

  const sessionBridge = {
    getAgent: vi.fn(() => mockAgent),
    getSessionId: vi.fn(() => 'ws:conn-1'),
    setBusy: vi.fn(),
    setOnToolUse: vi.fn(),
    isBusy: vi.fn(() => false),
    cancel: vi.fn(() => true),
  };

  const sessionManager = {
    getSessionConfig: vi.fn(),
    setSessionConfig: vi.fn(),
    deleteSessionConfig: vi.fn(),
  };

  const commandRegistry = new CommandRegistry();
  registerBuiltinCommands(commandRegistry);
  const mockAuditLogger = {
    userAction: vi.fn(),
    queuePreempt: vi.fn(),
  };

  const handlers = createAgentHandlers({
    sessionBridge: sessionBridge as unknown as AgentHandlerDeps['sessionBridge'],
    laneQueue: new LaneQueue(),
    sessionManager: sessionManager as unknown as AgentHandlerDeps['sessionManager'],
    commandRegistry,
  });

  beforeEach(() => {
    vi.clearAllMocks();
    initAuditLogger(mockAuditLogger as any);
    mockAgent.process.mockResolvedValue('agent response');
    mockAgent.compact.mockResolvedValue(null);
  });

  it('handles known commands without calling agent.process', async () => {
    const sent: OutboundMessage[] = [];
    const send = vi.fn((msg: OutboundMessage) => sent.push(msg));
    const req: GatewayRequest = {
      id: 1,
      method: 'agent.send',
      params: { message: '/usage', connectionId: 'conn-1' },
    };

    await handlers['agent.send'](req, send);

    expect(mockAgent.process).not.toHaveBeenCalled();
    expect(sent).toHaveLength(1);
    const event = sent[0] as GatewayEvent;
    expect(event.event).toBe('done');
    expect((event.data as { content: string }).content).toContain('Token Usage');
  });

  it('handles /context command via fast-path', async () => {
    const sent: OutboundMessage[] = [];
    const send = vi.fn((msg: OutboundMessage) => sent.push(msg));
    const req: GatewayRequest = {
      id: 8,
      method: 'agent.send',
      params: { message: '/context', connectionId: 'conn-1' },
    };

    await handlers['agent.send'](req, send);

    expect(mockAgent.process).not.toHaveBeenCalled();
    expect(sent).toHaveLength(1);
    const event = sent[0] as GatewayEvent;
    expect(event.event).toBe('done');
    expect((event.data as { content: string }).content).toContain('Context Usage');
    expect((event.data as { content: string }).content).toContain('Compaction threshold');
  });

  it('handles metadata commands via fast-path', async () => {
    const sent: OutboundMessage[] = [];
    const send = vi.fn((msg: OutboundMessage) => sent.push(msg));
    const req: GatewayRequest = {
      id: 2,
      method: 'agent.send',
      params: {
        message: '/reset',
        connectionId: 'conn-1',
        metadata: { isCommand: true, command: 'reset' },
      },
    };

    await handlers['agent.send'](req, send);

    expect(mockAgent.reset).toHaveBeenCalledOnce();
    expect(sessionManager.deleteSessionConfig).toHaveBeenCalledWith('ws', 'ws:conn-1', 'modelTier');
    expect(mockAgent.process).not.toHaveBeenCalled();
    expect(((sent[0] as GatewayEvent).data as { content: string }).content).toContain('Session reset.');
  });

  it('handles /model command via fast-path and persists session tier', async () => {
    const sent: OutboundMessage[] = [];
    const send = vi.fn((msg: OutboundMessage) => sent.push(msg));
    const req: GatewayRequest = {
      id: 4,
      method: 'agent.send',
      params: {
        message: '/model fast',
        connectionId: 'conn-1',
        metadata: { isCommand: true, command: 'model', commandArgs: 'fast' },
      },
    };

    await handlers['agent.send'](req, send);

    expect(mockAgent.setModelTier).toHaveBeenCalledWith('fast');
    expect(sessionManager.setSessionConfig).toHaveBeenCalledWith('ws', 'ws:conn-1', 'modelTier', 'fast');
    expect(mockAgent.process).not.toHaveBeenCalled();
    expect(((sent[0] as GatewayEvent).data as { content: string }).content).toContain('Switched to model tier: fast');
  });

  it('handles /model <tier> <provider/model> in gateway sessions', async () => {
    const sent: OutboundMessage[] = [];
    const send = vi.fn((msg: OutboundMessage) => sent.push(msg));
    const modelRouter = {
      setClient: vi.fn(),
      setTierStrict: vi.fn(),
    };
    const handlersWithRouter = createAgentHandlers({
      sessionBridge: sessionBridge as unknown as AgentHandlerDeps['sessionBridge'],
      laneQueue: new LaneQueue(),
      sessionManager: sessionManager as unknown as AgentHandlerDeps['sessionManager'],
      commandRegistry,
      modelRouter: modelRouter as unknown as AgentHandlerDeps['modelRouter'],
      runtimeConfig: {
        models: {
          default: { provider: 'anthropic', model: 'claude-sonnet-4' },
          fallback_chain: ['anthropic'],
        },
      } as unknown as AgentHandlerDeps['runtimeConfig'],
    });
    const req: GatewayRequest = {
      id: 9,
      method: 'agent.send',
      params: {
        message: '/model default github/gpt-5-mini',
        connectionId: 'conn-1',
        metadata: { isCommand: true, command: 'model', commandArgs: 'default github/gpt-5-mini' },
      },
    };

    await handlersWithRouter['agent.send'](req, send);

    expect(modelRouter.setClient).toHaveBeenCalledWith('default', expect.anything(), 'github/gpt-5-mini');
    expect(modelRouter.setTierStrict).toHaveBeenCalledWith('default', true);
    expect(mockAgent.setModelTier).toHaveBeenCalledWith('default');
    expect(((sent[0] as GatewayEvent).data as { content: string }).content).toContain('Set default to: github/gpt-5-mini');
  });

  it('handles /model <tier> reset in gateway sessions', async () => {
    const sent: OutboundMessage[] = [];
    const send = vi.fn((msg: OutboundMessage) => sent.push(msg));
    const modelRouter = {
      setClient: vi.fn(),
      setTierStrict: vi.fn(),
    };
    const handlersWithRouter = createAgentHandlers({
      sessionBridge: sessionBridge as unknown as AgentHandlerDeps['sessionBridge'],
      laneQueue: new LaneQueue(),
      sessionManager: sessionManager as unknown as AgentHandlerDeps['sessionManager'],
      commandRegistry,
      modelRouter: modelRouter as unknown as AgentHandlerDeps['modelRouter'],
      runtimeConfig: {
        models: {
          default: { provider: 'anthropic', model: 'claude-sonnet-4' },
          fallback_chain: ['anthropic'],
        },
      } as unknown as AgentHandlerDeps['runtimeConfig'],
    });
    const req: GatewayRequest = {
      id: 10,
      method: 'agent.send',
      params: {
        message: '/model default reset',
        connectionId: 'conn-1',
        metadata: { isCommand: true, command: 'model', commandArgs: 'default reset' },
      },
    };

    await handlersWithRouter['agent.send'](req, send);

    expect(modelRouter.setClient).toHaveBeenCalledWith('default', expect.anything(), 'anthropic/claude-sonnet-4');
    expect(modelRouter.setTierStrict).toHaveBeenCalledWith('default', false);
    expect(mockAgent.setModelTier).toHaveBeenCalledWith('default');
    expect(((sent[0] as GatewayEvent).data as { content: string }).content).toContain('Reset default to: anthropic/claude-sonnet-4');
  });

  it('handles /stop command via fast-path and requests cancellation', async () => {
    const sent: OutboundMessage[] = [];
    const send = vi.fn((msg: OutboundMessage) => sent.push(msg));
    const req: GatewayRequest = {
      id: 12,
      method: 'agent.send',
      params: {
        message: '/stop',
        connectionId: 'conn-1',
        metadata: { isCommand: true, command: 'stop' },
      },
    };

    await handlers['agent.send'](req, send);

    expect(sessionBridge.cancel).toHaveBeenCalledWith('conn-1');
    expect(mockAgent.process).not.toHaveBeenCalled();
    expect(((sent[0] as GatewayEvent).data as { content: string }).content).toContain('Cancellation requested');
  });

  it('falls through to agent.process for unknown commands', async () => {
    const sent: OutboundMessage[] = [];
    const send = vi.fn((msg: OutboundMessage) => sent.push(msg));
    const req: GatewayRequest = {
      id: 3,
      method: 'agent.send',
      params: { message: '/not-a-real-command', connectionId: 'conn-1' },
    };

    await handlers['agent.send'](req, send);

    expect(mockAgent.process).toHaveBeenCalledWith('/not-a-real-command', undefined);
    expect((sent[0] as GatewayEvent).event).toBe('done');
    expect(((sent[0] as GatewayEvent).data as { content: string }).content).toBe('agent response');
  });

  it('handles /approvals command via fast-path when hook engine is available', async () => {
    const sent: OutboundMessage[] = [];
    const send = vi.fn((msg: OutboundMessage) => sent.push(msg));
    const hookEngine = {
      getPendingConfirmations: vi.fn(() => []),
      resolveConfirmation: vi.fn(() => false),
    };
    const handlersWithHooks = createAgentHandlers({
      sessionBridge: sessionBridge as unknown as AgentHandlerDeps['sessionBridge'],
      laneQueue: new LaneQueue(),
      sessionManager: sessionManager as unknown as AgentHandlerDeps['sessionManager'],
      commandRegistry,
      hookEngine: hookEngine as unknown as AgentHandlerDeps['hookEngine'],
    });

    const req: GatewayRequest = {
      id: 11,
      method: 'agent.send',
      params: {
        message: '/approvals',
        connectionId: 'conn-1',
        metadata: { isCommand: true, command: 'approvals' },
      },
    };

    await handlersWithHooks['agent.send'](req, send);

    expect(mockAgent.process).not.toHaveBeenCalled();
    expect(hookEngine.getPendingConfirmations).toHaveBeenCalledWith({ sessionId: 'ws:conn-1' });
    expect(((sent[0] as GatewayEvent).data as { content: string }).content).toContain('No pending approvals');
  });

  it('emits user.action audit events for gateway requests', async () => {
    const sent: OutboundMessage[] = [];
    const send = vi.fn((msg: OutboundMessage) => sent.push(msg));
    const req: GatewayRequest = {
      id: 7,
      method: 'agent.send',
      params: { message: 'hello there', connectionId: 'conn-1' },
    };

    await handlers['agent.send'](req, send);

    expect(mockAuditLogger.userAction).toHaveBeenCalledWith(
      expect.objectContaining({
        source: 'gateway',
        action_type: 'message',
        sender: 'conn-1',
      }),
    );
  });

  it('handles /queue command via fast-path and persists queue session config', async () => {
    const sent: OutboundMessage[] = [];
    const send = vi.fn((msg: OutboundMessage) => sent.push(msg));
    const req: GatewayRequest = {
      id: 5,
      method: 'agent.send',
      params: {
        message: '/queue set mode followup',
        connectionId: 'conn-1',
        metadata: { isCommand: true, command: 'queue', commandArgs: 'set mode followup' },
      },
    };

    await handlers['agent.send'](req, send);

    expect(sessionManager.setSessionConfig).toHaveBeenCalledWith('ws', 'ws:conn-1', 'queue.mode', 'followup');
    expect(mockAgent.process).not.toHaveBeenCalled();
    expect(((sent[0] as GatewayEvent).data as { content: string }).content).toContain('Set queue.mode=followup');
  });

  it('emits context_warning event before done when orchestrator reports an alert', async () => {
    mockAgent.consumeContextAlert.mockReturnValueOnce({
      level: 'checkpoint',
      message: 'Context usage is 86.0% (172000/200000 estimated tokens).',
      budget: {
        estimatedTokens: 172000,
        contextWindow: 200000,
        remainingTokens: 28000,
        usagePct: 86,
        thresholdPct: 80,
        thresholdTokens: 160000,
        shouldCompact: true,
      },
      actions: {
        checkpointSaved: true,
        autoCompacted: false,
        checkpointNamespace: 'session/checkpoints/ws/conn-1',
      },
    });

    const sent: OutboundMessage[] = [];
    const send = vi.fn((msg: OutboundMessage) => sent.push(msg));

    await handlers['agent.send']({
      id: 6,
      method: 'agent.send',
      params: { message: 'hello', connectionId: 'conn-1' },
    }, send);

    expect(sent).toHaveLength(2);
    expect((sent[0] as GatewayEvent).event).toBe('context_warning');
    expect((sent[1] as GatewayEvent).event).toBe('done');
  });
});

describe('createAgentHandlers queue policy resolution', () => {
  const mockAuditLogger = {
    userAction: vi.fn(),
    queuePreempt: vi.fn(),
  };

  beforeEach(() => {
    vi.clearAllMocks();
    initAuditLogger(mockAuditLogger as any);
  });

  it('passes resolved per-request queue policy into lane enqueue', async () => {
    const mockAgent = {
      process: vi.fn(async () => 'ok'),
      consumeContextAlert: vi.fn(() => undefined),
      getContextBudget: vi.fn(() => ({
        estimatedTokens: 0,
        contextWindow: 128000,
        remainingTokens: 128000,
        usagePct: 0,
        thresholdPct: 80,
        thresholdTokens: 102400,
        shouldCompact: false,
      })),
      getUsage: vi.fn(() => ({
        primary: { inputTokens: 0, outputTokens: 0, calls: 0 },
        delegation: {},
        total: { inputTokens: 0, outputTokens: 0, calls: 0, estimatedCost: 0 },
      })),
      getModelTier: vi.fn(() => 'default'),
      setModelTier: vi.fn(),
      compact: vi.fn(async () => null),
      reset: vi.fn(),
      setOnToolUse: vi.fn(),
    };

    const sessionBridge = {
      getAgent: vi.fn(() => mockAgent),
      getSessionId: vi.fn(() => 'ws:s1'),
      setBusy: vi.fn(),
      setOnToolUse: vi.fn(),
      isBusy: vi.fn(() => false),
    };

    const laneQueue = {
      enqueue: vi.fn(async (_laneId: string, work: () => Promise<unknown>) => work()),
      cancel: vi.fn(),
    } as unknown as LaneQueue;

    const resolveQueuePolicy = vi.fn(() => ({ mode: 'steer_backlog' as const, cap: 3, debounceMs: 25 }));

    const handlers = createAgentHandlers({
      sessionBridge: sessionBridge as unknown as AgentHandlerDeps['sessionBridge'],
      laneQueue,
      resolveQueuePolicy,
    });

    const sent: OutboundMessage[] = [];
    const send = vi.fn((msg: OutboundMessage) => sent.push(msg));

    await handlers['agent.send']({
      id: 1,
      method: 'agent.send',
      params: { message: 'hello', connectionId: 'conn-1' },
    }, send);

    expect(resolveQueuePolicy).toHaveBeenCalledWith({
      laneId: 'ws:s1',
      sessionId: 'ws:s1',
      connectionId: 'conn-1',
      channel: 'ws',
    });
    expect((laneQueue.enqueue as unknown as ReturnType<typeof vi.fn>).mock.calls[0][2]).toEqual({
      mode: 'steer_backlog',
      cap: 3,
      debounceMs: 25,
    });
  });

  it('emits structured queue error events for lane rejections', async () => {
    const sessionBridge = {
      getAgent: vi.fn(() => ({
        process: vi.fn(async () => 'ok'),
        consumeContextAlert: vi.fn(() => undefined),
        getContextBudget: vi.fn(() => ({
          estimatedTokens: 0,
          contextWindow: 128000,
          remainingTokens: 128000,
          usagePct: 0,
          thresholdPct: 80,
          thresholdTokens: 102400,
          shouldCompact: false,
        })),
        getUsage: vi.fn(() => ({
          primary: { inputTokens: 0, outputTokens: 0, calls: 0 },
          delegation: {},
          total: { inputTokens: 0, outputTokens: 0, calls: 0, estimatedCost: 0 },
        })),
        getModelTier: vi.fn(() => 'default'),
        setModelTier: vi.fn(),
        compact: vi.fn(async () => null),
        reset: vi.fn(),
      })),
      getSessionId: vi.fn(() => 'ws:s1'),
      setBusy: vi.fn(),
      setOnToolUse: vi.fn(),
      isBusy: vi.fn(() => false),
    };

    const laneQueue = {
      enqueue: vi.fn(async () => {
        throw new LaneQueueRejectedError({
          code: 'overflow',
          laneId: 'ws:s1',
          mode: 'followup',
          overflow: 'drop_new',
          droppedCount: 1,
          message: 'Lane queue full (drop_new)',
        });
      }),
      cancel: vi.fn(),
    } as unknown as LaneQueue;

    const handlers = createAgentHandlers({
      sessionBridge: sessionBridge as unknown as AgentHandlerDeps['sessionBridge'],
      laneQueue,
    });

    const sent: OutboundMessage[] = [];
    const send = vi.fn((msg: OutboundMessage) => sent.push(msg));

    await handlers['agent.send']({
      id: 6,
      method: 'agent.send',
      params: { message: 'hello', connectionId: 'conn-1' },
    }, send);

    expect(sent).toHaveLength(1);
    const event = sent[0] as GatewayEvent;
    expect(event.event).toBe('error');
    expect((event.data as { code: number }).code).toBe(3);
    expect((event.data as { queue?: { code: string } }).queue?.code).toBe('overflow');
  });

  it('requests active-session cancellation when interrupt mode receives a new message', async () => {
    const mockAgent = {
      process: vi.fn(async () => 'ok'),
      consumeContextAlert: vi.fn(() => undefined),
      getContextBudget: vi.fn(() => ({
        estimatedTokens: 0,
        contextWindow: 128000,
        remainingTokens: 128000,
        usagePct: 0,
        thresholdPct: 80,
        thresholdTokens: 102400,
        shouldCompact: false,
      })),
      getUsage: vi.fn(() => ({
        primary: { inputTokens: 0, outputTokens: 0, calls: 0 },
        delegation: {},
        total: { inputTokens: 0, outputTokens: 0, calls: 0, estimatedCost: 0 },
      })),
      getModelTier: vi.fn(() => 'default'),
      setModelTier: vi.fn(),
      compact: vi.fn(async () => null),
      reset: vi.fn(),
    };

    const sessionBridge = {
      getAgent: vi.fn(() => mockAgent),
      getSessionId: vi.fn(() => 'ws:s1'),
      setBusy: vi.fn(),
      setOnToolUse: vi.fn(),
      isBusy: vi.fn(() => false),
      cancelSession: vi.fn(() => true),
      cancel: vi.fn(() => true),
    };

    const laneQueue = {
      enqueue: vi.fn(async (_laneId: string, work: () => Promise<unknown>) => work()),
      cancel: vi.fn(),
      isProcessing: vi.fn(() => true),
    } as unknown as LaneQueue;

    const handlers = createAgentHandlers({
      sessionBridge: sessionBridge as unknown as AgentHandlerDeps['sessionBridge'],
      laneQueue,
      resolveQueuePolicy: vi.fn(() => ({ mode: 'interrupt' as const })),
    });

    const sent: OutboundMessage[] = [];
    const send = vi.fn((msg: OutboundMessage) => sent.push(msg));

    await handlers['agent.send']({
      id: 7,
      method: 'agent.send',
      params: { message: 'newest', connectionId: 'conn-1' },
    }, send);

    expect(sessionBridge.cancelSession).toHaveBeenCalledWith('ws:s1');
    expect(sessionBridge.cancel).not.toHaveBeenCalled();
    expect(mockAuditLogger.queuePreempt).toHaveBeenCalledWith(expect.objectContaining({
      session_id: 'ws:s1',
      lane_id: 'ws:s1',
      request_id: '7',
      mode: 'interrupt',
      cancelled_active_run: true,
    }));
    expect((sent[0] as GatewayEvent).event).toBe('content');
    expect(((sent[0] as GatewayEvent).data as { text: string }).text).toContain('Interrupt mode');
    expect((sent[1] as GatewayEvent).event).toBe('done');
  });
});