dexorder/ai
2
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

data pipeline refactor and fix

Browse Source
This commit is contained in:
Tim Olson
2026-04-13 18:30:04 -04:00
parent 6418729b16
commit 326bf80846
96 changed files with 7107 additions and 1763 deletions
Download Patch File Download Diff File Expand all files Collapse all files

178
gateway/src/channels/websocket-handler.ts
View File

@@ -2,12 +2,14 @@ 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 } from '../services/ohlc-service.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,
@@ -42,6 +44,7 @@ export interface WebSocketHandlerConfig {
createHarness: HarnessFactory;
ohlcService?: OHLCService; // Optional for historical data support
symbolIndexService?: SymbolIndexService; // Optional for symbol search
conversationService?: ConversationService; // Optional for history replay on reconnect
}
/**
@@ -50,10 +53,18 @@ export interface WebSocketHandlerConfig {
* 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;
@@ -106,17 +117,22 @@ export class WebSocketHandler {
// If container is spinning up, wait for it to be ready before continuing
if (isSpinningUp) {
sendStatus(socket, 'spinning_up', 'Your workspace is starting up, please wait...');
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 }, 'Container failed to become ready within timeout');
socket.send(JSON.stringify({ type: 'error', message: 'Workspace failed to start. Please try again later.' }));
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 }, 'Container is ready, proceeding with session setup');
logger.info({ userId: authContext.userId }, 'Sandbox is ready, proceeding with session setup');
}
sendStatus(socket, 'initializing', 'Starting your workspace...');
@@ -241,6 +257,17 @@ export class WebSocketHandler {
})
);
// Replay conversation history so the UI pre-populates on reconnect
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 }));
}
}
// Handle messages
socket.on('message', async (data: Buffer) => {
try {
@@ -266,15 +293,45 @@ export class WebSocketHandler {
return;
}
// Chunks are streamed via channelAdapter.sendChunk() during handleMessage
try {
// Acknowledge receipt immediately so the client can show the seen indicator
socket.send(JSON.stringify({ type: 'agent_chunk', content: '', done: false }));
logger.info('Calling harness.handleMessage');
await harness.handleMessage(inboundMessage);
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;
}
}
// Send done marker after all chunks have been streamed
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({
@@ -332,6 +389,17 @@ export class WebSocketHandler {
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);
@@ -356,6 +424,7 @@ export class WebSocketHandler {
}, 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();
@@ -527,19 +596,92 @@ export class WebSocketHandler {
break;
}
case 'subscribe_bars':
case 'unsubscribe_bars':
// TODO: Implement real-time subscriptions
socket.send(
JSON.stringify({
type: `${payload.type}_response`,
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: 'Real-time subscriptions not yet implemented',
})
);
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