@kadi.build/core 0.7.2 → 0.9.0

package/src/client.ts

@@ -31,6 +31,7 @@ import type {
   BrokerState,
   RegisteredTool,
   ToolDefinition,
+  BrokerToolDefinition,
   ZodToolDefinition,
   ToolHandler,
   LoadedAbility,
@@ -52,6 +53,7 @@ import type {
   EmitOptions,
 } from './types.js';
 import { KadiError } from './errors.js';
+import * as protocol from './protocol.js';
 import { zodToJsonSchema, isZodSchema } from './zod.js';
 import { resolveAbilityEntry, resolveAbilityScript } from './lockfile.js';
 import { loadNativeTransport } from './transports/native.js';
@@ -156,8 +158,8 @@ function isMcpCallToolResult(result: unknown): boolean {
  * const client = new KadiClient({
  *   name: 'my-agent',
  *   brokers: {
- *     production: 'ws://broker-prod:8080',
- *     internal: 'ws://broker-internal:8080',
+ *     production: { url: 'ws://broker-prod:8080/kadi' },
+ *     internal: { url: 'ws://broker-internal:8080/kadi', networks: ['private'] },
  *   },
  *   defaultBroker: 'production',
  * });
@@ -241,17 +243,24 @@ export class KadiClient {
     this._agentId = crypto.createHash('sha256').update(this._publicKeyBase64).digest('hex').substring(0, 16);
 
     // Resolve configuration with defaults
-    // Auto-select first broker as default if not specified (matches Python behavior)
-    const brokers = config.brokers ?? {};
-    const firstBrokerName = Object.keys(brokers)[0];
+    const rawBrokers = config.brokers ?? {};
+    const normalizedBrokers: Record<string, { url: string; networks: string[] }> = {};
+
+    for (const [name, entry] of Object.entries(rawBrokers)) {
+      normalizedBrokers[name] = {
+        url: entry.url,
+        networks: entry.networks ?? ['global'],
+      };
+    }
+
+    const firstBrokerName = Object.keys(normalizedBrokers)[0];
 
     this.config = {
       name: config.name,
       version: config.version ?? '1.0.0',
       description: config.description ?? '',
-      brokers,
+      brokers: normalizedBrokers,
       defaultBroker: config.defaultBroker ?? firstBrokerName,
-      networks: config.networks ?? ['global'],
       heartbeatInterval: config.heartbeatInterval ?? DEFAULT_HEARTBEAT_INTERVAL,
       requestTimeout: config.requestTimeout ?? DEFAULT_REQUEST_TIMEOUT,
       autoReconnect: config.autoReconnect ?? DEFAULT_AUTO_RECONNECT,
@@ -445,8 +454,7 @@ export class KadiClient {
     const failures: string[] = [];
     for (const [i, result] of results.entries()) {
       if (result.status === 'rejected') {
-        const failedBroker = brokerNames[i];
-        if (failedBroker === undefined) continue;
+        const failedBroker = brokerNames[i]!;
         const reason = result.reason instanceof Error ? result.reason.message : String(result.reason);
         console.error(`[KADI] Failed to connect to broker "${failedBroker}": ${reason}`);
         failures.push(failedBroker);
@@ -485,8 +493,8 @@ export class KadiClient {
    * across all broker connections, ensuring consistent identity.
    */
   private async connectToBroker(brokerName: string): Promise<void> {
-    const url = this.config.brokers[brokerName];
-    if (!url) {
+    const brokerConfig = this.config.brokers[brokerName];
+    if (!brokerConfig) {
       throw new KadiError(`Broker "${brokerName}" not found in configuration`, 'UNKNOWN_BROKER', {
         broker: brokerName,
         available: Object.keys(this.config.brokers),
@@ -494,6 +502,8 @@ export class KadiClient {
       });
     }
 
+    const url = brokerConfig.url;
+
     // Check if already connected
     if (this.brokers.has(brokerName)) {
       const existing = this.brokers.get(brokerName)!;
@@ -506,6 +516,7 @@ export class KadiClient {
     const state: BrokerState = {
       name: brokerName,
       url,
+      networks: brokerConfig.networks,
       ws: null,
       heartbeatTimer: null,
       pendingRequests: new Map(),
@@ -530,12 +541,8 @@ export class KadiClient {
     // Register with broker (transitions to "ready" state)
     await this.registerWithBroker(state);
 
-    // Start heartbeat
-    state.heartbeatTimer = setInterval(() => {
-      this.sendHeartbeat(state);
-    }, this.config.heartbeatInterval);
-
-    state.status = 'connected';
+    // Finalize connection (heartbeat + status)
+    this.finalizeConnection(state);
   }
 
   // ─────────────────────────────────────────────────────────────
@@ -543,12 +550,7 @@ export class KadiClient {
   // ─────────────────────────────────────────────────────────────
 
   private buildHelloMessage(): JsonRpcRequest {
-    return {
-      jsonrpc: '2.0',
-      id: this.nextRequestId++,
-      method: 'kadi.session.hello',
-      params: { role: 'agent' },
-    };
+    return protocol.hello(this.nextRequestId++);
   }
 
   private buildAuthMessage(nonce: string): JsonRpcRequest {
@@ -560,38 +562,20 @@ export class KadiClient {
     });
     const signature = crypto.sign(null, Buffer.from(nonce, 'utf8'), privateKey);
 
-    return {
-      jsonrpc: '2.0',
-      id: this.nextRequestId++,
-      method: 'kadi.session.authenticate',
-      params: {
-        publicKey: this._publicKeyBase64,
-        signature: signature.toString('base64'),
-        nonce,
-      },
-    };
+    return protocol.authenticate(
+      this.nextRequestId++,
+      this._publicKeyBase64,
+      signature.toString('base64'),
+      nonce,
+    );
   }
 
-  private buildRegisterMessage(tools: ToolDefinition[], networks: string[]): JsonRpcRequest {
-    return {
-      jsonrpc: '2.0',
-      id: this.nextRequestId++,
-      method: 'kadi.agent.register',
-      params: {
-        tools,
-        networks,
-        displayName: this.config.name,
-      },
-    };
+  private buildRegisterMessage(tools: BrokerToolDefinition[], networks: string[]): JsonRpcRequest {
+    return protocol.register(this.nextRequestId++, tools, networks, this.config.name);
   }
 
   private buildHeartbeatMessage(): JsonRpcRequest {
-    return {
-      jsonrpc: '2.0',
-      id: this.nextRequestId++,
-      method: 'kadi.session.heartbeat',
-      params: { timestamp: Date.now() },
-    };
+    return protocol.heartbeat(this.nextRequestId++);
   }
 
   // ─────────────────────────────────────────────────────────────
@@ -718,10 +702,17 @@ export class KadiClient {
     // Note: If registration fails, sendRequest will reject with the error.
     // We don't need to check response.error here - errors are already
     // handled via Promise rejection in handleBrokerResponse.
-    await this.sendRequest<{ status: string }>(
+    const { droppedNetworks } = await this.sendRequest<{ status: string; droppedNetworks?: string[] }>(
       state,
-      this.buildRegisterMessage(tools, this.config.networks)
+      this.buildRegisterMessage(tools, state.networks)
     );
+
+    if (droppedNetworks && droppedNetworks.length > 0) {
+      console.warn(
+        `[KADI] Broker "${state.name}" dropped networks ${JSON.stringify(droppedNetworks)} ` +
+        `during registration. Tools scoped to those networks will not be discoverable.`
+      );
+    }
   }
 
   /**
@@ -1037,12 +1028,7 @@ export class KadiClient {
 
     // If this is the first handler for this pattern, subscribe on broker
     if (!state.subscribedPatterns.has(pattern)) {
-      await this.sendRequest(state, {
-        jsonrpc: '2.0',
-        id: this.nextRequestId++,
-        method: 'kadi.event.subscribe',
-        params: { pattern },
-      });
+      await this.sendRequest(state, protocol.eventSubscribe(this.nextRequestId++, pattern));
       state.subscribedPatterns.add(pattern);
     }
   }
@@ -1097,12 +1083,7 @@ export class KadiClient {
         // Unsubscribe from broker if we were subscribed
         if (state.subscribedPatterns.has(pattern) && state.status === 'connected') {
           try {
-            await this.sendRequest(state, {
-              jsonrpc: '2.0',
-              id: this.nextRequestId++,
-              method: 'kadi.event.unsubscribe',
-              params: { pattern },
-            });
+            await this.sendRequest(state, protocol.eventUnsubscribe(this.nextRequestId++, pattern));
           } catch {
             // Ignore errors during unsubscribe (broker might be disconnecting)
           }
@@ -1134,20 +1115,11 @@ export class KadiClient {
   async publish(channel: string, data: unknown, options: PublishOptions = {}): Promise<void> {
     const { state } = this.getConnectedBrokerState(options.broker);
 
-    // Resolve which network to publish to
-    const networkId = options.network ?? this.config.networks[0] ?? 'global';
+    // Resolve which network to publish to (prefer broker-specific, fallback to global)
+    const networkId = options.network ?? state.networks[0] ?? '';
 
     // Send publish request to broker
-    await this.sendRequest(state, {
-      jsonrpc: '2.0',
-      id: this.nextRequestId++,
-      method: 'kadi.event.publish',
-      params: {
-        channel,
-        data,
-        networkId,
-      },
-    });
+    await this.sendRequest(state, protocol.eventPublish(this.nextRequestId++, channel, data, networkId));
   }
 
   // ─────────────────────────────────────────────────────────────
@@ -1463,34 +1435,17 @@ export class KadiClient {
         responseResult = result;
       }
 
-      const response: JsonRpcResponse = {
-        jsonrpc: '2.0',
-        id: requestId,
-        result: responseResult,
-      };
-      state.ws?.send(JSON.stringify(response));
+      state.ws?.send(JSON.stringify(protocol.resultResponse(requestId, responseResult)));
     } catch (error) {
       const errorMessage = error instanceof Error ? error.message : String(error);
 
       // For MCP clients, wrap errors in CallToolResult format too
       if (context.callerProtocol === 'mcp') {
-        const response: JsonRpcResponse = {
-          jsonrpc: '2.0',
-          id: requestId,
-          result: { content: [{ type: 'text', text: errorMessage }], isError: true },
-        };
-        state.ws?.send(JSON.stringify(response));
+        const mcpResult = { content: [{ type: 'text', text: errorMessage }], isError: true };
+        state.ws?.send(JSON.stringify(protocol.resultResponse(requestId, mcpResult)));
       } else {
         // KADI clients receive JSON-RPC error
-        const response: JsonRpcResponse = {
-          jsonrpc: '2.0',
-          id: requestId,
-          error: {
-            code: -32000,
-            message: errorMessage,
-          },
-        };
-        state.ws?.send(JSON.stringify(response));
+        state.ws?.send(JSON.stringify(protocol.errorResponse(requestId, -32000, errorMessage)));
       }
     }
   }
@@ -1520,6 +1475,17 @@ export class KadiClient {
     }
   }
 
+  /**
+   * Finalize a broker connection: start heartbeat and mark as connected.
+   * Shared by initial connect and reconnect paths.
+   */
+  private finalizeConnection(state: BrokerState): void {
+    state.heartbeatTimer = setInterval(() => {
+      this.sendHeartbeat(state);
+    }, this.config.heartbeatInterval);
+    state.status = 'connected';
+  }
+
   /**
    * Send heartbeat to keep connection alive.
    */
@@ -1565,6 +1531,10 @@ export class KadiClient {
       );
     }
     state.pendingInvocations.clear();
+
+    // Clear broker-side subscription tracking — subscriptions are lost when
+    // the connection drops and must be re-established after reconnection.
+    state.subscribedPatterns.clear();
   }
 
   // ─────────────────────────────────────────────────────────────
@@ -1701,15 +1671,12 @@ export class KadiClient {
       // Re-register tools
       await this.registerWithBroker(state);
 
-      // Restart heartbeat
-      state.heartbeatTimer = setInterval(() => {
-        this.sendHeartbeat(state);
-      }, this.config.heartbeatInterval);
+      // Finalize connection (heartbeat + status)
+      this.finalizeConnection(state);
 
       // Success!
       console.error(`[KADI] Reconnected to broker "${state.name}" after ${state.reconnectAttempts} attempts`);
       state.reconnectAttempts = 0;
-      state.status = 'connected';
     } catch (error) {
       // Log the error and try again
       const message = error instanceof Error ? error.message : String(error);
@@ -1819,12 +1786,7 @@ export class KadiClient {
       this.eventHandler(event, data);
     } else if (this.isServingStdio) {
       // Stdio transport: write notification to stdout
-      const notification = {
-        jsonrpc: '2.0',
-        method: 'event',
-        params: { name: event, data },
-      };
-      const json = JSON.stringify(notification);
+      const json = JSON.stringify(protocol.eventNotification(event, data));
       process.stdout.write(`Content-Length: ${Buffer.byteLength(json)}\r\n\r\n${json}`);
     }
 
@@ -1861,8 +1823,12 @@ export class KadiClient {
    *   output: z.object({ result: z.number() }),
    * }, async ({ a, b }) => ({ result: a + b }));
    *
-   * // Register only on specific broker
-   * client.registerTool(def, handler, { brokers: ['internal'] });
+   * // Register on specific brokers with per-broker networks
+   * client.registerTool(def, handler, {
+   *   brokers: {
+   *     internal: { networks: ['private'] },
+   *   },
+   * });
    * ```
    */
   registerTool<TInput, TOutput>(
@@ -1888,12 +1854,44 @@ export class KadiClient {
         : undefined,
     };
 
+    // Validate per-broker network scoping (skip when no brokers configured)
+    const brokerScopes = options.brokers ?? {};
+    if (Object.keys(brokerScopes).length > 0 && Object.keys(this.config.brokers).length > 0) {
+      for (const [brokerName, scope] of Object.entries(brokerScopes)) {
+        const brokerConfig = this.config.brokers[brokerName];
+        if (!brokerConfig) {
+          throw new KadiError(
+            `Tool "${definition.name}" references broker "${brokerName}" which is not configured. Available brokers: ${JSON.stringify(Object.keys(this.config.brokers))}`,
+            'INVALID_CONFIG',
+            {
+              toolName: definition.name,
+              broker: brokerName,
+              availableBrokers: Object.keys(this.config.brokers),
+            }
+          );
+        }
+        const brokerNetworks = brokerConfig.networks;
+        const invalid = scope.networks.filter((n) => !brokerNetworks.includes(n));
+        if (invalid.length > 0) {
+          throw new KadiError(
+            `Tool "${definition.name}" has networks ${JSON.stringify(invalid)} not present on broker "${brokerName}" networks ${JSON.stringify(brokerNetworks)}. Per-tool networks must be a subset of the broker's networks — a tool can only be visible on networks the agent has joined.`,
+            'INVALID_CONFIG',
+            {
+              toolName: definition.name,
+              broker: brokerName,
+              invalidNetworks: invalid,
+              brokerNetworks,
+            }
+          );
+        }
+      }
+    }
+
     // Store registration
     const registered: RegisteredTool = {
       definition: jsonDefinition,
       handler: handler as ToolHandler,
-      registeredAt: new Date(),
-      targetBrokers: options.brokers ?? [],
+      brokerNetworks: brokerScopes,
     };
     this.tools.set(definition.name, registered);
   }
@@ -1902,21 +1900,36 @@ export class KadiClient {
    * Get tool definitions, optionally filtered for a specific broker.
    *
    * @param forBroker - If provided, only return tools targeted for this broker.
-   *                    Tools with empty targetBrokers are included for all brokers.
+   *                    Tools with empty brokerNetworks are included for all brokers.
    */
-  private getToolDefinitions(forBroker?: string): ToolDefinition[] {
-    return Array.from(this.tools.values())
-      .filter((t) => {
-        // If no broker specified, return all tools (e.g., for readAgentJson)
-        if (!forBroker) return true;
+  private getToolDefinitions(forBroker?: string): BrokerToolDefinition[] {
+    const results: BrokerToolDefinition[] = [];
+
+    for (const t of this.tools.values()) {
+      const hasBrokerNetworks = Object.keys(t.brokerNetworks).length > 0;
+
+      if (!forBroker) {
+        // No broker filter — return all tools (e.g., for readAgentJson)
+        results.push({ ...t.definition });
+        continue;
+      }
+
+      if (!hasBrokerNetworks) {
+        // Empty brokerNetworks = register with all brokers, no network field
+        results.push({ ...t.definition });
+        continue;
+      }
 
-        // Empty targetBrokers means "register with all brokers"
-        if (t.targetBrokers.length === 0) return true;
+      const scope = t.brokerNetworks[forBroker];
+      if (!scope) continue; // Tool not registered on this broker
 
-        // Otherwise, only include if this broker is in the target list
-        return t.targetBrokers.includes(forBroker);
-      })
-      .map((t) => t.definition);
+      results.push({
+        ...t.definition,
+        ...(scope.networks.length > 0 && { networks: scope.networks }),
+      });
+    }
+
+    return results;
   }
 
   /**
@@ -2090,7 +2103,7 @@ export class KadiClient {
     const ability = await loadBrokerTransport(name, {
       broker: state,
       requestTimeout: options.timeout ?? this.config.requestTimeout,
-      networks: options.networks,
+      networks: options.networks ?? state.networks,
       // Provide subscribe/unsubscribe for ability.on()/off() support
       subscribe: (pattern, handler) => this.subscribe(pattern, handler, { broker: brokerName }),
       unsubscribe: (pattern, handler) => this.unsubscribe(pattern, handler, { broker: brokerName }),
@@ -2184,13 +2197,8 @@ export class KadiClient {
     try {
       const pendingResult = await this.sendRequest<{ status: string; requestId: string }>(
         state,
-        {
-          jsonrpc: '2.0',
-          id: ++this.nextRequestId,
-          method: 'kadi.ability.request',
-          // timeout sent so the broker can enforce a matching server-side deadline
-          params: { toolName, toolInput: params, requestId, timeout },
-        }
+        // timeout sent so the broker can enforce a matching server-side deadline
+        protocol.abilityRequest(++this.nextRequestId, toolName, params, requestId, timeout)
       );
 
       // Validate the broker accepted the request
@@ -2365,12 +2373,7 @@ export class KadiClient {
    * Send a response via stdio.
    */
   private sendStdioResponse(id: string | number, result: unknown): void {
-    const response: JsonRpcResponse = {
-      jsonrpc: '2.0',
-      id,
-      result,
-    };
-    const json = JSON.stringify(response);
+    const json = JSON.stringify(protocol.resultResponse(id, result));
     process.stdout.write(`Content-Length: ${Buffer.byteLength(json)}\r\n\r\n${json}`);
   }
 
@@ -2378,12 +2381,7 @@ export class KadiClient {
    * Send an error response via stdio.
    */
   private sendStdioError(id: string | number, error: { code: number; message: string }): void {
-    const response: JsonRpcResponse = {
-      jsonrpc: '2.0',
-      id,
-      error,
-    };
-    const json = JSON.stringify(response);
+    const json = JSON.stringify(protocol.errorResponse(id, error.code, error.message));
     process.stdout.write(`Content-Length: ${Buffer.byteLength(json)}\r\n\r\n${json}`);
   }
 
@@ -2418,7 +2416,7 @@ export class KadiClient {
   /**
    * Get agent information (for readAgentJson protocol).
    */
-  readAgentJson(): { name: string; version: string; tools: ToolDefinition[] } {
+  readAgentJson(): { name: string; version: string; tools: BrokerToolDefinition[] } {
     return {
       name: this.config.name,
       version: this.config.version,

package/src/index.ts

@@ -10,7 +10,7 @@
  * const client = new KadiClient({
  *   name: 'my-agent',
  *   brokers: {
- *     production: 'ws://broker:8080',
+ *     production: { url: 'ws://broker:8080/kadi' },
  *   },
  *   defaultBroker: 'production',
  * });
@@ -39,10 +39,12 @@ export type { ErrorCode, ErrorContext } from './errors.js';
 export type {
   // Configuration
   ClientConfig,
+  BrokerEntry,
   ResolvedConfig,
 
   // Tools
   ToolDefinition,
+  BrokerToolDefinition,
   ZodToolDefinition,
   ToolHandler,
   RegisteredTool,
@@ -67,6 +69,7 @@ export type {
   BrokerState,
   BrokerStatus,
   PendingRequest,
+  PendingInvocation,
 
   // Broker Events (Pub/Sub)
   BrokerEvent,
@@ -109,6 +112,9 @@ export { loadStdioTransport } from './transports/stdio.js';
 export { loadBrokerTransport } from './transports/broker.js';
 export type { BrokerTransportOptions } from './transports/broker.js';
 
+// Protocol message builders
+export * as protocol from './protocol.js';
+
 // Crypto utilities (Ed25519 to X25519 conversion)
 export { convertToEncryptionKey, convertToEncryptionKeyPair } from './crypto.js';
 export type { EncryptionKeyPair } from './crypto.js';