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.
This commit is contained in:
@@ -0,0 +1,144 @@
|
||||
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(),
|
||||
};
|
||||
|
||||
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'],
|
||||
});
|
||||
}
|
||||
|
||||
describe('SessionBridge', () => {
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
});
|
||||
|
||||
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('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);
|
||||
});
|
||||
});
|
||||
Reference in New Issue
Block a user