@kadi.build/core 0.3.4 → 0.5.0

package/src/transports/broker.ts

@@ -0,0 +1,682 @@
+/**
+ * Broker transport for kadi-core v0.1.0
+ *
+ * Loads abilities that run on REMOTE agents connected to a broker.
+ * Communication happens via WebSocket using JSON-RPC.
+ *
+ * Protocol flow for invoking a remote tool:
+ * 1. Client sends kadi.ability.request → broker returns { status: 'pending', requestId }
+ * 2. Broker forwards request to provider agent
+ * 3. Provider executes tool and sends result back to broker
+ * 4. Broker sends kadi.ability.response notification with the actual result
+ *
+ * This transport doesn't manage the broker connection itself.
+ * It receives a BrokerState from KadiClient and uses it for communication.
+ *
+ * ═══════════════════════════════════════════════════════════════════════════════
+ * WORKAROUND NOTE: Ability Discovery
+ * ═══════════════════════════════════════════════════════════════════════════════
+ *
+ * The current broker is TOOL-CENTRIC (tracks tools, not abilities). There's no
+ * direct API to ask "What tools does agent X provide?"
+ *
+ * CURRENT WORKAROUND:
+ * 1. Call kadi.ability.list with includeProviders: true (fetches ALL tools)
+ * 2. Filter tools where provider.displayName matches the ability name
+ * 3. Use the provider's agentId for routing via _kadi hints
+ *
+ * FUTURE: When the broker implements kadi.agent.discover (see design doc at
+ * kadi-broker/docs/design/ABILITY_CENTRIC_DISCOVERY.md), replace the
+ * discoverAbilityWorkaround() function with a direct API call.
+ * ═══════════════════════════════════════════════════════════════════════════════
+ */
+
+import crypto from 'node:crypto';
+import type {
+  LoadedAbility,
+  InvokeOptions,
+  ToolDefinition,
+  BrokerState,
+  JsonRpcRequest,
+  EventHandler,
+  BrokerEventHandler,
+  BrokerEvent,
+} from '../types.js';
+import { KadiError } from '../errors.js';
+
+// ═══════════════════════════════════════════════════════════════════════════════
+// TYPES
+// ═══════════════════════════════════════════════════════════════════════════════
+
+/**
+ * Provider information returned by kadi.ability.list with includeProviders: true.
+ * This matches the ProviderSummary type from the broker.
+ */
+interface ProviderInfo {
+  sessionId: string;
+  agentId?: string;
+  displayName?: string;
+  source: 'kadi-agent' | 'mcp-upstream';
+  networks: string[];
+  tags?: string[];
+}
+
+/**
+ * Tool with provider information from kadi.ability.list.
+ */
+interface ToolWithProviders {
+  name: string;
+  description?: string;
+  inputSchema?: Record<string, unknown>;
+  tags?: string[];
+  providers?: ProviderInfo[];
+}
+
+/**
+ * Response from kadi.ability.list.
+ */
+interface AbilityListResponse {
+  tools: ToolWithProviders[];
+}
+
+/**
+ * Information about a discovered ability on the broker.
+ */
+interface DiscoveredAbility {
+  /** Agent ID for routing and event filtering (stable across reconnects) */
+  agentId: string;
+
+  /** Display name of the agent (what we searched for) */
+  name: string;
+
+  /** Tools available on this ability */
+  tools: ToolDefinition[];
+}
+
+/**
+ * Options for creating a broker transport.
+ */
+export interface BrokerTransportOptions {
+  /** The broker connection to use */
+  broker: BrokerState;
+
+  /** Timeout for requests in milliseconds (default: 600000 = 10 minutes) */
+  requestTimeout?: number;
+
+  /** Networks to filter by when discovering (default: all) */
+  networks?: string[];
+
+  /**
+   * Subscribe to broker events.
+   * Provided by KadiClient to enable ability.on() for broker transport.
+   */
+  subscribe?: (pattern: string, handler: BrokerEventHandler) => Promise<void>;
+
+  /**
+   * Unsubscribe from broker events.
+   * Provided by KadiClient to enable ability.off() for broker transport.
+   */
+  unsubscribe?: (pattern: string, handler: BrokerEventHandler) => Promise<void>;
+}
+
+// ═══════════════════════════════════════════════════════════════════════════════
+// REQUEST HELPERS
+// ═══════════════════════════════════════════════════════════════════════════════
+
+/**
+ * Send a JSON-RPC request over WebSocket and wait for response.
+ *
+ * This is a low-level helper used by both discovery and invoke.
+ * The broker connection's pendingRequests map is used to match responses.
+ */
+async function sendBrokerRequest(
+  broker: BrokerState,
+  method: string,
+  params: unknown,
+  timeoutMs: number
+): Promise<unknown> {
+  // Verify connection is active
+  if (!broker.ws || broker.status !== 'connected') {
+    throw new KadiError(
+      `Broker "${broker.name}" is not connected`,
+      'BROKER_NOT_CONNECTED',
+      {
+        broker: broker.name,
+        status: broker.status,
+        hint: 'Call client.connect() first',
+      }
+    );
+  }
+
+  // Generate unique request ID
+  const id = `${Date.now()}-${Math.random().toString(36).slice(2, 9)}`;
+
+  // Build JSON-RPC request
+  const request: JsonRpcRequest = {
+    jsonrpc: '2.0',
+    id,
+    method,
+    params,
+  };
+
+  // Create promise that will be resolved when response arrives
+  return new Promise((resolve, reject) => {
+    // Set up timeout
+    const timeout = setTimeout(() => {
+      broker.pendingRequests.delete(id);
+      reject(
+        new KadiError(`Request "${method}" timed out after ${timeoutMs}ms`, 'BROKER_TIMEOUT', {
+          broker: broker.name,
+          method,
+          timeout: timeoutMs,
+        })
+      );
+    }, timeoutMs);
+
+    // Store in pending requests map
+    // The message handler in client.ts will resolve this when response arrives
+    broker.pendingRequests.set(id, {
+      resolve: (result: unknown) => {
+        clearTimeout(timeout);
+        resolve(result);
+      },
+      reject: (error: Error) => {
+        clearTimeout(timeout);
+        reject(error);
+      },
+      timeout,
+      method,
+      sentAt: new Date(),
+    });
+
+    // Re-check connection before sending (could have disconnected during setup)
+    const ws = broker.ws;
+    if (!ws || broker.status !== 'connected') {
+      broker.pendingRequests.delete(id);
+      clearTimeout(timeout);
+      reject(
+        new KadiError(`Broker "${broker.name}" disconnected`, 'BROKER_NOT_CONNECTED', {
+          broker: broker.name,
+        })
+      );
+      return;
+    }
+
+    // Send the request
+    try {
+      ws.send(JSON.stringify(request));
+    } catch (error) {
+      broker.pendingRequests.delete(id);
+      clearTimeout(timeout);
+      reject(
+        new KadiError(`Failed to send request to broker "${broker.name}"`, 'BROKER_ERROR', {
+          broker: broker.name,
+          method,
+          reason: error instanceof Error ? error.message : String(error),
+        })
+      );
+    }
+  });
+}
+
+// ═══════════════════════════════════════════════════════════════════════════════
+// ABILITY DISCOVERY (WORKAROUND)
+// ═══════════════════════════════════════════════════════════════════════════════
+
+/**
+ * Discover an ability on the broker by name.
+ *
+ * ┌─────────────────────────────────────────────────────────────────────────────┐
+ * │ WORKAROUND: This function uses kadi.ability.list + client-side filtering   │
+ * │                                                                             │
+ * │ The broker is currently tool-centric and doesn't support querying by       │
+ * │ agent/ability name directly. This workaround:                               │
+ * │                                                                             │
+ * │ 1. Fetches ALL tools with provider information                             │
+ * │ 2. Filters to find tools where provider.displayName matches abilityName    │
+ * │ 3. Extracts the agentId for consistent routing                             │
+ * │                                                                             │
+ * │ TODO: Replace with kadi.agent.discover when broker supports it.            │
+ * │ See: kadi-broker/docs/design/ABILITY_CENTRIC_DISCOVERY.md                  │
+ * └─────────────────────────────────────────────────────────────────────────────┘
+ *
+ * @param broker - The broker connection to query
+ * @param abilityName - Name of the ability/agent to find (matches displayName)
+ * @param options - Discovery options (timeout)
+ * @returns Information about the discovered ability
+ * @throws KadiError if ability not found
+ */
+async function discoverAbilityWorkaround(
+  broker: BrokerState,
+  abilityName: string,
+  options: { timeoutMs: number }
+): Promise<DiscoveredAbility> {
+  // Step 1: Fetch all tools with provider information
+  // This is inefficient but necessary until the broker supports agent queries
+  const response = (await sendBrokerRequest(
+    broker,
+    'kadi.ability.list',
+    { includeProviders: true },
+    options.timeoutMs
+  )) as AbilityListResponse;
+
+  if (!response?.tools) {
+    throw new KadiError(
+      `Failed to list tools from broker "${broker.name}"`,
+      'BROKER_ERROR',
+      { broker: broker.name }
+    );
+  }
+
+  // Step 2: Filter tools to find those provided by the named agent
+  // We match on provider.displayName (the agent's name from registration)
+  const matchingTools: ToolDefinition[] = [];
+  let targetAgentId: string | undefined;
+  let targetDisplayName: string | undefined;
+
+  for (const tool of response.tools) {
+    if (!tool.providers) continue;
+
+    // Find a provider whose displayName matches our ability name
+    const matchingProvider = tool.providers.find(
+      (p) => p.displayName === abilityName
+    );
+
+    if (matchingProvider) {
+      // Extract the agentId (stable identifier for routing and event filtering)
+      // We use the first matching provider's agentId for all invocations
+      if (!targetAgentId && matchingProvider.agentId) {
+        targetAgentId = matchingProvider.agentId;
+        targetDisplayName = matchingProvider.displayName;
+      }
+
+      // Add tool to our list (convert to ToolDefinition format)
+      matchingTools.push({
+        name: tool.name,
+        description: tool.description ?? '',
+        inputSchema: tool.inputSchema ?? { type: 'object' },
+      });
+    }
+  }
+
+  // Step 3: Validate we found the ability
+  if (!targetAgentId || matchingTools.length === 0) {
+    // Build list of available agents for helpful error message
+    const availableAgents = new Set<string>();
+    for (const tool of response.tools) {
+      for (const provider of tool.providers ?? []) {
+        if (provider.displayName) {
+          availableAgents.add(provider.displayName);
+        }
+      }
+    }
+
+    const agentList = [...availableAgents];
+    const hint = agentList.length === 0
+      ? `No agents with tools found. Make sure the ability: (1) is connected via serve('broker'), (2) has at least one tool registered`
+      : `Available agents with tools: [${agentList.join(', ')}]. Is '${abilityName}' connected with tools?`;
+
+    throw new KadiError(
+      `Ability "${abilityName}" not found on broker "${broker.name}"`,
+      'ABILITY_NOT_FOUND',
+      {
+        abilityName,
+        broker: broker.name,
+        availableAgents: agentList,
+        hint,
+      }
+    );
+  }
+
+  return {
+    agentId: targetAgentId,
+    name: targetDisplayName ?? abilityName,
+    tools: matchingTools,
+  };
+}
+
+// ═══════════════════════════════════════════════════════════════════════════════
+// REMOTE INVOCATION
+// ═══════════════════════════════════════════════════════════════════════════════
+
+/**
+ * Invoke a tool on a remote agent via the broker.
+ *
+ * Uses the KADI async invocation pattern:
+ * 1. Generate requestId and set up response listener FIRST (avoids race condition)
+ * 2. Send kadi.ability.request with toolName, toolInput, and requestId
+ * 3. Broker returns { status: 'pending', requestId }
+ * 4. Wait for kadi.ability.response notification with the result
+ *
+ * For consistent routing to a specific agent, we inject _kadi routing hints
+ * into the toolInput. The broker strips these before forwarding to the provider.
+ *
+ * @param broker - The broker connection to use
+ * @param targetAgentId - Agent ID to route to (uses requireAgent hint)
+ * @param toolName - Name of the tool to invoke
+ * @param params - Parameters for the tool
+ * @param timeoutMs - Timeout in milliseconds
+ * @returns The tool's result
+ */
+async function invokeViaBroker<T>(
+  broker: BrokerState,
+  targetAgentId: string,
+  toolName: string,
+  params: unknown,
+  timeoutMs: number
+): Promise<T> {
+  // Fail fast if broker isn't properly initialized
+  if (!broker.pendingInvocations) {
+    throw new KadiError(
+      `Broker "${broker.name}" is missing pendingInvocations map`,
+      'BROKER_ERROR',
+      { broker: broker.name, hint: 'This is a bug in broker initialization' }
+    );
+  }
+
+  // Generate requestId FIRST, before any async operations.
+  // This allows us to set up the response listener before sending,
+  // eliminating the race condition where fast responses arrive
+  // before the listener exists.
+  const requestId = crypto.randomUUID();
+
+  // Set up result listener BEFORE sending the request.
+  // When the broker sends kadi.ability.response, the message handler
+  // will find this listener and resolve/reject the promise.
+  const resultPromise = new Promise<T>((resolve, reject) => {
+    const timeoutHandle = setTimeout(() => {
+      broker.pendingInvocations.delete(requestId);
+      reject(
+        new KadiError(`Tool "${toolName}" invocation timed out`, 'BROKER_TIMEOUT', {
+          broker: broker.name,
+          toolName,
+          requestId,
+          timeout: timeoutMs,
+        })
+      );
+    }, timeoutMs);
+
+    broker.pendingInvocations.set(requestId, {
+      resolve: (result: unknown) => {
+        clearTimeout(timeoutHandle);
+        resolve(result as T);
+      },
+      reject: (error: Error) => {
+        clearTimeout(timeoutHandle);
+        reject(error);
+      },
+      timeout: timeoutHandle,
+      toolName,
+      sentAt: new Date(),
+    });
+  });
+
+  // Inject routing hint to ensure the broker routes to our target agent
+  // The _kadi property is stripped by the broker before forwarding to the provider
+  const paramsWithRouting = {
+    ...(typeof params === 'object' && params !== null ? params : {}),
+    _kadi: { requireAgent: targetAgentId },
+  };
+
+  // Helper to clean up the pending invocation on failure
+  const cleanupPendingInvocation = () => {
+    const pending = broker.pendingInvocations.get(requestId);
+    if (pending) {
+      clearTimeout(pending.timeout);
+      broker.pendingInvocations.delete(requestId);
+    }
+  };
+
+  // NOW send the request (listener already exists, no race possible)
+  try {
+    const pendingResult = (await sendBrokerRequest(
+      broker,
+      'kadi.ability.request',
+      {
+        toolName,
+        toolInput: paramsWithRouting,
+        requestId,
+      },
+      timeoutMs
+    )) as { status: string; requestId?: string };
+
+    // Validate the broker accepted the request
+    if (pendingResult.status !== 'pending') {
+      cleanupPendingInvocation();
+      throw new KadiError(
+        'Unexpected response from broker: expected pending acknowledgment',
+        'BROKER_ERROR',
+        {
+          broker: broker.name,
+          toolName,
+          response: pendingResult,
+        }
+      );
+    }
+  } catch (error) {
+    // Clean up listener if request failed
+    cleanupPendingInvocation();
+    throw error;
+  }
+
+  return resultPromise;
+}
+
+// ═══════════════════════════════════════════════════════════════════════════════
+// BROKER TRANSPORT
+// ═══════════════════════════════════════════════════════════════════════════════
+
+/**
+ * Load a remote ability via broker.
+ *
+ * This discovers an agent providing the ability and returns a LoadedAbility
+ * interface that invokes tools via the broker, consistently routing to that
+ * specific agent.
+ *
+ * @param abilityName - Name of the ability to load (matches agent's displayName)
+ * @param options - Broker connection and options
+ * @returns LoadedAbility that can invoke tools remotely
+ *
+ * @example
+ * ```typescript
+ * // The broker state comes from KadiClient
+ * const brokerState = client.getBrokerState('production');
+ *
+ * // Load a remote ability
+ * const analyzer = await loadBrokerTransport('text-analyzer', {
+ *   broker: brokerState,
+ *   requestTimeout: 60000,
+ * });
+ *
+ * // Invoke tools on the remote agent
+ * const result = await analyzer.invoke('analyze', { text: 'hello' });
+ *
+ * // All invocations go to the same agent (consistent routing)
+ * const result2 = await analyzer.invoke('summarize', { text: 'world' });
+ * ```
+ */
+export async function loadBrokerTransport(
+  abilityName: string,
+  options: BrokerTransportOptions
+): Promise<LoadedAbility> {
+  const timeoutMs = options.requestTimeout ?? 600000; // 10 minutes default
+
+  // ─────────────────────────────────────────────────────────────────────────────
+  // Step 1: Discover the ability on the broker
+  // ─────────────────────────────────────────────────────────────────────────────
+  // TODO: Replace discoverAbilityWorkaround with kadi.agent.discover when
+  // the broker implements it. See: kadi-broker/docs/design/ABILITY_CENTRIC_DISCOVERY.md
+  const discovered = await discoverAbilityWorkaround(options.broker, abilityName, {
+    timeoutMs,
+  });
+
+  // ─────────────────────────────────────────────────────────────────────────────
+  // Step 2: Set up event subscription tracking
+  // ─────────────────────────────────────────────────────────────────────────────
+  // Maps user handlers to their wrapper handlers for cleanup in off()
+  // Structure: handler → (pattern → wrapper)
+  const handlerWrappers = new Map<EventHandler, Map<string, BrokerEventHandler>>();
+
+  // ─────────────────────────────────────────────────────────────────────────────
+  // Step 3: Return LoadedAbility interface
+  // ─────────────────────────────────────────────────────────────────────────────
+  // The agentId is used for routing hints, ensuring all invocations go to the
+  // same agent. This is stable across reconnects (unlike sessionId).
+  return {
+    name: discovered.name,
+    transport: 'broker',
+
+    /**
+     * Invoke a tool on the remote agent.
+     *
+     * Routes through the broker to the target agent using _kadi.requireAgent
+     * routing hint. This ensures all calls go to the same agent that was
+     * discovered during loadBrokerTransport().
+     *
+     * @param toolName - Name of the tool to invoke
+     * @param params - Input parameters for the tool
+     * @param invokeOptions - Optional settings (timeout override)
+     */
+    async invoke<T>(toolName: string, params: unknown, invokeOptions?: InvokeOptions): Promise<T> {
+      // Per-call timeout overrides the default timeout set when loading
+      const effectiveTimeout = invokeOptions?.timeout ?? timeoutMs;
+
+      try {
+        return await invokeViaBroker<T>(
+          options.broker,
+          discovered.agentId,
+          toolName,
+          params,
+          effectiveTimeout
+        );
+      } catch (error) {
+        // Re-throw KadiErrors as-is
+        if (error instanceof KadiError) {
+          throw error;
+        }
+        // Wrap other errors
+        throw new KadiError(`Remote tool "${toolName}" invocation failed`, 'TOOL_INVOCATION_FAILED', {
+          toolName,
+          ability: discovered.name,
+          targetAgent: discovered.agentId,
+          broker: options.broker.name,
+          reason: error instanceof Error ? error.message : String(error),
+        });
+      }
+    },
+
+    /**
+     * Get the list of tools this ability provides.
+     *
+     * Returns tools discovered from the agent during loadBrokerTransport().
+     * This is a snapshot from discovery time - if the agent registers new
+     * tools, you'd need to reload the ability to see them.
+     */
+    getTools(): ToolDefinition[] {
+      return discovered.tools;
+    },
+
+    /**
+     * Subscribe to events from this ability via the broker.
+     *
+     * Events are filtered to only include those from this specific ability
+     * (by matching the source agentId). The agentId is stable across reconnects.
+     *
+     * NOTE: Broker subscription is async but this method is sync for API consistency.
+     * If subscription fails, an error is logged but not thrown.
+     *
+     * @param pattern - Event pattern to subscribe to (e.g., 'order.*')
+     * @param handler - Handler function called with event data (not full BrokerEvent)
+     */
+    on(pattern: string, handler: EventHandler): void {
+      if (!options.subscribe) {
+        throw new KadiError(
+          'Event subscriptions not available: subscribe callback not provided',
+          'NOT_IMPLEMENTED',
+          {
+            transport: 'broker',
+            hint: 'Ensure loadBroker() passes subscribe/unsubscribe callbacks',
+          }
+        );
+      }
+
+      // Create wrapper that filters by source (agentId) and extracts just the data
+      const wrapper: BrokerEventHandler = (event: BrokerEvent) => {
+        // Filter: only process events from this ability's agent
+        // source now contains agentId (stable across reconnects)
+        if (event.source === discovered.agentId) {
+          // Call user handler with just the data (matching stdio/native behavior)
+          handler(event.data);
+        }
+      };
+
+      // Track the wrapper for cleanup in off()
+      let patternMap = handlerWrappers.get(handler);
+      if (!patternMap) {
+        patternMap = new Map();
+        handlerWrappers.set(handler, patternMap);
+      }
+      patternMap.set(pattern, wrapper);
+
+      // Subscribe (fire-and-forget, errors logged by subscribe implementation)
+      options.subscribe(pattern, wrapper).catch((err) => {
+        const message = err instanceof Error ? err.message : String(err);
+        console.error(`[KADI] ability.on() subscribe failed for "${pattern}": ${message}`);
+      });
+    },
+
+    /**
+     * Unsubscribe from events.
+     *
+     * @param pattern - Event pattern to unsubscribe from
+     * @param handler - The same handler function passed to on()
+     */
+    off(pattern: string, handler: EventHandler): void {
+      if (!options.unsubscribe) {
+        // Silently ignore - off() is cleanup, shouldn't throw during teardown
+        return;
+      }
+
+      // Find the wrapper for this handler/pattern combo
+      const patternMap = handlerWrappers.get(handler);
+      if (!patternMap) return;
+
+      const wrapper = patternMap.get(pattern);
+      if (!wrapper) return;
+
+      // Clean up tracking
+      patternMap.delete(pattern);
+      if (patternMap.size === 0) {
+        handlerWrappers.delete(handler);
+      }
+
+      // Unsubscribe (fire-and-forget)
+      options.unsubscribe(pattern, wrapper).catch((err) => {
+        const message = err instanceof Error ? err.message : String(err);
+        console.error(`[KADI] ability.off() unsubscribe failed for "${pattern}": ${message}`);
+      });
+    },
+
+    /**
+     * Disconnect and cleanup.
+     *
+     * For broker transport, cleans up event subscriptions. The broker connection
+     * itself is managed by KadiClient, not by individual abilities.
+     */
+    async disconnect(): Promise<void> {
+      // Clean up all event subscriptions
+      if (options.unsubscribe) {
+        for (const [_handler, patternMap] of handlerWrappers) {
+          for (const [pattern, wrapper] of patternMap) {
+            try {
+              await options.unsubscribe(pattern, wrapper);
+            } catch {
+              // Ignore errors during cleanup
+            }
+          }
+        }
+      }
+      handlerWrappers.clear();
+    },
+  };
+}