feat: add state bag and ServerCallContextBuilder to ServerCallContext

src/samples/authentication/server_call_context.ts

@@ -0,0 +1,192 @@
+/**
+ * Sample: ServerCallContext state headers
+ *
+ * Demonstrates two patterns for reading request headers inside an AgentExecutor:
+ *
+ * 1. AUTOMATIC (default builder) - `defaultServerCallContextBuilder` stores all
+ *    request headers in `context.state` under `STATE_HEADERS_KEY` with no extra
+ *    configuration needed.
+ *
+ * 2. CUSTOM BUILDER - supply a `contextBuilder` to `jsonRpcHandler` to extract
+ *    specific headers and store them in `state` under your own keys, so the
+ *    AgentExecutor receives clean, typed values without coupling to raw headers.
+ *
+ * Run:
+ *   cd src/samples && npx tsx authentication/server_call_context.ts
+ *
+ * Then send a request with a custom header:
+ *   curl -X POST http://localhost:41242 \
+ *     -H "Content-Type: application/json" \
+ *     -H "x-tenant-id: acme-corp" \
+ *     -d '{"jsonrpc":"2.0","id":"1","method":"SendMessage","params":{"message":{"messageId":"m1","role":"user","parts":[{"kind":"text","text":"hello"}]}}}'
+ */
+
+import express from 'express';
+import { v4 as uuidv4 } from 'uuid';
+import { AgentCard } from '../../index.js';
+import { Role } from '../../index.js';
+import {
+  AgentEvent,
+  AgentExecutor,
+  DefaultRequestHandler,
+  ExecutionEventBus,
+  InMemoryTaskStore,
+  RequestContext,
+  ServerCallContext,
+  ServerCallContextBuilder,
+  STATE_HEADERS_KEY,
+  RequestHeaders,
+  UnauthenticatedUser,
+} from '../../server/index.js';
+import { jsonRpcHandler } from '../../server/express/index.js';
+import { Message } from '../../index.js';
+
+// --- Custom state keys ---
+
+const STATE_TENANT_ID_KEY = 'tenantId';
+const STATE_REQUEST_ID_KEY = 'requestId';
+
+// --- Custom context builder ---
+
+/**
+ * Reads well-known headers and stores them as clean typed values in state,
+ * alongside the full raw headers stored automatically under STATE_HEADERS_KEY.
+ */
+const tenantContextBuilder: ServerCallContextBuilder = ({
+  extensions,
+  user,
+  headers,
+  requestedVersion,
+  tenant,
+}): ServerCallContext => {
+  const state = new Map<string, unknown>([
+    // Always include raw headers (mirrors defaultServerCallContextBuilder)
+    [STATE_HEADERS_KEY, headers],
+    // Extract specific headers into typed state entries
+    [STATE_TENANT_ID_KEY, headers['x-tenant-id'] ?? tenant ?? 'unknown'],
+    [STATE_REQUEST_ID_KEY, headers['x-request-id'] ?? uuidv4()],
+  ]);
+  return new ServerCallContext({
+    requestedExtensions: extensions,
+    user,
+    state,
+    requestedVersion,
+    tenant,
+  });
+};
+
+// --- AgentExecutor ---
+
+class StateHeadersAgentExecutor implements AgentExecutor {
+  public cancelTask = async (): Promise<void> => {};
+
+  async execute(requestContext: RequestContext, eventBus: ExecutionEventBus): Promise<void> {
+    const state = requestContext.context?.state;
+
+    // Pattern 1: read a typed value stored by the custom builder
+    const tenantId = state?.get(STATE_TENANT_ID_KEY) as string | undefined;
+    const requestId = state?.get(STATE_REQUEST_ID_KEY) as string | undefined;
+
+    // Pattern 2: read a specific header directly from the raw headers map
+    const rawHeaders = state?.get(STATE_HEADERS_KEY) as RequestHeaders | undefined;
+    const userAgent = rawHeaders?.['user-agent'];
+
+    const lines = [
+      `Tenant ID  : ${tenantId ?? '(not set)'}`,
+      `Request ID : ${requestId ?? '(not set)'}`,
+      `User-Agent : ${userAgent ?? '(not set)'}`,
+    ];
+
+    const finalMessage: Message = {
+      messageId: uuidv4(),
+      contextId: '',
+      taskId: '',
+      role: Role.ROLE_AGENT,
+      parts: [
+        {
+          content: { $case: 'text', value: lines.join('\n') },
+          metadata: undefined,
+          filename: '',
+          mediaType: '',
+        },
+      ],
+      metadata: undefined,
+      extensions: [],
+      referenceTaskIds: [],
+    };
+
+    eventBus.publish(AgentEvent.message(finalMessage));
+  }
+}
+
+// --- Server setup ---
+
+const agentCard: AgentCard = {
+  name: 'ServerCallContext State Headers Sample',
+  description: 'Demonstrates reading request headers from ServerCallContext.state',
+  supportedInterfaces: [
+    {
+      url: 'http://localhost:41242/',
+      protocolBinding: 'JSONRPC',
+      tenant: '',
+      protocolVersion: '0.3',
+    },
+  ],
+  provider: { organization: 'A2A Samples', url: 'https://example.com' },
+  version: '1.0.0',
+  documentationUrl: '',
+  capabilities: { streaming: false, pushNotifications: false, extensions: [] },
+  securitySchemes: {},
+  securityRequirements: [],
+  defaultInputModes: ['text'],
+  defaultOutputModes: ['text'],
+  signatures: [],
+  skills: [
+    {
+      id: 'echo_headers',
+      name: 'Echo Headers',
+      description: 'Echoes x-tenant-id, x-request-id and User-Agent from request headers.',
+      tags: ['sample'],
+      examples: ['hello'],
+      inputModes: ['text'],
+      outputModes: ['text'],
+      securityRequirements: [],
+    },
+  ],
+};
+
+async function main() {
+  const requestHandler = new DefaultRequestHandler(
+    agentCard,
+    new InMemoryTaskStore(),
+    new StateHeadersAgentExecutor()
+  );
+
+  const app = express();
+  app.use(express.json());
+  app.use(
+    jsonRpcHandler({
+      requestHandler,
+      userBuilder: async () => new UnauthenticatedUser(),
+      // Swap contextBuilder to see the difference between custom and default:
+      //   custom  → tenantId and requestId are extracted into typed state entries
+      //   default → only raw headers are stored under STATE_HEADERS_KEY
+      contextBuilder: tenantContextBuilder,
+    })
+  );
+
+  const PORT = 41242;
+  app.listen(PORT, () => {
+    console.log(`[StateHeadersSample] Listening on http://localhost:${PORT}`);
+    console.log(`[StateHeadersSample] Try:`);
+    console.log(
+      `  curl -X POST http://localhost:${PORT}` +
+        ` -H "Content-Type: application/json"` +
+        ` -H "x-tenant-id: acme-corp"` +
+        ` -H "x-request-id: req-123"` +
+        ` -d '{"jsonrpc":"2.0","id":"1","method":"SendMessage","params":{"message":{"messageId":"m1","role":"user","parts":[{"kind":"text","text":"hello"}]}}}'`
+    );
+  });
+}
+
+main().catch(console.error);

src/server/context.ts

@@ -5,30 +5,101 @@ import { User } from './authentication/user.js';
 // header as a v0.3 request.
 const ABSENT_HEADER_VERSION = '0.3';
 
+/**
+ * Transport-agnostic representation of request headers.
+ * Express passes `req.headers`; gRPC passes metadata converted to this shape.
+ */
+export type RequestHeaders = Record<string, string | string[] | undefined>;
+
+/**
+ * Options passed to a {@link ServerCallContextBuilder}.
+ */
+export interface ServerCallContextBuilderOptions {
+  /** Protocol extensions parsed from the request headers. */
+  extensions: Extensions | undefined;
+  /** Authenticated user extracted from the request. */
+  user: User | undefined;
+  /** Raw request headers (transport-agnostic). */
+  headers: RequestHeaders;
+  /** A2A protocol version from the A2A-Version header. Absent means '0.3'. */
+  requestedVersion?: string;
+  /** Tenant identifier extracted from the request path or metadata. */
+  tenant?: string;
+}
+
+/**
+ * Factory function type for creating {@link ServerCallContext} instances.
+ *
+ * Provide a custom implementation to inject additional state or produce a
+ * subclass of `ServerCallContext` (e.g. to mirror the Python A2A SDK's
+ * `state` pattern used by operator SDKs).
+ *
+ * @param options - All data available at request time.
+ * @returns A `ServerCallContext` (or subclass) for the current call.
+ */
+export type ServerCallContextBuilder = (
+  options: ServerCallContextBuilderOptions
+) => ServerCallContext;
+
+/**
+ * Key under which request headers are stored in {@link ServerCallContext.state}
+ * by the default builder. Mirrors Python SDK's `state['headers']`.
+ */
+export const STATE_HEADERS_KEY = 'headers';
+
 export interface ServerCallContextOptions {
   requestedExtensions?: Extensions;
   user?: User;
   tenant?: string;
-
   /**
    * The A2A protocol version requested by the client via the A2A-Version
    * service parameter. Defaults to `'0.3'` when the header is absent.
    */
   requestedVersion?: string;
+  /**
+   * Arbitrary key/value state bag for carrying custom data
+   * (e.g. request headers, tenant IDs) through the call pipeline.
+   */
+  state?: Map<string, unknown>;
 }
 
+/**
+ * The default {@link ServerCallContextBuilder}. Creates a `ServerCallContext`
+ * with the raw request headers pre-populated in {@link ServerCallContext.state}
+ * under the {@link STATE_HEADERS_KEY} key, mirroring the Python SDK's
+ * `DefaultCallContextBuilder`.
+ */
+export const defaultServerCallContextBuilder: ServerCallContextBuilder = ({
+  extensions,
+  user,
+  headers,
+  requestedVersion,
+  tenant,
+}: ServerCallContextBuilderOptions): ServerCallContext => {
+  const state = new Map<string, unknown>([[STATE_HEADERS_KEY, headers]]);
+  return new ServerCallContext({
+    requestedExtensions: extensions,
+    user,
+    state,
+    requestedVersion,
+    tenant,
+  });
+};
+
 export class ServerCallContext {
   private _requestedExtensions?: Extensions;
   private readonly _user?: User;
   private readonly _requestedVersion: string;
   private readonly _tenant?: string;
   private _activatedExtensions?: Extensions;
+  private readonly _state: Map<string, unknown>;
 
   constructor(options?: ServerCallContextOptions) {
     this._requestedExtensions = options?.requestedExtensions;
     this._user = options?.user;
     this._tenant = options?.tenant;
     this._requestedVersion = options?.requestedVersion || ABSENT_HEADER_VERSION;
+    this._state = options?.state ?? new Map();
   }
 
   get tenant(): string | undefined {
@@ -51,6 +122,15 @@ export class ServerCallContext {
     return this._requestedVersion;
   }
 
+  /**
+   * Arbitrary key/value state bag, equivalent to the `state` field on the
+   * Python A2A SDK's `ServerCallContext`. Use this to carry custom data
+   * (e.g. request headers, tenant IDs) through the call pipeline.
+   */
+  get state(): Map<string, unknown> {
+    return this._state;
+  }
+
   public addActivatedExtension(uri: string) {
     this._activatedExtensions = Extensions.createFrom(this._activatedExtensions, uri);
   }

src/server/express/json_rpc_handler.ts

@@ -9,7 +9,7 @@ import { JSONRPCErrorResponse } from '../../core.js';
 import { JSONRPCResponse } from '../transports/jsonrpc/jsonrpc_transport_handler.js';
 import { A2ARequestHandler } from '../request_handler/a2a_request_handler.js';
 import { JsonRpcTransportHandler } from '../transports/jsonrpc/jsonrpc_transport_handler.js';
-import { ServerCallContext } from '../context.js';
+import { ServerCallContextBuilder, defaultServerCallContextBuilder } from '../context.js';
 import { A2A_VERSION_HEADER, HTTP_EXTENSION_HEADER, JSON_CONTENT_TYPE } from '../../constants.js';
 import { UserBuilder, delegateAsyncIterator } from './common.js';
 import { SSE_HEADERS, formatSSEEvent, formatSSEErrorEvent } from '../../sse_utils.js';
@@ -38,6 +38,7 @@ export interface JsonRpcHandlerOptions {
    * as JSON-RPC `method not found` (-32601).
    */
   legacyCompat?: { enabled: boolean };
+  contextBuilder?: ServerCallContextBuilder;
 }
 
 /**
@@ -101,9 +102,11 @@ export function jsonRpcHandler(options: JsonRpcHandlerOptions): RequestHandler {
       const requestedExtensionsHeader = useLegacy
         ? (req.header(LEGACY_HTTP_EXTENSION_HEADER) ?? req.header(HTTP_EXTENSION_HEADER))
         : req.header(HTTP_EXTENSION_HEADER);
-      const context = new ServerCallContext({
-        requestedExtensions: Extensions.parseServiceParameter(requestedExtensionsHeader),
+      const ctxBuilder = options.contextBuilder ?? defaultServerCallContextBuilder;
+      const context = ctxBuilder({
+        extensions: Extensions.parseServiceParameter(requestedExtensionsHeader),
         user,
+        headers: req.headers,
         requestedVersion,
       });
       const agentCard = await options.requestHandler.getAgentCard();