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

package/src/loadAbility.js

@@ -0,0 +1,1086 @@
+import path from 'path';
+import fs from 'fs';
+import { spawn } from 'child_process';
+import { WebSocket } from 'ws';
+import { generateKeyPairSync, randomUUID } from 'crypto';
+import { fileURLToPath } from 'url';
+import EventEmitter from 'events';
+
+import { getProjectJSON, getAbilityJSON } from './utils/agentUtils.js';
+import { rootDir } from './utils/pathUtils.js';
+import { StdioFrameReader, Ipc } from './transport/IpcMessageBuilder.js';
+import {
+  Broker,
+  IdFactory,
+  toBase64Der
+} from './transport/BrokerMessageBuilder.js';
+
+import { createComponentLogger, formatObject } from './utils/logger.js';
+import { FRAME_HEADERS, FRAME_VALUES } from './transport/IpcMessageBuilder.js';
+
+const logger = createComponentLogger('loadAbility');
+
+const __filename = fileURLToPath(import.meta.url);
+
+// ===== ability loader cache (avoid re-spawning on repeated loads) =====
+// key: `${name}@${version}:protocol` -> proxy/module
+const __abilityCache = new Map();
+
+/**
+ * Load an ability by name and return a uniform interface (module or proxy).
+ *
+ * Supports three protocols via an optional manifest at:
+ *    abilities/<name>/<version>/ability.json
+ *
+ *  - protocol: "native"          -> in-process import
+ *  - protocol: "stdio"           -> spawn local process, JSON-RPC over stdio
+ *  - protocol: "broker"          -> request/response via the KADI Broker
+ */
+export async function loadAbility(abilityName, protocol) {
+  logger.lifecycle('start', `Loading ability: ${abilityName}`);
+  logger.info('loadAbility', `Protocol: ${protocol || 'auto'}`);
+
+  // ------------------------------
+  // 1) Resolve version to load (project or nested ability context)
+  // ------------------------------
+  let agentJSON = await getProjectJSON();
+  let abilityVersion = agentJSON?.abilities?.[abilityName] || null;
+
+  if (!abilityVersion) {
+    // Fallback when loadAbility() is called from inside an ability directory
+    logger.warn(
+      'config',
+      `Ability ${abilityName} not found in Project configuration`
+    );
+
+    // Identify parent ability name/version from current file path
+    const filePath = __filename;
+    const dirPath = path.dirname(filePath);
+    const pathParts = dirPath.split(path.sep);
+    const abilitiesIndex = pathParts.indexOf('abilities');
+
+    if (abilitiesIndex !== -1 && pathParts.length > abilitiesIndex + 2) {
+      const parentAbility = pathParts[abilitiesIndex + 1];
+      const parentVersion = pathParts[abilitiesIndex + 2];
+
+      logger.info(
+        'context',
+        `Loading from nested ability context: ${parentAbility}:${parentVersion}`
+      );
+
+      agentJSON = await getAbilityJSON(parentAbility, parentVersion);
+      abilityVersion = agentJSON?.abilities?.[abilityName] || null;
+
+      if (!abilityVersion) {
+        logger.error(
+          'config',
+          `Ability ${abilityName} not found in ${parentAbility}:${parentVersion} configuration`
+        );
+
+        throw new Error(
+          `Ability ${abilityName} not found in ${parentAbility}:${parentVersion} configuration`
+        );
+      }
+    } else {
+      logError('Ability not found in Project agent.json');
+      throw new Error('Ability not found in Project agent.json');
+    }
+  }
+
+  const abilityDir = path.join(
+    rootDir,
+    'abilities',
+    abilityName,
+    abilityVersion
+  );
+  logger.success(
+    'version',
+    `Resolved ${abilityName} version: ${abilityVersion}`
+  );
+  logger.trace('paths', `Ability directory: ${abilityDir}`);
+
+  // Check if ability directory exists
+  if (!fs.existsSync(abilityDir)) {
+    throw new Error(
+      `Ability '${abilityName}' version '${abilityVersion}' is not installed.\n` +
+        `Expected location: ${abilityDir}\n` +
+        `Please install the ability using: kadi install`
+    );
+  }
+
+  // ------------------------------
+  // 2) Read ability manifest (optional). If missing -> default to node-module
+  // ------------------------------
+  const manifestPath = path.join(abilityDir, 'agent.json');
+
+  logger.trace('manifest', `Looking for agent.json at: ${manifestPath}`);
+
+  const manifest = fs.existsSync(manifestPath)
+    ? JSON.parse(fs.readFileSync(manifestPath, 'utf8'))
+    : null;
+
+  logger.info('manifest', `Agent.json found: ${manifest ? 'yes' : 'no'}`);
+
+  if (!protocol) {
+    const firstKey = Object.keys(manifest?.interfaces)[0];
+    protocol = firstKey;
+    logger.info('protocol', `Auto-selected protocol: ${protocol}`);
+  } else {
+    protocol = protocol.toLowerCase();
+    logger.info('protocol', `Using specified protocol: ${protocol}`);
+  }
+
+  // Check cache with protocol-specific key
+  const cacheKey = `${abilityName}@${abilityVersion}:${protocol}`;
+  if (__abilityCache.has(cacheKey)) {
+    logger.info('cache', `Using cached ability: ${cacheKey}`);
+    return __abilityCache.get(cacheKey);
+  }
+
+  const startCmd = manifest?.scripts?.start || null;
+  if (!startCmd) {
+    logger.error('config', `Missing start command for ${abilityName}`);
+    throw new Error(
+      `Ability ${abilityName} missing start command (scripts.start)`
+    );
+  }
+
+  logger.trace('command', `Start command: ${startCmd}`);
+
+  // ------------------------------
+  // 3) Dispatch by protocol
+  // ------------------------------
+  let loaded;
+
+  if (protocol === 'native') {
+    logger.info('native', `Loading ${abilityName} as native JavaScript module`);
+
+    const entryFromManifest = manifest?.interfaces?.native?.entry;
+    let entryRelPath = entryFromManifest;
+
+    if (!entryRelPath) {
+      const moduleJSON = getAbilityJSON(abilityName, abilityVersion);
+      entryRelPath = moduleJSON.entry;
+    }
+
+    if (!entryRelPath) {
+      throw new Error(
+        `Ability '${abilityName}' is missing entry point configuration.\n` +
+          `Please specify 'entry' in the ability's agent.json or 'interface.entry' in the manifest.`
+      );
+    }
+
+    const modulePath = path.join(abilityDir, entryRelPath);
+
+    if (!fs.existsSync(modulePath)) {
+      throw new Error(
+        `Ability '${abilityName}' entry file not found: ${modulePath}\n` +
+          `Expected entry point: ${entryRelPath}`
+      );
+    }
+
+    console.error(`Loading module from: ${modulePath}`);
+    // const mod = await import(modulePath + `?v=${Date.now()}`);
+    const mod = await import(modulePath);
+    const rawModule = mod.default || mod;
+
+    // Detect "ability export" shape
+    const isAbilityExport =
+      rawModule &&
+      typeof rawModule.on === 'function' && // EventEmitter
+      typeof rawModule.publishEvent === 'function' &&
+      rawModule.methodHandlers instanceof Map;
+
+    const eventEmitter = new EventEmitter();
+
+    let functions = {};
+    let call; // (method: string, params: any) => Promise<any>
+
+    if (isAbilityExport) {
+      // 1) expose registered methods
+      for (const name of rawModule.methodHandlers.keys()) {
+        // value doesn't matter for buildAbilityProxy if it only
+        // needs the keys,
+        // but giving a callable makes it future-proof.
+        functions[name] = async (params) => {
+          const handler = rawModule.methodHandlers.get(name);
+          return await handler(params);
+        };
+      }
+
+      // 2) route calls through the handlers map
+      call = async (method, params) => {
+        const handler = rawModule.methodHandlers.get(method);
+        if (typeof handler !== 'function') {
+          throw new Error(
+            `Method '${method}' is not registered on KadiAbility`
+          );
+        }
+        return await handler(params);
+      };
+    } else {
+      // existing behavior for plain function modules
+      for (const [key, value] of Object.entries(rawModule)) {
+        if (typeof value === 'function' && !key.startsWith('_')) {
+          functions[key] = value;
+        }
+      }
+      call = async (method, params) => {
+        const func = rawModule[method];
+        if (typeof func !== 'function') {
+          throw new Error(
+            `Method '${method}' is not a function or doesn't exist`
+          );
+        }
+        return await func(params);
+      };
+    }
+
+    // Create uniform proxy
+    loaded = buildAbilityProxy({
+      call,
+      functions,
+      eventEmitter
+    });
+
+    // Wire native events (works for the KadiAbility instance)
+    setupNativeEventListener(rawModule, eventEmitter);
+  } else if (protocol === 'stdio') {
+    logger.info(
+      'stdio',
+      `Loading ${abilityName} with stdio JSON-RPC transport`
+    );
+
+    const timeoutMs = manifest?.interface?.timeoutMs ?? 15000;
+    const env = { ...process.env, KADI_PROTOCOL: protocol };
+    const logFile = path.join(abilityDir, `${abilityName}.log`);
+
+    // Use the enhanced spawnJSONRPCProcess that supports events
+    const rpc = spawnJSONRPCProcess({
+      command: startCmd,
+      cwd: abilityDir,
+      env,
+      timeoutMs,
+      logFile
+    });
+
+    // Create IPC helper bound to this rpc instance
+    const IPC = Ipc.with(rpc);
+
+    // Handshake: init
+    const initRest = await IPC.init({ api: '1.0' });
+
+    // Get static contract from the agent.json (if any)
+    let staticFunctions = null;
+    if (Array.isArray(manifest?.exports)) {
+      const tools = manifest.exports;
+      staticFunctions = Object.fromEntries(tools.map((t) => [t.name]));
+    }
+
+    // Always call discovery when enabled to get dynamic methods
+    let discoveredFunctions = null;
+    if (manifest?.interface?.discover !== false) {
+      try {
+        const disc = await IPC.discover();
+        discoveredFunctions = disc?.functions || null;
+      } catch (error) {
+        console.warn(`Discovery failed for ${abilityName}: ${error.message}`);
+        discoveredFunctions = null;
+      }
+    }
+
+    // Merge static and discovered functions
+    let functions = { ...staticFunctions };
+    if (discoveredFunctions) {
+      functions = { ...functions, ...discoveredFunctions };
+    }
+
+    // Create EventEmitter for events
+    const eventEmitter = new EventEmitter();
+
+    // Set up event listener for __kadi_event notifications
+    setupStdioEventListener(rpc, eventEmitter);
+
+    // Create proxy with event support
+    loaded = buildAbilityProxy({
+      call: (method, params) => IPC.call(method, params),
+      functions,
+      eventEmitter
+    });
+  } else if (protocol === 'broker') {
+    logger.info('broker', `Loading ${abilityName} with broker transport`);
+
+    const timeoutMs = manifest?.interface?.timeoutMs ?? 10000;
+    const scope = randomUUID();
+    logger.trace('broker', `Generated agent scope: ${scope}`);
+
+    const brokerUrl =
+      manifest?.brokers?.local ||
+      manifest?.brokers?.remote ||
+      (await getProjectJSON())?.brokers?.local;
+
+    logger.info('broker', `Broker URL: ${brokerUrl}`);
+
+    const serviceName =
+      manifest?.interface?.serviceName ||
+      `ability.${abilityName}.${abilityVersion.replace(/\./g, '_')}`;
+
+    if (startCmd) {
+      // Spawn the ability process
+      const env = {
+        ...process.env,
+        KADI_PROTOCOL: protocol,
+        KADI_BROKER_URL: brokerUrl,
+        KADI_SERVICE_NAME: serviceName,
+        KADI_AGENT_SCOPE: scope
+      };
+
+      logger.info('spawn', `Spawning ${abilityName} ability process`);
+
+      const child = spawn(startCmd, {
+        cwd: abilityDir,
+        env,
+        shell: true,
+        stdio: 'inherit',
+        detached: false
+      });
+
+      child.on('error', (error) => {
+        logger.error(
+          'spawn',
+          `Failed to spawn ${abilityName}: ${error.message}`
+        );
+      });
+
+      child.on('exit', (code, signal) => {
+        if (code !== 0) {
+          logger.warn(
+            'spawn',
+            `Ability process exited with code ${code}, signal ${signal}`
+          );
+        } else {
+          logger.success('spawn', 'Ability process exited cleanly');
+        }
+      });
+
+      // Give the ability time to connect to the broker
+      await new Promise((resolve) => setTimeout(resolve, 1000));
+    }
+
+    logger.info('broker', 'Creating broker RPC client');
+
+    let rpc;
+    try {
+      rpc = createBrokerRPCClient({
+        brokerUrl,
+        serviceName,
+        timeoutMs,
+        scope
+      });
+
+      await rpc.testConnection();
+      logger.success('broker', 'Agent connected to broker successfully');
+    } catch (error) {
+      logger.error(
+        'broker',
+        `Agent failed to connect to broker: ${error.message}`
+      );
+      throw new Error(
+        `Failed to connect to KADI Broker at ${brokerUrl}\n` +
+          `Original error: ${error.message}\n` +
+          `\n` +
+          `Common solutions:\n` +
+          `• Make sure the KADI Broker is running\n` +
+          `• Check if the broker URL is correct in agent.json\n` +
+          `• Verify network connectivity to the broker\n` +
+          `• Ensure the broker is accepting connections`
+      );
+    }
+
+    // Get static contract from the agent.json (if any)
+    let staticFunctions = null;
+    if (Array.isArray(manifest?.exports)) {
+      const tools = manifest.exports;
+      staticFunctions = Object.fromEntries(tools.map((t) => [t.name]));
+    }
+
+    // Get the full method list from the broker
+    let brokerFunctions = null;
+    if (manifest?.interface?.discover !== false) {
+      try {
+        logger.trace(
+          'discovery',
+          `Querying broker for methods with scopes: ${JSON.stringify([scope])}`
+        );
+
+        const brokerMethods = await rpc.getAvailableMethods([scope]);
+
+        logger.success(
+          'discovery',
+          `Broker returned ${brokerMethods.tools?.length || 0} methods`
+        );
+
+        if (logger.enabled) {
+          brokerMethods.tools?.forEach((tool) => {
+            logger.trace('discovery', `  - ${tool.name}`);
+          });
+        }
+
+        brokerFunctions = Object.fromEntries(
+          brokerMethods.tools.map((t) => [t.name])
+        );
+      } catch (error) {
+        logger.warn(
+          'discovery',
+          `Broker method discovery failed: ${error.message}`
+        );
+        brokerFunctions = null;
+      }
+    }
+
+    // Merge static and broker functions
+    let functions = { ...staticFunctions };
+    if (brokerFunctions) {
+      functions = { ...functions, ...brokerFunctions };
+    }
+
+    // Create EventEmitter for events
+    const eventEmitter = new EventEmitter();
+
+    // Set up event listener for kadi.event messages
+    setupBrokerEventListener(rpc, eventEmitter);
+
+    // Create proxy with event support
+    loaded = buildAbilityProxy({
+      call: (method, params) => rpc.call(method, params),
+      functions,
+      eventEmitter
+    });
+  } else {
+    throw new Error(`Unsupported ability interface protocol: ${protocol}`);
+  }
+
+  __abilityCache.set(cacheKey, loaded);
+
+  logger.success(
+    'complete',
+    `Successfully loaded ${abilityName} with ${Object.keys(loaded.__list?.() || {}).length} methods`
+  );
+
+  return loaded;
+}
+
+/* =========================================================================
+ * Helpers: small, dependency-free shims to keep the loader readable.
+ * ======================================================================= */
+
+/**
+ * Spawn a child process and speak newline-delimited JSON-RPC 2.0 over stdio.
+ * Enhanced to support event notifications via __kadi_event messages.
+ *
+ * Minimal reference implementation; production code should add heartbeats,
+ * backpressure control, and graceful shutdown handling.
+ * - write JSON-RPC requests to the child's stdin,
+ * - read/parse JSON-RPC responses from the child's stdout,
+ * - handle __kadi_event notifications for event publishing
+ * - enforce timeouts
+ * - logFile if present will see its content append with information echoed to stderr
+ *
+ * Reference: https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/
+ */
+function spawnJSONRPCProcess({ command, cwd, env, timeoutMs, logFile }) {
+  const ability = spawn(command, {
+    cwd,
+    env,
+    shell: true,
+    stdio: ['pipe', 'pipe', 'pipe'] // IMPORTANT: we need pipes, not 'inherit'
+    // 'pipe' for stdin, stdout, stderr
+  });
+
+  // ────────────────────────────────────
+  // Optional file logger for stderr
+  // ────────────────────────────────────
+  let logStream = null;
+  if (logFile) {
+    try {
+      fs.mkdirSync(path.dirname(logFile), { recursive: true });
+      logStream = fs.createWriteStream(logFile, { flags: 'a' });
+    } catch (e) {
+      console.error('Unable to open log file', logFile, e);
+    }
+  }
+
+  const pending = new Map(); // id -> { resolve, reject, timer }
+  let nextId = 1;
+
+  // Storage for event handlers
+  const eventHandlers = [];
+
+  // Create frame reader for parsing LSP-style messages from stdout
+  const frameReader = new StdioFrameReader(ability.stdout, {
+    maxBufferSize: 8 * 1024 * 1024 // 8MB buffer limit
+  });
+
+  // Set up frame processing with rate limiting to prevent event loop blocking
+  let processedFrames = 0;
+  const maxFramesPerChunk = 10; // Prevent infinite loops under extreme load
+  let pendingProcessing = false;
+
+  frameReader.onMessage((result) => {
+    // Rate limiting: if we've processed too many frames in this tick, defer to next tick
+    if (processedFrames >= maxFramesPerChunk && !pendingProcessing) {
+      pendingProcessing = true;
+      setImmediate(() => {
+        processedFrames = 0;
+        pendingProcessing = false;
+        // Process this result in the next tick
+        handleResult(result);
+      });
+      return;
+    }
+
+    handleResult(result);
+    processedFrames++;
+  });
+
+  function handleResult(result) {
+    if (!result.success) {
+      // Frame parsing failed - log detailed error and fail all pending calls
+      logger.error(`[rpc] Frame corruption detected: ${result.error}`);
+      logger.error(`[rpc] Corruption type: ${result.corruption_type}`);
+      logger.error(`[rpc] Details: ${result.message}`);
+
+      // Fail all pending calls since we can't trust the stream anymore
+      for (const { reject, timer } of pending.values()) {
+        clearTimeout(timer);
+        reject(
+          new Error(`Frame corruption: ${result.error} - ${result.message}`)
+        );
+      }
+      pending.clear();
+      return;
+    }
+
+    const msg = result.data;
+
+    // Handle event notifications (__kadi_event messages)
+    if (!msg.id && msg.method === '__kadi_event' && msg.params) {
+      // This is an event notification from the ability
+      const { eventName, data, timestamp } = msg.params;
+
+      // Log the event for debugging
+      if (logger && logger.enabled) {
+        logger.trace('event', `Received event notification: ${eventName}`);
+      }
+
+      // Call all registered event handlers
+      for (const handler of eventHandlers) {
+        try {
+          handler({
+            method: '__kadi_event',
+            params: { eventName, data, timestamp }
+          });
+        } catch (err) {
+          console.error(`[rpc] Error in event handler: ${err.message}`);
+        }
+      }
+
+      return; // Don't process further - events don't have responses
+    }
+
+    // Handle regular JSON-RPC responses
+    if (msg.id && pending.has(msg.id)) {
+      const { resolve, reject, timer } = pending.get(msg.id);
+      clearTimeout(timer);
+      pending.delete(msg.id);
+      if (msg.error) reject(new Error(msg.error.message || 'RPC error'));
+      else resolve(msg.result);
+    }
+  }
+
+  ability.stderr.on('data', (chunk) => {
+    // Optional: forward ability stderr to your logger
+    console.error(`[${cwd}]`, chunk.toString().trim());
+    if (logStream) logStream.write(chunk);
+  });
+
+  ability.on('close', (code) => {
+    // Fail all pending calls on exit
+    for (const { reject, timer } of pending.values()) {
+      clearTimeout(timer);
+      reject(new Error(`Ability process exited with code ${code}`));
+    }
+    pending.clear();
+    if (logStream) logStream.end();
+  });
+
+  function call(method, params) {
+    const id = nextId++;
+    const body = JSON.stringify({
+      jsonrpc: '2.0',
+      id,
+      method,
+      params: params ?? {}
+    });
+
+    // Compose LSP-style frame. We *could* omit Content-Type because UTF-8 JSON
+    // is the default, but including it keeps packet dumps self-describing.
+    const header =
+      `${FRAME_HEADERS.CONTENT_LENGTH}: ${Buffer.byteLength(body, 'utf8')}\r\n` +
+      `${FRAME_HEADERS.CONTENT_TYPE}: ${FRAME_VALUES.CONTENT_TYPE_VALUE}\r\n` +
+      `\r\n`;
+
+    return new Promise((resolve, reject) => {
+      const timer = setTimeout(() => {
+        pending.delete(id);
+        reject(new Error(`RPC timeout after ${timeoutMs}ms for ${method}`));
+      }, timeoutMs);
+
+      pending.set(id, { resolve, reject, timer });
+
+      try {
+        ability.stdin.write(header + body, 'utf8');
+      } catch (e) {
+        clearTimeout(timer);
+        pending.delete(id);
+        reject(e);
+      }
+    });
+  }
+
+  /**
+   * Register a handler for event notifications
+   * @param {Function} handler - Function to call when __kadi_event messages arrive
+   */
+  function onEvent(handler) {
+    if (typeof handler === 'function') {
+      eventHandlers.push(handler);
+    }
+  }
+
+  /**
+   * Remove an event handler
+   * @param {Function} handler - Handler to remove
+   */
+  function offEvent(handler) {
+    const index = eventHandlers.indexOf(handler);
+    if (index > -1) {
+      eventHandlers.splice(index, 1);
+    }
+  }
+
+  // Return the RPC interface with event support
+  return {
+    call,
+    onEvent, // Method to register event handlers
+    offEvent, // Method to unregister event handlers
+    // Expose the process for advanced use cases
+    _process: ability
+  };
+}
+
+/**
+ * Create a thin client that speaks request/response over the KADI Broker.
+ * Enhanced to support event notifications via kadi.event messages.
+ */
+function createBrokerRPCClient({ brokerUrl, serviceName, timeoutMs, scope }) {
+  const idFactory = new IdFactory();
+
+  // Storage for event handlers
+  const eventHandlers = [];
+
+  function rpc(ws, messageBuilder) {
+    return new Promise((resolve) => {
+      const message = messageBuilder.build();
+      const onMessage = (raw) => {
+        try {
+          const msg = JSON.parse(raw.toString());
+          if (msg.id === message.id) {
+            ws.off('message', onMessage);
+            resolve(msg);
+          }
+        } catch {}
+      };
+      ws.on('message', onMessage);
+      ws.send(JSON.stringify(message));
+    });
+  }
+
+  async function connect() {
+    return new Promise((resolve, reject) => {
+      const ws = new WebSocket(brokerUrl);
+
+      // Handle connection errors
+      ws.once('error', (error) => {
+        reject(new Error(`WebSocket connection failed: ${error.message}`));
+      });
+
+      ws.once('open', async () => {
+        try {
+          // Remove error listener since connection succeeded
+          ws.removeAllListeners('error');
+
+          // Add error handler for runtime errors
+          ws.on('error', (error) => {
+            console.error(`Broker WebSocket error: ${error.message}`);
+          });
+
+          ws.on('message', (raw) => {
+            try {
+              const msg = JSON.parse(raw.toString());
+
+              // Check if this is an kadi.event notification
+              if (msg.method === 'kadi.event' && msg.params) {
+                const { eventName, eventData, timestamp, from } = msg.params;
+
+                // Log the event for debugging
+                if (logger && logger.enabled) {
+                  logger.trace(
+                    'event',
+                    `Received broker event: ${eventName} from ${from}`
+                  );
+                }
+
+                // Call all registered event handlers
+                for (const handler of eventHandlers) {
+                  try {
+                    handler({
+                      method: 'kadi.event',
+                      params: { eventName, eventData, timestamp, from }
+                    });
+                  } catch (err) {
+                    console.error(
+                      `[broker-rpc] Error in event handler: ${err.message}`
+                    );
+                  }
+                }
+              }
+            } catch (err) {
+              // Ignore parsing errors for non-JSON messages
+            }
+          });
+
+          // hello/auth (ephemeral ed25519) using BrokerMessageBuilder
+          const helloMsg = Broker.hello({ role: 'agent', version: '0.1' }).id(
+            idFactory.next()
+          );
+          const hello = await rpc(ws, helloMsg);
+          const nonce = hello?.result?.nonce;
+
+          if (!nonce) {
+            throw new Error('No nonce received from broker hello');
+          }
+
+          const { publicKey, privateKey } = generateKeyPairSync('ed25519');
+          const publicKeyBase64Der = toBase64Der(publicKey);
+
+          const authMsg = Broker.authenticate({
+            publicKeyBase64Der,
+            privateKey,
+            nonce,
+            wantNewId: true
+          }).id(idFactory.next());
+
+          const auth = await rpc(ws, authMsg);
+          if (auth.error) {
+            throw new Error('Broker auth failed: ' + auth.error.message);
+          }
+
+          // Register capabilities (no tools since we are a client)
+          const registerMsg = Broker.registerCapabilities({
+            displayName: `client:${serviceName}`,
+            scopes: ['global', scope],
+            tools: []
+          }).id(idFactory.next());
+
+          await rpc(ws, registerMsg);
+
+          // Heartbeat using ping notification
+          setInterval(() => {
+            if (ws.readyState === ws.OPEN) {
+              ws.send(Broker.ping().toString());
+            }
+          }, 25_000);
+
+          resolve(ws);
+        } catch (error) {
+          reject(error);
+        }
+      });
+    });
+  }
+
+  const ready = connect();
+
+  return {
+    async testConnection() {
+      // This will throw if connection fails during the handshake
+      const ws = await ready;
+      if (!ws || ws.readyState !== ws.OPEN) {
+        throw new Error('WebSocket connection failed');
+      }
+    },
+
+    async call(toolName, args) {
+      const ws = await ready;
+
+      // Use BrokerMessageBuilder for callAbility
+      const callMsg = Broker.callAbility({
+        toolName,
+        args
+      }).id(idFactory.next());
+
+      const ack = await rpc(ws, callMsg);
+      if (ack.error) {
+        throw new Error(ack.error.message || 'agent.callAbility failed');
+      }
+
+      const expectedRequestId = ack?.result?.requestId;
+      return await new Promise((resolve, reject) => {
+        const timer = setTimeout(() => {
+          ws.off('message', onMessage);
+          reject(new Error('broker call timeout'));
+        }, timeoutMs || 10000);
+
+        const onMessage = (raw) => {
+          try {
+            const msg = JSON.parse(raw.toString());
+            if (msg.method === 'ability.result') {
+              const { requestId, result, error } = msg.params || {};
+              if (!expectedRequestId || requestId === expectedRequestId) {
+                clearTimeout(timer);
+                ws.off('message', onMessage);
+                if (error) reject(new Error(error.message || 'ability error'));
+                else resolve(result);
+              }
+            }
+          } catch {}
+        };
+        ws.on('message', onMessage);
+      });
+    },
+
+    async getAvailableMethods(scopes = ['global']) {
+      const ws = await ready;
+      const msg = Broker.listTools({ scopes }).id(idFactory.next());
+      const response = await rpc(ws, msg);
+      if (response.error) {
+        throw new Error(response.error.message || 'listTools failed');
+      }
+      return response.result;
+    },
+
+    /**
+     * Register a handler for event notifications
+     * @param {Function} handler - Function to call when kadi.event messages
+     * arrive
+     */
+    onEvent(handler) {
+      if (typeof handler === 'function') {
+        eventHandlers.push(handler);
+      }
+    },
+
+    /**
+     * Remove an event handler
+     * @param {Function} handler - Handler to remove
+     */
+    offEvent(handler) {
+      const index = eventHandlers.indexOf(handler);
+      if (index > -1) {
+        eventHandlers.splice(index, 1);
+      }
+    }
+  };
+}
+
+/**
+ * Set up event listener for stdio transport
+ * Listens for __kadi_event notifications and emits them on the EventEmitter
+ *
+ * @param {Object} rpc - The RPC object returned from spawnJSONRPCProcess
+ * @param {EventEmitter} eventEmitter - The EventEmitter to emit events on
+ */
+function setupStdioEventListener(rpc, eventEmitter) {
+  if (!rpc.onEvent) {
+    logger.warn('event', 'RPC object does not support event notifications');
+    return;
+  }
+
+  // Register handler for __kadi_event messages
+  rpc.onEvent((message) => {
+    if (message.method === '__kadi_event' && message.params) {
+      const { eventName, data, timestamp } = message.params;
+
+      logger.trace('event', `Received stdio event: ${eventName}`);
+      logger.trace('event-data', { eventName, data, timestamp });
+
+      // Emit the event on the EventEmitter
+      // This allows agents to do: ability.events.on('eventName', handler)
+      eventEmitter.emit(eventName, data);
+
+      // Also emit a generic 'event' event for catch-all handlers
+      eventEmitter.emit('event', { eventName, data, timestamp });
+    }
+  });
+
+  logger.trace('event', 'Stdio event listener configured');
+}
+
+/**
+ * Set up event listener for broker transport
+ * Listens for kadi.event messages and emits them on the EventEmitter
+ *
+ * @param {Object} rpc - The RPC client returned from createBrokerRPCClient
+ * @param {EventEmitter} eventEmitter - The EventEmitter to emit events on
+ */
+function setupBrokerEventListener(rpc, eventEmitter) {
+  if (!rpc.onEvent) {
+    logger.warn('event', 'RPC client does not support event notifications');
+    return;
+  }
+
+  // Register handler for kadi.event messages
+  rpc.onEvent((message) => {
+    if (message.method === 'kadi.event' && message.params) {
+      const { eventName, eventData, timestamp, from } = message.params;
+
+      logger.trace('event', `Received broker event: ${eventName} from ${from}`);
+      logger.trace('event-data', { eventName, eventData, timestamp, from });
+
+      // Emit the event on the EventEmitter
+      // Use eventData (broker's field name) as the data payload
+      eventEmitter.emit(eventName, eventData);
+
+      // Also emit a generic 'event' event for catch-all handlers
+      eventEmitter.emit('event', {
+        eventName,
+        data: eventData,
+        timestamp,
+        source: from
+      });
+    }
+  });
+
+  logger.trace('event', 'Broker event listener configured');
+}
+
+/**
+ * Set up event listener for native transport
+ * For native abilities that are KadiAbility instances
+ *
+ * @param {any} rawModule - The imported module (might be a KadiAbility instance)
+ * @param {EventEmitter} eventEmitter - The EventEmitter to emit events on
+ */
+function setupNativeEventListener(rawModule, eventEmitter) {
+  // Check if the module is an EventEmitter (like KadiAbility)
+  if (rawModule && typeof rawModule.on === 'function') {
+    logger.trace('event', 'Native module appears to be an EventEmitter');
+
+    // Remove any existing ability:event listeners first
+    rawModule.removeAllListeners('ability:event');
+
+    // Listen for ability:event which KadiAbility emits for native protocol
+    rawModule.on('ability:event', ({ eventName, data }) => {
+      logger.trace('event', `Received native event: ${eventName}`);
+      logger.trace('event-data', { eventName, data });
+
+      // Emit the event on the proxy's EventEmitter
+      eventEmitter.emit(eventName, data);
+
+      // Also emit a generic 'event' event
+      eventEmitter.emit('event', { eventName, data, timestamp: Date.now() });
+    });
+
+    logger.trace('event', 'Native event listener configured');
+  } else {
+    logger.trace(
+      'event',
+      'Native module is not an EventEmitter, skipping event setup'
+    );
+  }
+}
+
+/**
+ * Enhanced buildAbilityProxy with event support
+ * Replace the existing buildAbilityProxy function with this version
+ */
+function buildAbilityProxy({ call, functions, eventEmitter }) {
+  // Create EventEmitter if not provided
+  const events = eventEmitter || new EventEmitter();
+
+  // Set max listeners to avoid warnings when many handlers are attached
+  events.setMaxListeners(100);
+
+  const base = {
+    call: async (method, params) => {
+      try {
+        return await call(method, params);
+      } catch (error) {
+        // Enhance error messages for better user experience
+        if (functions && !functions[method]) {
+          const availableMethods = Object.keys(functions);
+          throw new Error(
+            `Method '${method}' is not exposed by this ability.\n` +
+              `Available methods: ${availableMethods.length > 0 ? availableMethods.join(', ') : 'none'}\n` +
+              `Use ability.__list() to see all available methods.`
+          );
+        }
+        throw error;
+      }
+    },
+    __list: () => (functions ? Object.keys(functions) : []),
+    events // Expose the EventEmitter for event subscriptions
+  };
+
+  if (!functions) {
+    // No discovery -> return generic proxy that routes any property call to RPC .call(name)
+    return new Proxy(base, {
+      get(target, prop) {
+        if (prop in target) return target[prop];
+        if (prop === 'then') return undefined; // Avoid being treated as a Promise/thenable
+        if (typeof prop !== 'string') return undefined;
+        return async (params) => {
+          try {
+            return await call(prop, params);
+          } catch (error) {
+            // For abilities without discovery, provide a generic helpful message
+            throw new Error(
+              `Failed to call method '${prop}' on ability.\n` +
+                `Original error: ${error.message}\n` +
+                `Tip: Use ability.__list() to see available methods or check if the method name is correct.`
+            );
+          }
+        };
+      }
+    });
+  }
+
+  // Discovery available → define concrete methods plus dynamic fallback
+  for (const fnName of Object.keys(functions)) {
+    base[fnName] = async (params) => {
+      try {
+        return await call(fnName, params);
+      } catch (error) {
+        throw new Error(`Error calling method '${fnName}': ${error.message}`);
+      }
+    };
+  }
+
+  return new Proxy(base, {
+    get(target, prop) {
+      if (prop in target) return target[prop];
+      if (prop === 'then') return undefined; // Avoid thenable detection
+      if (typeof prop !== 'string') return undefined;
+      // Unknown method name → provide helpful error
+      return async (params) => {
+        const availableMethods = Object.keys(functions);
+        throw new Error(
+          `Method '${prop}' is not exposed by this ability.\n` +
+            `Available methods: ${availableMethods.length > 0 ? availableMethods.join(', ') : 'none'}\n` +
+            `Use ability.__list() to see all available methods.`
+        );
+      };
+    }
+  });
+}