sandbox connected and streaming
This commit is contained in:
@@ -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();
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user