feat(gateway): add WebSocket gateway with JSON-RPC protocol and auth

Phase 2 of the Flynn roadmap. Adds a WebSocket gateway server that starts alongside the Telegram bot, providing real-time API access to the agent, sessions, and tools. Protocol: JSON-RPC-like (request/response/event) over WebSocket. 8 methods: agent.send, agent.cancel, sessions.list, sessions.history, sessions.create, tools.list, tools.invoke, system.health. Auth: Bearer token + Tailscale identity header support. Session bridge: per-connection agent instances with shared model router. New files: src/gateway/ (protocol, router, server, auth, session-bridge, handlers for agent/sessions/tools/system). 57 new tests (181 total), typecheck clean.
2026-02-05 19:11:25 -08:00
parent ad7fc241f1
commit f30a8bc318
21 changed files with 1878 additions and 2 deletions
@@ -0,0 +1,271 @@
+import { describe, it, expect, vi, beforeEach } from 'vitest';
+import { createSystemHandlers } from './system.js';
+import { createSessionHandlers } from './sessions.js';
+import { createToolHandlers } from './tools.js';
+import { createAgentHandlers } from './agent.js';
+import { ErrorCode } from '../protocol.js';
+import type { GatewayRequest, GatewayResponse, GatewayError, GatewayEvent, OutboundMessage } from '../protocol.js';
+
+describe('system handlers', () => {
+  const deps = {
+    startTime: Date.now() - 60_000,
+    version: '0.1.0',
+    getSessionCount: () => 3,
+    getToolCount: () => 6,
+    getConnectionCount: () => 2,
+  };
+  const handlers = createSystemHandlers(deps);
+
+  it('system.health returns status info', async () => {
+    const req: GatewayRequest = { id: 1, method: 'system.health' };
+    const result = await handlers['system.health'](req) as GatewayResponse;
+
+    expect(result.id).toBe(1);
+    const r = result.result as Record<string, unknown>;
+    expect(r.status).toBe('ok');
+    expect(r.version).toBe('0.1.0');
+    expect(r.sessions).toBe(3);
+    expect(r.tools).toBe(6);
+    expect(r.connections).toBe(2);
+    expect(typeof r.uptime).toBe('number');
+    expect(r.uptime).toBeGreaterThanOrEqual(59);
+  });
+});
+
+describe('session handlers', () => {
+  const mockHistory = [
+    { role: 'user' as const, content: 'hello' },
+    { role: 'assistant' as const, content: 'hi' },
+  ];
+
+  const mockSession = {
+    id: 'ws:test',
+    addMessage: vi.fn(),
+    getHistory: vi.fn(() => mockHistory),
+    clear: vi.fn(),
+  };
+
+  const mockSessionManager = {
+    listSessions: vi.fn(() => ['ws:test']),
+    getSession: vi.fn(() => mockSession),
+    transferSession: vi.fn(),
+    closeSession: vi.fn(),
+  };
+
+  const handlers = createSessionHandlers({
+    sessionManager: mockSessionManager as any,
+  });
+
+  beforeEach(() => {
+    vi.clearAllMocks();
+    mockSessionManager.listSessions.mockReturnValue(['ws:test']);
+    mockSessionManager.getSession.mockReturnValue(mockSession);
+    mockSession.getHistory.mockReturnValue(mockHistory);
+  });
+
+  it('sessions.list returns session list with message counts', async () => {
+    const req: GatewayRequest = { id: 1, method: 'sessions.list' };
+    const result = await handlers['sessions.list'](req) as GatewayResponse;
+
+    expect(result.id).toBe(1);
+    const r = result.result as { sessions: Array<{ id: string; messageCount: number }> };
+    expect(r.sessions).toHaveLength(1);
+    expect(r.sessions[0].id).toBe('ws:test');
+    expect(r.sessions[0].messageCount).toBe(2);
+  });
+
+  it('sessions.history returns messages with pagination', async () => {
+    const req: GatewayRequest = { id: 2, method: 'sessions.history', params: { sessionId: 'ws:test', limit: 1, offset: 0 } };
+    const result = await handlers['sessions.history'](req) as GatewayResponse;
+
+    const r = result.result as { messages: unknown[]; total: number };
+    expect(r.messages).toHaveLength(1);
+    expect(r.total).toBe(2);
+  });
+
+  it('sessions.history requires sessionId', async () => {
+    const req: GatewayRequest = { id: 3, method: 'sessions.history', params: {} };
+    const result = await handlers['sessions.history'](req) as GatewayError;
+
+    expect(result.error.code).toBe(ErrorCode.InvalidRequest);
+  });
+
+  it('sessions.create creates a new session', async () => {
+    const req: GatewayRequest = { id: 4, method: 'sessions.create', params: { sessionId: 'ws:new' } };
+    const result = await handlers['sessions.create'](req) as GatewayResponse;
+
+    const r = result.result as { sessionId: string };
+    expect(r.sessionId).toBe('ws:new');
+    expect(mockSessionManager.getSession).toHaveBeenCalledWith('ws', 'new');
+  });
+
+  it('sessions.create auto-generates session ID', async () => {
+    const req: GatewayRequest = { id: 5, method: 'sessions.create' };
+    const result = await handlers['sessions.create'](req) as GatewayResponse;
+
+    const r = result.result as { sessionId: string };
+    expect(r.sessionId).toMatch(/^ws:\d+$/);
+  });
+});
+
+describe('tool handlers', () => {
+  const mockTool = {
+    name: 'test.tool',
+    description: 'A test tool',
+    inputSchema: { type: 'object' as const, properties: {} },
+    execute: vi.fn(),
+  };
+
+  const mockRegistry = {
+    list: vi.fn(() => [mockTool]),
+    get: vi.fn((name: string) => (name === 'test.tool' ? mockTool : undefined)),
+    register: vi.fn(),
+    toAnthropicFormat: vi.fn(),
+    toOpenAIFormat: vi.fn(),
+  };
+
+  const mockExecutor = {
+    execute: vi.fn(async () => ({ success: true, output: 'done' })),
+  };
+
+  const handlers = createToolHandlers({
+    toolRegistry: mockRegistry as any,
+    toolExecutor: mockExecutor as any,
+  });
+
+  beforeEach(() => {
+    vi.clearAllMocks();
+    mockRegistry.list.mockReturnValue([mockTool]);
+    mockRegistry.get.mockImplementation((name: string) => (name === 'test.tool' ? mockTool : undefined));
+    mockExecutor.execute.mockResolvedValue({ success: true, output: 'done' });
+  });
+
+  it('tools.list returns tool definitions', async () => {
+    const req: GatewayRequest = { id: 1, method: 'tools.list' };
+    const result = await handlers['tools.list'](req) as GatewayResponse;
+
+    const r = result.result as { tools: Array<{ name: string }> };
+    expect(r.tools).toHaveLength(1);
+    expect(r.tools[0].name).toBe('test.tool');
+  });
+
+  it('tools.invoke executes a tool', async () => {
+    const req: GatewayRequest = { id: 2, method: 'tools.invoke', params: { tool: 'test.tool', args: {} } };
+    const result = await handlers['tools.invoke'](req) as GatewayResponse;
+
+    expect(result.result).toEqual({ success: true, output: 'done' });
+    expect(mockExecutor.execute).toHaveBeenCalledWith('test.tool', {});
+  });
+
+  it('tools.invoke errors on missing tool name', async () => {
+    const req: GatewayRequest = { id: 3, method: 'tools.invoke', params: {} };
+    const result = await handlers['tools.invoke'](req) as GatewayError;
+    expect(result.error.code).toBe(ErrorCode.InvalidRequest);
+  });
+
+  it('tools.invoke errors on unknown tool', async () => {
+    const req: GatewayRequest = { id: 4, method: 'tools.invoke', params: { tool: 'unknown' } };
+    const result = await handlers['tools.invoke'](req) as GatewayError;
+    expect(result.error.code).toBe(ErrorCode.ToolNotFound);
+  });
+});
+
+describe('agent handlers', () => {
+  const mockAgent = {
+    process: vi.fn(async () => 'response text'),
+    setOnToolUse: vi.fn(),
+  };
+
+  const mockBridge = {
+    getAgent: vi.fn(() => mockAgent),
+    getSessionId: vi.fn(() => 'ws:conn-1'),
+    isBusy: vi.fn(() => false),
+    setBusy: vi.fn(),
+    setOnToolUse: vi.fn(),
+  };
+
+  const handlers = createAgentHandlers({
+    sessionBridge: mockBridge as any,
+  });
+
+  beforeEach(() => {
+    vi.clearAllMocks();
+    mockBridge.isBusy.mockReturnValue(false);
+    mockBridge.getAgent.mockReturnValue(mockAgent);
+    mockAgent.process.mockResolvedValue('response text');
+  });
+
+  it('agent.send processes message and sends done event', async () => {
+    const req: GatewayRequest = { id: 1, method: 'agent.send', params: { message: 'hello', connectionId: 'conn-1' } };
+    const sent: OutboundMessage[] = [];
+    const send = vi.fn((msg: OutboundMessage) => sent.push(msg));
+
+    await handlers['agent.send'](req, send);
+
+    expect(mockAgent.process).toHaveBeenCalledWith('hello');
+    expect(sent).toHaveLength(1);
+    const doneEvent = sent[0] as GatewayEvent;
+    expect(doneEvent.event).toBe('done');
+    expect((doneEvent.data as any).content).toBe('response text');
+  });
+
+  it('agent.send requires message', async () => {
+    const req: GatewayRequest = { id: 2, method: 'agent.send', params: { connectionId: 'conn-1' } };
+    const send = vi.fn();
+    const result = await handlers['agent.send'](req, send) as GatewayError;
+
+    expect(result.error.code).toBe(ErrorCode.InvalidRequest);
+    expect(result.error.message).toContain('message');
+  });
+
+  it('agent.send rejects when busy', async () => {
+    mockBridge.isBusy.mockReturnValue(true);
+    const req: GatewayRequest = { id: 3, method: 'agent.send', params: { message: 'hi', connectionId: 'conn-1' } };
+    const send = vi.fn();
+    const result = await handlers['agent.send'](req, send) as GatewayError;
+
+    expect(result.error.code).toBe(ErrorCode.AgentBusy);
+  });
+
+  it('agent.send handles errors gracefully', async () => {
+    mockAgent.process.mockRejectedValue(new Error('model failed'));
+    const req: GatewayRequest = { id: 4, method: 'agent.send', params: { message: 'hi', connectionId: 'conn-1' } };
+    const sent: OutboundMessage[] = [];
+    const send = vi.fn((msg: OutboundMessage) => sent.push(msg));
+
+    await handlers['agent.send'](req, send);
+
+    const errorEvent = sent[0] as GatewayEvent;
+    expect(errorEvent.event).toBe('error');
+    expect((errorEvent.data as any).message).toBe('model failed');
+  });
+
+  it('agent.send sets and cleans up tool use callback', async () => {
+    const req: GatewayRequest = { id: 5, method: 'agent.send', params: { message: 'hi', connectionId: 'conn-1' } };
+    const send = vi.fn();
+
+    await handlers['agent.send'](req, send);
+
+    // setOnToolUse called twice: once to set callback, once to clear it
+    expect(mockBridge.setOnToolUse).toHaveBeenCalledTimes(2);
+    expect(mockBridge.setOnToolUse).toHaveBeenLastCalledWith('conn-1', undefined);
+  });
+
+  it('agent.send sets busy state correctly', async () => {
+    const req: GatewayRequest = { id: 6, method: 'agent.send', params: { message: 'hi', connectionId: 'conn-1' } };
+    const send = vi.fn();
+
+    await handlers['agent.send'](req, send);
+
+    expect(mockBridge.setBusy).toHaveBeenCalledWith('conn-1', true);
+    expect(mockBridge.setBusy).toHaveBeenCalledWith('conn-1', false);
+  });
+
+  it('agent.cancel returns cancelled state', async () => {
+    mockBridge.isBusy.mockReturnValue(true);
+    const req: GatewayRequest = { id: 7, method: 'agent.cancel', params: { connectionId: 'conn-1' } };
+    const result = await handlers['agent.cancel'](req) as GatewayResponse;
+
+    expect((result.result as any).cancelled).toBe(true);
+  });
+});