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,135 @@
|
||||
import { randomUUID } from 'crypto';
|
||||
import type { SessionManager } from '../session/manager.js';
|
||||
import type { Session } from '../session/manager.js';
|
||||
import type { ModelClient } from '../models/types.js';
|
||||
import type { ModelRouter } from '../models/router.js';
|
||||
import type { ToolRegistry } from '../tools/registry.js';
|
||||
import type { ToolExecutor } from '../tools/executor.js';
|
||||
import { NativeAgent } from '../backends/native/agent.js';
|
||||
import type { ToolUseEvent } from '../backends/native/agent.js';
|
||||
|
||||
export interface SessionBridgeConfig {
|
||||
sessionManager: SessionManager;
|
||||
modelClient: ModelClient | ModelRouter;
|
||||
systemPrompt: string;
|
||||
toolRegistry: ToolRegistry;
|
||||
toolExecutor: ToolExecutor;
|
||||
}
|
||||
|
||||
interface ClientEntry {
|
||||
connectionId: string;
|
||||
sessionId: string;
|
||||
agent: NativeAgent;
|
||||
busy: boolean;
|
||||
}
|
||||
|
||||
export class SessionBridge {
|
||||
private clients: Map<string, ClientEntry> = new Map();
|
||||
private agents: Map<string, NativeAgent> = new Map();
|
||||
private config: SessionBridgeConfig;
|
||||
|
||||
constructor(config: SessionBridgeConfig) {
|
||||
this.config = config;
|
||||
}
|
||||
|
||||
/** Register a new WS connection. Returns the assigned connection ID. */
|
||||
connect(connectionId?: string): string {
|
||||
const id = connectionId ?? randomUUID();
|
||||
const sessionId = `ws:${id}`;
|
||||
const agent = this.getOrCreateAgent(sessionId);
|
||||
|
||||
this.clients.set(id, {
|
||||
connectionId: id,
|
||||
sessionId,
|
||||
agent,
|
||||
busy: false,
|
||||
});
|
||||
|
||||
return id;
|
||||
}
|
||||
|
||||
/** Remove a WS connection. Does NOT destroy the session (persists in SQLite). */
|
||||
disconnect(connectionId: string): void {
|
||||
const client = this.clients.get(connectionId);
|
||||
if (client) {
|
||||
// Only remove the agent if no other clients share the session
|
||||
const otherClients = Array.from(this.clients.values())
|
||||
.filter(c => c.sessionId === client.sessionId && c.connectionId !== connectionId);
|
||||
if (otherClients.length === 0) {
|
||||
this.agents.delete(client.sessionId);
|
||||
}
|
||||
this.clients.delete(connectionId);
|
||||
}
|
||||
}
|
||||
|
||||
/** Switch a connection to a different session (e.g. resuming an old session). */
|
||||
switchSession(connectionId: string, sessionId: string): void {
|
||||
const client = this.clients.get(connectionId);
|
||||
if (!client) throw new Error(`Unknown connection: ${connectionId}`);
|
||||
if (client.busy) throw new Error('Cannot switch session while agent is busy');
|
||||
|
||||
const agent = this.getOrCreateAgent(sessionId);
|
||||
client.sessionId = sessionId;
|
||||
client.agent = agent;
|
||||
}
|
||||
|
||||
/** Get the NativeAgent for a connection. */
|
||||
getAgent(connectionId: string): NativeAgent | undefined {
|
||||
return this.clients.get(connectionId)?.agent;
|
||||
}
|
||||
|
||||
/** Get the session ID for a connection. */
|
||||
getSessionId(connectionId: string): string | undefined {
|
||||
return this.clients.get(connectionId)?.sessionId;
|
||||
}
|
||||
|
||||
/** Check if a connection's agent is busy. */
|
||||
isBusy(connectionId: string): boolean {
|
||||
return this.clients.get(connectionId)?.busy ?? false;
|
||||
}
|
||||
|
||||
/** Mark a connection's agent as busy/idle. */
|
||||
setBusy(connectionId: string, busy: boolean): void {
|
||||
const client = this.clients.get(connectionId);
|
||||
if (client) client.busy = busy;
|
||||
}
|
||||
|
||||
/** Set onToolUse callback for a connection's agent. */
|
||||
setOnToolUse(connectionId: string, callback: ((event: ToolUseEvent) => void) | undefined): void {
|
||||
const client = this.clients.get(connectionId);
|
||||
if (client) client.agent.setOnToolUse(callback);
|
||||
}
|
||||
|
||||
/** List all active sessions with connection counts. */
|
||||
listSessions(): Array<{ sessionId: string; connections: number }> {
|
||||
const sessionMap = new Map<string, number>();
|
||||
for (const client of this.clients.values()) {
|
||||
sessionMap.set(client.sessionId, (sessionMap.get(client.sessionId) ?? 0) + 1);
|
||||
}
|
||||
return Array.from(sessionMap.entries()).map(([sessionId, connections]) => ({
|
||||
sessionId,
|
||||
connections,
|
||||
}));
|
||||
}
|
||||
|
||||
/** Get count of active connections. */
|
||||
get connectionCount(): number {
|
||||
return this.clients.size;
|
||||
}
|
||||
|
||||
private getOrCreateAgent(sessionId: string): NativeAgent {
|
||||
let agent = this.agents.get(sessionId);
|
||||
if (!agent) {
|
||||
const session = this.config.sessionManager.getSession('ws', sessionId);
|
||||
agent = new NativeAgent({
|
||||
modelClient: this.config.modelClient,
|
||||
systemPrompt: this.config.systemPrompt,
|
||||
session,
|
||||
toolRegistry: this.config.toolRegistry,
|
||||
toolExecutor: this.config.toolExecutor,
|
||||
});
|
||||
this.agents.set(sessionId, agent);
|
||||
}
|
||||
return agent;
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user