omnara 0.1.3 → 0.1.4

package/dist/relay/relay-bridge.d.ts

@@ -1,6 +1,6 @@
 import EventEmitter from 'events';
-import { AgentOrchestrator } from '../orchestrator/agent-orchestrator';
-import { CredentialsStore } from '../auth/credentials-store';
+import { AgentOrchestrator } from '../orchestrator/agent-orchestrator.js';
+import { CredentialsStore } from '../auth/credentials-store.js';
 export interface RelayBridgeOptions {
     store?: CredentialsStore;
     logger?: Console;
@@ -20,15 +20,36 @@ export declare class RelayBridge extends EventEmitter {
     private pingTimer;
     private heartbeatIntervalMs;
     private authProviderPromise;
-    private currentSessionId;
+    private userSessionId;
+    private currentAgentSessionId;
     constructor(orchestrator: AgentOrchestrator, options: RelayBridgeOptions);
     start(): Promise<void>;
     stop(): Promise<void>;
     private establishConnection;
     private getAuthProvider;
+    private ensureAgentSessionExists;
+    private ensureUserSession;
+    private sendPrimarySessionInit;
+    private sendFrame;
     private handleWebsocketLifecycle;
     private handleOutbound;
     private handleInbound;
+    private toRelayFrame;
+    private tryHandleControlMessage;
+    private tryHandleMessageEnvelope;
+    private handleSessionBootstrap;
+    private handleSessionInitAck;
+    private handleSessionAnnounce;
+    private handleMessageEnvelope;
+    private isTextContent;
+    private buildTextEnvelope;
+    private buildToolCallEnvelope;
+    private normalizeToolArguments;
+    private buildToolResultEnvelope;
+    private buildAgentCompleteEnvelope;
+    private buildProgressEnvelope;
+    private buildEnvelope;
+    private isMessageEnvelope;
     private enqueue;
     private flushQueue;
     private startPingTimer;

package/dist/relay/relay-bridge.js

@@ -1,47 +1,35 @@
-"use strict";
-var __importDefault = (this && this.__importDefault) || function (mod) {
-    return (mod && mod.__esModule) ? mod : { "default": mod };
-};
-Object.defineProperty(exports, "__esModule", { value: true });
-exports.RelayBridge = void 0;
-const events_1 = __importDefault(require("events"));
-const ws_1 = __importDefault(require("ws"));
-const types_1 = require("../orchestrator/types");
-const credentials_store_1 = require("../auth/credentials-store");
-const api_client_1 = require("../http/api-client");
-const relay_service_1 = require("./relay-service");
-const auth_provider_1 = require("../auth/auth-provider");
+import { randomUUID } from 'crypto';
+import EventEmitter from 'events';
+import WebSocket from 'ws';
+import { AgentOrchestrator, } from '../orchestrator/agent-orchestrator.js';
+import { CredentialsStore } from '../auth/credentials-store.js';
+import { ApiClient, normalizeBaseUrl } from '../http/api-client.js';
+import { registerRelaySession } from './relay-service.js';
+import { createAuthProvider } from '../auth/auth-provider.js';
+import { createUserSession } from '../http/user-sessions.js';
+import { validateControlMessage, validateMessageEnvelope, } from '../schemas/index.js';
 const MAX_QUEUE_SIZE = 500;
-const OUTBOUND_MESSAGE_TYPES = new Set([
-    'session_created',
-    'session_joined',
-    'session_init',
-    'session_initialized',
-    'user_text',
-    'assistant_text',
-    'assistant_complete',
-    'tool_call',
-    'tool_result',
-    'system',
-    'error',
-    'worker_update',
-    'pong',
-]);
-class RelayBridge extends events_1.default {
+export class RelayBridge extends EventEmitter {
+    orchestrator;
+    store;
+    logger;
+    backendUrl;
+    frontendUrl;
+    connectionHandle = null;
+    ws = null;
+    stopped = false;
+    pendingOutbound = [];
+    pingTimer = null;
+    heartbeatIntervalMs = 30_000;
+    authProviderPromise = null;
+    userSessionId = null;
+    currentAgentSessionId = null;
     constructor(orchestrator, options) {
         super();
-        this.connectionHandle = null;
-        this.ws = null;
-        this.stopped = false;
-        this.pendingOutbound = [];
-        this.pingTimer = null;
-        this.heartbeatIntervalMs = 30000;
-        this.authProviderPromise = null;
-        this.currentSessionId = null;
         this.orchestrator = orchestrator;
-        this.store = options.store ?? new credentials_store_1.CredentialsStore();
+        this.store = options.store ?? new CredentialsStore();
         this.logger = options.logger ?? console;
-        this.backendUrl = (0, api_client_1.normalizeBaseUrl)(options.backendUrl);
+        this.backendUrl = normalizeBaseUrl(options.backendUrl);
         this.frontendUrl = options.frontendUrl;
     }
     async start() {
@@ -49,16 +37,17 @@ class RelayBridge extends events_1.default {
             throw new Error('Relay bridge already started');
         }
         this.connectionHandle = this.orchestrator.addClient({
-            canSend: () => this.ws?.readyState === ws_1.default.OPEN,
+            canSend: () => this.ws?.readyState === WebSocket.OPEN,
             send: (message) => {
                 this.handleOutbound(message);
             },
         });
-        let backoffMs = 1000;
+        this.ensureAgentSessionExists();
+        let backoffMs = 1_000;
         while (!this.stopped) {
             try {
                 await this.establishConnection();
-                backoffMs = 1000;
+                backoffMs = 1_000;
             }
             catch (error) {
                 const aborted = isAbortError(error);
@@ -70,13 +59,13 @@ class RelayBridge extends events_1.default {
                 }
                 this.logger.error(`Relay connection failed: ${extractErrorMessage(error)}. Retrying in ${Math.round(backoffMs / 1000)}s.`);
                 await delay(backoffMs);
-                backoffMs = Math.min(backoffMs * 2, 30000);
+                backoffMs = Math.min(backoffMs * 2, 30_000);
             }
         }
     }
     async stop() {
         this.stopped = true;
-        this.currentSessionId = null;
+        this.currentAgentSessionId = null;
         if (this.pingTimer) {
             clearInterval(this.pingTimer);
             this.pingTimer = null;
@@ -85,36 +74,41 @@ class RelayBridge extends events_1.default {
             this.connectionHandle.disconnect();
             this.connectionHandle = null;
         }
-        if (this.ws && this.ws.readyState === ws_1.default.OPEN) {
+        if (this.ws && this.ws.readyState === WebSocket.OPEN) {
             this.ws.close(1000, 'CLI shutting down');
         }
         this.ws = null;
     }
     async establishConnection() {
         const authProvider = await this.getAuthProvider();
-        const client = new api_client_1.ApiClient({ baseUrl: this.backendUrl });
+        const client = new ApiClient({ baseUrl: this.backendUrl });
         const accessToken = await authProvider.getAccessToken();
-        const registration = await (0, relay_service_1.registerRelaySession)(client, accessToken, this.currentSessionId ? { sessionId: this.currentSessionId } : undefined);
-        this.currentSessionId = registration.sessionId;
-        this.orchestrator.ensureSessionExists(this.currentSessionId);
+        const userSessionId = await this.ensureUserSession(client, accessToken);
+        const registration = await registerRelaySession(client, accessToken, {
+            userSessionId,
+            clientKind: 'cli',
+        });
+        this.userSessionId = registration.userSessionId;
+        const agentSessionId = this.ensureAgentSessionExists();
         const url = buildRelayWebsocketUrl(this.backendUrl, registration.relayWsPath);
-        const ws = new ws_1.default(url, {
+        const ws = new WebSocket(url, {
             headers: {
                 authorization: `Bearer ${accessToken}`,
                 'x-relay-token': registration.relayToken,
             },
         });
         await waitForOpen(ws);
-        this.logger.log(`Connected to relay session ${registration.sessionId}`);
+        this.logger.log(`Connected to relay (user_session=${registration.userSessionId}, agent_session=${agentSessionId})`);
         this.ws = ws;
         this.heartbeatIntervalMs = registration.relayHeartbeatIntervalSec * 1000;
         this.startPingTimer();
+        await this.sendPrimarySessionInit(ws);
         this.flushQueue();
         await this.handleWebsocketLifecycle(ws);
     }
     async getAuthProvider() {
         if (!this.authProviderPromise) {
-            this.authProviderPromise = (0, auth_provider_1.createAuthProvider)({
+            this.authProviderPromise = createAuthProvider({
                 store: this.store,
                 backendUrl: this.backendUrl,
                 frontendUrl: this.frontendUrl,
@@ -123,6 +117,64 @@ class RelayBridge extends events_1.default {
         }
         return this.authProviderPromise;
     }
+    ensureAgentSessionExists() {
+        if (!this.currentAgentSessionId) {
+            this.currentAgentSessionId = randomUUID();
+        }
+        this.orchestrator.ensureSessionExists(this.currentAgentSessionId);
+        return this.currentAgentSessionId;
+    }
+    async ensureUserSession(client, accessToken) {
+        if (this.userSessionId) {
+            return this.userSessionId;
+        }
+        const descriptor = await createUserSession(client, accessToken, {
+            metadata: { source: 'cli' },
+        });
+        this.userSessionId = descriptor.session_id;
+        this.logger.log(`Created user session ${this.userSessionId} for CLI orchestrator`);
+        return this.userSessionId;
+    }
+    async sendPrimarySessionInit(ws) {
+        if (!this.userSessionId) {
+            throw new Error('Cannot send session_init before user session is available.');
+        }
+        const agentSessionId = this.ensureAgentSessionExists();
+        const initFrame = {
+            type: 'session_init',
+            session_id: agentSessionId,
+            user_session_id: this.userSessionId,
+            session_type: 'CODE',
+            parent_session_id: null,
+        };
+        this.sendFrame(initFrame, ws);
+    }
+    sendFrame(frame, socket) {
+        if (this.isMessageEnvelope(frame)) {
+            if (!validateMessageEnvelope(frame)) {
+                this.logger.error('[RelayBridge] Message envelope validation failed before sending', {
+                    frame: JSON.stringify(frame),
+                    errors: validateMessageEnvelope.errors,
+                });
+                return;
+            }
+        }
+        const serialized = JSON.stringify(frame);
+        const target = socket ?? this.ws;
+        if (target && target.readyState === WebSocket.OPEN) {
+            try {
+                target.send(serialized);
+                return;
+            }
+            catch (error) {
+                this.logger.error(`Failed to send relay frame: ${extractErrorMessage(error)}`);
+                if (target === this.ws) {
+                    this.closeCurrentSocket();
+                }
+            }
+        }
+        this.enqueue(serialized);
+    }
     async handleWebsocketLifecycle(ws) {
         await new Promise((resolve) => {
             const onMessage = (data) => {
@@ -160,20 +212,11 @@ class RelayBridge extends events_1.default {
         });
     }
     handleOutbound(message) {
-        const payload = JSON.stringify(message);
-        if (this.ws && this.ws.readyState === ws_1.default.OPEN) {
-            try {
-                this.ws.send(payload);
-            }
-            catch (error) {
-                this.logger.error(`Failed to send relay message: ${extractErrorMessage(error)}`);
-                this.enqueue(payload);
-                this.closeCurrentSocket();
-            }
-        }
-        else {
-            this.enqueue(payload);
+        const frame = this.toRelayFrame(message);
+        if (!frame) {
+            return;
         }
+        this.sendFrame(frame);
     }
     handleInbound(data) {
         let parsed;
@@ -185,16 +228,273 @@ class RelayBridge extends events_1.default {
             this.logger.error(`Received invalid JSON from relay: ${extractErrorMessage(error)}`);
             return;
         }
-        const result = types_1.inboundMessageSchema.safeParse(parsed);
-        if (result.success) {
-            this.connectionHandle?.dispatch(result.data);
+        if (this.tryHandleControlMessage(parsed)) {
             return;
         }
-        if (isOutboundMessage(parsed)) {
-            // Relay echoed a message originally produced by the orchestrator; ignore.
+        if (this.tryHandleMessageEnvelope(parsed)) {
             return;
         }
-        this.logger.error('Received invalid relay message payload.', parsed);
+        this.logger.error('Received invalid relay message payload.', { payload: parsed });
+    }
+    toRelayFrame(message) {
+        switch (message.type) {
+            case 'session_init': {
+                if (!this.userSessionId) {
+                    this.logger.warn('Dropping session_init frame: user session is not initialised yet.');
+                    return null;
+                }
+                const frame = {
+                    type: 'session_init',
+                    session_id: message.sessionId,
+                    user_session_id: this.userSessionId,
+                    session_type: 'CODE',
+                    parent_session_id: message.parentSessionId ?? null,
+                };
+                if (message.metadata) {
+                    frame.metadata = message.metadata;
+                }
+                return frame;
+            }
+            case 'assistant_text':
+                return this.buildTextEnvelope(message);
+            case 'tool_call':
+                return this.buildToolCallEnvelope(message);
+            case 'tool_result':
+                return this.buildToolResultEnvelope(message);
+            case 'agent_complete':
+                return this.buildAgentCompleteEnvelope(message);
+            case 'agent_progress':
+                return this.buildProgressEnvelope(message);
+            default:
+                return null;
+        }
+    }
+    tryHandleControlMessage(payload) {
+        if (!validateControlMessage(payload)) {
+            return false;
+        }
+        const control = payload;
+        const type = control.type;
+        switch (type) {
+            case 'session_bootstrap':
+                this.handleSessionBootstrap(control);
+                break;
+            case 'session_init_ack':
+                this.handleSessionInitAck(control);
+                break;
+            case 'session_announce':
+                this.handleSessionAnnounce(control);
+                break;
+            case 'session_closed':
+                this.logger.info(`Session closed: ${control.session_id} -> status ${control.status}`);
+                break;
+            case 'heartbeat':
+                if (control.mode === 'ping') {
+                    const response = {
+                        type: 'heartbeat',
+                        mode: 'pong',
+                        sent_at: new Date().toISOString(),
+                    };
+                    this.sendFrame(response);
+                }
+                break;
+            default:
+                this.logger.debug(`Ignoring control message ${type ?? 'unknown'}`);
+        }
+        return true;
+    }
+    tryHandleMessageEnvelope(payload) {
+        if (!validateMessageEnvelope(payload)) {
+            return false;
+        }
+        const envelope = payload;
+        this.handleMessageEnvelope(envelope);
+        return true;
+    }
+    handleSessionBootstrap(control) {
+        const userSessionId = control.root.session_id;
+        if (!this.userSessionId) {
+            this.userSessionId = userSessionId;
+        }
+        if (this.userSessionId !== userSessionId) {
+            this.logger.warn(`Bootstrap user session mismatch (expected ${this.userSessionId}, got ${userSessionId})`);
+        }
+        for (const agent of control.agent_sessions ?? []) {
+            this.orchestrator.ensureSessionExists(agent.session_id);
+        }
+    }
+    handleSessionInitAck(control) {
+        const sessionId = control.session_id;
+        if (!this.currentAgentSessionId) {
+            this.currentAgentSessionId = sessionId;
+        }
+        if (this.currentAgentSessionId !== sessionId) {
+            this.logger.warn(`Session init ack for unexpected session: ${sessionId} (current=${this.currentAgentSessionId})`);
+        }
+    }
+    handleSessionAnnounce(control) {
+        if (control.session.kind === 'agent_session') {
+            this.orchestrator.ensureSessionExists(control.session.session_id);
+        }
+    }
+    handleMessageEnvelope(envelope) {
+        const content = envelope.payload.content;
+        if (this.isTextContent(content) && envelope.sender.kind === 'user') {
+            this.connectionHandle?.dispatch({
+                type: 'chat',
+                sessionId: envelope.session_id,
+                text: content.text,
+            });
+            return;
+        }
+        const label = content.type ?? 'unknown';
+        this.logger.debug(`Ignoring inbound envelope with content type ${label}`);
+    }
+    isTextContent(content) {
+        return typeof content.text === 'string';
+    }
+    buildTextEnvelope(message) {
+        const content = {
+            type: 'text',
+            text: message.text,
+        };
+        const payloadMetadata = message.workerId || message.workerTask
+            ? {
+                worker_id: message.workerId,
+                worker_task: message.workerTask,
+            }
+            : undefined;
+        const senderSessionId = message.workerId ?? message.sessionId;
+        return this.buildEnvelope({
+            sessionId: message.sessionId,
+            senderSessionId,
+            content,
+            payloadMetadata,
+            messageId: message.messageId,
+        });
+    }
+    buildToolCallEnvelope(message) {
+        const args = this.normalizeToolArguments(message.input);
+        const content = {
+            type: 'tool_call',
+            name: message.toolName,
+            tool_use_id: message.toolUseId,
+            arguments: args,
+        };
+        const payloadMetadata = message.workerId || message.workerTask
+            ? {
+                worker_id: message.workerId,
+                worker_task: message.workerTask,
+            }
+            : undefined;
+        const senderSessionId = message.workerId ?? message.sessionId;
+        return this.buildEnvelope({
+            sessionId: message.sessionId,
+            senderSessionId,
+            content,
+            payloadMetadata,
+            messageId: message.messageId,
+        });
+    }
+    normalizeToolArguments(input) {
+        if (input && typeof input === 'object' && !Array.isArray(input)) {
+            return input;
+        }
+        return { value: input };
+    }
+    buildToolResultEnvelope(message) {
+        const output = typeof message.content === 'object' && message.content !== null
+            ? message.content
+            : { value: message.content };
+        const content = {
+            type: 'tool_result',
+            tool_use_id: message.toolUseId,
+            output,
+            is_error: Boolean(message.isError),
+        };
+        const payloadMetadata = message.workerId || message.workerTask
+            ? {
+                worker_id: message.workerId,
+                worker_task: message.workerTask,
+            }
+            : undefined;
+        const senderSessionId = message.workerId ?? message.sessionId;
+        return this.buildEnvelope({
+            sessionId: message.sessionId,
+            senderSessionId,
+            content,
+            payloadMetadata,
+            messageId: message.messageId,
+        });
+    }
+    buildAgentCompleteEnvelope(message) {
+        const content = {
+            type: 'agent_complete',
+            success: message.success,
+            duration_ms: message.durationMs ?? null,
+            cost_usd: message.costUsd ?? null,
+            message: message.message ?? null,
+        };
+        if (message.metadata) {
+            content.metadata = message.metadata;
+        }
+        const metadataAgent = typeof message.metadata?.agent_session_id === 'string'
+            ? message.metadata.agent_session_id
+            : undefined;
+        const senderSessionId = metadataAgent ?? this.currentAgentSessionId ?? message.sessionId;
+        return this.buildEnvelope({
+            sessionId: message.sessionId,
+            senderSessionId,
+            content,
+            messageId: message.messageId,
+        });
+    }
+    buildProgressEnvelope(message) {
+        const content = {
+            type: 'agent_progress',
+            status: message.status,
+        };
+        if (message.detail) {
+            content.detail = message.detail;
+        }
+        if (message.metadata) {
+            content.metadata = message.metadata;
+        }
+        const metadataAgent = typeof message.metadata?.agent_session_id === 'string'
+            ? message.metadata.agent_session_id
+            : undefined;
+        const sessionId = message.sessionId ?? this.currentAgentSessionId ?? this.ensureAgentSessionExists();
+        const senderSessionId = metadataAgent ?? this.currentAgentSessionId ?? sessionId;
+        return this.buildEnvelope({
+            sessionId,
+            senderSessionId,
+            content,
+            messageId: message.messageId,
+        });
+    }
+    buildEnvelope(options) {
+        const payload = {
+            version: 0,
+            content: options.content,
+        };
+        if (options.payloadMetadata && Object.keys(options.payloadMetadata).length > 0) {
+            payload.metadata = options.payloadMetadata;
+        }
+        const sender = {
+            kind: 'agent_session',
+            id: options.senderSessionId ?? options.sessionId,
+        };
+        const envelope = {
+            message_id: options.messageId ?? randomUUID(),
+            session_id: options.sessionId,
+            sender,
+            payload,
+            created_at: new Date().toISOString(),
+        };
+        return envelope;
+    }
+    isMessageEnvelope(frame) {
+        return frame.payload !== undefined;
     }
     enqueue(serialized) {
         if (this.pendingOutbound.length >= MAX_QUEUE_SIZE) {
@@ -203,7 +503,7 @@ class RelayBridge extends events_1.default {
         this.pendingOutbound.push(serialized);
     }
     flushQueue() {
-        if (!this.ws || this.ws.readyState !== ws_1.default.OPEN) {
+        if (!this.ws || this.ws.readyState !== WebSocket.OPEN) {
             return;
         }
         while (this.pendingOutbound.length > 0) {
@@ -226,27 +526,26 @@ class RelayBridge extends events_1.default {
         if (this.pingTimer) {
             clearInterval(this.pingTimer);
         }
-        const interval = Math.max(this.heartbeatIntervalMs, 5000);
+        const interval = Math.max(this.heartbeatIntervalMs, 5_000);
         this.pingTimer = setInterval(() => {
-            if (!this.ws || this.ws.readyState !== ws_1.default.OPEN) {
+            const socket = this.ws;
+            if (!socket || socket.readyState !== WebSocket.OPEN) {
                 return;
             }
-            try {
-                this.ws.send(JSON.stringify({ type: 'ping' }));
-            }
-            catch (error) {
-                this.logger.error(`Failed to send relay ping: ${extractErrorMessage(error)}`);
-                this.closeCurrentSocket();
-            }
+            const frame = {
+                type: 'heartbeat',
+                mode: 'ping',
+                sent_at: new Date().toISOString(),
+            };
+            this.sendFrame(frame, socket);
         }, interval);
     }
     closeCurrentSocket() {
-        if (this.ws && this.ws.readyState === ws_1.default.OPEN) {
+        if (this.ws && this.ws.readyState === WebSocket.OPEN) {
             this.ws.close(1012, 'Reconnecting');
         }
     }
 }
-exports.RelayBridge = RelayBridge;
 function buildRelayWebsocketUrl(baseUrl, path) {
     const url = new URL(baseUrl);
     url.protocol = url.protocol === 'https:' ? 'wss:' : 'ws:';
@@ -270,13 +569,6 @@ function extractErrorMessage(error) {
     }
     return String(error);
 }
-function isOutboundMessage(value) {
-    if (!value || typeof value !== 'object') {
-        return false;
-    }
-    const type = value.type;
-    return typeof type === 'string' && OUTBOUND_MESSAGE_TYPES.has(type);
-}
 function isAbortError(error) {
     if (!(error instanceof Error)) {
         return false;