Skip to content

Commit 3c295e4

Browse files
cbugwadia32cbugwadia32
committed
make card modifier and extended card modifier async
1 parent ffe31e2 commit 3c295e4

9 files changed

Lines changed: 76 additions & 59 deletions

File tree

‎src/a2a/server/apps/jsonrpc/fastapi_app.py‎

Lines changed: 6 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,7 @@
11
import logging
2-
3-
from collections.abc import Callable
2+
from collections.abc import Awaitable, Callable
43
from typing import TYPE_CHECKING, Any
54

6-
75
if TYPE_CHECKING:
86
from fastapi import FastAPI
97

@@ -32,7 +30,6 @@
3230
PREV_AGENT_CARD_WELL_KNOWN_PATH,
3331
)
3432

35-
3633
logger = logging.getLogger(__name__)
3734

3835

@@ -72,9 +69,10 @@ def __init__( # noqa: PLR0913
7269
http_handler: RequestHandler,
7370
extended_agent_card: AgentCard | None = None,
7471
context_builder: CallContextBuilder | None = None,
75-
card_modifier: Callable[[AgentCard], AgentCard] | None = None,
72+
card_modifier: Callable[[AgentCard], Awaitable[AgentCard]]
73+
| None = None,
7674
extended_card_modifier: Callable[
77-
[AgentCard, ServerCallContext], AgentCard
75+
[AgentCard, ServerCallContext], Awaitable[AgentCard]
7876
]
7977
| None = None,
8078
max_content_length: int | None = 10 * 1024 * 1024, # 10MB
@@ -90,9 +88,9 @@ def __init__( # noqa: PLR0913
9088
context_builder: The CallContextBuilder used to construct the
9189
ServerCallContext passed to the http_handler. If None, no
9290
ServerCallContext is passed.
93-
card_modifier: An optional callback to dynamically modify the public
91+
card_modifier: An optional async callback to dynamically modify the public
9492
agent card before it is served.
95-
extended_card_modifier: An optional callback to dynamically modify
93+
extended_card_modifier: An optional async callback to dynamically modify
9694
the extended agent card before it is served. It receives the
9795
call context.
9896
max_content_length: The maximum allowed content length for incoming

‎src/a2a/server/apps/jsonrpc/jsonrpc_app.py‎

Lines changed: 10 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -2,9 +2,8 @@
22
import json
33
import logging
44
import traceback
5-
65
from abc import ABC, abstractmethod
7-
from collections.abc import AsyncGenerator, Callable
6+
from collections.abc import AsyncGenerator, Awaitable, Callable
87
from typing import TYPE_CHECKING, Any
98

109
from pydantic import ValidationError
@@ -52,7 +51,6 @@
5251
)
5352
from a2a.utils.errors import MethodNotImplementedError
5453

55-
5654
logger = logging.getLogger(__name__)
5755

5856
if TYPE_CHECKING:
@@ -178,9 +176,10 @@ def __init__( # noqa: PLR0913
178176
http_handler: RequestHandler,
179177
extended_agent_card: AgentCard | None = None,
180178
context_builder: CallContextBuilder | None = None,
181-
card_modifier: Callable[[AgentCard], AgentCard] | None = None,
179+
card_modifier: Callable[[AgentCard], Awaitable[AgentCard]]
180+
| None = None,
182181
extended_card_modifier: Callable[
183-
[AgentCard, ServerCallContext], AgentCard
182+
[AgentCard, ServerCallContext], Awaitable[AgentCard]
184183
]
185184
| None = None,
186185
max_content_length: int | None = 10 * 1024 * 1024, # 10MB
@@ -196,9 +195,9 @@ def __init__( # noqa: PLR0913
196195
context_builder: The CallContextBuilder used to construct the
197196
ServerCallContext passed to the http_handler. If None, no
198197
ServerCallContext is passed.
199-
card_modifier: An optional callback to dynamically modify the public
198+
card_modifier: An optional async callback to dynamically modify the public
200199
agent card before it is served.
201-
extended_card_modifier: An optional callback to dynamically modify
200+
extended_card_modifier: An optional async callback to dynamically modify
202201
the extended agent card before it is served. It receives the
203202
call context.
204203
max_content_length: The maximum allowed content length for incoming
@@ -576,7 +575,7 @@ async def _handle_get_agent_card(self, request: Request) -> JSONResponse:
576575

577576
card_to_serve = self.agent_card
578577
if self.card_modifier:
579-
card_to_serve = self.card_modifier(card_to_serve)
578+
card_to_serve = await self.card_modifier(card_to_serve)
580579

581580
return JSONResponse(
582581
card_to_serve.model_dump(
@@ -605,7 +604,9 @@ async def _handle_get_authenticated_extended_agent_card(
605604
context = self._context_builder.build(request)
606605
# If no base extended card is provided, pass the public card to the modifier
607606
base_card = card_to_serve if card_to_serve else self.agent_card
608-
card_to_serve = self.extended_card_modifier(base_card, context)
607+
card_to_serve = await self.extended_card_modifier(
608+
base_card, context
609+
)
609610

610611
if card_to_serve:
611612
return JSONResponse(

‎src/a2a/server/apps/jsonrpc/starlette_app.py‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
import logging
22

3-
from collections.abc import Callable
3+
from collections.abc import Awaitable, Callable
44
from typing import TYPE_CHECKING, Any
55

66

@@ -54,9 +54,9 @@ def __init__( # noqa: PLR0913
5454
http_handler: RequestHandler,
5555
extended_agent_card: AgentCard | None = None,
5656
context_builder: CallContextBuilder | None = None,
57-
card_modifier: Callable[[AgentCard], AgentCard] | None = None,
57+
card_modifier: Callable[[AgentCard], Awaitable[AgentCard]] | None = None,
5858
extended_card_modifier: Callable[
59-
[AgentCard, ServerCallContext], AgentCard
59+
[AgentCard, ServerCallContext], Awaitable[AgentCard]
6060
]
6161
| None = None,
6262
max_content_length: int | None = 10 * 1024 * 1024, # 10MB
@@ -72,9 +72,9 @@ def __init__( # noqa: PLR0913
7272
context_builder: The CallContextBuilder used to construct the
7373
ServerCallContext passed to the http_handler. If None, no
7474
ServerCallContext is passed.
75-
card_modifier: An optional callback to dynamically modify the public
75+
card_modifier: An optional async callback to dynamically modify the public
7676
agent card before it is served.
77-
extended_card_modifier: An optional callback to dynamically modify
77+
extended_card_modifier: An optional async callback to dynamically modify
7878
the extended agent card before it is served. It receives the
7979
call context.
8080
max_content_length: The maximum allowed content length for incoming

‎src/a2a/server/apps/rest/fastapi_app.py‎

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
import logging
22

3-
from collections.abc import Callable
3+
from collections.abc import Awaitable, Callable
44
from typing import TYPE_CHECKING, Any
55

66

@@ -49,9 +49,10 @@ def __init__( # noqa: PLR0913
4949
http_handler: RequestHandler,
5050
extended_agent_card: AgentCard | None = None,
5151
context_builder: CallContextBuilder | None = None,
52-
card_modifier: Callable[[AgentCard], AgentCard] | None = None,
52+
card_modifier: Callable[[AgentCard], Awaitable[AgentCard]]
53+
| None = None,
5354
extended_card_modifier: Callable[
54-
[AgentCard, ServerCallContext], AgentCard
55+
[AgentCard, ServerCallContext], Awaitable[AgentCard]
5556
]
5657
| None = None,
5758
):
@@ -66,9 +67,9 @@ def __init__( # noqa: PLR0913
6667
context_builder: The CallContextBuilder used to construct the
6768
ServerCallContext passed to the http_handler. If None, no
6869
ServerCallContext is passed.
69-
card_modifier: An optional callback to dynamically modify the public
70+
card_modifier: An optional async callback to dynamically modify the public
7071
agent card before it is served.
71-
extended_card_modifier: An optional callback to dynamically modify
72+
extended_card_modifier: An optional async callback to dynamically modify
7273
the extended agent card before it is served. It receives the
7374
call context.
7475
"""

‎src/a2a/server/apps/rest/rest_adapter.py‎

Lines changed: 10 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -58,9 +58,10 @@ def __init__( # noqa: PLR0913
5858
http_handler: RequestHandler,
5959
extended_agent_card: AgentCard | None = None,
6060
context_builder: CallContextBuilder | None = None,
61-
card_modifier: Callable[[AgentCard], AgentCard] | None = None,
61+
card_modifier: Callable[[AgentCard], Awaitable[AgentCard]]
62+
| None = None,
6263
extended_card_modifier: Callable[
63-
[AgentCard, ServerCallContext], AgentCard
64+
[AgentCard, ServerCallContext], Awaitable[AgentCard]
6465
]
6566
| None = None,
6667
):
@@ -75,9 +76,9 @@ def __init__( # noqa: PLR0913
7576
context_builder: The CallContextBuilder used to construct the
7677
ServerCallContext passed to the http_handler. If None, no
7778
ServerCallContext is passed.
78-
card_modifier: An optional callback to dynamically modify the public
79+
card_modifier: An optional async callback to dynamically modify the public
7980
agent card before it is served.
80-
extended_card_modifier: An optional callback to dynamically modify
81+
extended_card_modifier: An optional async callback to dynamically modify
8182
the extended agent card before it is served. It receives the
8283
call context.
8384
"""
@@ -150,7 +151,7 @@ async def handle_get_agent_card(
150151
"""
151152
card_to_serve = self.agent_card
152153
if self.card_modifier:
153-
card_to_serve = self.card_modifier(card_to_serve)
154+
card_to_serve = await self.card_modifier(card_to_serve)
154155

155156
return card_to_serve.model_dump(mode='json', exclude_none=True)
156157

@@ -182,9 +183,11 @@ async def handle_authenticated_agent_card(
182183

183184
if self.extended_card_modifier:
184185
context = self._context_builder.build(request)
185-
card_to_serve = self.extended_card_modifier(card_to_serve, context)
186+
card_to_serve = await self.extended_card_modifier(
187+
card_to_serve, context
188+
)
186189
elif self.card_modifier:
187-
card_to_serve = self.card_modifier(card_to_serve)
190+
card_to_serve = await self.card_modifier(card_to_serve)
188191

189192
return card_to_serve.model_dump(mode='json', exclude_none=True)
190193

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

Lines changed: 2 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,15 +1,12 @@
11
# ruff: noqa: N802
22
import contextlib
33
import logging
4-
54
from abc import ABC, abstractmethod
65
from collections.abc import AsyncIterable, Sequence
76

8-
97
try:
108
import grpc
119
import grpc.aio
12-
1310
from grpc.aio import Metadata
1411
except ImportError as e:
1512
raise ImportError(
@@ -21,7 +18,6 @@
2118
from collections.abc import Callable
2219

2320
import a2a.grpc.a2a_pb2_grpc as a2a_grpc
24-
2521
from a2a import types
2622
from a2a.auth.user import UnauthenticatedUser
2723
from a2a.extensions.common import (
@@ -36,7 +32,6 @@
3632
from a2a.utils.errors import ServerError
3733
from a2a.utils.helpers import validate, validate_async_generator
3834

39-
4035
logger = logging.getLogger(__name__)
4136

4237
# For now we use a trivial wrapper on the grpc context object
@@ -99,7 +94,7 @@ def __init__(
9994
delegate requests to.
10095
context_builder: The CallContextBuilder object. If none the
10196
DefaultCallContextBuilder is used.
102-
card_modifier: An optional callback to dynamically modify the public
97+
card_modifier: An optional async callback to dynamically modify the public
10398
agent card before it is served.
10499
"""
105100
self.agent_card = agent_card
@@ -339,7 +334,7 @@ async def GetAgentCard(
339334
"""Get the agent card for the agent served."""
340335
card_to_serve = self.agent_card
341336
if self.card_modifier:
342-
card_to_serve = self.card_modifier(card_to_serve)
337+
card_to_serve = await self.card_modifier(card_to_serve)
343338
return proto_utils.ToProto.agent_card(card_to_serve)
344339

345340
async def abort_context(

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

Lines changed: 10 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,5 @@
11
import logging
2-
3-
from collections.abc import AsyncIterable, Callable
2+
from collections.abc import AsyncIterable, Awaitable, Callable
43

54
from a2a.server.context import ServerCallContext
65
from a2a.server.request_handlers.request_handler import RequestHandler
@@ -49,7 +48,6 @@
4948
from a2a.utils.helpers import validate
5049
from a2a.utils.telemetry import SpanKind, trace_class
5150

52-
5351
logger = logging.getLogger(__name__)
5452

5553

@@ -63,21 +61,22 @@ def __init__(
6361
request_handler: RequestHandler,
6462
extended_agent_card: AgentCard | None = None,
6563
extended_card_modifier: Callable[
66-
[AgentCard, ServerCallContext], AgentCard
64+
[AgentCard, ServerCallContext], Awaitable[AgentCard]
6765
]
6866
| None = None,
69-
card_modifier: Callable[[AgentCard], AgentCard] | None = None,
67+
card_modifier: Callable[[AgentCard], Awaitable[AgentCard]]
68+
| None = None,
7069
):
7170
"""Initializes the JSONRPCHandler.
7271
7372
Args:
7473
agent_card: The AgentCard describing the agent's capabilities.
7574
request_handler: The underlying `RequestHandler` instance to delegate requests to.
7675
extended_agent_card: An optional, distinct Extended AgentCard to be served
77-
extended_card_modifier: An optional callback to dynamically modify
76+
extended_card_modifier: An optional async callback to dynamically modify
7877
the extended agent card before it is served. It receives the
7978
call context.
80-
card_modifier: An optional callback to dynamically modify the public
79+
card_modifier: An optional async callback to dynamically modify the public
8180
agent card before it is served.
8281
"""
8382
self.agent_card = agent_card
@@ -450,9 +449,11 @@ async def get_authenticated_extended_card(
450449

451450
card_to_serve = base_card
452451
if self.extended_card_modifier and context:
453-
card_to_serve = self.extended_card_modifier(base_card, context)
452+
card_to_serve = await self.extended_card_modifier(
453+
base_card, context
454+
)
454455
elif self.card_modifier:
455-
card_to_serve = self.card_modifier(base_card)
456+
card_to_serve = await self.card_modifier(base_card)
456457

457458
return GetAuthenticatedExtendedCardResponse(
458459
root=GetAuthenticatedExtendedCardSuccessResponse(

‎src/a2a/utils/signing.py‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
import json
22

3-
from collections.abc import Callable
3+
from collections.abc import Awaitable, Callable
44
from typing import Any, TypedDict
55

66
from a2a.utils.helpers import canonicalize_agent_card
@@ -54,7 +54,7 @@ def create_agent_card_signer(
5454
signing_key: PyJWK | str | bytes,
5555
protected_header: ProtectedHeader,
5656
header: dict[str, Any] | None = None,
57-
) -> Callable[[AgentCard], AgentCard]:
57+
) -> Callable[[AgentCard], Awaitable[AgentCard]]:
5858
"""Creates a function that signs an AgentCard and adds the signature.
5959
6060
Args:
@@ -63,10 +63,10 @@ def create_agent_card_signer(
6363
header: Unprotected header parameters.
6464
6565
Returns:
66-
A callable that takes an AgentCard and returns the modified AgentCard with a signature.
66+
A callable that takes an AgentCard and returns the awaitable modified AgentCard with a signature.
6767
"""
6868

69-
def agent_card_signer(agent_card: AgentCard) -> AgentCard:
69+
async def agent_card_signer(agent_card: AgentCard) -> AgentCard:
7070
"""Signs agent card."""
7171
canonical_payload = canonicalize_agent_card(agent_card)
7272
payload_dict = json.loads(canonical_payload)

0 commit comments

Comments
 (0)