@google/gemini-cli-core 0.34.0-preview.3 → 0.35.0-nightly.20260314.3038fdce2

package/dist/src/agents/a2a-client-manager.js

@@ -3,41 +3,51 @@
  * Copyright 2025 Google LLC
  * SPDX-License-Identifier: Apache-2.0
  */
-import { ClientFactory, ClientFactoryOptions, DefaultAgentCardResolver, RestTransportFactory, JsonRpcTransportFactory, createAuthenticatingFetchWithRetry, } from '@a2a-js/sdk/client';
+import { ClientFactory, ClientFactoryOptions, DefaultAgentCardResolver, JsonRpcTransportFactory, RestTransportFactory, createAuthenticatingFetchWithRetry, } from '@a2a-js/sdk/client';
+import { GrpcTransportFactory } from '@a2a-js/sdk/client/grpc';
+import * as grpc from '@grpc/grpc-js';
 import { v4 as uuidv4 } from 'uuid';
-import { Agent as UndiciAgent } from 'undici';
+import { Agent as UndiciAgent, ProxyAgent } from 'undici';
+import { normalizeAgentCard } from './a2aUtils.js';
 import { debugLogger } from '../utils/debugLogger.js';
-import { safeLookup } from '../utils/fetch.js';
 import { classifyAgentError } from './a2a-errors.js';
 // Remote agents can take 10+ minutes (e.g. Deep Research).
 // Use a dedicated dispatcher so the global 5-min timeout isn't affected.
 const A2A_TIMEOUT = 1800000; // 30 minutes
-const a2aDispatcher = new UndiciAgent({
-    headersTimeout: A2A_TIMEOUT,
-    bodyTimeout: A2A_TIMEOUT,
-    connect: {
-        lookup: safeLookup, // SSRF protection at connection level
-    },
-});
-const a2aFetch = (input, init) => 
-// eslint-disable-next-line no-restricted-syntax -- TODO: Migrate to safeFetch for SSRF protection
-fetch(input, { ...init, dispatcher: a2aDispatcher });
 /**
- * Manages A2A clients and caches loaded agent information.
- * Follows a singleton pattern to ensure a single client instance.
+ * Orchestrates communication with remote A2A agents.
+ * Manages protocol negotiation, authentication, and transport selection.
  */
 export class A2AClientManager {
     static instance;
     // Each agent should manage their own context/taskIds/card/etc
     clients = new Map();
     agentCards = new Map();
-    constructor() { }
+    a2aDispatcher;
+    a2aFetch;
+    constructor(config) {
+        const proxyUrl = config?.getProxy();
+        const agentOptions = {
+            headersTimeout: A2A_TIMEOUT,
+            bodyTimeout: A2A_TIMEOUT,
+        };
+        if (proxyUrl) {
+            this.a2aDispatcher = new ProxyAgent({
+                uri: proxyUrl,
+                ...agentOptions,
+            });
+        }
+        else {
+            this.a2aDispatcher = new UndiciAgent(agentOptions);
+        }
+        this.a2aFetch = (input, init) => fetch(input, { ...init, dispatcher: this.a2aDispatcher });
+    }
     /**
      * Gets the singleton instance of the A2AClientManager.
      */
-    static getInstance() {
+    static getInstance(config) {
         if (!A2AClientManager.instance) {
-            A2AClientManager.instance = new A2AClientManager();
+            A2AClientManager.instance = new A2AClientManager(config);
         }
         return A2AClientManager.instance;
     }
@@ -61,15 +71,15 @@ export class A2AClientManager {
             throw new Error(`Agent with name '${name}' is already loaded.`);
         }
         // Authenticated fetch for API calls (transports).
-        let authFetch = a2aFetch;
+        let authFetch = this.a2aFetch;
         if (authHandler) {
-            authFetch = createAuthenticatingFetchWithRetry(a2aFetch, authHandler);
+            authFetch = createAuthenticatingFetchWithRetry(this.a2aFetch, authHandler);
         }
         // Use unauthenticated fetch for the agent card unless explicitly required.
         // Some servers reject unexpected auth headers on the card endpoint (e.g. 400).
         const cardFetch = async (input, init) => {
             // Try without auth first
-            const response = await a2aFetch(input, init);
+            const response = await this.a2aFetch(input, init);
             // Retry with auth if we hit a 401/403
             if ((response.status === 401 || response.status === 403) && authFetch) {
                 return authFetch(input, init);
@@ -77,17 +87,28 @@ export class A2AClientManager {
             return response;
         };
         const resolver = new DefaultAgentCardResolver({ fetchImpl: cardFetch });
-        const options = ClientFactoryOptions.createFrom(ClientFactoryOptions.default, {
+        const rawCard = await resolver.resolve(agentCardUrl, '');
+        // TODO: Remove normalizeAgentCard once @a2a-js/sdk handles
+        // proto field name aliases (supportedInterfaces → additionalInterfaces,
+        // protocolBinding → transport).
+        const agentCard = normalizeAgentCard(rawCard);
+        const grpcUrl = agentCard.additionalInterfaces?.find((i) => i.transport === 'GRPC')
+            ?.url ?? agentCard.url;
+        const clientOptions = ClientFactoryOptions.createFrom(ClientFactoryOptions.default, {
             transports: [
                 new RestTransportFactory({ fetchImpl: authFetch }),
                 new JsonRpcTransportFactory({ fetchImpl: authFetch }),
+                new GrpcTransportFactory({
+                    grpcChannelCredentials: grpcUrl.startsWith('https://')
+                        ? grpc.credentials.createSsl()
+                        : grpc.credentials.createInsecure(),
+                }),
             ],
             cardResolver: resolver,
         });
         try {
-            const factory = new ClientFactory(options);
-            const client = await factory.createFromUrl(agentCardUrl, '');
-            const agentCard = await client.getAgentCard();
+            const factory = new ClientFactory(clientOptions);
+            const client = await factory.createFromAgentCard(agentCard);
             this.clients.set(name, client);
             this.agentCards.set(name, agentCard);
             debugLogger.debug(`[A2AClientManager] Loaded agent '${name}' from ${agentCardUrl}`);
@@ -115,9 +136,8 @@ export class A2AClientManager {
      */
     async *sendMessageStream(agentName, message, options) {
         const client = this.clients.get(agentName);
-        if (!client) {
+        if (!client)
             throw new Error(`Agent '${agentName}' not found.`);
-        }
         const messageParams = {
             message: {
                 kind: 'message',
@@ -128,9 +148,18 @@ export class A2AClientManager {
                 taskId: options?.taskId,
             },
         };
-        yield* client.sendMessageStream(messageParams, {
-            signal: options?.signal,
-        });
+        try {
+            yield* client.sendMessageStream(messageParams, {
+                signal: options?.signal,
+            });
+        }
+        catch (error) {
+            const prefix = `[A2AClientManager] sendMessageStream Error [${agentName}]`;
+            if (error instanceof Error) {
+                throw new Error(`${prefix}: ${error.message}`, { cause: error });
+            }
+            throw new Error(`${prefix}: Unexpected error during sendMessageStream: ${String(error)}`);
+        }
     }
     /**
      * Retrieves a loaded agent card.
@@ -156,9 +185,8 @@ export class A2AClientManager {
      */
     async getTask(agentName, taskId) {
         const client = this.clients.get(agentName);
-        if (!client) {
+        if (!client)
             throw new Error(`Agent '${agentName}' not found.`);
-        }
         try {
             return await client.getTask({ id: taskId });
         }
@@ -178,9 +206,8 @@ export class A2AClientManager {
      */
     async cancelTask(agentName, taskId) {
         const client = this.clients.get(agentName);
-        if (!client) {
+        if (!client)
             throw new Error(`Agent '${agentName}' not found.`);
-        }
         try {
             return await client.cancelTask({ id: taskId });
         }

package/dist/src/agents/a2a-client-manager.js.map

@@ -1 +1 @@
-{"version":3,"file":"a2a-client-manager.js","sourceRoot":"","sources":["../../../src/agents/a2a-client-manager.ts"],"names":[],"mappings":"AAAA;;;;GAIG;AAUH,OAAO,EAEL,aAAa,EACb,oBAAoB,EACpB,wBAAwB,EACxB,oBAAoB,EACpB,uBAAuB,EAEvB,kCAAkC,GACnC,MAAM,oBAAoB,CAAC;AAC5B,OAAO,EAAE,EAAE,IAAI,MAAM,EAAE,MAAM,MAAM,CAAC;AACpC,OAAO,EAAE,KAAK,IAAI,WAAW,EAAE,MAAM,QAAQ,CAAC;AAC9C,OAAO,EAAE,WAAW,EAAE,MAAM,yBAAyB,CAAC;AACtD,OAAO,EAAE,UAAU,EAAE,MAAM,mBAAmB,CAAC;AAC/C,OAAO,EAAE,kBAAkB,EAAE,MAAM,iBAAiB,CAAC;AAErD,2DAA2D;AAC3D,yEAAyE;AACzE,MAAM,WAAW,GAAG,OAAO,CAAC,CAAC,aAAa;AAC1C,MAAM,aAAa,GAAG,IAAI,WAAW,CAAC;IACpC,cAAc,EAAE,WAAW;IAC3B,WAAW,EAAE,WAAW;IACxB,OAAO,EAAE;QACP,MAAM,EAAE,UAAU,EAAE,sCAAsC;KAC3D;CACF,CAAC,CAAC;AACH,MAAM,QAAQ,GAAiB,CAAC,KAAK,EAAE,IAAI,EAAE,EAAE;AAC7C,kGAAkG;AAClG,KAAK,CAAC,KAAK,EAAE,EAAE,GAAG,IAAI,EAAE,UAAU,EAAE,aAAa,EAAiB,CAAC,CAAC;AAQtE;;;GAGG;AACH,MAAM,OAAO,gBAAgB;IACnB,MAAM,CAAC,QAAQ,CAAmB;IAE1C,8DAA8D;IACtD,OAAO,GAAG,IAAI,GAAG,EAAkB,CAAC;IACpC,UAAU,GAAG,IAAI,GAAG,EAAqB,CAAC;IAElD,gBAAuB,CAAC;IAExB;;OAEG;IACH,MAAM,CAAC,WAAW;QAChB,IAAI,CAAC,gBAAgB,CAAC,QAAQ,EAAE,CAAC;YAC/B,gBAAgB,CAAC,QAAQ,GAAG,IAAI,gBAAgB,EAAE,CAAC;QACrD,CAAC;QACD,OAAO,gBAAgB,CAAC,QAAQ,CAAC;IACnC,CAAC;IAED;;;OAGG;IACH,MAAM,CAAC,uBAAuB;QAC5B,qDAAqD;QACrD,gBAAgB,CAAC,QAAQ,GAAG,SAAS,CAAC;IACxC,CAAC;IAED;;;;;;OAMG;IACH,KAAK,CAAC,SAAS,CACb,IAAY,EACZ,YAAoB,EACpB,WAAmC;QAEnC,IAAI,IAAI,CAAC,OAAO,CAAC,GAAG,CAAC,IAAI,CAAC,IAAI,IAAI,CAAC,UAAU,CAAC,GAAG,CAAC,IAAI,CAAC,EAAE,CAAC;YACxD,MAAM,IAAI,KAAK,CAAC,oBAAoB,IAAI,sBAAsB,CAAC,CAAC;QAClE,CAAC;QAED,kDAAkD;QAClD,IAAI,SAAS,GAAiB,QAAQ,CAAC;QACvC,IAAI,WAAW,EAAE,CAAC;YAChB,SAAS,GAAG,kCAAkC,CAAC,QAAQ,EAAE,WAAW,CAAC,CAAC;QACxE,CAAC;QAED,2EAA2E;QAC3E,+EAA+E;QAC/E,MAAM,SAAS,GAAG,KAAK,EACrB,KAAwB,EACxB,IAAkB,EACC,EAAE;YACrB,yBAAyB;YACzB,MAAM,QAAQ,GAAG,MAAM,QAAQ,CAAC,KAAK,EAAE,IAAI,CAAC,CAAC;YAE7C,sCAAsC;YACtC,IAAI,CAAC,QAAQ,CAAC,MAAM,KAAK,GAAG,IAAI,QAAQ,CAAC,MAAM,KAAK,GAAG,CAAC,IAAI,SAAS,EAAE,CAAC;gBACtE,OAAO,SAAS,CAAC,KAAK,EAAE,IAAI,CAAC,CAAC;YAChC,CAAC;YAED,OAAO,QAAQ,CAAC;QAClB,CAAC,CAAC;QAEF,MAAM,QAAQ,GAAG,IAAI,wBAAwB,CAAC,EAAE,SAAS,EAAE,SAAS,EAAE,CAAC,CAAC;QAExE,MAAM,OAAO,GAAG,oBAAoB,CAAC,UAAU,CAC7C,oBAAoB,CAAC,OAAO,EAC5B;YACE,UAAU,EAAE;gBACV,IAAI,oBAAoB,CAAC,EAAE,SAAS,EAAE,SAAS,EAAE,CAAC;gBAClD,IAAI,uBAAuB,CAAC,EAAE,SAAS,EAAE,SAAS,EAAE,CAAC;aACtD;YACD,YAAY,EAAE,QAAQ;SACvB,CACF,CAAC;QAEF,IAAI,CAAC;YACH,MAAM,OAAO,GAAG,IAAI,aAAa,CAAC,OAAO,CAAC,CAAC;YAC3C,MAAM,MAAM,GAAG,MAAM,OAAO,CAAC,aAAa,CAAC,YAAY,EAAE,EAAE,CAAC,CAAC;YAC7D,MAAM,SAAS,GAAG,MAAM,MAAM,CAAC,YAAY,EAAE,CAAC;YAE9C,IAAI,CAAC,OAAO,CAAC,GAAG,CAAC,IAAI,EAAE,MAAM,CAAC,CAAC;YAC/B,IAAI,CAAC,UAAU,CAAC,GAAG,CAAC,IAAI,EAAE,SAAS,CAAC,CAAC;YAErC,WAAW,CAAC,KAAK,CACf,oCAAoC,IAAI,UAAU,YAAY,EAAE,CACjE,CAAC;YAEF,OAAO,SAAS,CAAC;QACnB,CAAC;QAAC,OAAO,KAAc,EAAE,CAAC;YACxB,MAAM,kBAAkB,CAAC,IAAI,EAAE,YAAY,EAAE,KAAK,CAAC,CAAC;QACtD,CAAC;IACH,CAAC;IAED;;OAEG;IACH,UAAU;QACR,IAAI,CAAC,OAAO,CAAC,KAAK,EAAE,CAAC;QACrB,IAAI,CAAC,UAAU,CAAC,KAAK,EAAE,CAAC;QACxB,WAAW,CAAC,KAAK,CAAC,mCAAmC,CAAC,CAAC;IACzD,CAAC;IAED;;;;;;;OAOG;IACH,KAAK,CAAC,CAAC,iBAAiB,CACtB,SAAiB,EACjB,OAAe,EACf,OAAuE;QAEvE,MAAM,MAAM,GAAG,IAAI,CAAC,OAAO,CAAC,GAAG,CAAC,SAAS,CAAC,CAAC;QAC3C,IAAI,CAAC,MAAM,EAAE,CAAC;YACZ,MAAM,IAAI,KAAK,CAAC,UAAU,SAAS,cAAc,CAAC,CAAC;QACrD,CAAC;QAED,MAAM,aAAa,GAAsB;YACvC,OAAO,EAAE;gBACP,IAAI,EAAE,SAAS;gBACf,IAAI,EAAE,MAAM;gBACZ,SAAS,EAAE,MAAM,EAAE;gBACnB,KAAK,EAAE,CAAC,EAAE,IAAI,EAAE,MAAM,EAAE,IAAI,EAAE,OAAO,EAAE,CAAC;gBACxC,SAAS,EAAE,OAAO,EAAE,SAAS;gBAC7B,MAAM,EAAE,OAAO,EAAE,MAAM;aACxB;SACF,CAAC;QAEF,KAAK,CAAC,CAAC,MAAM,CAAC,iBAAiB,CAAC,aAAa,EAAE;YAC7C,MAAM,EAAE,OAAO,EAAE,MAAM;SACxB,CAAC,CAAC;IACL,CAAC;IAED;;;;OAIG;IACH,YAAY,CAAC,IAAY;QACvB,OAAO,IAAI,CAAC,UAAU,CAAC,GAAG,CAAC,IAAI,CAAC,CAAC;IACnC,CAAC;IAED;;;;OAIG;IACH,SAAS,CAAC,IAAY;QACpB,OAAO,IAAI,CAAC,OAAO,CAAC,GAAG,CAAC,IAAI,CAAC,CAAC;IAChC,CAAC;IAED;;;;;OAKG;IACH,KAAK,CAAC,OAAO,CAAC,SAAiB,EAAE,MAAc;QAC7C,MAAM,MAAM,GAAG,IAAI,CAAC,OAAO,CAAC,GAAG,CAAC,SAAS,CAAC,CAAC;QAC3C,IAAI,CAAC,MAAM,EAAE,CAAC;YACZ,MAAM,IAAI,KAAK,CAAC,UAAU,SAAS,cAAc,CAAC,CAAC;QACrD,CAAC;QACD,IAAI,CAAC;YACH,OAAO,MAAM,MAAM,CAAC,OAAO,CAAC,EAAE,EAAE,EAAE,MAAM,EAAE,CAAC,CAAC;QAC9C,CAAC;QAAC,OAAO,KAAc,EAAE,CAAC;YACxB,MAAM,MAAM,GAAG,4BAA4B,SAAS,GAAG,CAAC;YACxD,IAAI,KAAK,YAAY,KAAK,EAAE,CAAC;gBAC3B,MAAM,IAAI,KAAK,CAAC,GAAG,MAAM,KAAK,KAAK,CAAC,OAAO,EAAE,EAAE,EAAE,KAAK,EAAE,KAAK,EAAE,CAAC,CAAC;YACnE,CAAC;YACD,MAAM,IAAI,KAAK,CAAC,GAAG,MAAM,uBAAuB,MAAM,CAAC,KAAK,CAAC,EAAE,CAAC,CAAC;QACnE,CAAC;IACH,CAAC;IAED;;;;;OAKG;IACH,KAAK,CAAC,UAAU,CAAC,SAAiB,EAAE,MAAc;QAChD,MAAM,MAAM,GAAG,IAAI,CAAC,OAAO,CAAC,GAAG,CAAC,SAAS,CAAC,CAAC;QAC3C,IAAI,CAAC,MAAM,EAAE,CAAC;YACZ,MAAM,IAAI,KAAK,CAAC,UAAU,SAAS,cAAc,CAAC,CAAC;QACrD,CAAC;QACD,IAAI,CAAC;YACH,OAAO,MAAM,MAAM,CAAC,UAAU,CAAC,EAAE,EAAE,EAAE,MAAM,EAAE,CAAC,CAAC;QACjD,CAAC;QAAC,OAAO,KAAc,EAAE,CAAC;YACxB,MAAM,MAAM,GAAG,+BAA+B,SAAS,GAAG,CAAC;YAC3D,IAAI,KAAK,YAAY,KAAK,EAAE,CAAC;gBAC3B,MAAM,IAAI,KAAK,CAAC,GAAG,MAAM,KAAK,KAAK,CAAC,OAAO,EAAE,EAAE,EAAE,KAAK,EAAE,KAAK,EAAE,CAAC,CAAC;YACnE,CAAC;YACD,MAAM,IAAI,KAAK,CAAC,GAAG,MAAM,uBAAuB,MAAM,CAAC,KAAK,CAAC,EAAE,CAAC,CAAC;QACnE,CAAC;IACH,CAAC;CACF"}
+{"version":3,"file":"a2a-client-manager.js","sourceRoot":"","sources":["../../../src/agents/a2a-client-manager.ts"],"names":[],"mappings":"AAAA;;;;GAIG;AAWH,OAAO,EACL,aAAa,EACb,oBAAoB,EACpB,wBAAwB,EACxB,uBAAuB,EACvB,oBAAoB,EACpB,kCAAkC,GACnC,MAAM,oBAAoB,CAAC;AAC5B,OAAO,EAAE,oBAAoB,EAAE,MAAM,yBAAyB,CAAC;AAC/D,OAAO,KAAK,IAAI,MAAM,eAAe,CAAC;AACtC,OAAO,EAAE,EAAE,IAAI,MAAM,EAAE,MAAM,MAAM,CAAC;AACpC,OAAO,EAAE,KAAK,IAAI,WAAW,EAAE,UAAU,EAAE,MAAM,QAAQ,CAAC;AAC1D,OAAO,EAAE,kBAAkB,EAAE,MAAM,eAAe,CAAC;AAEnD,OAAO,EAAE,WAAW,EAAE,MAAM,yBAAyB,CAAC;AACtD,OAAO,EAAE,kBAAkB,EAAE,MAAM,iBAAiB,CAAC;AAYrD,2DAA2D;AAC3D,yEAAyE;AACzE,MAAM,WAAW,GAAG,OAAO,CAAC,CAAC,aAAa;AAE1C;;;GAGG;AACH,MAAM,OAAO,gBAAgB;IACnB,MAAM,CAAC,QAAQ,CAAmB;IAE1C,8DAA8D;IACtD,OAAO,GAAG,IAAI,GAAG,EAAkB,CAAC;IACpC,UAAU,GAAG,IAAI,GAAG,EAAqB,CAAC;IAE1C,aAAa,CAA2B;IACxC,QAAQ,CAAe;IAE/B,YAAoB,MAAe;QACjC,MAAM,QAAQ,GAAG,MAAM,EAAE,QAAQ,EAAE,CAAC;QACpC,MAAM,YAAY,GAAG;YACnB,cAAc,EAAE,WAAW;YAC3B,WAAW,EAAE,WAAW;SACzB,CAAC;QAEF,IAAI,QAAQ,EAAE,CAAC;YACb,IAAI,CAAC,aAAa,GAAG,IAAI,UAAU,CAAC;gBAClC,GAAG,EAAE,QAAQ;gBACb,GAAG,YAAY;aAChB,CAAC,CAAC;QACL,CAAC;aAAM,CAAC;YACN,IAAI,CAAC,aAAa,GAAG,IAAI,WAAW,CAAC,YAAY,CAAC,CAAC;QACrD,CAAC;QAED,IAAI,CAAC,QAAQ,GAAG,CAAC,KAAK,EAAE,IAAI,EAAE,EAAE,CAC9B,KAAK,CAAC,KAAK,EAAE,EAAE,GAAG,IAAI,EAAE,UAAU,EAAE,IAAI,CAAC,aAAa,EAAiB,CAAC,CAAC;IAC7E,CAAC;IAED;;OAEG;IACH,MAAM,CAAC,WAAW,CAAC,MAAe;QAChC,IAAI,CAAC,gBAAgB,CAAC,QAAQ,EAAE,CAAC;YAC/B,gBAAgB,CAAC,QAAQ,GAAG,IAAI,gBAAgB,CAAC,MAAM,CAAC,CAAC;QAC3D,CAAC;QACD,OAAO,gBAAgB,CAAC,QAAQ,CAAC;IACnC,CAAC;IAED;;;OAGG;IACH,MAAM,CAAC,uBAAuB;QAC5B,qDAAqD;QACrD,gBAAgB,CAAC,QAAQ,GAAG,SAAS,CAAC;IACxC,CAAC;IAED;;;;;;OAMG;IACH,KAAK,CAAC,SAAS,CACb,IAAY,EACZ,YAAoB,EACpB,WAAmC;QAEnC,IAAI,IAAI,CAAC,OAAO,CAAC,GAAG,CAAC,IAAI,CAAC,IAAI,IAAI,CAAC,UAAU,CAAC,GAAG,CAAC,IAAI,CAAC,EAAE,CAAC;YACxD,MAAM,IAAI,KAAK,CAAC,oBAAoB,IAAI,sBAAsB,CAAC,CAAC;QAClE,CAAC;QAED,kDAAkD;QAClD,IAAI,SAAS,GAAiB,IAAI,CAAC,QAAQ,CAAC;QAC5C,IAAI,WAAW,EAAE,CAAC;YAChB,SAAS,GAAG,kCAAkC,CAC5C,IAAI,CAAC,QAAQ,EACb,WAAW,CACZ,CAAC;QACJ,CAAC;QAED,2EAA2E;QAC3E,+EAA+E;QAC/E,MAAM,SAAS,GAAG,KAAK,EACrB,KAAwB,EACxB,IAAkB,EACC,EAAE;YACrB,yBAAyB;YACzB,MAAM,QAAQ,GAAG,MAAM,IAAI,CAAC,QAAQ,CAAC,KAAK,EAAE,IAAI,CAAC,CAAC;YAElD,sCAAsC;YACtC,IAAI,CAAC,QAAQ,CAAC,MAAM,KAAK,GAAG,IAAI,QAAQ,CAAC,MAAM,KAAK,GAAG,CAAC,IAAI,SAAS,EAAE,CAAC;gBACtE,OAAO,SAAS,CAAC,KAAK,EAAE,IAAI,CAAC,CAAC;YAChC,CAAC;YAED,OAAO,QAAQ,CAAC;QAClB,CAAC,CAAC;QAEF,MAAM,QAAQ,GAAG,IAAI,wBAAwB,CAAC,EAAE,SAAS,EAAE,SAAS,EAAE,CAAC,CAAC;QACxE,MAAM,OAAO,GAAG,MAAM,QAAQ,CAAC,OAAO,CAAC,YAAY,EAAE,EAAE,CAAC,CAAC;QACzD,2DAA2D;QAC3D,wEAAwE;QACxE,gCAAgC;QAChC,MAAM,SAAS,GAAG,kBAAkB,CAAC,OAAO,CAAC,CAAC;QAE9C,MAAM,OAAO,GACX,SAAS,CAAC,oBAAoB,EAAE,IAAI,CAAC,CAAC,CAAC,EAAE,EAAE,CAAC,CAAC,CAAC,SAAS,KAAK,MAAM,CAAC;YACjE,EAAE,GAAG,IAAI,SAAS,CAAC,GAAG,CAAC;QAE3B,MAAM,aAAa,GAAG,oBAAoB,CAAC,UAAU,CACnD,oBAAoB,CAAC,OAAO,EAC5B;YACE,UAAU,EAAE;gBACV,IAAI,oBAAoB,CAAC,EAAE,SAAS,EAAE,SAAS,EAAE,CAAC;gBAClD,IAAI,uBAAuB,CAAC,EAAE,SAAS,EAAE,SAAS,EAAE,CAAC;gBACrD,IAAI,oBAAoB,CAAC;oBACvB,sBAAsB,EAAE,OAAO,CAAC,UAAU,CAAC,UAAU,CAAC;wBACpD,CAAC,CAAC,IAAI,CAAC,WAAW,CAAC,SAAS,EAAE;wBAC9B,CAAC,CAAC,IAAI,CAAC,WAAW,CAAC,cAAc,EAAE;iBACtC,CAAC;aACH;YACD,YAAY,EAAE,QAAQ;SACvB,CACF,CAAC;QAEF,IAAI,CAAC;YACH,MAAM,OAAO,GAAG,IAAI,aAAa,CAAC,aAAa,CAAC,CAAC;YACjD,MAAM,MAAM,GAAG,MAAM,OAAO,CAAC,mBAAmB,CAAC,SAAS,CAAC,CAAC;YAE5D,IAAI,CAAC,OAAO,CAAC,GAAG,CAAC,IAAI,EAAE,MAAM,CAAC,CAAC;YAC/B,IAAI,CAAC,UAAU,CAAC,GAAG,CAAC,IAAI,EAAE,SAAS,CAAC,CAAC;YAErC,WAAW,CAAC,KAAK,CACf,oCAAoC,IAAI,UAAU,YAAY,EAAE,CACjE,CAAC;YAEF,OAAO,SAAS,CAAC;QACnB,CAAC;QAAC,OAAO,KAAc,EAAE,CAAC;YACxB,MAAM,kBAAkB,CAAC,IAAI,EAAE,YAAY,EAAE,KAAK,CAAC,CAAC;QACtD,CAAC;IACH,CAAC;IAED;;OAEG;IACH,UAAU;QACR,IAAI,CAAC,OAAO,CAAC,KAAK,EAAE,CAAC;QACrB,IAAI,CAAC,UAAU,CAAC,KAAK,EAAE,CAAC;QACxB,WAAW,CAAC,KAAK,CAAC,mCAAmC,CAAC,CAAC;IACzD,CAAC;IAED;;;;;;;OAOG;IACH,KAAK,CAAC,CAAC,iBAAiB,CACtB,SAAiB,EACjB,OAAe,EACf,OAAuE;QAEvE,MAAM,MAAM,GAAG,IAAI,CAAC,OAAO,CAAC,GAAG,CAAC,SAAS,CAAC,CAAC;QAC3C,IAAI,CAAC,MAAM;YAAE,MAAM,IAAI,KAAK,CAAC,UAAU,SAAS,cAAc,CAAC,CAAC;QAEhE,MAAM,aAAa,GAAsB;YACvC,OAAO,EAAE;gBACP,IAAI,EAAE,SAAS;gBACf,IAAI,EAAE,MAAM;gBACZ,SAAS,EAAE,MAAM,EAAE;gBACnB,KAAK,EAAE,CAAC,EAAE,IAAI,EAAE,MAAM,EAAE,IAAI,EAAE,OAAO,EAAE,CAAC;gBACxC,SAAS,EAAE,OAAO,EAAE,SAAS;gBAC7B,MAAM,EAAE,OAAO,EAAE,MAAM;aACxB;SACF,CAAC;QAEF,IAAI,CAAC;YACH,KAAK,CAAC,CAAC,MAAM,CAAC,iBAAiB,CAAC,aAAa,EAAE;gBAC7C,MAAM,EAAE,OAAO,EAAE,MAAM;aACxB,CAAC,CAAC;QACL,CAAC;QAAC,OAAO,KAAc,EAAE,CAAC;YACxB,MAAM,MAAM,GAAG,+CAA+C,SAAS,GAAG,CAAC;YAC3E,IAAI,KAAK,YAAY,KAAK,EAAE,CAAC;gBAC3B,MAAM,IAAI,KAAK,CAAC,GAAG,MAAM,KAAK,KAAK,CAAC,OAAO,EAAE,EAAE,EAAE,KAAK,EAAE,KAAK,EAAE,CAAC,CAAC;YACnE,CAAC;YACD,MAAM,IAAI,KAAK,CACb,GAAG,MAAM,gDAAgD,MAAM,CAAC,KAAK,CAAC,EAAE,CACzE,CAAC;QACJ,CAAC;IACH,CAAC;IAED;;;;OAIG;IACH,YAAY,CAAC,IAAY;QACvB,OAAO,IAAI,CAAC,UAAU,CAAC,GAAG,CAAC,IAAI,CAAC,CAAC;IACnC,CAAC;IAED;;;;OAIG;IACH,SAAS,CAAC,IAAY;QACpB,OAAO,IAAI,CAAC,OAAO,CAAC,GAAG,CAAC,IAAI,CAAC,CAAC;IAChC,CAAC;IAED;;;;;OAKG;IACH,KAAK,CAAC,OAAO,CAAC,SAAiB,EAAE,MAAc;QAC7C,MAAM,MAAM,GAAG,IAAI,CAAC,OAAO,CAAC,GAAG,CAAC,SAAS,CAAC,CAAC;QAC3C,IAAI,CAAC,MAAM;YAAE,MAAM,IAAI,KAAK,CAAC,UAAU,SAAS,cAAc,CAAC,CAAC;QAChE,IAAI,CAAC;YACH,OAAO,MAAM,MAAM,CAAC,OAAO,CAAC,EAAE,EAAE,EAAE,MAAM,EAAE,CAAC,CAAC;QAC9C,CAAC;QAAC,OAAO,KAAc,EAAE,CAAC;YACxB,MAAM,MAAM,GAAG,4BAA4B,SAAS,GAAG,CAAC;YACxD,IAAI,KAAK,YAAY,KAAK,EAAE,CAAC;gBAC3B,MAAM,IAAI,KAAK,CAAC,GAAG,MAAM,KAAK,KAAK,CAAC,OAAO,EAAE,EAAE,EAAE,KAAK,EAAE,KAAK,EAAE,CAAC,CAAC;YACnE,CAAC;YACD,MAAM,IAAI,KAAK,CAAC,GAAG,MAAM,uBAAuB,MAAM,CAAC,KAAK,CAAC,EAAE,CAAC,CAAC;QACnE,CAAC;IACH,CAAC;IAED;;;;;OAKG;IACH,KAAK,CAAC,UAAU,CAAC,SAAiB,EAAE,MAAc;QAChD,MAAM,MAAM,GAAG,IAAI,CAAC,OAAO,CAAC,GAAG,CAAC,SAAS,CAAC,CAAC;QAC3C,IAAI,CAAC,MAAM;YAAE,MAAM,IAAI,KAAK,CAAC,UAAU,SAAS,cAAc,CAAC,CAAC;QAChE,IAAI,CAAC;YACH,OAAO,MAAM,MAAM,CAAC,UAAU,CAAC,EAAE,EAAE,EAAE,MAAM,EAAE,CAAC,CAAC;QACjD,CAAC;QAAC,OAAO,KAAc,EAAE,CAAC;YACxB,MAAM,MAAM,GAAG,+BAA+B,SAAS,GAAG,CAAC;YAC3D,IAAI,KAAK,YAAY,KAAK,EAAE,CAAC;gBAC3B,MAAM,IAAI,KAAK,CAAC,GAAG,MAAM,KAAK,KAAK,CAAC,OAAO,EAAE,EAAE,EAAE,KAAK,EAAE,KAAK,EAAE,CAAC,CAAC;YACnE,CAAC;YACD,MAAM,IAAI,KAAK,CAAC,GAAG,MAAM,uBAAuB,MAAM,CAAC,KAAK,CAAC,EAAE,CAAC,CAAC;QACnE,CAAC;IACH,CAAC;CACF"}

package/dist/src/agents/a2a-client-manager.test.js

@@ -4,66 +4,72 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 import { vi, describe, it, expect, beforeEach, afterEach } from 'vitest';
-import { A2AClientManager, } from './a2a-client-manager.js';
+import { A2AClientManager } from './a2a-client-manager.js';
 import { ClientFactory, DefaultAgentCardResolver, createAuthenticatingFetchWithRetry, ClientFactoryOptions, } from '@a2a-js/sdk/client';
+import { Agent as UndiciAgent, ProxyAgent } from 'undici';
 import { debugLogger } from '../utils/debugLogger.js';
+vi.mock('@a2a-js/sdk/client', async (importOriginal) => {
+    const actual = await importOriginal();
+    return {
+        ...actual,
+        createAuthenticatingFetchWithRetry: vi.fn(),
+        ClientFactory: vi.fn(),
+        DefaultAgentCardResolver: vi.fn(),
+        ClientFactoryOptions: {
+            createFrom: vi.fn(),
+            default: {},
+        },
+    };
+});
 vi.mock('../utils/debugLogger.js', () => ({
     debugLogger: {
         debug: vi.fn(),
     },
 }));
-vi.mock('@a2a-js/sdk/client', () => {
-    const ClientFactory = vi.fn();
-    const DefaultAgentCardResolver = vi.fn();
-    const RestTransportFactory = vi.fn();
-    const JsonRpcTransportFactory = vi.fn();
-    const ClientFactoryOptions = {
-        default: {},
-        createFrom: vi.fn(),
-    };
-    const createAuthenticatingFetchWithRetry = vi.fn();
-    DefaultAgentCardResolver.prototype.resolve = vi.fn();
-    ClientFactory.prototype.createFromUrl = vi.fn();
-    return {
-        ClientFactory,
-        ClientFactoryOptions,
-        DefaultAgentCardResolver,
-        RestTransportFactory,
-        JsonRpcTransportFactory,
-        createAuthenticatingFetchWithRetry,
-    };
-});
 describe('A2AClientManager', () => {
     let manager;
-    // Stable mocks initialized once
-    const sendMessageStreamMock = vi.fn();
-    const getTaskMock = vi.fn();
-    const cancelTaskMock = vi.fn();
-    const getAgentCardMock = vi.fn();
-    const authFetchMock = vi.fn();
+    const mockAgentCard = {
+        name: 'test-agent',
+        description: 'A test agent',
+        url: 'http://test.agent',
+        version: '1.0.0',
+        protocolVersion: '0.1.0',
+        capabilities: {},
+        skills: [],
+        defaultInputModes: [],
+        defaultOutputModes: [],
+    };
     const mockClient = {
-        sendMessageStream: sendMessageStreamMock,
-        getTask: getTaskMock,
-        cancelTask: cancelTaskMock,
-        getAgentCard: getAgentCardMock,
+        sendMessageStream: vi.fn(),
+        getTask: vi.fn(),
+        cancelTask: vi.fn(),
     };
-    const mockAgentCard = { name: 'TestAgent' };
+    const authFetchMock = vi.fn();
     beforeEach(() => {
         vi.clearAllMocks();
         A2AClientManager.resetInstanceForTesting();
         manager = A2AClientManager.getInstance();
-        // Default mock implementations
-        getAgentCardMock.mockResolvedValue({
+        // Re-create the instances as plain objects that can be spied on
+        const factoryInstance = {
+            createFromUrl: vi.fn(),
+            createFromAgentCard: vi.fn(),
+        };
+        const resolverInstance = {
+            resolve: vi.fn(),
+        };
+        vi.mocked(ClientFactory).mockReturnValue(factoryInstance);
+        vi.mocked(DefaultAgentCardResolver).mockReturnValue(resolverInstance);
+        vi.spyOn(factoryInstance, 'createFromUrl').mockResolvedValue(mockClient);
+        vi.spyOn(factoryInstance, 'createFromAgentCard').mockResolvedValue(mockClient);
+        vi.spyOn(resolverInstance, 'resolve').mockResolvedValue({
             ...mockAgentCard,
             url: 'http://test.agent/real/endpoint',
         });
-        vi.mocked(ClientFactory.prototype.createFromUrl).mockResolvedValue(mockClient);
-        vi.mocked(DefaultAgentCardResolver.prototype.resolve).mockResolvedValue({
-            ...mockAgentCard,
-            url: 'http://test.agent/real/endpoint',
-        });
-        vi.mocked(ClientFactoryOptions.createFrom).mockImplementation((_defaults, overrides) => overrides);
-        vi.mocked(createAuthenticatingFetchWithRetry).mockReturnValue(authFetchMock);
+        vi.spyOn(ClientFactoryOptions, 'createFrom').mockImplementation((_defaults, overrides) => overrides);
+        vi.mocked(createAuthenticatingFetchWithRetry).mockImplementation(() => authFetchMock.mockResolvedValue({
+            ok: true,
+            json: async () => ({}),
+        }));
         vi.stubGlobal('fetch', vi.fn().mockResolvedValue({
             ok: true,
             json: async () => ({}),
@@ -78,16 +84,51 @@ describe('A2AClientManager', () => {
         const instance2 = A2AClientManager.getInstance();
         expect(instance1).toBe(instance2);
     });
+    describe('getInstance / dispatcher initialization', () => {
+        it('should use UndiciAgent when no proxy is configured', async () => {
+            await manager.loadAgent('TestAgent', 'http://test.agent/card');
+            const resolverOptions = vi.mocked(DefaultAgentCardResolver).mock
+                .calls[0][0];
+            const cardFetch = resolverOptions?.fetchImpl;
+            await cardFetch('http://test.agent/card');
+            const fetchCall = vi
+                .mocked(fetch)
+                .mock.calls.find((call) => call[0] === 'http://test.agent/card');
+            expect(fetchCall).toBeDefined();
+            expect(fetchCall[1]?.dispatcher).toBeInstanceOf(UndiciAgent);
+            expect(fetchCall[1]?.dispatcher).not.toBeInstanceOf(ProxyAgent);
+        });
+        it('should use ProxyAgent when a proxy is configured via Config', async () => {
+            A2AClientManager.resetInstanceForTesting();
+            const mockConfig = {
+                getProxy: () => 'http://my-proxy:8080',
+            };
+            manager = A2AClientManager.getInstance(mockConfig);
+            await manager.loadAgent('TestProxyAgent', 'http://test.proxy.agent/card');
+            const resolverOptions = vi.mocked(DefaultAgentCardResolver).mock
+                .calls[0][0];
+            const cardFetch = resolverOptions?.fetchImpl;
+            await cardFetch('http://test.proxy.agent/card');
+            const fetchCall = vi
+                .mocked(fetch)
+                .mock.calls.find((call) => call[0] === 'http://test.proxy.agent/card');
+            expect(fetchCall).toBeDefined();
+            expect(fetchCall[1]?.dispatcher).toBeInstanceOf(ProxyAgent);
+        });
+    });
     describe('loadAgent', () => {
         it('should create and cache an A2AClient', async () => {
             const agentCard = await manager.loadAgent('TestAgent', 'http://test.agent/card');
-            expect(agentCard).toMatchObject(mockAgentCard);
             expect(manager.getAgentCard('TestAgent')).toBe(agentCard);
             expect(manager.getClient('TestAgent')).toBeDefined();
         });
+        it('should configure ClientFactory with REST, JSON-RPC, and gRPC transports', async () => {
+            await manager.loadAgent('TestAgent', 'http://test.agent/card');
+            expect(ClientFactoryOptions.createFrom).toHaveBeenCalled();
+        });
         it('should throw an error if an agent with the same name is already loaded', async () => {
             await manager.loadAgent('TestAgent', 'http://test.agent/card');
-            await expect(manager.loadAgent('TestAgent', 'http://another.agent/card')).rejects.toThrow("Agent with name 'TestAgent' is already loaded.");
+            await expect(manager.loadAgent('TestAgent', 'http://test.agent/card')).rejects.toThrow("Agent with name 'TestAgent' is already loaded.");
         });
         it('should use native fetch by default', async () => {
             await manager.loadAgent('TestAgent', 'http://test.agent/card');
@@ -98,12 +139,8 @@ describe('A2AClientManager', () => {
                 headers: vi.fn(),
                 shouldRetryWithHeaders: vi.fn(),
             };
-            await manager.loadAgent('CustomAuthAgent', 'http://custom.agent/card', customAuthHandler);
-            expect(createAuthenticatingFetchWithRetry).toHaveBeenCalledWith(expect.anything(), customAuthHandler);
+            await manager.loadAgent('TestAgent', 'http://test.agent/card', customAuthHandler);
             // Card resolver should NOT use the authenticated fetch by default.
-            const resolverInstance = vi.mocked(DefaultAgentCardResolver).mock
-                .instances[0];
-            expect(resolverInstance).toBeDefined();
             const resolverOptions = vi.mocked(DefaultAgentCardResolver).mock
                 .calls[0][0];
             expect(resolverOptions?.fetchImpl).not.toBe(authFetchMock);
@@ -143,103 +180,134 @@ describe('A2AClientManager', () => {
         });
         it('should log a debug message upon loading an agent', async () => {
             await manager.loadAgent('TestAgent', 'http://test.agent/card');
-            expect(debugLogger.debug).toHaveBeenCalledWith("[A2AClientManager] Loaded agent 'TestAgent' from http://test.agent/card");
+            expect(debugLogger.debug).toHaveBeenCalledWith(expect.stringContaining("Loaded agent 'TestAgent'"));
         });
         it('should clear the cache', async () => {
             await manager.loadAgent('TestAgent', 'http://test.agent/card');
-            expect(manager.getAgentCard('TestAgent')).toBeDefined();
-            expect(manager.getClient('TestAgent')).toBeDefined();
             manager.clearCache();
             expect(manager.getAgentCard('TestAgent')).toBeUndefined();
             expect(manager.getClient('TestAgent')).toBeUndefined();
-            expect(debugLogger.debug).toHaveBeenCalledWith('[A2AClientManager] Cache cleared.');
+        });
+        it('should throw if resolveAgentCard fails', async () => {
+            const resolverInstance = {
+                resolve: vi.fn().mockRejectedValue(new Error('Resolution failed')),
+            };
+            vi.mocked(DefaultAgentCardResolver).mockReturnValue(resolverInstance);
+            await expect(manager.loadAgent('FailAgent', 'http://fail.agent')).rejects.toThrow('Resolution failed');
+        });
+        it('should throw if factory.createFromAgentCard fails', async () => {
+            const factoryInstance = {
+                createFromAgentCard: vi
+                    .fn()
+                    .mockRejectedValue(new Error('Factory failed')),
+            };
+            vi.mocked(ClientFactory).mockReturnValue(factoryInstance);
+            await expect(manager.loadAgent('FailAgent', 'http://fail.agent')).rejects.toThrow('Factory failed');
+        });
+    });
+    describe('getAgentCard and getClient', () => {
+        it('should return undefined if agent is not found', () => {
+            expect(manager.getAgentCard('Unknown')).toBeUndefined();
+            expect(manager.getClient('Unknown')).toBeUndefined();
         });
     });
     describe('sendMessageStream', () => {
         beforeEach(async () => {
-            await manager.loadAgent('TestAgent', 'http://test.agent');
+            await manager.loadAgent('TestAgent', 'http://test.agent/card');
         });
         it('should send a message and return a stream', async () => {
-            const mockResult = {
-                kind: 'message',
-                messageId: 'a',
-                parts: [],
-                role: 'agent',
-            };
-            sendMessageStreamMock.mockReturnValue((async function* () {
-                yield mockResult;
+            mockClient.sendMessageStream.mockReturnValue((async function* () {
+                yield { kind: 'message' };
             })());
             const stream = manager.sendMessageStream('TestAgent', 'Hello');
             const results = [];
-            for await (const res of stream) {
-                results.push(res);
+            for await (const result of stream) {
+                results.push(result);
             }
-            expect(results).toEqual([mockResult]);
-            expect(sendMessageStreamMock).toHaveBeenCalledWith(expect.objectContaining({
-                message: expect.anything(),
-            }), expect.any(Object));
+            expect(results).toHaveLength(1);
+            expect(mockClient.sendMessageStream).toHaveBeenCalled();
         });
         it('should use contextId and taskId when provided', async () => {
-            sendMessageStreamMock.mockReturnValue((async function* () {
-                yield {
-                    kind: 'message',
-                    messageId: 'a',
-                    parts: [],
-                    role: 'agent',
-                };
+            mockClient.sendMessageStream.mockReturnValue((async function* () {
+                yield { kind: 'message' };
             })());
-            const expectedContextId = 'user-context-id';
-            const expectedTaskId = 'user-task-id';
             const stream = manager.sendMessageStream('TestAgent', 'Hello', {
-                contextId: expectedContextId,
-                taskId: expectedTaskId,
+                contextId: 'ctx123',
+                taskId: 'task456',
             });
+            // trigger execution
             for await (const _ of stream) {
-                // consume stream
+                break;
             }
-            const call = sendMessageStreamMock.mock.calls[0][0];
-            expect(call.message.contextId).toBe(expectedContextId);
-            expect(call.message.taskId).toBe(expectedTaskId);
+            expect(mockClient.sendMessageStream).toHaveBeenCalledWith(expect.objectContaining({
+                message: expect.objectContaining({
+                    contextId: 'ctx123',
+                    taskId: 'task456',
+                }),
+            }), expect.any(Object));
         });
-        it('should propagate the original error on failure', async () => {
-            sendMessageStreamMock.mockImplementationOnce(() => {
-                throw new Error('Network error');
+        it('should correctly propagate AbortSignal to the stream', async () => {
+            mockClient.sendMessageStream.mockReturnValue((async function* () {
+                yield { kind: 'message' };
+            })());
+            const controller = new AbortController();
+            const stream = manager.sendMessageStream('TestAgent', 'Hello', {
+                signal: controller.signal,
+            });
+            // trigger execution
+            for await (const _ of stream) {
+                break;
+            }
+            expect(mockClient.sendMessageStream).toHaveBeenCalledWith(expect.any(Object), expect.objectContaining({ signal: controller.signal }));
+        });
+        it('should handle a multi-chunk stream with different event types', async () => {
+            mockClient.sendMessageStream.mockReturnValue((async function* () {
+                yield { kind: 'message', messageId: 'm1' };
+                yield { kind: 'status-update', taskId: 't1' };
+            })());
+            const stream = manager.sendMessageStream('TestAgent', 'Hello');
+            const results = [];
+            for await (const result of stream) {
+                results.push(result);
+            }
+            expect(results).toHaveLength(2);
+            expect(results[0].kind).toBe('message');
+            expect(results[1].kind).toBe('status-update');
+        });
+        it('should throw prefixed error on failure', async () => {
+            mockClient.sendMessageStream.mockImplementation(() => {
+                throw new Error('Network failure');
             });
             const stream = manager.sendMessageStream('TestAgent', 'Hello');
             await expect(async () => {
                 for await (const _ of stream) {
-                    // consume
+                    // empty
                 }
-            }).rejects.toThrow('Network error');
+            }).rejects.toThrow('[A2AClientManager] sendMessageStream Error [TestAgent]: Network failure');
         });
         it('should throw an error if the agent is not found', async () => {
             const stream = manager.sendMessageStream('NonExistentAgent', 'Hello');
             await expect(async () => {
                 for await (const _ of stream) {
-                    // consume
+                    // empty
                 }
             }).rejects.toThrow("Agent 'NonExistentAgent' not found.");
         });
     });
     describe('getTask', () => {
         beforeEach(async () => {
-            await manager.loadAgent('TestAgent', 'http://test.agent');
+            await manager.loadAgent('TestAgent', 'http://test.agent/card');
         });
         it('should get a task from the correct agent', async () => {
-            getTaskMock.mockResolvedValue({
-                id: 'task123',
-                contextId: 'a',
-                kind: 'task',
-                status: { state: 'completed' },
-            });
-            await manager.getTask('TestAgent', 'task123');
-            expect(getTaskMock).toHaveBeenCalledWith({
-                id: 'task123',
-            });
+            const mockTask = { id: 'task123', kind: 'task' };
+            mockClient.getTask.mockResolvedValue(mockTask);
+            const result = await manager.getTask('TestAgent', 'task123');
+            expect(result).toBe(mockTask);
+            expect(mockClient.getTask).toHaveBeenCalledWith({ id: 'task123' });
         });
         it('should throw prefixed error on failure', async () => {
-            getTaskMock.mockRejectedValueOnce(new Error('Network error'));
-            await expect(manager.getTask('TestAgent', 'task123')).rejects.toThrow('A2AClient getTask Error [TestAgent]: Network error');
+            mockClient.getTask.mockRejectedValue(new Error('Not found'));
+            await expect(manager.getTask('TestAgent', 'task123')).rejects.toThrow('A2AClient getTask Error [TestAgent]: Not found');
         });
         it('should throw an error if the agent is not found', async () => {
             await expect(manager.getTask('NonExistentAgent', 'task123')).rejects.toThrow("Agent 'NonExistentAgent' not found.");
@@ -247,23 +315,18 @@ describe('A2AClientManager', () => {
     });
     describe('cancelTask', () => {
         beforeEach(async () => {
-            await manager.loadAgent('TestAgent', 'http://test.agent');
+            await manager.loadAgent('TestAgent', 'http://test.agent/card');
         });
         it('should cancel a task on the correct agent', async () => {
-            cancelTaskMock.mockResolvedValue({
-                id: 'task123',
-                contextId: 'a',
-                kind: 'task',
-                status: { state: 'canceled' },
-            });
-            await manager.cancelTask('TestAgent', 'task123');
-            expect(cancelTaskMock).toHaveBeenCalledWith({
-                id: 'task123',
-            });
+            const mockTask = { id: 'task123', kind: 'task' };
+            mockClient.cancelTask.mockResolvedValue(mockTask);
+            const result = await manager.cancelTask('TestAgent', 'task123');
+            expect(result).toBe(mockTask);
+            expect(mockClient.cancelTask).toHaveBeenCalledWith({ id: 'task123' });
         });
         it('should throw prefixed error on failure', async () => {
-            cancelTaskMock.mockRejectedValueOnce(new Error('Network error'));
-            await expect(manager.cancelTask('TestAgent', 'task123')).rejects.toThrow('A2AClient cancelTask Error [TestAgent]: Network error');
+            mockClient.cancelTask.mockRejectedValue(new Error('Cannot cancel'));
+            await expect(manager.cancelTask('TestAgent', 'task123')).rejects.toThrow('A2AClient cancelTask Error [TestAgent]: Cannot cancel');
         });
         it('should throw an error if the agent is not found', async () => {
             await expect(manager.cancelTask('NonExistentAgent', 'task123')).rejects.toThrow("Agent 'NonExistentAgent' not found.");