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,204 @@
|
||||
import { WebSocketServer, WebSocket } from 'ws';
|
||||
import { randomUUID } from 'crypto';
|
||||
import type { IncomingMessage } from 'http';
|
||||
import { Router } from './router.js';
|
||||
import { SessionBridge } from './session-bridge.js';
|
||||
import type { SessionBridgeConfig } from './session-bridge.js';
|
||||
import { authenticateRequest } from './auth.js';
|
||||
import type { AuthConfig } from './auth.js';
|
||||
import {
|
||||
parseMessage,
|
||||
makeError,
|
||||
ErrorCode,
|
||||
type OutboundMessage,
|
||||
} from './protocol.js';
|
||||
import {
|
||||
createSystemHandlers,
|
||||
createSessionHandlers,
|
||||
createToolHandlers,
|
||||
createAgentHandlers,
|
||||
} from './handlers/index.js';
|
||||
import type { SessionManager } from '../session/manager.js';
|
||||
import type { ToolRegistry } from '../tools/registry.js';
|
||||
import type { ToolExecutor } from '../tools/executor.js';
|
||||
|
||||
export interface GatewayServerConfig {
|
||||
port: number;
|
||||
host?: string;
|
||||
sessionManager: SessionManager;
|
||||
modelClient: SessionBridgeConfig['modelClient'];
|
||||
systemPrompt: string;
|
||||
toolRegistry: ToolRegistry;
|
||||
toolExecutor: ToolExecutor;
|
||||
version?: string;
|
||||
auth?: AuthConfig;
|
||||
}
|
||||
|
||||
export class GatewayServer {
|
||||
private wss: WebSocketServer | null = null;
|
||||
private router: Router;
|
||||
private sessionBridge: SessionBridge;
|
||||
private connectionMap: Map<WebSocket, string> = new Map();
|
||||
private config: GatewayServerConfig;
|
||||
private startTime: number = Date.now();
|
||||
|
||||
constructor(config: GatewayServerConfig) {
|
||||
this.config = config;
|
||||
|
||||
this.sessionBridge = new SessionBridge({
|
||||
sessionManager: config.sessionManager,
|
||||
modelClient: config.modelClient,
|
||||
systemPrompt: config.systemPrompt,
|
||||
toolRegistry: config.toolRegistry,
|
||||
toolExecutor: config.toolExecutor,
|
||||
});
|
||||
|
||||
this.router = new Router();
|
||||
this.registerHandlers();
|
||||
}
|
||||
|
||||
private registerHandlers(): void {
|
||||
const systemHandlers = createSystemHandlers({
|
||||
startTime: this.startTime,
|
||||
version: this.config.version ?? '0.1.0',
|
||||
getSessionCount: () => this.sessionBridge.listSessions().length,
|
||||
getToolCount: () => this.config.toolRegistry.list().length,
|
||||
getConnectionCount: () => this.sessionBridge.connectionCount,
|
||||
});
|
||||
|
||||
const sessionHandlers = createSessionHandlers({
|
||||
sessionManager: this.config.sessionManager,
|
||||
});
|
||||
|
||||
const toolHandlers = createToolHandlers({
|
||||
toolRegistry: this.config.toolRegistry,
|
||||
toolExecutor: this.config.toolExecutor,
|
||||
});
|
||||
|
||||
const agentHandlers = createAgentHandlers({
|
||||
sessionBridge: this.sessionBridge,
|
||||
});
|
||||
|
||||
// Register all methods
|
||||
for (const [method, handler] of Object.entries(systemHandlers)) {
|
||||
this.router.register(method, handler);
|
||||
}
|
||||
for (const [method, handler] of Object.entries(sessionHandlers)) {
|
||||
this.router.register(method, handler);
|
||||
}
|
||||
for (const [method, handler] of Object.entries(toolHandlers)) {
|
||||
this.router.register(method, handler);
|
||||
}
|
||||
for (const [method, handler] of Object.entries(agentHandlers)) {
|
||||
this.router.register(method, handler);
|
||||
}
|
||||
}
|
||||
|
||||
async start(): Promise<void> {
|
||||
return new Promise((resolve) => {
|
||||
this.wss = new WebSocketServer({
|
||||
port: this.config.port,
|
||||
host: this.config.host ?? '127.0.0.1',
|
||||
});
|
||||
|
||||
this.wss.on('connection', (ws: WebSocket, req: IncomingMessage) => {
|
||||
// Auth check on upgrade
|
||||
const authResult = authenticateRequest(req, this.config.auth ?? {});
|
||||
if (!authResult.authenticated) {
|
||||
ws.close(4001, authResult.error ?? 'Authentication failed');
|
||||
return;
|
||||
}
|
||||
this.handleConnection(ws, authResult.identity);
|
||||
});
|
||||
|
||||
this.wss.on('listening', () => {
|
||||
const addr = this.wss?.address();
|
||||
const portStr = typeof addr === 'object' && addr ? `:${addr.port}` : '';
|
||||
console.log(`Gateway WebSocket server listening on ${this.config.host ?? '127.0.0.1'}${portStr}`);
|
||||
resolve();
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
async stop(): Promise<void> {
|
||||
return new Promise((resolve) => {
|
||||
if (!this.wss) {
|
||||
resolve();
|
||||
return;
|
||||
}
|
||||
|
||||
// Close all connections
|
||||
for (const [ws, connectionId] of this.connectionMap) {
|
||||
this.sessionBridge.disconnect(connectionId);
|
||||
ws.close(1001, 'Server shutting down');
|
||||
}
|
||||
this.connectionMap.clear();
|
||||
|
||||
this.wss.close(() => {
|
||||
this.wss = null;
|
||||
resolve();
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
private handleConnection(ws: WebSocket, identity?: string): void {
|
||||
const connectionId = randomUUID();
|
||||
this.sessionBridge.connect(connectionId);
|
||||
this.connectionMap.set(ws, connectionId);
|
||||
|
||||
ws.on('message', async (data) => {
|
||||
const raw = data.toString();
|
||||
await this.handleMessage(ws, connectionId, raw);
|
||||
});
|
||||
|
||||
ws.on('close', () => {
|
||||
this.sessionBridge.disconnect(connectionId);
|
||||
this.connectionMap.delete(ws);
|
||||
});
|
||||
|
||||
ws.on('error', (err) => {
|
||||
console.error(`WebSocket error (${connectionId}):`, err.message);
|
||||
});
|
||||
}
|
||||
|
||||
private async handleMessage(ws: WebSocket, connectionId: string, raw: string): Promise<void> {
|
||||
const request = parseMessage(raw);
|
||||
|
||||
if (!request) {
|
||||
this.send(ws, makeError(0, ErrorCode.ParseError, 'Invalid JSON or missing required fields'));
|
||||
return;
|
||||
}
|
||||
|
||||
// Inject connectionId into params so handlers can identify the client
|
||||
if (!request.params) request.params = {};
|
||||
request.params.connectionId = connectionId;
|
||||
|
||||
const send = (msg: OutboundMessage) => this.send(ws, msg);
|
||||
const response = await this.router.dispatch(request, send);
|
||||
|
||||
if (response) {
|
||||
this.send(ws, response);
|
||||
}
|
||||
}
|
||||
|
||||
private send(ws: WebSocket, msg: OutboundMessage): void {
|
||||
if (ws.readyState === WebSocket.OPEN) {
|
||||
ws.send(JSON.stringify(msg));
|
||||
}
|
||||
}
|
||||
|
||||
/** Get the underlying WebSocketServer (for testing). */
|
||||
getWss(): WebSocketServer | null {
|
||||
return this.wss;
|
||||
}
|
||||
|
||||
/** Get the session bridge (for testing/debugging). */
|
||||
getSessionBridge(): SessionBridge {
|
||||
return this.sessionBridge;
|
||||
}
|
||||
|
||||
/** Get list of registered methods. */
|
||||
getMethods(): string[] {
|
||||
return this.router.listMethods();
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user