wip

Author: jezekra1

agents/chat/uv.lock

@@ -22,6 +22,8 @@ dependencies = [
     "httpx",                                          # version determined by a2a-sdk
     "mcp>=1.12.3",
     "fastapi>=0.116.1",
+    "authlib>=1.3.0",
+    "async-lru>=2.0.4",
 ]
 
 [dependency-groups]
@@ -75,3 +77,4 @@ addopts = "-v"
 ignore = ["tests/**", "examples/cli.py"]
 venvPath = "."
 venv = ".venv"
+reportUnusedCallResult = "none"

apps/agentstack-sdk-py/pyproject.toml

@@ -8,10 +8,13 @@
 from typing import TYPE_CHECKING, Any, Self
 from urllib.parse import parse_qs
 
-import a2a.types
 import pydantic
+from a2a.server.agent_execution import RequestContext
+from a2a.types import Message as A2AMessage
+from a2a.types import Role, TextPart
 from mcp.client.auth import OAuthClientProvider
 from mcp.shared.auth import OAuthClientMetadata
+from typing_extensions import override
 
 from agentstack_sdk.a2a.extensions.auth.oauth.storage import MemoryTokenStorageFactory, TokenStorageFactory
 from agentstack_sdk.a2a.extensions.base import BaseExtensionClient, BaseExtensionServer, BaseExtensionSpec
@@ -58,13 +61,17 @@ class OAuthExtensionMetadata(pydantic.BaseModel):
 
 
 class OAuthExtensionServer(BaseExtensionServer[OAuthExtensionSpec, OAuthExtensionMetadata]):
+    context: RunContext
+    token_storage_factory: TokenStorageFactory
+
     def __init__(self, spec: OAuthExtensionSpec, token_storage_factory: TokenStorageFactory | None = None) -> None:
         super().__init__(spec)
         self.token_storage_factory = token_storage_factory or MemoryTokenStorageFactory()
 
-    def handle_incoming_message(self, message: a2a.types.Message, context: RunContext):
-        super().handle_incoming_message(message, context)
-        self.context = context
+    @override
+    def handle_incoming_message(self, message: A2AMessage, run_context: RunContext, request_context: RequestContext):
+        super().handle_incoming_message(message, run_context, request_context)
+        self.context = run_context
 
     def _get_fulfillment_for_resource(self, resource_url: pydantic.AnyUrl):
         if not self.data:
@@ -117,7 +124,7 @@ def create_auth_request(self, *, authorization_endpoint_url: pydantic.AnyUrl):
         data = AuthRequest(authorization_endpoint_url=authorization_endpoint_url)
         return AgentMessage(text="Authorization required", metadata={self.spec.URI: data.model_dump(mode="json")})
 
-    def parse_auth_response(self, *, message: a2a.types.Message):
+    def parse_auth_response(self, *, message: A2AMessage):
         if not message or not message.metadata or not (data := message.metadata.get(self.spec.URI)):
             raise RuntimeError("Invalid auth response")
         return AuthResponse.model_validate(data)
@@ -127,18 +134,18 @@ class OAuthExtensionClient(BaseExtensionClient[OAuthExtensionSpec, NoneType]):
     def fulfillment_metadata(self, *, oauth_fulfillments: dict[str, Any]) -> dict[str, Any]:
         return {self.spec.URI: OAuthExtensionMetadata(oauth_fulfillments=oauth_fulfillments).model_dump(mode="json")}
 
-    def parse_auth_request(self, *, message: a2a.types.Message):
+    def parse_auth_request(self, *, message: A2AMessage):
         if not message or not message.metadata or not (data := message.metadata.get(self.spec.URI)):
             raise ValueError("Invalid auth request")
         return AuthRequest.model_validate(data)
 
     def create_auth_response(self, *, task_id: str, redirect_uri: pydantic.AnyUrl):
         data = AuthResponse(redirect_uri=redirect_uri)
 
-        return a2a.types.Message(
+        return A2AMessage(
             message_id=str(uuid.uuid4()),
-            role=a2a.types.Role.user,
-            parts=[a2a.types.TextPart(text="Authorization completed")],  # type: ignore
+            role=Role.user,
+            parts=[TextPart(text="Authorization completed")],  # type: ignore
             task_id=task_id,
             metadata={self.spec.URI: data.model_dump(mode="json")},
         )

apps/agentstack-sdk-py/src/agentstack_sdk/a2a/extensions/auth/oauth/oauth.py

@@ -1,11 +1,14 @@
 # Copyright 2025 © BeeAI a Series of LF Projects, LLC
 # SPDX-License-Identifier: Apache-2.0
 
-import typing
-from typing import TYPE_CHECKING
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Self
 
 import pydantic
+from a2a.server.agent_execution.context import RequestContext
 from a2a.types import Message as A2AMessage
+from typing_extensions import override
 
 from agentstack_sdk.a2a.extensions.base import BaseExtensionClient, BaseExtensionServer, BaseExtensionSpec
 from agentstack_sdk.a2a.types import AgentMessage, AuthRequired
@@ -35,7 +38,7 @@ class SecretsExtensionSpec(BaseExtensionSpec[SecretsServiceExtensionParams | Non
     URI: str = "https://a2a-extensions.agentstack.beeai.dev/auth/secrets/v1"
 
     @classmethod
-    def single_demand(cls, name: str, key: str | None = None, description: str | None = None) -> typing.Self:
+    def single_demand(cls, name: str, key: str | None = None, description: str | None = None) -> Self:
         return cls(
             params=SecretsServiceExtensionParams(
                 secret_demands={key or "default": SecretDemand(description=description, name=name)}
@@ -44,9 +47,12 @@ def single_demand(cls, name: str, key: str | None = None, description: str | Non
 
 
 class SecretsExtensionServer(BaseExtensionServer[SecretsExtensionSpec, SecretsServiceExtensionMetadata]):
-    def handle_incoming_message(self, message: A2AMessage, context: "RunContext"):
-        super().handle_incoming_message(message, context)
-        self.context = context
+    context: RunContext
+
+    @override
+    def handle_incoming_message(self, message: A2AMessage, run_context: RunContext, request_context: RequestContext):
+        super().handle_incoming_message(message, run_context, request_context)
+        self.context = run_context
 
     def parse_secret_response(self, message: A2AMessage) -> SecretsServiceExtensionMetadata:
         if not message or not message.metadata or not (data := message.metadata.get(self.spec.URI)):

apps/agentstack-sdk-py/src/agentstack_sdk/a2a/extensions/auth/secrets/secrets.py

@@ -9,8 +9,11 @@
 from contextlib import asynccontextmanager
 from types import NoneType
 
-import a2a.types
 import pydantic
+from a2a.server.agent_execution.context import RequestContext
+from a2a.types import AgentCard, AgentExtension
+from a2a.types import Message as A2AMessage
+from typing_extensions import override
 
 ParamsT = typing.TypeVar("ParamsT")
 MetadataFromClientT = typing.TypeVar("MetadataFromClientT")
@@ -19,6 +22,7 @@
 
 if typing.TYPE_CHECKING:
     from agentstack_sdk.server.context import RunContext
+    from agentstack_sdk.server.dependencies import Dependency
 
 
 def _get_generic_args(cls: type, base_class: type) -> tuple[typing.Any, ...]:
@@ -68,7 +72,7 @@ def __init__(self, params: ParamsT) -> None:
         self.params = params
 
     @classmethod
-    def from_agent_card(cls, agent: a2a.types.AgentCard) -> typing.Self | None:
+    def from_agent_card(cls, agent: AgentCard) -> typing.Self | None:
         """
         Client should construct an extension instance using this classmethod.
         """
@@ -81,14 +85,14 @@ def from_agent_card(cls, agent: a2a.types.AgentCard) -> typing.Self | None:
         except StopIteration:
             return None
 
-    def to_agent_card_extensions(self, *, required: bool = False) -> list[a2a.types.AgentExtension]:
+    def to_agent_card_extensions(self, *, required: bool = False) -> list[AgentExtension]:
         """
         Agent should use this method to obtain extension definitions to advertise on the agent card.
         This returns a list, as it's possible to support multiple A2A extensions within a single class.
         (Usually, that would be different versions of the extension spec.)
         """
         return [
-            a2a.types.AgentExtension(
+            AgentExtension(
                 uri=self.URI,
                 description=self.DESCRIPTION,
                 params=typing.cast(
@@ -105,7 +109,8 @@ def __init__(self):
         super().__init__(None)
 
     @classmethod
-    def from_agent_card(cls, agent: a2a.types.AgentCard) -> typing.Self | None:
+    @override
+    def from_agent_card(cls, agent: AgentCard) -> typing.Self | None:
         if any(e.uri == cls.URI for e in agent.capabilities.extensions or []):
             return cls()
         return None
@@ -125,7 +130,7 @@ def __init_subclass__(cls, **kwargs):
         cls.MetadataFromClient = _get_generic_args(cls, BaseExtensionServer)[1]
 
     _metadata_from_client: MetadataFromClientT | None = None
-    _dependencies: dict
+    _dependencies: dict[str, Dependency] = {}
 
     @property
     def data(self):
@@ -139,7 +144,7 @@ def __init__(self, spec: ExtensionSpecT, *args, **kwargs) -> None:
         self._args = args
         self._kwargs = kwargs
 
-    def parse_client_metadata(self, message: a2a.types.Message) -> MetadataFromClientT | None:
+    def parse_client_metadata(self, message: A2AMessage) -> MetadataFromClientT | None:
         """
         Server should use this method to retrieve extension-associated metadata from a message.
         """
@@ -149,7 +154,7 @@ def parse_client_metadata(self, message: a2a.types.Message) -> MetadataFromClien
             else pydantic.TypeAdapter(self.MetadataFromClient).validate_python(message.metadata[self.spec.URI])
         )
 
-    def handle_incoming_message(self, message: a2a.types.Message, context: RunContext):
+    def handle_incoming_message(self, message: A2AMessage, run_context: RunContext, request_context: RequestContext):
         if self._metadata_from_client is None:
             self._metadata_from_client = self.parse_client_metadata(message)
 
@@ -158,12 +163,16 @@ def _fork(self) -> typing.Self:
         return type(self)(self.spec, *self._args, **self._kwargs)
 
     def __call__(
-        self, message: a2a.types.Message, context: RunContext, dependencies: dict[str, typing.Any]
+        self,
+        message: A2AMessage,
+        run_context: RunContext,
+        request_context: RequestContext,
+        dependencies: dict[str, Dependency],
     ) -> typing.Self:
         """Works as a dependency constructor - create a private instance for the request"""
         instance = self._fork()
         instance._dependencies = dependencies
-        instance.handle_incoming_message(message, context)
+        instance.handle_incoming_message(message, run_context, request_context)
         return instance
 
     @asynccontextmanager
@@ -185,7 +194,7 @@ def __init_subclass__(cls, **kwargs):
     def __init__(self, spec: ExtensionSpecT) -> None:
         self.spec = spec
 
-    def parse_server_metadata(self, message: a2a.types.Message) -> MetadataFromServerT | None:
+    def parse_server_metadata(self, message: A2AMessage) -> MetadataFromServerT | None:
         """
         Client should use this method to retrieve extension-associated metadata from a message.
         """

apps/agentstack-sdk-py/src/agentstack_sdk/a2a/extensions/base.py

@@ -5,12 +5,18 @@
 
 import re
 from types import NoneType
-from typing import Any, Self
+from typing import TYPE_CHECKING, Any, Self
 
 import pydantic
+from a2a.server.agent_execution.context import RequestContext
+from a2a.types import Message as A2AMessage
+from typing_extensions import override
 
 from agentstack_sdk.a2a.extensions.base import BaseExtensionClient, BaseExtensionServer, BaseExtensionSpec
 
+if TYPE_CHECKING:
+    from agentstack_sdk.server.context import RunContext
+
 
 class EmbeddingFulfillment(pydantic.BaseModel):
     identifier: str | None = None
@@ -78,10 +84,11 @@ class EmbeddingServiceExtensionMetadata(pydantic.BaseModel):
 class EmbeddingServiceExtensionServer(
     BaseExtensionServer[EmbeddingServiceExtensionSpec, EmbeddingServiceExtensionMetadata]
 ):
-    def handle_incoming_message(self, message, context):
+    @override
+    def handle_incoming_message(self, message: A2AMessage, run_context: RunContext, request_context: RequestContext):
         from agentstack_sdk.platform import get_platform_client
 
-        super().handle_incoming_message(message, context)
+        super().handle_incoming_message(message, run_context, request_context)
         if not self.data:
             return
 

apps/agentstack-sdk-py/src/agentstack_sdk/a2a/extensions/services/embedding.py

@@ -8,12 +8,13 @@
 from typing import TYPE_CHECKING, Any, Self
 
 import pydantic
+from a2a.server.agent_execution.context import RequestContext
+from a2a.types import Message as A2AMessage
+from typing_extensions import override
 
 from agentstack_sdk.a2a.extensions.base import BaseExtensionClient, BaseExtensionServer, BaseExtensionSpec
 
 if TYPE_CHECKING:
-    from a2a.types import Message
-
     from agentstack_sdk.server.context import RunContext
 
 
@@ -81,10 +82,11 @@ class LLMServiceExtensionMetadata(pydantic.BaseModel):
 
 
 class LLMServiceExtensionServer(BaseExtensionServer[LLMServiceExtensionSpec, LLMServiceExtensionMetadata]):
-    def handle_incoming_message(self, message: Message, context: RunContext):
+    @override
+    def handle_incoming_message(self, message: A2AMessage, run_context: RunContext, request_context: RequestContext):
         from agentstack_sdk.platform import get_platform_client
 
-        super().handle_incoming_message(message, context)
+        super().handle_incoming_message(message, run_context, request_context)
         if not self.data:
             return
 

apps/agentstack-sdk-py/src/agentstack_sdk/a2a/extensions/services/llm.py

@@ -8,10 +8,12 @@
 from types import NoneType
 from typing import TYPE_CHECKING, Annotated, Any, Literal, Self
 
-import a2a.types
 import pydantic
+from a2a.server.agent_execution.context import RequestContext
+from a2a.types import Message as A2AMessage
 from mcp.client.stdio import StdioServerParameters, stdio_client
 from mcp.client.streamable_http import streamablehttp_client
+from typing_extensions import override
 
 from agentstack_sdk.a2a.extensions.auth.oauth.oauth import OAuthExtensionServer
 from agentstack_sdk.a2a.extensions.base import BaseExtensionClient, BaseExtensionServer, BaseExtensionSpec
@@ -102,8 +104,9 @@ class MCPServiceExtensionMetadata(pydantic.BaseModel):
 
 
 class MCPServiceExtensionServer(BaseExtensionServer[MCPServiceExtensionSpec, MCPServiceExtensionMetadata]):
-    def handle_incoming_message(self, message: a2a.types.Message, context: RunContext):
-        super().handle_incoming_message(message, context)
+    @override
+    def handle_incoming_message(self, message: A2AMessage, run_context: RunContext, request_context: RequestContext):
+        super().handle_incoming_message(message, run_context, request_context)
         if not self.data:
             return
 
@@ -115,7 +118,8 @@ def handle_incoming_message(self, message: a2a.types.Message, context: RunContex
                 except Exception:
                     logger.warning("Platform URL substitution failed", exc_info=True)
 
-    def parse_client_metadata(self, message: a2a.types.Message) -> MCPServiceExtensionMetadata | None:
+    @override
+    def parse_client_metadata(self, message: A2AMessage) -> MCPServiceExtensionMetadata | None:
         metadata = super().parse_client_metadata(message)
         if metadata:
             for name, demand in self.spec.params.mcp_demands.items():