agentstack-sdk 0.5.0rc5__py3-none-any.whl → 0.5.1rc2__py3-none-any.whl

agentstack_sdk/a2a/extensions/__init__.py

@@ -2,6 +2,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from .auth import *
+from .interactions import *
 from .services import *
 from .tools import *
 from .ui import *

agentstack_sdk/a2a/extensions/auth/oauth/oauth.py

@@ -8,10 +8,13 @@ from types import NoneType
 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 @@ class OAuthExtensionServer(BaseExtensionServer[OAuthExtensionSpec, OAuthExtensio
         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,7 +134,7 @@ 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)
@@ -135,10 +142,10 @@ class OAuthExtensionClient(BaseExtensionClient[OAuthExtensionSpec, NoneType]):
     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")},
         )

agentstack_sdk/a2a/extensions/auth/secrets/secrets.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 @@ class SecretsExtensionSpec(BaseExtensionSpec[SecretsServiceExtensionParams | 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)):

agentstack_sdk/a2a/extensions/base.py

@@ -9,8 +9,11 @@ from collections.abc import AsyncIterator
 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 @@ MetadataFromServerT = typing.TypeVar("MetadataFromServerT")
 
 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 @@ class BaseExtensionSpec(abc.ABC, typing.Generic[ParamsT]):
         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 @@ class BaseExtensionSpec(abc.ABC, typing.Generic[ParamsT]):
         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 @@ class NoParamsBaseExtensionSpec(BaseExtensionSpec[NoneType]):
         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 @@ class BaseExtensionServer(abc.ABC, typing.Generic[ExtensionSpecT, MetadataFromCl
         cls.MetadataFromClient = _get_generic_args(cls, BaseExtensionServer)[1]
 
     _metadata_from_client: MetadataFromClientT | None = None
-    _dependencies: dict
+    _dependencies: dict[str, Dependency] = {}  # noqa: RUF012
 
     @property
     def data(self):
@@ -139,7 +144,7 @@ class BaseExtensionServer(abc.ABC, typing.Generic[ExtensionSpecT, MetadataFromCl
         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 @@ class BaseExtensionServer(abc.ABC, typing.Generic[ExtensionSpecT, MetadataFromCl
             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 @@ class BaseExtensionServer(abc.ABC, typing.Generic[ExtensionSpecT, MetadataFromCl
         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 @@ class BaseExtensionClient(abc.ABC, typing.Generic[ExtensionSpecT, MetadataFromSe
     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.
         """

agentstack_sdk/a2a/extensions/interactions/__init__.py

@@ -0,0 +1,4 @@
+# Copyright 2025 © BeeAI a Series of LF Projects, LLC
+# SPDX-License-Identifier: Apache-2.0
+
+from .approval import *

agentstack_sdk/a2a/extensions/interactions/approval.py

@@ -0,0 +1,125 @@
+# Copyright 2025 © BeeAI a Series of LF Projects, LLC
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import uuid
+from types import NoneType
+from typing import TYPE_CHECKING, Annotated, Any, Literal
+
+import a2a.types
+from mcp import Implementation, Tool
+from pydantic import BaseModel, Discriminator, Field, TypeAdapter
+
+from agentstack_sdk.a2a.extensions.base import BaseExtensionClient, BaseExtensionServer, BaseExtensionSpec
+from agentstack_sdk.a2a.types import AgentMessage, InputRequired
+
+if TYPE_CHECKING:
+    from agentstack_sdk.server.context import RunContext
+
+
+class ApprovalRejectionError(RuntimeError):
+    pass
+
+
+class GenericApprovalRequest(BaseModel):
+    action: Literal["generic"] = "generic"
+
+    title: str | None = Field(None, description="A human-readable title for the action being approved.")
+    description: str | None = Field(None, description="A human-readable description of the action being approved.")
+
+
+class ToolCallServer(BaseModel):
+    name: str = Field(description="The programmatic name of the server.")
+    title: str | None = Field(description="A human-readable title for the server.")
+    version: str = Field(description="The version of the server.")
+
+
+class ToolCallApprovalRequest(BaseModel):
+    action: Literal["tool-call"] = "tool-call"
+
+    title: str | None = Field(None, description="A human-readable title for the tool call being approved.")
+    description: str | None = Field(None, description="A human-readable description of the tool call being approved.")
+    name: str = Field(description="The programmatic name of the tool.")
+    input: dict[str, Any] | None = Field(description="The input for the tool.")
+    server: ToolCallServer | None = Field(None, description="The server executing the tool.")
+
+    @staticmethod
+    def from_mcp_tool(
+        tool: Tool, input: dict[str, Any] | None, server: Implementation | None = None
+    ) -> ToolCallApprovalRequest:
+        return ToolCallApprovalRequest(
+            name=tool.name,
+            title=tool.annotations.title if tool.annotations else None,
+            description=tool.description,
+            input=input,
+            server=ToolCallServer(name=server.name, title=server.title, version=server.version) if server else None,
+        )
+
+
+ApprovalRequest = Annotated[GenericApprovalRequest | ToolCallApprovalRequest, Discriminator("action")]
+
+
+class ApprovalResponse(BaseModel):
+    decision: Literal["approve", "reject"]
+
+    @property
+    def approved(self) -> bool:
+        return self.decision == "approve"
+
+    def raise_on_rejection(self) -> None:
+        if self.decision == "reject":
+            raise ApprovalRejectionError("Approval request has been rejected")
+
+
+class ApprovalExtensionParams(BaseModel):
+    pass
+
+
+class ApprovalExtensionSpec(BaseExtensionSpec[ApprovalExtensionParams]):
+    URI: str = "https://a2a-extensions.agentstack.beeai.dev/interactions/approval/v1"
+
+
+class ApprovalExtensionMetadata(BaseModel):
+    pass
+
+
+class ApprovalExtensionServer(BaseExtensionServer[ApprovalExtensionSpec, ApprovalExtensionMetadata]):
+    def create_request_message(self, *, request: ApprovalRequest):
+        return AgentMessage(text="Approval requested", metadata={self.spec.URI: request.model_dump(mode="json")})
+
+    def parse_response(self, *, message: a2a.types.Message):
+        if not message.metadata or not (data := message.metadata.get(self.spec.URI)):
+            raise ValueError("Approval response data is missing")
+        return ApprovalResponse.model_validate(data)
+
+    async def request_approval(
+        self,
+        request: ApprovalRequest,
+        *,
+        context: RunContext,
+    ) -> ApprovalResponse:
+        message = self.create_request_message(request=request)
+        message = await context.yield_async(InputRequired(message=message))
+        if not message:
+            raise RuntimeError("Yield did not return a message")
+        return self.parse_response(message=message)
+
+
+class ApprovalExtensionClient(BaseExtensionClient[ApprovalExtensionSpec, NoneType]):
+    def create_response_message(self, *, response: ApprovalResponse, task_id: str | None):
+        return a2a.types.Message(
+            message_id=str(uuid.uuid4()),
+            role=a2a.types.Role.user,
+            parts=[],
+            task_id=task_id,
+            metadata={self.spec.URI: response.model_dump(mode="json")},
+        )
+
+    def parse_request(self, *, message: a2a.types.Message):
+        if not message.metadata or not (data := message.metadata.get(self.spec.URI)):
+            raise ValueError("Approval request data is missing")
+        return TypeAdapter(ApprovalRequest).validate_python(data)
+
+    def metadata(self) -> dict[str, Any]:
+        return {self.spec.URI: ApprovalExtensionMetadata().model_dump(mode="json")}

agentstack_sdk/a2a/extensions/services/embedding.py

@@ -5,12 +5,18 @@ from __future__ import annotations
 
 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
 

agentstack_sdk/a2a/extensions/services/llm.py

@@ -8,12 +8,13 @@ from types import NoneType
 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
 

agentstack_sdk/a2a/extensions/services/mcp.py

@@ -8,10 +8,12 @@ from contextlib import asynccontextmanager
 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 @@ class MCPServiceExtensionServer(BaseExtensionServer[MCPServiceExtensionSpec, MCP
                 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():

agentstack_sdk/a2a/extensions/services/platform.py

@@ -9,9 +9,12 @@ from contextlib import asynccontextmanager
 from types import NoneType
 from typing import TYPE_CHECKING
 
-import a2a.types
 import pydantic
+from a2a.server.agent_execution.context import RequestContext
+from a2a.types import Message as A2AMessage
+from fastapi.security.utils import get_authorization_scheme_param
 from pydantic.networks import HttpUrl
+from typing_extensions import override
 
 from agentstack_sdk.a2a.extensions.base import (
     BaseExtensionClient,
@@ -21,6 +24,7 @@ from agentstack_sdk.a2a.extensions.base import (
 from agentstack_sdk.a2a.extensions.exceptions import ExtensionError
 from agentstack_sdk.platform import use_platform_client
 from agentstack_sdk.platform.client import PlatformClient
+from agentstack_sdk.server.middleware.platform_auth_backend import PlatformAuthenticatedUser
 from agentstack_sdk.util.httpx import BearerAuth
 
 if TYPE_CHECKING:
@@ -29,7 +33,7 @@ if TYPE_CHECKING:
 
 class PlatformApiExtensionMetadata(pydantic.BaseModel):
     base_url: HttpUrl | None = None
-    auth_token: pydantic.Secret[str]
+    auth_token: pydantic.Secret[str] | None = None
     expires_at: pydantic.AwareDatetime | None = None
 
 
@@ -53,13 +57,8 @@ class PlatformApiExtensionSpec(BaseExtensionSpec[PlatformApiExtensionParams]):
 class PlatformApiExtensionServer(BaseExtensionServer[PlatformApiExtensionSpec, PlatformApiExtensionMetadata]):
     context_id: str | None = None
 
-    def parse_client_metadata(self, message: a2a.types.Message) -> PlatformApiExtensionMetadata | None:
-        self.context_id = message.context_id
-        # we assume that the context id is the same ID as the platform context id
-        # if different IDs are passed, api requests to platform using this token will fail
-        return super().parse_client_metadata(message)
-
     @asynccontextmanager
+    @override
     async def lifespan(self) -> AsyncIterator[None]:
         """Called when entering the agent context after the first message was parsed (__call__ was already called)"""
         if self.data and self.spec.params.auto_use:
@@ -68,25 +67,44 @@ class PlatformApiExtensionServer(BaseExtensionServer[PlatformApiExtensionSpec, P
         else:
             yield
 
-    def handle_incoming_message(self, message: a2a.types.Message, context: RunContext):
-        super().handle_incoming_message(message, context)
-        if self.data:
-            self.data.base_url = self.data.base_url or HttpUrl(os.getenv("PLATFORM_URL", "http://127.0.0.1:8333"))
+    def _get_header_token(self, request_context: RequestContext) -> pydantic.Secret[str] | None:
+        header_token = None
+        call_context = request_context.call_context
+        assert call_context
+        if isinstance(call_context.user, PlatformAuthenticatedUser):
+            header_token = call_context.user.auth_token.get_secret_value()
+        elif auth_header := call_context.state.get("headers", {}).get("authorization", None):
+            _scheme, header_token = get_authorization_scheme_param(auth_header)
+        return pydantic.Secret(header_token) if header_token else None
+
+    @override
+    def handle_incoming_message(self, message: A2AMessage, run_context: RunContext, request_context: RequestContext):
+        super().handle_incoming_message(message, run_context, request_context)
+        # we assume that request context id is the same ID as the platform context id
+        # if different IDs are passed, api requests to platform using this token will fail
+        self.context_id = request_context.context_id
+
+        self._metadata_from_client = self._metadata_from_client or PlatformApiExtensionMetadata()
+        data = self._metadata_from_client
+        data.base_url = data.base_url or HttpUrl(os.getenv("PLATFORM_URL", "http://127.0.0.1:8333"))
+        data.auth_token = data.auth_token or self._get_header_token(request_context)
+
+        if not data.auth_token:
+            raise ExtensionError(self.spec, "Platform extension metadata was not provided")
 
     @asynccontextmanager
     async def use_client(self) -> AsyncIterator[PlatformClient]:
-        if not self.data:
+        if not self.data or not self.data.auth_token:
             raise ExtensionError(self.spec, "Platform extension metadata was not provided")
-        auth_token = self.data.auth_token.get_secret_value()
         async with use_platform_client(
             context_id=self.context_id,
             base_url=str(self.data.base_url),
-            auth_token=auth_token,
+            auth_token=self.data.auth_token.get_secret_value(),
         ) as client:
             yield client
 
     async def create_httpx_auth(self) -> BearerAuth:
-        if not self.data:
+        if not self.data or not self.data.auth_token:
             raise ExtensionError(self.spec, "Platform extension metadata was not provided")
         return BearerAuth(token=self.data.auth_token.get_secret_value())
 

agentstack_sdk/a2a/extensions/ui/__init__.py

@@ -2,6 +2,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from .agent_detail import *
+from .canvas import *
 from .citation import *
 from .error import *
 from .form_request import *

agentstack_sdk/a2a/extensions/ui/canvas.py

@@ -6,8 +6,10 @@ from __future__ import annotations
 from typing import TYPE_CHECKING
 
 import pydantic
+from a2a.server.agent_execution.context import RequestContext
 from a2a.types import Artifact, TextPart
 from a2a.types import Message as A2AMessage
+from typing_extensions import override
 
 if TYPE_CHECKING:
     from agentstack_sdk.server.context import RunContext
@@ -37,12 +39,13 @@ class CanvasExtensionSpec(NoParamsBaseExtensionSpec):
 
 
 class CanvasExtensionServer(BaseExtensionServer[CanvasExtensionSpec, CanvasEditRequestMetadata]):
-    def handle_incoming_message(self, message: A2AMessage, context: RunContext):
+    @override
+    def handle_incoming_message(self, message: A2AMessage, run_context: RunContext, request_context: RequestContext):
         if message.metadata and self.spec.URI in message.metadata and message.parts:
             message.parts = [part for part in message.parts if not isinstance(part.root, TextPart)]
 
-        super().handle_incoming_message(message, context)
-        self.context = context
+        super().handle_incoming_message(message, run_context, request_context)
+        self.context = run_context
 
     async def parse_canvas_edit_request(self, *, message: A2AMessage) -> CanvasEditRequest | None:
         if not message or not message.metadata or not (data := message.metadata.get(self.spec.URI)):

agentstack_sdk/a2a/extensions/ui/error.py

@@ -18,7 +18,8 @@ from agentstack_sdk.a2a.extensions.base import (
     BaseExtensionServer,
     BaseExtensionSpec,
 )
-from agentstack_sdk.a2a.types import AgentMessage, JsonDict, Metadata
+from agentstack_sdk.a2a.types import AgentMessage, Metadata
+from agentstack_sdk.types import JsonValue
 from agentstack_sdk.util import resource_context
 
 logger = logging.getLogger(__name__)
@@ -68,7 +69,7 @@ class ErrorMetadata(pydantic.BaseModel):
 
     error: Error | ErrorGroup
     stack_trace: str | None = None
-    context: JsonDict | None = None
+    context: JsonValue | None = None
 
 
 class ErrorExtensionParams(pydantic.BaseModel):
@@ -133,7 +134,7 @@ class ErrorExtensionServer(BaseExtensionServer[ErrorExtensionSpec, NoneType]):
             yield
 
     @property
-    def context(self) -> JsonDict:
+    def context(self) -> JsonValue:
         """Get the current request's error context."""
         try:
             return get_error_extension_context().context
@@ -214,7 +215,7 @@ DEFAULT_ERROR_EXTENSION: Final = ErrorExtensionServer(ErrorExtensionSpec(ErrorEx
 
 class ErrorContext(pydantic.BaseModel, arbitrary_types_allowed=True):
     server: ErrorExtensionServer = pydantic.Field(default=DEFAULT_ERROR_EXTENSION)
-    context: JsonDict = pydantic.Field(default_factory=dict)
+    context: JsonValue = pydantic.Field(default_factory=dict)
 
 
 get_error_extension_context, use_error_extension_context = resource_context(

agentstack_sdk/a2a/extensions/ui/form_request.py

@@ -5,8 +5,10 @@ from __future__ import annotations
 
 from typing import TYPE_CHECKING, TypeVar, cast
 
+from a2a.server.agent_execution.context import RequestContext
 from a2a.types import Message as A2AMessage
 from pydantic import TypeAdapter
+from typing_extensions import override
 
 from agentstack_sdk.a2a.extensions.base import (
     BaseExtensionClient,
@@ -27,9 +29,10 @@ class FormRequestExtensionSpec(NoParamsBaseExtensionSpec):
 
 
 class FormRequestExtensionServer(BaseExtensionServer[FormRequestExtensionSpec, FormResponse]):
-    def handle_incoming_message(self, message: A2AMessage, 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
 
     async def request_form(self, *, form: FormRender, model: type[T] = FormResponse) -> T | None:
         message = await self.context.yield_async(