sandbox connected and streaming

This commit is contained in:
2026-03-30 23:29:03 -04:00
parent c3a8fae132
commit 998f69fa1a
130 changed files with 7416 additions and 2123 deletions

View File

@@ -1,11 +1,9 @@
import type { FastifyInstance, FastifyRequest } from 'fastify';
import type { WebSocket } from '@fastify/websocket';
import type { Authenticator } from '../auth/authenticator.js';
import { AgentHarness } from '../harness/agent-harness.js';
import type { AgentHarness, HarnessFactory } from '../harness/agent-harness.js';
import type { InboundMessage } from '../types/messages.js';
import { randomUUID } from 'crypto';
import type { ProviderConfig } from '../llm/provider.js';
import type { SessionRegistry, EventSubscriber, Session } from '../events/index.js';
import type { OHLCService } from '../services/ohlc-service.js';
import type { SymbolIndexService } from '../services/symbol-index-service.js';
@@ -29,12 +27,18 @@ function jsonStringifySafe(obj: any): string {
);
}
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;
providerConfig: ProviderConfig;
sessionRegistry: SessionRegistry;
eventSubscriber: EventSubscriber;
createHarness: HarnessFactory;
ohlcService?: OHLCService; // Optional for historical data support
symbolIndexService?: SymbolIndexService; // Optional for symbol search
}
@@ -78,13 +82,7 @@ export class WebSocketHandler {
const logger = app.log;
// Send initial connecting message
socket.send(
JSON.stringify({
type: 'status',
status: 'authenticating',
message: 'Authenticating...',
})
);
sendStatus(socket, 'authenticating', 'Authenticating...');
// Authenticate (returns immediately if container is spinning up)
const { authContext, isSpinningUp } = await this.config.authenticator.authenticateWebSocket(request);
@@ -105,33 +103,23 @@ export class WebSocketHandler {
'WebSocket connection authenticated'
);
// If container is spinning up, send status and start background polling
// If container is spinning up, wait for it to be ready before continuing
if (isSpinningUp) {
socket.send(
JSON.stringify({
type: 'status',
status: 'spinning_up',
message: 'Your workspace is starting up, please wait...',
})
);
sendStatus(socket, 'spinning_up', 'Your workspace is starting up, please wait...');
// Start background polling for container readiness
this.pollContainerReadiness(socket, authContext, app).catch((error) => {
logger.error({ error, userId: authContext.userId }, 'Error polling container readiness');
});
const ready = await this.config.containerManager.waitForContainerReady(authContext.userId, 120000);
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.' }));
socket.close(1011, 'Container startup timeout');
return;
}
// Don't return - continue with session setup so we can receive messages once ready
} else {
// Send workspace starting message
socket.send(
JSON.stringify({
type: 'status',
status: 'initializing',
message: 'Starting your workspace...',
})
);
logger.info({ userId: authContext.userId }, 'Container is ready, proceeding with session setup');
}
sendStatus(socket, 'initializing', 'Starting your workspace...');
// Create workspace manager for this session
const workspace = new WorkspaceManager({
userId: authContext.userId,
@@ -149,6 +137,34 @@ export class WebSocketHandler {
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,
@@ -167,14 +183,17 @@ export class WebSocketHandler {
workspace.setAdapter(wsAdapter);
this.workspaces.set(authContext.sessionId, workspace);
// Create agent harness with workspace manager
harness = new AgentHarness({
// Create agent harness via factory (storage deps injected by factory)
harness = this.config.createHarness({
userId: authContext.userId,
sessionId: authContext.sessionId,
license: authContext.license,
providerConfig: this.config.providerConfig,
mcpServerUrl: authContext.mcpServerUrl,
logger,
workspaceManager: workspace,
channelAdapter: wsAdapter,
channelType: authContext.channelType,
channelUserId: authContext.channelUserId,
});
await harness.initialize();
@@ -182,7 +201,7 @@ export class WebSocketHandler {
// Register session for event system
// Container endpoint is derived from the MCP server URL (same container, different port)
const containerEventEndpoint = this.getContainerEventEndpoint(authContext.license.mcpServerUrl);
const containerEventEndpoint = this.getContainerEventEndpoint(authContext.mcpServerUrl);
const session: Session = {
userId: authContext.userId,
@@ -203,18 +222,16 @@ export class WebSocketHandler {
'Session registered for events'
);
// Send connected message (only if not spinning up - otherwise sent by pollContainerReadiness)
if (!isSpinningUp) {
socket.send(
JSON.stringify({
type: 'connected',
sessionId: authContext.sessionId,
userId: authContext.userId,
licenseType: authContext.license.licenseType,
message: 'Connected to Dexorder AI',
})
);
}
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',
})
);
// Handle messages
socket.on('message', async (data: Buffer) => {
@@ -241,19 +258,16 @@ export class WebSocketHandler {
return;
}
// Stream response chunks to client
// Chunks are streamed via channelAdapter.sendChunk() during handleMessage
try {
for await (const chunk of harness.streamMessage(inboundMessage)) {
socket.send(
JSON.stringify({
type: 'agent_chunk',
content: chunk,
done: false,
})
);
}
// Acknowledge receipt immediately so the client can show the seen indicator
socket.send(JSON.stringify({ type: 'agent_chunk', content: '', done: false }));
// Send final chunk with done flag
logger.info('Calling harness.handleMessage');
await harness.handleMessage(inboundMessage);
// Send done marker after all chunks have been streamed
logger.debug('Sending done marker to client');
socket.send(
JSON.stringify({
type: 'agent_chunk',
@@ -331,73 +345,11 @@ export class WebSocketHandler {
}
}
/**
* Poll for container readiness in the background
* Sends notification to client when container is ready
*/
private async pollContainerReadiness(
socket: WebSocket,
authContext: any,
app: FastifyInstance
): Promise<void> {
const logger = app.log;
const userId = authContext.userId;
logger.info({ userId }, 'Starting background poll for container readiness');
try {
// Wait for container to become ready (2 minute timeout)
const ready = await this.config.containerManager.waitForContainerReady(userId, 120000);
if (ready) {
logger.info({ userId }, 'Container is now ready, notifying client');
// Send ready notification
socket.send(
JSON.stringify({
type: 'status',
status: 'ready',
message: 'Your workspace is ready!',
})
);
// Also send the 'connected' message
socket.send(
JSON.stringify({
type: 'connected',
sessionId: authContext.sessionId,
userId: authContext.userId,
licenseType: authContext.license.licenseType,
message: 'Connected to Dexorder AI',
})
);
} else {
logger.warn({ userId }, 'Container failed to become ready within timeout');
socket.send(
JSON.stringify({
type: 'error',
message: 'Workspace failed to start. Please try again later.',
})
);
}
} catch (error) {
logger.error({ error, userId }, 'Error waiting for container readiness');
socket.send(
JSON.stringify({
type: 'error',
message: 'Error starting workspace. Please try again later.',
})
);
}
}
/**
* Derive the container's XPUB event endpoint from the MCP server URL.
*
* MCP URL format: http://agent-user-abc123.dexorder-agents.svc.cluster.local:3000
* Event endpoint: tcp://agent-user-abc123.dexorder-agents.svc.cluster.local:5570
* MCP URL format: http://sandbox-user-abc123.dexorder-sandboxes.svc.cluster.local:3000
* Event endpoint: tcp://sandbox-user-abc123.dexorder-sandboxes.svc.cluster.local:5570
*/
private getContainerEventEndpoint(mcpServerUrl: string): string {
try {
@@ -578,4 +530,14 @@ export class WebSocketHandler {
);
}
}
/**
* 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();
}
}