feat: Add support for extension headers on client side (#227)

# Description

This PR implements the support for extensions on client side

Fixes #170 🦕

Release-As: 0.3.6

---------

Co-authored-by: Ivan Shymko <ishymko@google.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

Author: 3 people

src/client/index.ts

@@ -25,4 +25,9 @@ export {
   ClientCallInput,
   ClientCallResult,
 } from './interceptors.js';
+export {
+  ServiceParameters,
+  ServiceParametersUpdate,
+  withA2AExtensions,
+} from './service-parameters.js';
 export { ClientCallContext, ContextUpdate, ClientCallContextKey } from './context.js';

src/client/multitransport-client.ts

@@ -19,6 +19,7 @@ import {
   ClientCallResult,
   ClientCallInput,
 } from './interceptors.js';
+import { ServiceParameters } from './service-parameters.js';
 import { Transport } from './transports/transport.js';
 
 export interface ClientConfig {
@@ -51,7 +52,11 @@ export interface RequestOptions {
    */
   signal?: AbortSignal;
 
-  // TODO: propagate extensions
+  /**
+   * A key-value map for passing horizontally applicable context or parameters.
+   * All parameters are passed to the server via underlying transports (e.g. In JsonRPC via Headers).
+   */
+  serviceParameters?: ServiceParameters;
 
   /**
    * Arbitrary data available to interceptors and transport implementation.

src/client/service-parameters.ts

@@ -0,0 +1,28 @@
+import { HTTP_EXTENSION_HEADER } from '../constants.js';
+
+export type ServiceParameters = Record<string, string>;
+
+export type ServiceParametersUpdate = (parameters: ServiceParameters) => void;
+
+export const ServiceParameters = {
+  create(...updates: ServiceParametersUpdate[]): ServiceParameters {
+    return ServiceParameters.createFrom(undefined, ...updates);
+  },
+
+  createFrom: (
+    serviceParameters: ServiceParameters | undefined,
+    ...updates: ServiceParametersUpdate[]
+  ): ServiceParameters => {
+    const result = serviceParameters ? { ...serviceParameters } : {};
+    for (const update of updates) {
+      update(result);
+    }
+    return result;
+  },
+};
+
+export function withA2AExtensions(...extensions: string[]): ServiceParametersUpdate {
+  return (parameters: ServiceParameters) => {
+    parameters[HTTP_EXTENSION_HEADER] = extensions.join(',');
+  };
+}

src/client/transports/json_rpc_transport.ts

@@ -197,7 +197,7 @@ export class JsonRpcTransport implements Transport {
       id: requestId,
     };
 
-    const httpResponse = await this._fetchRpc(rpcRequest, 'application/json', options?.signal);
+    const httpResponse = await this._fetchRpc(rpcRequest, 'application/json', options);
 
     if (!httpResponse.ok) {
       let errorBodyText = '(empty or non-JSON response)';
@@ -237,16 +237,17 @@ export class JsonRpcTransport implements Transport {
   private async _fetchRpc(
     rpcRequest: JSONRPCRequest,
     acceptHeader: string = 'application/json',
-    signal?: AbortSignal
+    options?: RequestOptions
   ): Promise<Response> {
     const requestInit: RequestInit = {
       method: 'POST',
       headers: {
+        ...options?.serviceParameters,
         'Content-Type': 'application/json',
         Accept: acceptHeader,
       },
       body: JSON.stringify(rpcRequest),
-      signal,
+      signal: options?.signal,
     };
     return this._fetch(this.endpoint, requestInit);
   }
@@ -264,7 +265,7 @@ export class JsonRpcTransport implements Transport {
       id: clientRequestId,
     };
 
-    const response = await this._fetchRpc(rpcRequest, 'text/event-stream', options?.signal);
+    const response = await this._fetchRpc(rpcRequest, 'text/event-stream', options);
 
     if (!response.ok) {
       let errorBody = '';

test/client/transports/json_rpc_transport.test.ts

@@ -0,0 +1,56 @@
+import { JsonRpcTransport } from '../../../src/client/transports/json_rpc_transport.js';
+import { expect } from 'chai';
+import sinon from 'sinon';
+import { describe, it, beforeEach } from 'mocha';
+import { MessageSendParams, TextPart } from '../../../src/types.js';
+import { RequestOptions } from '../../../src/client/multitransport-client.js';
+import { HTTP_EXTENSION_HEADER } from '../../../src/constants.js';
+import { ServiceParameters, withA2AExtensions } from '../../../src/client/service-parameters.js';
+
+describe('JsonRpcTransport', () => {
+  let transport: JsonRpcTransport;
+  let mockFetch: sinon.SinonStubbedFunction<typeof fetch>;
+  const endpoint = 'https://test.endpoint/api';
+
+  beforeEach(() => {
+    mockFetch = sinon.stub();
+    transport = new JsonRpcTransport({
+      endpoint,
+      fetchImpl: mockFetch,
+    });
+  });
+
+  describe('sendMessage', () => {
+    it('should correctly add the extension headers', async () => {
+      const messageParams: MessageSendParams = {
+        message: {
+          kind: 'message',
+          messageId: 'test-msg-1',
+          role: 'user',
+          parts: [
+            {
+              kind: 'text',
+              text: 'Hello, agent!',
+            } as TextPart,
+          ],
+        },
+      };
+
+      const expectedExtensions = 'extension1,extension2';
+      const serviceParameters = ServiceParameters.create(withA2AExtensions(expectedExtensions));
+      const options: RequestOptions = {
+        serviceParameters,
+      };
+
+      mockFetch.resolves(
+        new Response(JSON.stringify({ jsonrpc: '2.0', result: {}, id: 1 }), {
+          status: 200,
+        })
+      );
+      await transport.sendMessage(messageParams, options);
+      const fetchArgs = mockFetch.firstCall.args[1];
+      const headers = fetchArgs.headers;
+      expect((headers as any)[HTTP_EXTENSION_HEADER]).to.deep.equal(expectedExtensions);
+    });
+  });
+});

test/server/a2a_express_app.spec.ts

@@ -8,7 +8,7 @@ import { A2AExpressApp } from '../../src/server/express/a2a_express_app.js';
 import { A2ARequestHandler } from '../../src/server/request_handler/a2a_request_handler.js';
 import { JsonRpcTransportHandler } from '../../src/server/transports/jsonrpc_transport_handler.js';
 import { AgentCard, JSONRPCSuccessResponse, JSONRPCErrorResponse } from '../../src/index.js';
-import { AGENT_CARD_PATH } from '../../src/constants.js';
+import { AGENT_CARD_PATH, HTTP_EXTENSION_HEADER } from '../../src/constants.js';
 import { A2AError } from '../../src/server/error.js';
 import { ServerCallContext } from '../../src/server/context.js';
 import { User, UnauthenticatedUser } from '../../src/server/authentication/user.js';
@@ -262,7 +262,7 @@ describe('A2AExpressApp', () => {
 
       await request(expressApp)
         .post('/')
-        .set('X-A2A-Extensions', uriExtensionsValues)
+        .set(HTTP_EXTENSION_HEADER, uriExtensionsValues)
         .set('Not-Relevant-Header', 'unused-value')
         .send(requestBody)
         .expect(200);
@@ -294,12 +294,12 @@ describe('A2AExpressApp', () => {
       });
       const response = await request(expressApp)
         .post('/')
-        .set('X-A2A-Extensions', uriExtensionsValues)
+        .set(HTTP_EXTENSION_HEADER, uriExtensionsValues)
         .set('Not-Relevant-Header', 'unused-value')
         .send(requestBody)
         .expect(200);
 
-      expect(response.get('X-A2A-Extensions')).to.equal('activated-extension');
+      expect(response.get(HTTP_EXTENSION_HEADER)).to.equal('activated-extension');
     });
   });