@kadi.build/core 0.0.1-alpha.3 → 0.0.1-alpha.5

package/src/servers/BrokerRpcServer.js

@@ -1,776 +0,0 @@
-import { BaseRpcServer } from './BaseRpcServer.js';
-import { Broker, IdFactory } from '../transport/BrokerMessageBuilder.js';
-import { WebSocket } from 'ws';
-import crypto from 'node:crypto';
-import { createComponentLogger } from '../utils/logger.js';
-/**
- * RPC Server for broker transport using WebSocket
- *
- * This server connects abilities to the Kadi broker via WebSocket and speaks
- * the broker's native protocol directly (agent.message/ability.result),
- * providing distributed ability execution across the network.
- */
-export class BrokerRpcServer extends BaseRpcServer {
-  /**
-   * Create a new BrokerRpcServer instance
-   *
-   * @param {Object} options - Configuration options
-   * @param {string} options.brokerUrl - Broker WebSocket URL
-   * @param {string} options.serviceName - Service name for broker registration
-   * @param {number} options.connectTimeoutMs - Connection timeout
-   * @param {number} options.heartbeatIntervalMs - Heartbeat interval
-   */
-  constructor(options = {}) {
-    super({ ...options, protocol: 'broker' });
-
-    this.brokerUrl =
-      options.brokerUrl || process.env.KADI_BROKER_URL || 'ws://localhost:8080';
-    this.serviceName =
-      options.serviceName || process.env.KADI_SERVICE_NAME || 'unnamed-ability';
-    this.connectTimeoutMs = options.connectTimeoutMs || 10000;
-    this.heartbeatIntervalMs = options.heartbeatIntervalMs || 25000;
-
-    this.logger = createComponentLogger('BrokerRpcServer');
-    this.logger.lifecycle('constructor', 'BrokerRpcServer initialized');
-    this.logger.trace('constructor', `Broker URL: ${this.brokerUrl}`);
-    this.logger.trace('constructor', `Service name: ${this.serviceName}`);
-    this.logger.trace(
-      'constructor',
-      `Connect timeout: ${this.connectTimeoutMs}ms`
-    );
-
-    this._ws = null;
-    this._isConnected = false;
-    this._idFactory = new IdFactory();
-    this._requestMetadata = new Map(); // Store broker metadata by request ID
-    this._heartbeatTimer = null;
-    this._reconnectAttempts = 0;
-    this._maxReconnectAttempts = options.maxReconnectAttempts || 5;
-    this._reconnectDelayMs = options.reconnectDelayMs || 1000;
-
-    this.ability = null;
-    this.resolve = null;
-    this.reject = null;
-  }
-
-  /**
-   * Start serving the ability via broker
-   *
-   * @param {KadiAbility} ability - The ability instance to serve
-   * @returns {Promise<void>} - Promise that resolves when server stops
-   */
-  async serve(ability) {
-    if (this.isServing) {
-      this.logger.error('serve', 'BrokerRpcServer is already serving');
-      throw new Error('BrokerRpcServer is already serving');
-    }
-
-    this.ability = ability;
-    this.isServing = true;
-
-    this.logger.lifecycle(
-      'serve',
-      `Starting ${ability.name || 'unnamed ability'} on broker`
-    );
-    this.logger.info('serve', `Broker URL: ${this.brokerUrl}`);
-    this.logger.info('serve', `Service name: ${this.serviceName}`);
-    this.logger.info(
-      'serve',
-      `Available methods: ${ability.getMethodNames().join(', ')}`
-    );
-
-    try {
-      // Connect to broker and complete handshake
-      await this._connect();
-      await this._handshake();
-
-      // Emit start event
-      this.emit('start', {
-        name: ability.name,
-        version: ability.version,
-        methods: ability.getMethodNames()
-      });
-
-      this.logger.success(
-        'serve',
-        'Broker service started, ready to receive tool calls'
-      );
-
-      // Keep the process alive and handle messages
-      return new Promise((resolve, reject) => {
-        this.resolve = resolve;
-        this.reject = reject;
-      });
-    } catch (error) {
-      this.isServing = false;
-      this.logger.error(
-        'serve',
-        `Failed to start broker service: ${error.message}`
-      );
-      this.emit('error', error);
-      throw error;
-    }
-  }
-
-  /**
-   * Publish an event to connected agents via the broker
-   *
-   * Events are sent as 'kadi.event' messages through the WebSocket connection.
-   * The broker will route these to subscribed agents.
-   *
-   * @param {string} eventName - Name of the event to publish
-   * @param {any} data - Event data payload (must be JSON-serializable)
-   */
-  async publishEvent(eventName, data = {}) {
-    this.logger.trace(
-      'publishEvent',
-      `Publishing event via broker: ${eventName}`
-    );
-
-    if (!this._ws || !this._isConnected) {
-      this.logger.warn(
-        'publishEvent',
-        `Cannot publish event ${eventName}: not connected to broker`
-      );
-      return;
-    }
-
-    try {
-      // Send event as a broker notification (no id field)
-      const eventMessage = {
-        jsonrpc: '2.0',
-        method: 'kadi.event',
-        params: {
-          eventName,
-          eventData: data,
-          timestamp: Date.now(),
-          from: this.serviceName // Include source for routing
-        }
-      };
-
-      this._ws.send(JSON.stringify(eventMessage));
-
-      this.logger.trace('publishEvent', `Event ${eventName} sent to broker`);
-    } catch (error) {
-      // Events are best-effort - log but don't throw
-      this.logger.warn(
-        'publishEvent',
-        `Failed to publish event ${eventName}: ${error.message}`
-      );
-    }
-  }
-
-  /**
-   * Handle requests in broker-native format
-   *
-   * This implementation handles incoming tool calls from the broker and
-   * converts them to internal method calls. Broker requests come as
-   * agent.message notifications with tool execution parameters.
-   *
-   * @param {Object} message - The broker message object
-   */
-  async _handleBrokerMessage(message) {
-    try {
-      this.logger.trace(
-        'message',
-        `Received broker message: ${message.method}`
-      );
-      this.logger.trace('message-detail', JSON.stringify(message, null, 2));
-
-      // Handle incoming tool calls
-      if (message.method === 'agent.message' && message.params) {
-        await this._handleToolCall(message);
-      }
-      // Handle ability results (responses to our tool calls)
-      else if (message.method === 'ability.result') {
-        this.logger.trace('message', 'Received ability.result (not handling)');
-      }
-      // Handle other broker messages
-      else {
-        this.logger.warn(
-          'message',
-          `Unhandled broker message: ${message.method}`
-        );
-        this.logger.trace('message-detail', JSON.stringify(message, null, 2));
-      }
-    } catch (error) {
-      this.logger.error(
-        'message',
-        `Failed to handle broker message: ${error.message}`
-      );
-      this.logError('Failed to handle broker message:', error);
-      this.emit('error', error);
-    }
-  }
-
-  /**
-   * Handle incoming tool calls from the broker
-   *
-   * @param {Object} message - Broker message containing tool call
-   */
-  async _handleToolCall(message) {
-    const { toolName, args, requestId, from } = message.params;
-
-    this.logger.request(requestId, toolName, 'Handling tool call');
-    this.logger.trace(
-      'broker-tool-call',
-      `From: ${from}, Args: ${JSON.stringify(args)}`
-    );
-
-    try {
-      // Convert broker message to internal request format
-      const request = {
-        id: requestId,
-        method: toolName,
-        params: args || {}
-      };
-
-      // Store broker metadata for response routing
-      this._requestMetadata.set(requestId, { requestId, from });
-
-      // Process the request using base class logic
-      const response = await this.handleRequest(request);
-
-      // Send response back to broker
-      await this._sendBrokerResponse(requestId, from, response);
-    } catch (error) {
-      this.logger.error(
-        'tool-call',
-        `Error handling tool call ${toolName}: ${error.message}`
-      );
-
-      this.logError(`Error handling tool call ${toolName}:`, error);
-
-      // Send error response to broker
-      await this._sendBrokerError(requestId, from, error.message);
-    }
-  }
-
-  /**
-   * Send a successful result back to the broker
-   *
-   * @param {string} requestId - Original request ID
-   * @param {string} toSessionId - Target session ID
-   * @param {Object} response - Response object
-   */
-  async _sendBrokerResponse(requestId, toSessionId, response) {
-    if (!response) {
-      // Notification - no response needed
-      this.logger.trace(
-        'response',
-        `No response needed for notification ${requestId}`
-      );
-      return;
-    }
-
-    const brokerResponse = {
-      jsonrpc: '2.0',
-      method: 'ability.result',
-      params: {
-        requestId,
-        toSessionId,
-        result: response.result
-      }
-    };
-
-    this.logger.response(requestId, 'success', 'Sending broker response');
-    this.logger.trace(
-      'response-detail',
-      JSON.stringify(brokerResponse, null, 2)
-    );
-
-    this._ws.send(JSON.stringify(brokerResponse));
-
-    // Clean up metadata
-    this._requestMetadata.delete(requestId);
-  }
-
-  /**
-   * Send an error result back to the broker
-   *
-   * @param {string} requestId - Original request ID
-   * @param {string} toSessionId - Target session ID
-   * @param {string} errorMessage - Error message
-   */
-  async _sendBrokerError(requestId, toSessionId, errorMessage) {
-    const brokerResponse = {
-      jsonrpc: '2.0',
-      method: 'ability.result',
-      params: {
-        requestId,
-        toSessionId,
-        error: errorMessage
-      }
-    };
-
-    this.logger.response(
-      requestId,
-      'error',
-      `Sending broker error: ${errorMessage}`
-    );
-    this.logger.trace(
-      'broker-response-detail',
-      JSON.stringify(brokerResponse, null, 2)
-    );
-
-    this._ws.send(JSON.stringify(brokerResponse));
-
-    // Clean up metadata
-    this._requestMetadata.delete(requestId);
-  }
-
-  /**
-   * Connect to the broker WebSocket
-   */
-  async _connect() {
-    this.logger.info('connect', `Connecting to broker at ${this.brokerUrl}`);
-
-    return new Promise((resolve, reject) => {
-      const timeout = setTimeout(() => {
-        this.logger.error(
-          'connect',
-          `Connection timeout after ${this.connectTimeoutMs}ms`
-        );
-
-        reject(
-          new Error(`Connection timeout after ${this.connectTimeoutMs}ms`)
-        );
-      }, this.connectTimeoutMs);
-
-      this._ws = new WebSocket(this.brokerUrl);
-
-      this._ws.on('open', () => {
-        clearTimeout(timeout);
-        this._isConnected = true;
-        this._reconnectAttempts = 0;
-        this.logger.success(
-          'connect',
-          `Connected to broker at ${this.brokerUrl}`
-        );
-        resolve();
-      });
-
-      this._ws.on('error', (error) => {
-        clearTimeout(timeout);
-        this.logger.error(
-          'connect',
-          `WebSocket connection error: ${error.message || error}`
-        );
-        this.logError('WebSocket connection error:', error.message);
-
-        reject(new Error(`Failed to connect to broker: ${error.message}`));
-      });
-
-      this._ws.on('close', (code, reason) => {
-        this._isConnected = false;
-        this.logger.warn(
-          'connect',
-          `Disconnected from broker (${code}): ${reason}`
-        );
-
-        if (
-          this.isServing &&
-          this._reconnectAttempts < this._maxReconnectAttempts
-        ) {
-          this._attemptReconnect();
-        } else {
-          this.shutdown('broker_disconnected');
-        }
-      });
-
-      this._ws.on('message', (data) => {
-        try {
-          const message = JSON.parse(data.toString());
-          this._handleBrokerMessage(message);
-        } catch (error) {
-          this.logError('Failed to parse broker message:', error);
-        }
-      });
-    });
-  }
-
-  /**
-   * Attempt to reconnect to the broker
-   */
-  async _attemptReconnect() {
-    this._reconnectAttempts++;
-    const delay =
-      this._reconnectDelayMs * Math.pow(2, this._reconnectAttempts - 1);
-
-    this.logger.info(
-      'reconnect',
-      `Attempting to reconnect (${this._reconnectAttempts}/${this._maxReconnectAttempts}) in ${delay}ms`
-    );
-
-    setTimeout(async () => {
-      try {
-        await this._connect();
-        await this._handshake();
-        this.logger.success('reconnect', 'Successfully reconnected to broker');
-      } catch (error) {
-        this.logger.error('reconnect', `Reconnection failed: ${error.message}`);
-        if (this._reconnectAttempts >= this._maxReconnectAttempts) {
-          this.shutdown('max_reconnect_attempts_reached');
-        }
-      }
-    }, delay);
-  }
-
-  /**
-   * Complete the broker handshake
-   */
-  async _handshake() {
-    this.logger.info('handshake', 'Starting broker handshake');
-
-    // Step 1: Send hello as agent
-    const hello = Broker.hello({ role: 'agent' })
-      .id(this._idFactory.next())
-      .build();
-    this._ws.send(JSON.stringify(hello));
-
-    // Wait for hello response with nonce
-    const helloResponse = await this._waitForResponse(hello.id);
-    if (!helloResponse.result || !helloResponse.result.nonce) {
-      this.logger.trace('handshake', 'Received nonce from broker');
-
-      throw new Error('Invalid hello response from broker');
-    }
-
-    const nonce = helloResponse.result.nonce;
-    this.log(`Received nonce from broker`);
-
-    // Step 2: Generate ephemeral keys and authenticate
-    const { publicKey, privateKey } = crypto.generateKeyPairSync('ed25519');
-    const publicKeyBase64 = publicKey
-      .export({ format: 'der', type: 'spki' })
-      .toString('base64');
-
-    const authenticate = Broker.authenticate({
-      publicKeyBase64Der: publicKeyBase64,
-      privateKey,
-      nonce,
-      wantNewId: true
-    })
-      .id(this._idFactory.next())
-      .build();
-
-    this.logger.trace('handshake', 'Sending authentication to broker');
-    this._ws.send(JSON.stringify(authenticate));
-
-    // Wait for authentication response
-    const authResponse = await this._waitForResponse(authenticate.id);
-    if (!authResponse.result || !authResponse.result.agentId) {
-      this.logger.error(
-        'handshake',
-        `Authentication failed: ${JSON.stringify(authResponse.error)}`
-      );
-
-      throw new Error(
-        `Authentication failed: ${JSON.stringify(authResponse.error)}`
-      );
-    }
-
-    this.logger.success(
-      'handshake',
-      `Authenticated as agent: ${authResponse.result.agentId}`
-    );
-
-    // Step 3: Register capabilities
-    await this._registerCapabilities();
-
-    // Step 4: Start heartbeat
-    this._startHeartbeat();
-
-    this.logger.success('handshake', 'Handshake completed successfully');
-  }
-
-  /**
-   * Register capabilities with the broker
-   */
-  async _registerCapabilities() {
-    if (
-      !this.ability ||
-      typeof this.ability.extractToolsForBroker !== 'function'
-    ) {
-      this.logger.warn(
-        'capabilities',
-        'No ability reference or extractToolsForBroker method not found'
-      );
-
-      return;
-    }
-
-    try {
-      this.logger.info('capabilities', 'Extracting tools from ability');
-      const tools = await this.ability.extractToolsForBroker();
-      // this.log(`Extracted tools:`, JSON.stringify(tools, null, 2));
-
-      if (tools.length > 0) {
-        const registerCapabilities = Broker.registerCapabilities({
-          displayName: this.serviceName,
-          tools,
-          mailboxMode: 'persistent',
-          scopes: [this.ability.scope] // We want to register the tools in the same scope as the agent
-        })
-          .id(this._idFactory.next())
-          .build();
-
-        this._ws.send(JSON.stringify(registerCapabilities));
-
-        // Wait for registration response
-        const registrationResponse = await this._waitForResponse(
-          registerCapabilities.id
-        );
-        this.logger.trace(
-          'capabilities',
-          `Registration response: ${JSON.stringify(registrationResponse, null, 2)}`
-        );
-        this.logger.success(
-          'capabilities',
-          `Registered ${tools.length} capabilities with broker`
-        );
-      } else {
-        this.logger.warn(
-          'capabilities',
-          'No tools found, skipping capability registration'
-        );
-      }
-    } catch (error) {
-      this.logger.error(
-        'capabilities',
-        `Failed to register capabilities: ${error.message}`
-      );
-      throw error;
-    }
-  }
-
-  /**
-   * Start heartbeat to keep connection alive
-   */
-  _startHeartbeat() {
-    this.logger.trace(
-      'heartbeat',
-      `Starting heartbeat every ${this.heartbeatIntervalMs}ms`
-    );
-
-    this._heartbeatTimer = setInterval(() => {
-      if (this._isConnected && this._ws.readyState === WebSocket.OPEN) {
-        if (this.logger.enabled) {
-          this.logger.trace('heartbeat', 'Sending ping to broker');
-        }
-        this._ws.send(
-          JSON.stringify({
-            jsonrpc: '2.0',
-            method: 'kadi.ping'
-          })
-        );
-      }
-    }, this.heartbeatIntervalMs);
-  }
-
-  /**
-   * Wait for a response to a specific message
-   *
-   * @param {string|number} messageId - Message ID to wait for
-   * @param {number} timeoutMs - Timeout in milliseconds
-   * @returns {Promise<Object>} - Response message
-   */
-  async _waitForResponse(messageId, timeoutMs = 5000) {
-    this.logger.trace(
-      'wait-response',
-      `Waiting for response to message ${messageId} (timeout: ${timeoutMs}ms)`
-    );
-
-    return new Promise((resolve, reject) => {
-      const timeout = setTimeout(() => {
-        this._ws.off('message', handler);
-        this.logger.error(
-          'wait-response',
-          `Timeout waiting for response to message ${messageId}`
-        );
-
-        reject(
-          new Error(`Timeout waiting for response to message ${messageId}`)
-        );
-      }, timeoutMs);
-
-      const handler = (data) => {
-        try {
-          const message = JSON.parse(data.toString());
-          if (message.id === messageId) {
-            clearTimeout(timeout);
-            this._ws.off('message', handler);
-            this.logger.trace(
-              'wait-response',
-              `Received response for message ${messageId}`
-            );
-
-            resolve(message);
-          }
-        } catch (error) {
-          // Ignore parsing errors for non-matching messages
-        }
-      };
-
-      this._ws.on('message', handler);
-    });
-  }
-
-  /**
-   * Gracefully shutdown the broker server
-   *
-   * @param {string} reason - Reason for shutdown
-   */
-  async shutdown(reason = 'unknown') {
-    this.logger.lifecycle(
-      'shutdown',
-      `Shutting down broker server, reason: ${reason}`
-    );
-
-    await super.shutdown(reason);
-
-    // Stop heartbeat
-    if (this._heartbeatTimer) {
-      this.logger.trace('shutdown', 'Stopping heartbeat timer');
-
-      clearInterval(this._heartbeatTimer);
-      this._heartbeatTimer = null;
-    }
-
-    // Close WebSocket connection
-    if (this._ws) {
-      this._isConnected = false;
-      this.logger.trace('shutdown', 'Closing WebSocket connection');
-
-      this._ws.close();
-      this._ws = null;
-    }
-
-    // Clear request metadata
-    this._requestMetadata.clear();
-
-    // Resolve the serve promise
-    if (this.resolve) {
-      this.logger.trace('shutdown', 'Resolving serve promise');
-      this.resolve();
-    }
-  }
-
-  /**
-   * Send a message to the broker
-   *
-   * @param {Object} message - Message to send
-   */
-  async sendMessage(message) {
-    if (!this._ws || !this._isConnected) {
-      throw new Error('Not connected to broker');
-    }
-
-    this._ws.send(JSON.stringify(message));
-  }
-
-  /**
-   * Call another ability through the broker
-   *
-   * @param {string} toolName - Name of the tool/ability to call
-   * @param {Object} args - Arguments to pass
-   * @returns {Promise<any>} - Result from the ability
-   */
-  async callAbility(toolName, args = {}) {
-    this.logger.info('call-ability', `Calling ability: ${toolName}`);
-    this.logger.trace('call-ability', `Args: ${JSON.stringify(args)}`);
-
-    const callMessage = Broker.callAbility({ toolName, args })
-      .id(this._idFactory.next())
-      .build();
-
-    const ack = await this._sendAndWaitForResponse(callMessage);
-    if (ack.error) {
-      this.logger.error(
-        'call-ability',
-        `Call failed: ${ack.error.message || 'agent.callAbility failed'}`
-      );
-
-      throw new Error(ack.error.message || 'agent.callAbility failed');
-    }
-
-    const expectedRequestId = ack.result?.requestId;
-    this.logger.trace(
-      'call-ability',
-      `Waiting for ability result, requestId: ${expectedRequestId}`
-    );
-
-    return new Promise((resolve, reject) => {
-      const timeout = setTimeout(() => {
-        this._ws.off('message', handler);
-        this.logger.error(
-          'call-ability',
-          `Ability call timeout for ${toolName}`
-        );
-
-        reject(new Error('Ability call timeout'));
-      }, this.timeoutMs);
-
-      const handler = (data) => {
-        try {
-          const message = JSON.parse(data.toString());
-          if (message.method === 'ability.result') {
-            const { requestId, result, error } = message.params || {};
-            if (requestId === expectedRequestId) {
-              clearTimeout(timeout);
-              this._ws.off('message', handler);
-              if (error) {
-                this.logger.error(
-                  'call-ability',
-                  `Ability ${toolName} returned error: ${error.message || error}`
-                );
-
-                reject(new Error(error.message || 'Ability error'));
-              } else {
-                this.logger.success(
-                  'call-ability',
-                  `Ability ${toolName} completed successfully`
-                );
-
-                resolve(result);
-              }
-            }
-          }
-        } catch (parseError) {
-          // Ignore parsing errors
-        }
-      };
-
-      this._ws.on('message', handler);
-    });
-  }
-
-  /**
-   * Send a message and wait for its response
-   *
-   * @param {Object} message - Message to send
-   * @returns {Promise<Object>} - Response message
-   */
-  async _sendAndWaitForResponse(message) {
-    this.logger.trace(
-      'send-wait',
-      `Sending message and waiting for response: ${message.method || 'unknown'}`
-    );
-
-    this._ws.send(JSON.stringify(message));
-    const response = await this._waitForResponse(message.id);
-
-    if (response.error) {
-      this.logger.warn(
-        'send-wait',
-        `Response contained error: ${JSON.stringify(response.error)}`
-      );
-    } else {
-      this.logger.trace('send-wait', 'Response received successfully');
-    }
-
-    return response;
-  }
-}
-
-export default BrokerRpcServer;