Skip to content

Commit f8b0b8c

Browse files
guglielmo-sanguglielmo-san
committed
fix changes
1 parent 30466c4 commit f8b0b8c

11 files changed

Lines changed: 81 additions & 84 deletions

File tree

‎src/a2a/compat/v0_3/grpc_handler.py‎

Lines changed: 6 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -23,9 +23,8 @@
2323
from a2a.server.context import ServerCallContext
2424
from a2a.server.request_handlers.grpc_handler import (
2525
_ERROR_CODE_MAP,
26+
DefaultGrpcContextBuilder,
2627
GrpcContextBuilder,
27-
build_grpc_server_call_context,
28-
default_grpc_user_builder,
2928
)
3029
from a2a.server.request_handlers.request_handler import RequestHandler
3130
from a2a.types.a2a_pb2 import AgentCard
@@ -45,7 +44,7 @@ def __init__(
4544
self,
4645
agent_card: AgentCard,
4746
request_handler: RequestHandler,
48-
context_builder: GrpcContextBuilder | None = None,
47+
context_builder: GrpcContextBuilder | None = None,
4948
card_modifier: Callable[[AgentCard], Awaitable[AgentCard] | AgentCard]
5049
| None = None,
5150
):
@@ -55,14 +54,14 @@ def __init__(
5554
agent_card: The AgentCard describing the agent's capabilities (v1.0).
5655
request_handler: The underlying `RequestHandler` instance to
5756
delegate requests to.
58-
context_builder: Optional custom user builder to extract user from the
57+
context_builder: Optional custom user builder to extract user from the
5958
gRPC context.
6059
card_modifier: An optional callback to dynamically modify the public
6160
agent card before it is served.
6261
"""
6362
self.agent_card = agent_card
6463
self.handler03 = RequestHandler03(request_handler=request_handler)
65-
self.user_builder =context_builder or default_grpc_user_builder
64+
self._context_builder = context_builder or DefaultGrpcContextBuilder()
6665
self.card_modifier = card_modifier
6766

6867
async def _handle_unary(
@@ -73,9 +72,7 @@ async def _handle_unary(
7372
) -> TResponse:
7473
"""Centralized error handling and context management for unary calls."""
7574
try:
76-
server_context = build_grpc_server_call_context(
77-
context, self.user_builder
78-
)
75+
server_context = self._context_builder.build(context)
7976
result = await handler_func(server_context)
8077
self._set_extension_metadata(context, server_context)
8178
except A2AError as e:
@@ -91,9 +88,7 @@ async def _handle_stream(
9188
) -> AsyncIterable[TResponse]:
9289
"""Centralized error handling and context management for streaming calls."""
9390
try:
94-
server_context = build_grpc_server_call_context(
95-
context, self.user_builder
96-
)
91+
server_context = self._context_builder.build(context)
9792
async for item in handler_func(server_context):
9893
yield item
9994
self._set_extension_metadata(context, server_context)

‎src/a2a/compat/v0_3/rest_adapter.py‎

Lines changed: 5 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -36,8 +36,7 @@
3636
from a2a.server.context import ServerCallContext
3737
from a2a.server.routes.common import (
3838
ContextBuilder,
39-
build_server_call_context,
40-
default_user_builder,
39+
DefaultContextBuilder,
4140
)
4241
from a2a.utils.error_handlers import (
4342
rest_error_handler,
@@ -64,7 +63,7 @@ def __init__( # noqa: PLR0913
6463
agent_card: 'AgentCard',
6564
http_handler: 'RequestHandler',
6665
extended_agent_card: 'AgentCard | None' = None,
67-
context_builder: 'ContextBuilder | None' = None,
66+
context_builder: 'ContextBuilder | None' = None,
6867
card_modifier: 'Callable[[AgentCard], Awaitable[AgentCard] | AgentCard] | None' = None,
6968
extended_card_modifier: 'Callable[[AgentCard, ServerCallContext], Awaitable[AgentCard] | AgentCard] | None' = None,
7069
):
@@ -75,15 +74,15 @@ def __init__( # noqa: PLR0913
7574
self.handler = REST03Handler(
7675
agent_card=agent_card, request_handler=http_handler
7776
)
78-
self._user_builder =context_builder or default_user_builder
77+
self._context_builder = context_builder or DefaultContextBuilder()
7978

8079
@rest_error_handler
8180
async def _handle_request(
8281
self,
8382
method: 'Callable[[Request, ServerCallContext], Awaitable[Any]]',
8483
request: Request,
8584
) -> Response:
86-
call_context = build_server_call_context(request, self._user_builder)
85+
call_context = self._context_builder.build(request)
8786
response = await method(request, call_context)
8887
return JSONResponse(content=response)
8988

@@ -100,7 +99,7 @@ async def _handle_streaming_request(
10099
message=f'Failed to pre-consume request body: {e}'
101100
) from e
102101

103-
call_context = build_server_call_context(request, self._user_builder)
102+
call_context = self._context_builder.build(request)
104103

105104
async def event_generator(
106105
stream: AsyncIterable[Any],

‎src/a2a/server/request_handlers/grpc_handler.py‎

Lines changed: 34 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
# ruff: noqa: N802
22
import logging
33

4+
from abc import ABC, abstractmethod
45
from collections.abc import AsyncIterable, Awaitable, Callable
56
from typing import TypeVar
67

@@ -40,12 +41,37 @@
4041

4142
logger = logging.getLogger(__name__)
4243

43-
GrpcContextBuilder = Callable[[grpc.aio.ServicerContext], User]
4444

45+
class GrpcContextBuilder(ABC):
46+
"""Interface for building ServerCallContext from gRPC context."""
4547

46-
def default_grpc_user_builder(context: grpc.aio.ServicerContext) -> User:
47-
"""Default strategy for creating a User from a gRPC context."""
48-
return UnauthenticatedUser()
48+
@abstractmethod
49+
def build(self, context: grpc.aio.ServicerContext) -> ServerCallContext:
50+
"""Builds a ServerCallContext from a gRPC ServicerContext."""
51+
52+
@abstractmethod
53+
def build_user(self, context: grpc.aio.ServicerContext) -> User:
54+
"""Builds a User from a gRPC ServicerContext."""
55+
56+
57+
class DefaultGrpcContextBuilder(GrpcContextBuilder):
58+
"""Default implementation of GrpcContextBuilder."""
59+
60+
def build(self, context: grpc.aio.ServicerContext) -> ServerCallContext:
61+
"""Builds a ServerCallContext from a gRPC ServicerContext."""
62+
user = self.build_user(context)
63+
state = {'grpc_context': context}
64+
return ServerCallContext(
65+
user=user,
66+
state=state,
67+
requested_extensions=get_requested_extensions(
68+
_get_metadata_value(context, HTTP_EXTENSION_HEADER)
69+
),
70+
)
71+
72+
def build_user(self, context: grpc.aio.ServicerContext) -> User:
73+
"""Builds a User from a gRPC ServicerContext."""
74+
return UnauthenticatedUser()
4975

5076

5177
def _get_metadata_value(
@@ -62,22 +88,6 @@ def _get_metadata_value(
6288
if k.lower() == lower_key
6389
]
6490

65-
66-
def build_grpc_server_call_context(
67-
context: grpc.aio.ServicerContext,context_builder: GrpcContextBuilder
68-
) -> ServerCallContext:
69-
"""Builds a ServerCallContext from a gRPC ServicerContext."""
70-
user =context_builder(context)
71-
state = {'grpc_context': context}
72-
return ServerCallContext(
73-
user=user,
74-
state=state,
75-
requested_extensions=get_requested_extensions(
76-
_get_metadata_value(context, HTTP_EXTENSION_HEADER)
77-
),
78-
)
79-
80-
8191
_ERROR_CODE_MAP = {
8292
types.InvalidRequestError: grpc.StatusCode.INVALID_ARGUMENT,
8393
types.MethodNotFoundError: grpc.StatusCode.NOT_FOUND,
@@ -105,7 +115,7 @@ def __init__(
105115
self,
106116
agent_card: AgentCard,
107117
request_handler: RequestHandler,
108-
context_builder: GrpcContextBuilder | None = None,
118+
context_builder: GrpcContextBuilder | None = None,
109119
card_modifier: Callable[[AgentCard], Awaitable[AgentCard] | AgentCard]
110120
| None = None,
111121
):
@@ -115,14 +125,14 @@ def __init__(
115125
agent_card: The AgentCard describing the agent's capabilities.
116126
request_handler: The underlying `RequestHandler` instance to
117127
delegate requests to.
118-
context_builder: Optional custom user builder to extract user from the
128+
context_builder: Optional custom user builder to extract user from the
119129
gRPC context.
120130
card_modifier: An optional callback to dynamically modify the public
121131
agent card before it is served.
122132
"""
123133
self.agent_card = agent_card
124134
self.request_handler = request_handler
125-
self.user_builder =context_builder or default_grpc_user_builder
135+
self._context_builder = context_builder or DefaultGrpcContextBuilder()
126136
self.card_modifier = card_modifier
127137

128138
async def _handle_unary(
@@ -446,8 +456,6 @@ def _build_call_context(
446456
context: grpc.aio.ServicerContext,
447457
request: message.Message,
448458
) -> ServerCallContext:
449-
server_context = build_grpc_server_call_context(
450-
context, self.user_builder
451-
)
459+
server_context = self._context_builder.build(context)
452460
server_context.tenant = getattr(request, 'tenant', '')
453461
return server_context

‎src/a2a/server/routes/__init__.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,14 @@
11
"""A2A Routes."""
22

33
from a2a.server.routes.agent_card_routes import create_agent_card_routes
4-
from a2a.server.routes.common import ContextBuilder
4+
from a2a.server.routes.common import ContextBuilder, DefaultContextBuilder
55
from a2a.server.routes.jsonrpc_routes import create_jsonrpc_routes
66
from a2a.server.routes.rest_routes import create_rest_routes
77

88

99
__all__ = [
1010
'ContextBuilder',
11+
'DefaultContextBuilder',
1112
'create_agent_card_routes',
1213
'create_jsonrpc_routes',
1314
'create_rest_routes',

‎src/a2a/server/routes/common.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,4 @@
11
from abc import ABC, abstractmethod
2-
from collections.abc import Callable
32
from typing import TYPE_CHECKING, Any
43

54

@@ -21,6 +20,7 @@
2120
)
2221
from a2a.server.context import ServerCallContext
2322

23+
2424
class StarletteUser(User):
2525
"""Adapts a Starlette BaseUser to the A2A User interface."""
2626

‎src/a2a/server/routes/jsonrpc_dispatcher.py‎

Lines changed: 5 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,7 @@
2929
)
3030
from a2a.server.routes.common import (
3131
ContextBuilder,
32-
32+
DefaultContextBuilder,
3333
)
3434
from a2a.types import A2ARequest
3535
from a2a.types.a2a_pb2 import (
@@ -144,7 +144,7 @@ def __init__( # noqa: PLR0913
144144
agent_card: AgentCard,
145145
request_handler: RequestHandler,
146146
extended_agent_card: AgentCard | None = None,
147-
context_builder: ContextBuilder | None = None,
147+
context_builder: ContextBuilder | None = None,
148148
card_modifier: Callable[[AgentCard], Awaitable[AgentCard] | AgentCard]
149149
| None = None,
150150
extended_card_modifier: Callable[
@@ -161,7 +161,7 @@ def __init__( # noqa: PLR0913
161161
requests via http.
162162
extended_agent_card: An optional, distinct AgentCard to be served
163163
at the authenticated extended card endpoint.
164-
context_builder: Optional custom user builder to extract user from the
164+
context_builder: Optional custom user builder to extract user from the
165165
request.
166166
card_modifier: An optional callback to dynamically modify the public
167167
agent card before it is served.
@@ -191,7 +191,7 @@ def __init__( # noqa: PLR0913
191191
agent_card=agent_card,
192192
http_handler=request_handler,
193193
extended_agent_card=extended_agent_card,
194-
context_builder=self._context_builder,
194+
context_builder=self._context_builder,
195195
card_modifier=card_modifier,
196196
extended_card_modifier=extended_card_modifier,
197197
)
@@ -333,9 +333,7 @@ async def handle_requests(self, request: Request) -> Response: # noqa: PLR0911,
333333
)
334334

335335
# 3) Build call context and wrap the request for downstream handling
336-
call_context = build_server_call_context(
337-
request, self._user_builder
338-
)
336+
call_context = self._context_builder.build(request)
339337
call_context.tenant = getattr(specific_request, 'tenant', '')
340338
call_context.state['method'] = method
341339
call_context.state['request_id'] = request_id

‎src/a2a/server/routes/jsonrpc_routes.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,7 @@ def create_jsonrpc_routes( # noqa: PLR0913
2929
request_handler: RequestHandler,
3030
rpc_url: str,
3131
extended_agent_card: AgentCard | None = None,
32-
context_builder: ContextBuilder | None = None,
32+
context_builder: ContextBuilder | None = None,
3333
card_modifier: Callable[[AgentCard], Awaitable[AgentCard] | AgentCard]
3434
| None = None,
3535
extended_card_modifier: Callable[
@@ -51,7 +51,7 @@ def create_jsonrpc_routes( # noqa: PLR0913
5151
rpc_url: The URL prefix for the RPC endpoints.
5252
extended_agent_card: An optional, distinct AgentCard to be served
5353
at the authenticated extended card endpoint.
54-
context_builder: Optional custom context builder to extract user from the
54+
context_builder: Optional custom context builder to extract user from the
5555
request.
5656
card_modifier: An optional callback to dynamically modify the public
5757
agent card before it is served.
@@ -71,7 +71,7 @@ def create_jsonrpc_routes( # noqa: PLR0913
7171
agent_card=agent_card,
7272
request_handler=request_handler,
7373
extended_agent_card=extended_agent_card,
74-
context_builder=context_builder,
74+
context_builder=context_builder,
7575
card_modifier=card_modifier,
7676
extended_card_modifier=extended_card_modifier,
7777
enable_v0_3_compat=enable_v0_3_compat,

‎src/a2a/server/routes/rest_dispatcher.py‎

Lines changed: 5 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -10,8 +10,7 @@
1010
from a2a.server.request_handlers.request_handler import RequestHandler
1111
from a2a.server.routes.common import (
1212
ContextBuilder,
13-
build_server_call_context,
14-
default_user_builder,
13+
DefaultContextBuilder,
1514
)
1615
from a2a.types import a2a_pb2
1716
from a2a.types.a2a_pb2 import (
@@ -72,7 +71,7 @@ def __init__( # noqa: PLR0913
7271
agent_card: AgentCard,
7372
request_handler: RequestHandler,
7473
extended_agent_card: AgentCard | None = None,
75-
context_builder: ContextBuilder | None = None,
74+
context_builder: ContextBuilder | None = None,
7675
card_modifier: Callable[[AgentCard], Awaitable[AgentCard] | AgentCard]
7776
| None = None,
7877
extended_card_modifier: Callable[
@@ -87,7 +86,7 @@ def __init__( # noqa: PLR0913
8786
request_handler: The underlying `RequestHandler` instance to delegate requests to.
8887
extended_agent_card: An optional, distinct AgentCard to be served
8988
at the authenticated extended card endpoint.
90-
context_builder: Optional custom user builder to extract user from the
89+
context_builder: Optional custom user builder to extract user from the
9190
request.
9291
card_modifier: An optional callback to dynamically modify the public
9392
agent card before it is served.
@@ -106,11 +105,11 @@ def __init__( # noqa: PLR0913
106105
self.extended_agent_card = extended_agent_card
107106
self.card_modifier = card_modifier
108107
self.extended_card_modifier = extended_card_modifier
109-
self._user_builder =context_builder or default_user_builder
108+
self._context_builder = context_builder or DefaultContextBuilder()
110109
self.request_handler = request_handler
111110

112111
def _build_call_context(self, request: Request) -> ServerCallContext:
113-
call_context = build_server_call_context(request, self._user_builder)
112+
call_context = self._context_builder.build(request)
114113
if 'tenant' in request.path_params:
115114
call_context.tenant = request.path_params['tenant']
116115
return call_context

‎src/a2a/server/routes/rest_routes.py‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -46,7 +46,7 @@ def create_rest_routes( # noqa: PLR0913
4646
agent_card: AgentCard,
4747
request_handler: RequestHandler,
4848
extended_agent_card: AgentCard | None = None,
49-
context_builder: ContextBuilder | None = None,
49+
context_builder: ContextBuilder | None = None,
5050
card_modifier: Callable[[AgentCard], Awaitable[AgentCard] | AgentCard]
5151
| None = None,
5252
extended_card_modifier: Callable[
@@ -64,7 +64,7 @@ def create_rest_routes( # noqa: PLR0913
6464
requests via http.
6565
extended_agent_card: An optional, distinct AgentCard to be served
6666
at the authenticated extended card endpoint.
67-
context_builder: Optional custom user builder to extract user from the
67+
context_builder: Optional custom user builder to extract user from the
6868
request.
6969
card_modifier: An optional callback to dynamically modify the public
7070
agent card before it is served.
@@ -86,7 +86,7 @@ def create_rest_routes( # noqa: PLR0913
8686
agent_card=agent_card,
8787
request_handler=request_handler,
8888
extended_agent_card=extended_agent_card,
89-
context_builder=user_builder,
89+
context_builder=context_builder,
9090
card_modifier=card_modifier,
9191
extended_card_modifier=extended_card_modifier,
9292
)
@@ -97,7 +97,7 @@ def create_rest_routes( # noqa: PLR0913
9797
agent_card=agent_card,
9898
http_handler=request_handler,
9999
extended_agent_card=extended_agent_card,
100-
context_builder=user_builder,
100+
context_builder=context_builder,
101101
card_modifier=card_modifier,
102102
extended_card_modifier=extended_card_modifier,
103103
)

0 commit comments

Comments
 (0)