Skip to content

Commit 477e394

Browse files
ikubickiikubicki
authored
feat: add state bag and ServerCallContextBuilder to ServerCallContext (#364)
# Description ## Motivation The [Python A2A SDK](https://github.com/a2aproject/a2a-python) exposes a `state` property on [`ServerCallContext`](https://github.com/a2aproject/a2a-python/blob/main/src/a2a/server/context.py#L22) — an arbitrary `MutableMapping[str, Any]` used to pass metadata through the call pipeline. The Python app layer actively uses it (e.g. [`call_context.state['method'] = method`](https://github.com/a2aproject/a2a-python/blob/main/src/a2a/server/apps/jsonrpc/jsonrpc_app.py#L149-L153)), and op-SDK implementations rely on it to carry request-scoped data such as tenant IDs, auth tokens, or raw headers. The TypeScript SDK was missing this capability, making it impossible to faithfully port Python-based agent implementations to TypeScript without architectural workarounds. ## Changes ### `ServerCallContext` - Added a `state` property — a `Map<string, unknown>` key/value bag, directly mirroring Python's `state: MutableMapping[str, Any]`. - Added `withRequestedExtensions()` method that returns a new context with updated extensions while preserving `user`, `state`, and `activatedExtensions`. - Extended the constructor to accept an optional `state` map. ### `ServerCallContextBuilder` - Added `ServerCallContextBuilder` factory function type, mirroring Python's abstract `CallContextBuilder.build(request)` pattern. - Added `defaultServerCallContextBuilder` — the default implementation that pre-populates `state` with raw request headers under the `headers` key, mirroring Python's `DefaultCallContextBuilder`. - Added `STATE_HEADERS_KEY` constant for the headers key. - Added `RequestHeaders` type — transport-agnostic representation of request headers. ### Express handlers (`jsonRpcHandler`, `restHandler`) - Added optional `contextBuilder?: ServerCallContextBuilder` to `JsonRpcHandlerOptions` and `RestHandlerOptions`, falling back to `defaultServerCallContextBuilder`. - Both handlers now pass raw `req.headers` into the context builder. ### gRPC handler - Updated `grpc_service.ts` to pass request metadata as headers into the context builder. ### Exports - `ServerCallContextBuilder`, `RequestHeaders`, `defaultServerCallContextBuilder`, and `STATE_HEADERS_KEY` are now exported from `@a2a-js/sdk/server`. ### Tests & samples - Added unit tests for the new `ServerCallContext` API in `test/server/context.spec.ts`. - Added `src/samples/authentication/server_call_context.ts` demonstrating custom context usage. ## Impact This change enables direct migration of Python A2A agent implementations to TypeScript — any code relying on `context.state` in the Python SDK will have a direct equivalent in the TypeScript SDK with the same semantics. ## Checklist - [x] Follow the [`CONTRIBUTING` Guide](https://github.com/google-a2a/a2a-js/blob/main/CONTRIBUTING.md). - [x] Make your Pull Request title in the https://www.conventionalcommits.org/ specification. - [x] Ensure the tests and linter pass - [x] Appropriate docs were updated (if necessary)
1 parent 37ec204 commit 477e394

7 files changed

Lines changed: 508 additions & 55 deletions

File tree

‎src/samples/authentication/server_call_context.ts‎

Lines changed: 192 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,192 @@
1+
/**
2+
* Sample: ServerCallContext state headers
3+
*
4+
* Demonstrates two patterns for reading request headers inside an AgentExecutor:
5+
*
6+
* 1. AUTOMATIC (default builder) - `defaultServerCallContextBuilder` stores all
7+
* request headers in `context.state` under `STATE_HEADERS_KEY` with no extra
8+
* configuration needed.
9+
*
10+
* 2. CUSTOM BUILDER - supply a `contextBuilder` to `jsonRpcHandler` to extract
11+
* specific headers and store them in `state` under your own keys, so the
12+
* AgentExecutor receives clean, typed values without coupling to raw headers.
13+
*
14+
* Run:
15+
* cd src/samples && npx tsx authentication/server_call_context.ts
16+
*
17+
* Then send a request with a custom header:
18+
* curl -X POST http://localhost:41242 \
19+
* -H "Content-Type: application/json" \
20+
* -H "x-tenant-id: acme-corp" \
21+
* -d '{"jsonrpc":"2.0","id":"1","method":"SendMessage","params":{"message":{"messageId":"m1","role":"user","parts":[{"kind":"text","text":"hello"}]}}}'
22+
*/
23+
24+
import express from 'express';
25+
import { v4 as uuidv4 } from 'uuid';
26+
import { AgentCard } from '../../index.js';
27+
import { Role } from '../../index.js';
28+
import {
29+
AgentEvent,
30+
AgentExecutor,
31+
DefaultRequestHandler,
32+
ExecutionEventBus,
33+
InMemoryTaskStore,
34+
RequestContext,
35+
ServerCallContext,
36+
ServerCallContextBuilder,
37+
STATE_HEADERS_KEY,
38+
RequestHeaders,
39+
UnauthenticatedUser,
40+
} from '../../server/index.js';
41+
import { jsonRpcHandler } from '../../server/express/index.js';
42+
import { Message } from '../../index.js';
43+
44+
// --- Custom state keys ---
45+
46+
const STATE_TENANT_ID_KEY = 'tenantId';
47+
const STATE_REQUEST_ID_KEY = 'requestId';
48+
49+
// --- Custom context builder ---
50+
51+
/**
52+
* Reads well-known headers and stores them as clean typed values in state,
53+
* alongside the full raw headers stored automatically under STATE_HEADERS_KEY.
54+
*/
55+
const tenantContextBuilder: ServerCallContextBuilder = ({
56+
extensions,
57+
user,
58+
headers,
59+
requestedVersion,
60+
tenant,
61+
}): ServerCallContext => {
62+
const state = new Map<string, unknown>([
63+
// Always include raw headers (mirrors defaultServerCallContextBuilder)
64+
[STATE_HEADERS_KEY, headers],
65+
// Extract specific headers into typed state entries
66+
[STATE_TENANT_ID_KEY, headers['x-tenant-id'] ?? tenant ?? 'unknown'],
67+
[STATE_REQUEST_ID_KEY, headers['x-request-id'] ?? uuidv4()],
68+
]);
69+
return new ServerCallContext({
70+
requestedExtensions: extensions,
71+
user,
72+
state,
73+
requestedVersion,
74+
tenant,
75+
});
76+
};
77+
78+
// --- AgentExecutor ---
79+
80+
class StateHeadersAgentExecutor implements AgentExecutor {
81+
public cancelTask = async (): Promise<void> => {};
82+
83+
async execute(requestContext: RequestContext, eventBus: ExecutionEventBus): Promise<void> {
84+
const state = requestContext.context?.state;
85+
86+
// Pattern 1: read a typed value stored by the custom builder
87+
const tenantId = state?.get(STATE_TENANT_ID_KEY) as string | undefined;
88+
const requestId = state?.get(STATE_REQUEST_ID_KEY) as string | undefined;
89+
90+
// Pattern 2: read a specific header directly from the raw headers map
91+
const rawHeaders = state?.get(STATE_HEADERS_KEY) as RequestHeaders | undefined;
92+
const userAgent = rawHeaders?.['user-agent'];
93+
94+
const lines = [
95+
`Tenant ID : ${tenantId ?? '(not set)'}`,
96+
`Request ID : ${requestId ?? '(not set)'}`,
97+
`User-Agent : ${userAgent ?? '(not set)'}`,
98+
];
99+
100+
const finalMessage: Message = {
101+
messageId: uuidv4(),
102+
contextId: '',
103+
taskId: '',
104+
role: Role.ROLE_AGENT,
105+
parts: [
106+
{
107+
content: { $case: 'text', value: lines.join('\n') },
108+
metadata: undefined,
109+
filename: '',
110+
mediaType: '',
111+
},
112+
],
113+
metadata: undefined,
114+
extensions: [],
115+
referenceTaskIds: [],
116+
};
117+
118+
eventBus.publish(AgentEvent.message(finalMessage));
119+
}
120+
}
121+
122+
// --- Server setup ---
123+
124+
const agentCard: AgentCard = {
125+
name: 'ServerCallContext State Headers Sample',
126+
description: 'Demonstrates reading request headers from ServerCallContext.state',
127+
supportedInterfaces: [
128+
{
129+
url: 'http://localhost:41242/',
130+
protocolBinding: 'JSONRPC',
131+
tenant: '',
132+
protocolVersion: '0.3',
133+
},
134+
],
135+
provider: { organization: 'A2A Samples', url: 'https://example.com' },
136+
version: '1.0.0',
137+
documentationUrl: '',
138+
capabilities: { streaming: false, pushNotifications: false, extensions: [] },
139+
securitySchemes: {},
140+
securityRequirements: [],
141+
defaultInputModes: ['text'],
142+
defaultOutputModes: ['text'],
143+
signatures: [],
144+
skills: [
145+
{
146+
id: 'echo_headers',
147+
name: 'Echo Headers',
148+
description: 'Echoes x-tenant-id, x-request-id and User-Agent from request headers.',
149+
tags: ['sample'],
150+
examples: ['hello'],
151+
inputModes: ['text'],
152+
outputModes: ['text'],
153+
securityRequirements: [],
154+
},
155+
],
156+
};
157+
158+
async function main() {
159+
const requestHandler = new DefaultRequestHandler(
160+
agentCard,
161+
new InMemoryTaskStore(),
162+
new StateHeadersAgentExecutor()
163+
);
164+
165+
const app = express();
166+
app.use(express.json());
167+
app.use(
168+
jsonRpcHandler({
169+
requestHandler,
170+
userBuilder: async () => new UnauthenticatedUser(),
171+
// Swap contextBuilder to see the difference between custom and default:
172+
// custom → tenantId and requestId are extracted into typed state entries
173+
// default → only raw headers are stored under STATE_HEADERS_KEY
174+
contextBuilder: tenantContextBuilder,
175+
})
176+
);
177+
178+
const PORT = 41242;
179+
app.listen(PORT, () => {
180+
console.log(`[StateHeadersSample] Listening on http://localhost:${PORT}`);
181+
console.log(`[StateHeadersSample] Try:`);
182+
console.log(
183+
` curl -X POST http://localhost:${PORT}` +
184+
` -H "Content-Type: application/json"` +
185+
` -H "x-tenant-id: acme-corp"` +
186+
` -H "x-request-id: req-123"` +
187+
` -d '{"jsonrpc":"2.0","id":"1","method":"SendMessage","params":{"message":{"messageId":"m1","role":"user","parts":[{"kind":"text","text":"hello"}]}}}'`
188+
);
189+
});
190+
}
191+
192+
main().catch(console.error);

‎src/server/context.ts‎

Lines changed: 81 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,30 +5,101 @@ import { User } from './authentication/user.js';
55
// header as a v0.3 request.
66
const ABSENT_HEADER_VERSION = '0.3';
77

8+
/**
9+
* Transport-agnostic representation of request headers.
10+
* Express passes `req.headers`; gRPC passes metadata converted to this shape.
11+
*/
12+
export type RequestHeaders = Record<string, string | string[] | undefined>;
13+
14+
/**
15+
* Options passed to a {@link ServerCallContextBuilder}.
16+
*/
17+
export interface ServerCallContextBuilderOptions {
18+
/** Protocol extensions parsed from the request headers. */
19+
extensions: Extensions | undefined;
20+
/** Authenticated user extracted from the request. */
21+
user: User | undefined;
22+
/** Raw request headers (transport-agnostic). */
23+
headers: RequestHeaders;
24+
/** A2A protocol version from the A2A-Version header. Absent means '0.3'. */
25+
requestedVersion?: string;
26+
/** Tenant identifier extracted from the request path or metadata. */
27+
tenant?: string;
28+
}
29+
30+
/**
31+
* Factory function type for creating {@link ServerCallContext} instances.
32+
*
33+
* Provide a custom implementation to inject additional state or produce a
34+
* subclass of `ServerCallContext` (e.g. to mirror the Python A2A SDK's
35+
* `state` pattern used by operator SDKs).
36+
*
37+
* @param options - All data available at request time.
38+
* @returns A `ServerCallContext` (or subclass) for the current call.
39+
*/
40+
export type ServerCallContextBuilder = (
41+
options: ServerCallContextBuilderOptions
42+
) => ServerCallContext;
43+
44+
/**
45+
* Key under which request headers are stored in {@link ServerCallContext.state}
46+
* by the default builder. Mirrors Python SDK's `state['headers']`.
47+
*/
48+
export const STATE_HEADERS_KEY = 'headers';
49+
850
export interface ServerCallContextOptions {
951
requestedExtensions?: Extensions;
1052
user?: User;
1153
tenant?: string;
12-
1354
/**
1455
* The A2A protocol version requested by the client via the A2A-Version
1556
* service parameter. Defaults to `'0.3'` when the header is absent.
1657
*/
1758
requestedVersion?: string;
59+
/**
60+
* Arbitrary key/value state bag for carrying custom data
61+
* (e.g. request headers, tenant IDs) through the call pipeline.
62+
*/
63+
state?: Map<string, unknown>;
1864
}
1965

66+
/**
67+
* The default {@link ServerCallContextBuilder}. Creates a `ServerCallContext`
68+
* with the raw request headers pre-populated in {@link ServerCallContext.state}
69+
* under the {@link STATE_HEADERS_KEY} key, mirroring the Python SDK's
70+
* `DefaultCallContextBuilder`.
71+
*/
72+
export const defaultServerCallContextBuilder: ServerCallContextBuilder = ({
73+
extensions,
74+
user,
75+
headers,
76+
requestedVersion,
77+
tenant,
78+
}: ServerCallContextBuilderOptions): ServerCallContext => {
79+
const state = new Map<string, unknown>([[STATE_HEADERS_KEY, headers]]);
80+
return new ServerCallContext({
81+
requestedExtensions: extensions,
82+
user,
83+
state,
84+
requestedVersion,
85+
tenant,
86+
});
87+
};
88+
2089
export class ServerCallContext {
2190
private _requestedExtensions?: Extensions;
2291
private readonly _user?: User;
2392
private readonly _requestedVersion: string;
2493
private readonly _tenant?: string;
2594
private _activatedExtensions?: Extensions;
95+
private readonly _state: Map<string, unknown>;
2696

2797
constructor(options?: ServerCallContextOptions) {
2898
this._requestedExtensions = options?.requestedExtensions;
2999
this._user = options?.user;
30100
this._tenant = options?.tenant;
31101
this._requestedVersion = options?.requestedVersion || ABSENT_HEADER_VERSION;
102+
this._state = options?.state ?? new Map();
32103
}
33104

34105
get tenant(): string | undefined {
@@ -51,6 +122,15 @@ export class ServerCallContext {
51122
return this._requestedVersion;
52123
}
53124

125+
/**
126+
* Arbitrary key/value state bag, equivalent to the `state` field on the
127+
* Python A2A SDK's `ServerCallContext`. Use this to carry custom data
128+
* (e.g. request headers, tenant IDs) through the call pipeline.
129+
*/
130+
get state(): Map<string, unknown> {
131+
return this._state;
132+
}
133+
54134
public addActivatedExtension(uri: string) {
55135
this._activatedExtensions = Extensions.createFrom(this._activatedExtensions, uri);
56136
}

‎src/server/express/json_rpc_handler.ts‎

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@ import { JSONRPCErrorResponse } from '../../core.js';
99
import { JSONRPCResponse } from '../transports/jsonrpc/jsonrpc_transport_handler.js';
1010
import { A2ARequestHandler } from '../request_handler/a2a_request_handler.js';
1111
import { JsonRpcTransportHandler } from '../transports/jsonrpc/jsonrpc_transport_handler.js';
12-
import { ServerCallContext } from '../context.js';
12+
import { ServerCallContextBuilder, defaultServerCallContextBuilder } from '../context.js';
1313
import { A2A_VERSION_HEADER, HTTP_EXTENSION_HEADER, JSON_CONTENT_TYPE } from '../../constants.js';
1414
import { UserBuilder, delegateAsyncIterator } from './common.js';
1515
import { SSE_HEADERS, formatSSEEvent, formatSSEErrorEvent } from '../../sse_utils.js';
@@ -38,6 +38,7 @@ export interface JsonRpcHandlerOptions {
3838
* as JSON-RPC `method not found` (-32601).
3939
*/
4040
legacyCompat?: { enabled: boolean };
41+
contextBuilder?: ServerCallContextBuilder;
4142
}
4243

4344
/**
@@ -101,9 +102,11 @@ export function jsonRpcHandler(options: JsonRpcHandlerOptions): RequestHandler {
101102
const requestedExtensionsHeader = useLegacy
102103
? (req.header(LEGACY_HTTP_EXTENSION_HEADER) ?? req.header(HTTP_EXTENSION_HEADER))
103104
: req.header(HTTP_EXTENSION_HEADER);
104-
const context = new ServerCallContext({
105-
requestedExtensions: Extensions.parseServiceParameter(requestedExtensionsHeader),
105+
const ctxBuilder = options.contextBuilder ?? defaultServerCallContextBuilder;
106+
const context = ctxBuilder({
107+
extensions: Extensions.parseServiceParameter(requestedExtensionsHeader),
106108
user,
109+
headers: req.headers,
107110
requestedVersion,
108111
});
109112
const agentCard = await options.requestHandler.getAgentCard();

0 commit comments

Comments
 (0)