Files
ai/gateway/src/channels/websocket-handler.ts
2026-04-17 17:15:33 -04:00

804 lines
30 KiB
TypeScript

import type { FastifyInstance, FastifyRequest } from 'fastify';
import type { WebSocket } from '@fastify/websocket';
import type { Authenticator } from '../auth/authenticator.js';
import type { AgentHarness, HarnessFactory } from '../harness/agent-harness.js';
import type { HarnessEvent } from '../harness/harness-events.js';
import type { InboundMessage } from '../types/messages.js';
import { randomUUID } from 'crypto';
import type { SessionRegistry, EventSubscriber, Session } from '../events/index.js';
import type { OHLCService, BarUpdateCallback } from '../services/ohlc-service.js';
import type { SymbolIndexService } from '../services/symbol-index-service.js';
import type { ContainerManager } from '../k8s/container-manager.js';
import type { ConversationService } from '../services/conversation-service.js';
import {
WorkspaceManager,
ContainerSync,
DEFAULT_STORES,
type ChannelAdapter,
type ChannelCapabilities,
type SnapshotMessage,
type PatchMessage,
} from '../workspace/index.js';
/**
* Safe JSON stringifier that handles BigInt values
* Converts BigInt to Number (safe for timestamps and other integer values)
*/
function jsonStringifySafe(obj: any): string {
return JSON.stringify(obj, (_key, value) =>
typeof value === 'bigint' ? Number(value) : value
);
}
export type SessionStatus = 'authenticating' | 'spinning_up' | 'initializing' | 'ready' | 'error'
function sendStatus(socket: WebSocket, status: SessionStatus, message: string): void {
socket.send(JSON.stringify({ type: 'status', status, message }))
}
export interface WebSocketHandlerConfig {
authenticator: Authenticator;
containerManager: ContainerManager;
sessionRegistry: SessionRegistry;
eventSubscriber: EventSubscriber;
createHarness: HarnessFactory;
ohlcService?: OHLCService; // Optional for historical data support
symbolIndexService?: SymbolIndexService; // Optional for symbol search
conversationService?: ConversationService; // Optional for history replay on reconnect
}
/**
* WebSocket channel handler
*
* Handles WebSocket connections for chat and integrates with the event system
* for container-to-client notifications.
*/
interface BarSubscription {
ticker: string;
periodSeconds: number;
callback: BarUpdateCallback;
}
export class WebSocketHandler {
private config: WebSocketHandlerConfig;
private harnesses = new Map<string, AgentHarness>();
private workspaces = new Map<string, WorkspaceManager>();
/** Per-session realtime bar subscriptions for cleanup on disconnect */
private barSubscriptions = new Map<string, BarSubscription[]>();
constructor(config: WebSocketHandlerConfig) {
this.config = config;
}
/**
* Register WebSocket routes
*/
register(app: FastifyInstance): void {
app.get(
'/ws/chat',
{ websocket: true },
async (socket: WebSocket, request: FastifyRequest) => {
await this.handleConnection(socket, request, app);
}
);
}
/**
* Handle WebSocket connection
*/
private async handleConnection(
socket: WebSocket,
request: FastifyRequest,
app: FastifyInstance
): Promise<void> {
const logger = app.log;
// Send initial connecting message
sendStatus(socket, 'authenticating', 'Authenticating...');
// Authenticate (returns immediately if container is spinning up)
const { authContext, isSpinningUp } = await this.config.authenticator.authenticateWebSocket(request);
if (!authContext) {
logger.warn('WebSocket authentication failed');
socket.send(
JSON.stringify({
type: 'error',
message: 'Authentication failed',
})
);
socket.close(1008, 'Authentication failed');
return;
}
logger.info(
{ userId: authContext.userId, sessionId: authContext.sessionId, isSpinningUp },
'WebSocket connection authenticated'
);
// If container is spinning up, wait for it to be ready before continuing
if (isSpinningUp) {
sendStatus(socket, 'spinning_up', 'Your personal agent is starting up, please wait...');
const startupPingInterval = setInterval(() => {
if (socket.readyState === 1) socket.ping();
}, 10000);
const ready = await this.config.containerManager.waitForContainerReady(authContext.userId, 120000);
clearInterval(startupPingInterval);
if (!ready) {
logger.warn({ userId: authContext.userId }, 'Sandbox failed to become ready within timeout');
socket.send(JSON.stringify({ type: 'error', message: 'Agent workspace failed to start. Please try again later.' }));
socket.close(1011, 'Container startup timeout');
return;
}
logger.info({ userId: authContext.userId }, 'Sandbox is ready, proceeding with session setup');
}
sendStatus(socket, 'initializing', 'Starting your workspace...');
// Create WebSocket channel adapter
const wsAdapter: ChannelAdapter = {
sendSnapshot: (msg: SnapshotMessage) => {
socket.send(JSON.stringify(msg));
},
sendPatch: (msg: PatchMessage) => {
socket.send(JSON.stringify(msg));
},
sendText: (msg) => {
socket.send(JSON.stringify({
type: 'text',
text: msg.text,
}));
},
sendChunk: (content) => {
socket.send(JSON.stringify({
type: 'agent_chunk',
content,
done: false,
}));
},
sendImage: (msg) => {
socket.send(JSON.stringify({
type: 'image',
data: msg.data,
mimeType: msg.mimeType,
caption: msg.caption,
}));
},
sendToolCall: (toolName, label) => {
socket.send(JSON.stringify({
type: 'agent_tool_call',
toolName,
label: label ?? toolName,
}));
},
getCapabilities: (): ChannelCapabilities => ({
supportsSync: true,
supportsImages: true,
supportsMarkdown: true,
supportsStreaming: true,
supportsTradingViewEmbed: true,
}),
};
// Declare harness and workspace outside try block so they're available in catch
let harness: AgentHarness | undefined;
let workspace: WorkspaceManager | undefined;
try {
// Create and connect harness first so MCP client is available for ContainerSync
harness = this.config.createHarness({
userId: authContext.userId,
sessionId: authContext.sessionId,
license: authContext.license,
mcpServerUrl: authContext.mcpServerUrl,
logger,
channelAdapter: wsAdapter,
channelType: authContext.channelType,
channelUserId: authContext.channelUserId,
});
await harness.initialize();
// Wire ContainerSync now that MCP client is connected, then initialize workspace
const containerSync = new ContainerSync(harness.getMcpClient(), logger);
workspace = new WorkspaceManager({
userId: authContext.userId,
sessionId: authContext.sessionId,
stores: DEFAULT_STORES,
containerSync,
logger,
});
await workspace.initialize();
workspace.setAdapter(wsAdapter);
harness.setWorkspaceManager(workspace);
this.workspaces.set(authContext.sessionId, workspace);
this.harnesses.set(authContext.sessionId, harness);
// Push all store snapshots to the client now, before 'connected'.
// Empty seqs force full snapshots for every store, so the browser's
// message queue has the current workspace state (including persistent
// stores loaded from the container) before TradingView initializes.
await workspace.handleHello({});
// Register session for event system
// Container endpoint is derived from the MCP server URL (same container, different port)
const containerEventEndpoint = this.getContainerEventEndpoint(authContext.mcpServerUrl);
const session: Session = {
userId: authContext.userId,
sessionId: authContext.sessionId,
socket,
channelType: 'websocket',
containerEndpoint: containerEventEndpoint,
connectedAt: new Date(),
};
this.config.sessionRegistry.register(session);
// Subscribe to informational events from user's container
await this.config.eventSubscriber.onSessionConnect(session);
logger.info(
{ userId: authContext.userId, containerEndpoint: containerEventEndpoint },
'Session registered for events'
);
sendStatus(socket, 'ready', 'Your workspace is ready!');
socket.send(
JSON.stringify({
type: 'connected',
sessionId: authContext.sessionId,
userId: authContext.userId,
licenseType: authContext.license.licenseType,
message: 'Connected to Dexorder AI',
})
);
// Replay conversation history so the UI pre-populates on reconnect;
// greet new users on their first conversation
if (this.config.conversationService) {
const history = await this.config.conversationService.getHistory(
authContext.userId,
authContext.sessionId
);
if (history.length > 0) {
socket.send(JSON.stringify({ type: 'conversation_history', messages: history }));
} else {
// First conversation — auto-send greeting prompt and stream the response
socket.send(JSON.stringify({ type: 'agent_chunk', content: '', done: false }));
for await (const event of harness!.streamGreeting()) {
const e = event as HarnessEvent;
switch (e.type) {
case 'chunk':
socket.send(JSON.stringify({ type: 'agent_chunk', content: e.content, done: false }));
break;
case 'tool_call':
socket.send(JSON.stringify({ type: 'agent_tool_call', toolName: e.toolName, label: e.label }));
break;
case 'image':
socket.send(JSON.stringify({ type: 'image', data: e.data, mimeType: e.mimeType, caption: e.caption }));
break;
case 'error':
socket.send(JSON.stringify({ type: 'text', text: `An error occurred during greeting.` }));
break;
case 'done':
break;
}
}
socket.send(JSON.stringify({ type: 'agent_chunk', content: '', done: true }));
}
}
// Handle messages
socket.on('message', async (data: Buffer) => {
try {
logger.info({ rawMessage: data.toString().substring(0, 500) }, 'WebSocket message received');
const payload = JSON.parse(data.toString());
logger.info({ type: payload.type, request_id: payload.request_id }, 'WebSocket message parsed');
// Route based on message type
if (payload.type === 'message' || payload.type === 'agent_user_message') {
// Chat message - send to agent harness with streaming
const inboundMessage: InboundMessage = {
messageId: randomUUID(),
userId: authContext.userId,
sessionId: authContext.sessionId,
content: payload.content,
attachments: payload.attachments,
timestamp: new Date(),
};
if (!harness) {
logger.error('Harness not initialized');
socket.send(JSON.stringify({ type: 'error', message: 'Session not ready' }));
return;
}
try {
// Acknowledge receipt immediately so the client can show the seen indicator
socket.send(JSON.stringify({ type: 'agent_chunk', content: '', done: false }));
logger.info('Streaming harness response');
let fatalError = false;
for await (const event of harness.streamMessage(inboundMessage)) {
const e = event as HarnessEvent;
switch (e.type) {
case 'chunk':
socket.send(JSON.stringify({ type: 'agent_chunk', content: e.content, done: false }));
break;
case 'tool_call':
socket.send(JSON.stringify({ type: 'agent_tool_call', toolName: e.toolName, label: e.label }));
break;
case 'subagent_tool_call':
socket.send(JSON.stringify({ type: 'subagent_tool_call', agentName: e.agentName, toolName: e.toolName, label: e.label }));
break;
case 'subagent_chunk':
socket.send(JSON.stringify({ type: 'subagent_chunk', agentName: e.agentName, content: e.content }));
break;
case 'image':
socket.send(JSON.stringify({ type: 'image', data: e.data, mimeType: e.mimeType, caption: e.caption }));
break;
case 'error':
socket.send(JSON.stringify({ type: 'text', text: `An unrecoverable error occurred in the ${e.source}.` }));
if (e.fatal) fatalError = true;
break;
case 'done':
break;
}
}
if (fatalError) {
socket.close(1011, 'Fatal error');
return;
}
// Send done marker after all events have been streamed
logger.debug('Sending done marker to client');
socket.send(
JSON.stringify({
type: 'agent_chunk',
content: '',
done: true,
})
);
} catch (error) {
logger.error({ error }, 'Error streaming response');
socket.send(
JSON.stringify({
type: 'error',
message: 'Failed to generate response',
})
);
}
} else if (payload.type === 'hello') {
// Workspace sync: hello message
logger.debug({ seqs: payload.seqs }, 'Handling workspace hello');
await workspace!.handleHello(payload.seqs || {});
} else if (payload.type === 'patch') {
// Workspace sync: patch message
logger.debug({ store: payload.store, seq: payload.seq }, 'Handling workspace patch');
await workspace!.handlePatch(payload.store, payload.seq, payload.patch || []);
} else if (payload.type === 'agent_stop') {
logger.info('Agent stop requested');
harness?.interrupt();
} else if (this.isDatafeedMessage(payload)) {
// Historical data request - send to OHLC service
logger.info({ type: payload.type }, 'Routing to datafeed handler');
await this.handleDatafeedMessage(socket, payload, logger, authContext);
} else {
logger.warn({ type: payload.type }, 'Unknown message type received');
}
} catch (error) {
logger.error({ error }, 'Error handling WebSocket message');
socket.send(
JSON.stringify({
type: 'error',
message: 'Failed to process message',
})
);
}
});
// Handle disconnection
socket.on('close', async (code: number, reason: Buffer) => {
clearInterval(pingInterval);
logger.info({ sessionId: authContext.sessionId, code, reason: reason?.toString() }, 'WebSocket disconnected');
// Unregister from event system
const removedSession = this.config.sessionRegistry.unregister(authContext.sessionId);
if (removedSession) {
await this.config.eventSubscriber.onSessionDisconnect(removedSession);
}
// Cleanup realtime bar subscriptions
const sessionId = authContext.sessionId;
const subs = this.barSubscriptions.get(sessionId);
if (subs && this.config.ohlcService) {
for (const { ticker, periodSeconds, callback } of subs) {
this.config.ohlcService.unsubscribeFromTicker(ticker, periodSeconds, callback);
}
this.barSubscriptions.delete(sessionId);
logger.info({ sessionId, count: subs.length }, 'Cleaned up realtime bar subscriptions');
}
// Cleanup workspace
await workspace!.shutdown();
this.workspaces.delete(authContext.sessionId);
// Cleanup harness
if (harness) {
await harness.cleanup();
this.harnesses.delete(authContext.sessionId);
}
});
socket.on('error', (error: any) => {
logger.error({ error, sessionId: authContext.sessionId }, 'WebSocket error');
});
// Ping every 30 seconds to keep the connection alive through CloudFlare proxy.
// CloudFlare drops idle WebSocket connections after ~100 seconds.
const pingInterval = setInterval(() => {
if (socket.readyState === 1) { // OPEN
socket.ping();
}
}, 30000);
} catch (error) {
logger.error({ error }, 'Failed to initialize session');
socket.send(JSON.stringify({ type: 'text', text: 'An unrecoverable error occurred in the agent harness.' }));
socket.close(1011, 'Internal server error');
if (workspace) {
await workspace.shutdown();
this.workspaces.delete(authContext.sessionId);
}
if (harness) {
await harness.cleanup();
}
}
}
/**
* Derive the container's XPUB event endpoint from the MCP server URL.
*
* MCP URL format: http://sandbox-user-abc123.sandbox.svc.cluster.local:3000
* Event endpoint: tcp://sandbox-user-abc123.sandbox.svc.cluster.local:5570
*/
private getContainerEventEndpoint(mcpServerUrl: string): string {
try {
const url = new URL(mcpServerUrl);
// Replace protocol and port
return `tcp://${url.hostname}:5570`;
} catch {
// Fallback if URL parsing fails
return mcpServerUrl.replace('http://', 'tcp://').replace(':3000', ':5570');
}
}
/**
* Check if message is a datafeed message (TradingView protocol)
*/
private isDatafeedMessage(payload: any): boolean {
const datafeedTypes = [
'get_config',
'search_symbols',
'resolve_symbol',
'get_bars',
'subscribe_bars',
'unsubscribe_bars',
'evaluate_indicator',
];
return datafeedTypes.includes(payload.type);
}
/**
* Handle datafeed messages (TradingView protocol)
*/
private async handleDatafeedMessage(
socket: WebSocket,
payload: any,
logger: any,
authContext?: any
): Promise<void> {
logger.info({ type: payload.type, payload }, 'handleDatafeedMessage called');
const ohlcService = this.config.ohlcService;
const symbolIndexService = this.config.symbolIndexService;
logger.info({
hasOhlcService: !!ohlcService,
hasSymbolIndexService: !!symbolIndexService
}, 'Service availability');
const requestId = payload.request_id || randomUUID();
if (!ohlcService && !symbolIndexService) {
logger.warn({ requestId }, 'No datafeed services available yet');
socket.send(JSON.stringify({
type: 'error',
request_id: requestId,
error_message: 'Services initializing, please retry',
}));
return;
}
try {
switch (payload.type) {
case 'get_config': {
const config = ohlcService ? await ohlcService.getConfig() : { supported_resolutions: ['1', '5', '15', '60', '1D'] };
socket.send(
JSON.stringify({
type: 'get_config_response',
request_id: requestId,
config,
})
);
break;
}
case 'search_symbols': {
logger.info({ query: payload.query, limit: payload.limit }, 'Handling search_symbols');
// Use SymbolIndexService if available, otherwise fallback to OHLCService stub
const symbolIndexService = this.config.symbolIndexService;
logger.info({ hasSymbolIndexService: !!symbolIndexService }, 'Service check for search');
const results = symbolIndexService
? await symbolIndexService.search(payload.query, payload.limit || 30)
: (ohlcService ? await ohlcService.searchSymbols(
payload.query,
payload.symbol_type,
payload.exchange,
payload.limit || 30
) : []);
logger.info({ resultsCount: results.length }, 'Search complete');
socket.send(
JSON.stringify({
type: 'search_symbols_response',
request_id: requestId,
results,
})
);
break;
}
case 'resolve_symbol': {
logger.info({ symbol: payload.symbol }, 'Handling resolve_symbol');
// Use SymbolIndexService if available, otherwise fallback to OHLCService stub
const symbolIndexService = this.config.symbolIndexService;
logger.info({ hasSymbolIndexService: !!symbolIndexService }, 'Service check for resolve');
const symbolInfo = symbolIndexService
? await symbolIndexService.resolveSymbol(payload.symbol)
: (ohlcService ? await ohlcService.resolveSymbol(payload.symbol) : null);
logger.info({ found: !!symbolInfo }, 'Symbol resolution complete');
if (!symbolInfo) {
logger.warn({ symbol: payload.symbol }, 'Symbol not found');
socket.send(
JSON.stringify({
type: 'error',
request_id: requestId,
error_message: `Symbol not found: ${payload.symbol}`,
})
);
} else {
logger.info({ symbolInfo }, 'Sending symbol_info response');
socket.send(
JSON.stringify({
type: 'resolve_symbol_response',
request_id: requestId,
symbol_info: symbolInfo,
})
);
}
break;
}
case 'get_bars': {
if (!ohlcService) {
socket.send(JSON.stringify({
type: 'get_bars_response',
request_id: requestId,
error: 'OHLC service not available',
}));
break;
}
try {
const history = await ohlcService.fetchOHLC(
payload.symbol,
payload.period_seconds,
payload.from_time,
payload.to_time,
payload.countback
);
logger.info({ requestId, barCount: history.bars?.length ?? 0, noData: history.noData, socketState: socket.readyState }, 'Sending get_bars_response');
socket.send(jsonStringifySafe({ type: 'get_bars_response', request_id: requestId, history }));
logger.info({ requestId }, 'get_bars_response sent');
} catch (err: any) {
const errorMessage = err?.message ?? String(err);
logger.error({ requestId, ticker: payload.symbol, errorMessage }, 'get_bars failed');
socket.send(JSON.stringify({ type: 'get_bars_response', request_id: requestId, error: errorMessage }));
}
break;
}
case 'subscribe_bars': {
if (!ohlcService || !authContext) {
socket.send(JSON.stringify({
type: 'subscribe_bars_response',
request_id: requestId,
subscription_id: payload.subscription_id,
success: false,
message: 'Realtime service not available',
}));
break;
}
const subTicker: string = payload.symbol;
const subPeriod: number = payload.period_seconds ?? payload.resolution ?? 60;
const sessionId = authContext.sessionId;
// Create a per-subscription callback that forwards bars to this socket
const barCallback: BarUpdateCallback = (bar) => {
if (socket.readyState !== 1 /* OPEN */) return;
socket.send(JSON.stringify({
type: 'bar_update',
subscription_id: payload.subscription_id,
ticker: bar.ticker,
period_seconds: bar.periodSeconds,
bar: {
// Convert nanoseconds → seconds for client compatibility
time: Number(bar.timestamp / 1_000_000_000n),
open: bar.open,
high: bar.high,
low: bar.low,
close: bar.close,
volume: bar.volume,
},
}));
};
ohlcService.subscribeToTicker(subTicker, subPeriod, barCallback);
// Track for cleanup on disconnect
if (!this.barSubscriptions.has(sessionId)) {
this.barSubscriptions.set(sessionId, []);
}
this.barSubscriptions.get(sessionId)!.push({
ticker: subTicker,
periodSeconds: subPeriod,
callback: barCallback,
});
logger.info({ sessionId, ticker: subTicker, period: subPeriod }, 'Subscribed to realtime bars');
socket.send(JSON.stringify({
type: 'subscribe_bars_response',
request_id: requestId,
subscription_id: payload.subscription_id,
success: true,
}));
break;
}
case 'unsubscribe_bars': {
if (!ohlcService || !authContext) break;
const unsubTicker: string = payload.symbol;
const unsubPeriod: number = payload.period_seconds ?? payload.resolution ?? 60;
const sessionId = authContext.sessionId;
const subs = this.barSubscriptions.get(sessionId);
if (subs) {
const idx = subs.findIndex(
s => s.ticker === unsubTicker && s.periodSeconds === unsubPeriod
);
if (idx >= 0) {
const [removed] = subs.splice(idx, 1);
ohlcService.unsubscribeFromTicker(unsubTicker, unsubPeriod, removed.callback);
logger.info({ sessionId, ticker: unsubTicker, period: unsubPeriod }, 'Unsubscribed from realtime bars');
}
}
socket.send(JSON.stringify({
type: 'unsubscribe_bars_response',
request_id: requestId,
subscription_id: payload.subscription_id,
success: true,
}));
break;
}
case 'evaluate_indicator': {
// Direct MCP call — bypasses the agent/LLM for performance
const harness = this.harnesses.get(authContext.sessionId);
if (!harness) {
socket.send(JSON.stringify({
type: 'evaluate_indicator_result',
request_id: requestId,
error: 'Session not initialized',
}));
break;
}
try {
const mcpResult = await harness.callMcpTool('evaluate_indicator', {
symbol: payload.symbol,
from_time: payload.from_time,
to_time: payload.to_time,
period_seconds: payload.period_seconds,
pandas_ta_name: payload.pandas_ta_name,
parameters: payload.parameters ?? {},
}) as any;
// MCP returns { content: [{type: 'text', text: '...json...'}] }
// When the tool raises an exception, the MCP framework sets isError: true
// and puts the raw exception text in content[0].text (not JSON-wrapped).
const rawText = mcpResult?.content?.[0]?.text ?? mcpResult?.[0]?.text;
if (mcpResult?.isError || rawText == null) {
const errMsg = rawText ?? 'evaluate_indicator returned no content';
logger.error({ pandas_ta_name: payload.pandas_ta_name, rawText }, 'evaluate_indicator sandbox error');
socket.send(JSON.stringify({
type: 'evaluate_indicator_result',
request_id: requestId,
error: errMsg,
}));
break;
}
let data: any;
try {
data = JSON.parse(rawText);
} catch {
// Sandbox returned non-JSON (e.g. bare exception text)
logger.error({ pandas_ta_name: payload.pandas_ta_name, rawText }, 'evaluate_indicator returned non-JSON');
socket.send(JSON.stringify({
type: 'evaluate_indicator_result',
request_id: requestId,
error: rawText,
}));
break;
}
socket.send(JSON.stringify({
type: 'evaluate_indicator_result',
request_id: requestId,
...data,
}));
} catch (err: any) {
logger.error({ err: err?.message, pandas_ta_name: payload.pandas_ta_name }, 'evaluate_indicator handler error');
socket.send(JSON.stringify({
type: 'evaluate_indicator_result',
request_id: requestId,
error: err?.message ?? String(err),
}));
}
break;
}
default:
logger.warn({ type: payload.type }, 'Unknown datafeed message type');
}
} catch (error: any) {
logger.error({ error, type: payload.type }, 'Error handling datafeed message');
socket.send(
jsonStringifySafe({
type: 'error',
request_id: requestId,
error_code: 'INTERNAL_ERROR',
error_message: error.message || 'Internal server error',
})
);
}
}
/**
* Flush and clean up all active sessions.
* Called during graceful shutdown to ensure conversations are persisted.
*/
async endAllSessions(): Promise<void> {
const cleanups = Array.from(this.harnesses.values()).map(h => h.cleanup());
await Promise.allSettled(cleanups);
this.harnesses.clear();
}
}