agentstack-sdk 0.5.2rc2__py3-none-any.whl

agentstack_sdk/__init__.py

@@ -0,0 +1,6 @@
+# Copyright 2025 © BeeAI a Series of LF Projects, LLC
+# SPDX-License-Identifier: Apache-2.0
+
+from importlib.metadata import version
+
+__version__ = version("agentstack-sdk")

agentstack_sdk/a2a/__init__.py

@@ -0,0 +1,2 @@
+# Copyright 2025 © BeeAI a Series of LF Projects, LLC
+# SPDX-License-Identifier: Apache-2.0

agentstack_sdk/a2a/extensions/__init__.py

@@ -0,0 +1,8 @@
+# Copyright 2025 © BeeAI a Series of LF Projects, LLC
+# 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/__init__.py

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

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

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

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

@@ -0,0 +1,151 @@
+# 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, Any, Self
+from urllib.parse import parse_qs
+
+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
+from agentstack_sdk.a2a.types import AgentMessage, AuthRequired, RunYieldResume
+
+if TYPE_CHECKING:
+    from agentstack_sdk.server.context import RunContext
+
+_DEFAULT_DEMAND_NAME = "default"
+
+
+class AuthRequest(pydantic.BaseModel):
+    authorization_endpoint_url: pydantic.AnyUrl
+
+
+class AuthResponse(pydantic.BaseModel):
+    redirect_uri: pydantic.AnyUrl
+
+
+class OAuthFulfillment(pydantic.BaseModel):
+    redirect_uri: pydantic.AnyUrl
+
+
+class OAuthDemand(pydantic.BaseModel):
+    redirect_uri: bool = True
+
+
+class OAuthExtensionParams(pydantic.BaseModel):
+    oauth_demands: dict[str, OAuthDemand]
+    """Server requests that the agent requires to be provided by the client."""
+
+
+class OAuthExtensionSpec(BaseExtensionSpec[OAuthExtensionParams]):
+    URI: str = "https://a2a-extensions.agentstack.beeai.dev/auth/oauth/v1"
+
+    @classmethod
+    def single_demand(cls, name: str = _DEFAULT_DEMAND_NAME) -> Self:
+        return cls(params=OAuthExtensionParams(oauth_demands={name: OAuthDemand()}))
+
+
+class OAuthExtensionMetadata(pydantic.BaseModel):
+    oauth_fulfillments: dict[str, OAuthFulfillment] = {}
+    """Provided servers corresponding to the server requests."""
+
+
+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()
+
+    @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:
+            raise RuntimeError("No fulfillments found")
+
+        fulfillment = self.data.oauth_fulfillments.get(str(resource_url)) or self.data.oauth_fulfillments.get(
+            _DEFAULT_DEMAND_NAME
+        )
+        if fulfillment:
+            return fulfillment
+
+        raise RuntimeError("Fulfillment not found")
+
+    async def create_httpx_auth(self, *, resource_url: pydantic.AnyUrl):
+        fulfillment = self._get_fulfillment_for_resource(resource_url=resource_url)
+
+        resume: RunYieldResume = None
+
+        async def handle_redirect(auth_url: str) -> None:
+            nonlocal resume
+            if resume:
+                raise RuntimeError("Another redirect is already pending")
+            message = self.create_auth_request(authorization_endpoint_url=pydantic.AnyUrl(auth_url))
+            resume = await self.context.yield_async(AuthRequired(message=message))
+
+        async def handle_callback() -> tuple[str, str | None]:
+            nonlocal resume
+            try:
+                if not resume:
+                    raise ValueError("Missing resume data")
+                response = self.parse_auth_response(message=resume)
+                params = parse_qs(response.redirect_uri.query)
+                return params["code"][0], params.get("state", [None])[0]
+            finally:
+                resume = None
+
+        # A2A Client is responsible for catching the redirect and forwarding it over the A2A connection
+        oauth_auth = OAuthClientProvider(
+            server_url=str(resource_url),
+            client_metadata=OAuthClientMetadata(
+                redirect_uris=[fulfillment.redirect_uri],
+            ),
+            storage=await self.token_storage_factory.create_storage(),
+            redirect_handler=handle_redirect,
+            callback_handler=handle_callback,
+        )
+        return oauth_auth
+
+    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: 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)
+
+
+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: 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 A2AMessage(
+            message_id=str(uuid.uuid4()),
+            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/oauth/storage/__init__.py

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

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

@@ -0,0 +1,11 @@
+# Copyright 2025 © BeeAI a Series of LF Projects, LLC
+# SPDX-License-Identifier: Apache-2.0
+
+import abc
+
+from mcp.client.auth import TokenStorage
+
+
+class TokenStorageFactory(abc.ABC):
+    @abc.abstractmethod
+    async def create_storage(self) -> TokenStorage: ...

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

@@ -0,0 +1,38 @@
+# Copyright 2025 © BeeAI a Series of LF Projects, LLC
+# SPDX-License-Identifier: Apache-2.0
+
+
+from mcp.client.auth import TokenStorage
+from mcp.shared.auth import OAuthClientInformationFull, OAuthToken
+
+from .base import TokenStorageFactory
+
+
+class MemoryTokenStorage(TokenStorage):
+    def __init__(self):
+        self.tokens: OAuthToken | None = None
+        self.client_info: OAuthClientInformationFull | None = None
+
+    async def get_tokens(self) -> OAuthToken | None:
+        return self.tokens
+
+    async def set_tokens(self, tokens: OAuthToken) -> None:
+        self.tokens = tokens
+
+    async def get_client_info(self) -> OAuthClientInformationFull | None:
+        return self.client_info
+
+    async def set_client_info(self, client_info: OAuthClientInformationFull) -> None:
+        self.client_info = client_info
+
+
+class MemoryTokenStorageFactory(TokenStorageFactory):
+    def __init__(self, *, client_info: OAuthClientInformationFull | None = None):
+        super().__init__()
+        self._client_info = client_info
+
+    async def create_storage(self) -> TokenStorage:
+        storage = MemoryTokenStorage()
+        if self._client_info:
+            await storage.set_client_info(self._client_info)
+        return storage

agentstack_sdk/a2a/extensions/auth/secrets/__init__.py

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

agentstack_sdk/a2a/extensions/auth/secrets/secrets.py

@@ -0,0 +1,77 @@
+# Copyright 2025 © BeeAI a Series of LF Projects, LLC
+# SPDX-License-Identifier: Apache-2.0
+
+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
+
+if TYPE_CHECKING:
+    from agentstack_sdk.server.context import RunContext
+
+
+class SecretDemand(pydantic.BaseModel):
+    name: str
+    description: str | None = None
+
+
+class SecretFulfillment(pydantic.BaseModel):
+    secret: str
+
+
+class SecretsServiceExtensionParams(pydantic.BaseModel):
+    secret_demands: dict[str, SecretDemand]
+
+
+class SecretsServiceExtensionMetadata(pydantic.BaseModel):
+    secret_fulfillments: dict[str, SecretFulfillment] = {}
+
+
+class SecretsExtensionSpec(BaseExtensionSpec[SecretsServiceExtensionParams | None]):
+    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) -> Self:
+        return cls(
+            params=SecretsServiceExtensionParams(
+                secret_demands={key or "default": SecretDemand(description=description, name=name)}
+            )
+        )
+
+
+class SecretsExtensionServer(BaseExtensionServer[SecretsExtensionSpec, SecretsServiceExtensionMetadata]):
+    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)):
+            raise ValueError("Secrets has not been provided in response.")
+
+        return SecretsServiceExtensionMetadata.model_validate(data)
+
+    async def request_secrets(self, params: SecretsServiceExtensionParams) -> SecretsServiceExtensionMetadata:
+        resume = await self.context.yield_async(
+            AuthRequired(
+                message=AgentMessage(
+                    metadata={self.spec.URI: params.model_dump(mode="json")},
+                )
+            )
+        )
+        if isinstance(resume, A2AMessage):
+            return self.parse_secret_response(message=resume)
+        else:
+            raise ValueError("Secrets has not been provided in response.")
+
+
+class SecretsExtensionClient(BaseExtensionClient[SecretsExtensionSpec, SecretsServiceExtensionParams]): ...

agentstack_sdk/a2a/extensions/base.py

@@ -0,0 +1,205 @@
+# Copyright 2025 © BeeAI a Series of LF Projects, LLC
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import abc
+import typing
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
+from types import NoneType
+
+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")
+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, ...]:
+    for base in getattr(cls, "__orig_bases__", ()):
+        if typing.get_origin(base) is base_class and (args := typing.get_args(base)):
+            return args
+    raise TypeError(f"Missing Params type for {cls.__name__}")
+
+
+class BaseExtensionSpec(abc.ABC, typing.Generic[ParamsT]):
+    """
+    Base class for an A2A extension handler.
+
+    The base implementations assume a single URI. More complex extension
+    handlers (e.g. serving multiple versions of an extension spec) may override
+    the appropriate methods.
+    """
+
+    URI: str
+    """
+    URI of the extension spec, or the preferred one if there are multiple supported.
+    """
+
+    DESCRIPTION: str | None = None
+    """
+    Description to be attached with the extension spec.
+    """
+
+    Params: type[ParamsT]
+    """
+    Type of the extension params, attached to the agent card.
+    """
+
+    def __init_subclass__(cls, **kwargs):
+        super().__init_subclass__(**kwargs)
+        cls.Params = _get_generic_args(cls, BaseExtensionSpec)[0]
+
+    params: ParamsT
+    """
+    Params from the agent card.
+    """
+
+    def __init__(self, params: ParamsT) -> None:
+        """
+        Agent should construct an extension instance using the constructor.
+        """
+        self.params = params
+
+    @classmethod
+    def from_agent_card(cls, agent: AgentCard) -> typing.Self | None:
+        """
+        Client should construct an extension instance using this classmethod.
+        """
+        try:
+            return cls(
+                params=pydantic.TypeAdapter(cls.Params).validate_python(
+                    next(x for x in agent.capabilities.extensions or [] if x.uri == cls.URI).params
+                ),
+            )
+        except StopIteration:
+            return None
+
+    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 [
+            AgentExtension(
+                uri=self.URI,
+                description=self.DESCRIPTION,
+                params=typing.cast(
+                    dict[str, typing.Any] | None,
+                    pydantic.TypeAdapter(self.Params).dump_python(self.params, mode="json"),
+                ),
+                required=required,
+            )
+        ]
+
+
+class NoParamsBaseExtensionSpec(BaseExtensionSpec[NoneType]):
+    def __init__(self):
+        super().__init__(None)
+
+    @classmethod
+    @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
+
+
+ExtensionSpecT = typing.TypeVar("ExtensionSpecT", bound=BaseExtensionSpec[typing.Any])
+
+
+class BaseExtensionServer(abc.ABC, typing.Generic[ExtensionSpecT, MetadataFromClientT]):
+    MetadataFromClient: type[MetadataFromClientT]
+    """
+    Type of the extension metadata, attached to messages.
+    """
+
+    def __init_subclass__(cls, **kwargs):
+        super().__init_subclass__(**kwargs)
+        cls.MetadataFromClient = _get_generic_args(cls, BaseExtensionServer)[1]
+
+    _metadata_from_client: MetadataFromClientT | None = None
+    _dependencies: dict[str, Dependency] = {}  # noqa: RUF012
+
+    @property
+    def data(self):
+        return self._metadata_from_client
+
+    def __bool__(self):
+        return bool(self.data)
+
+    def __init__(self, spec: ExtensionSpecT, *args, **kwargs) -> None:
+        self.spec = spec
+        self._args = args
+        self._kwargs = kwargs
+
+    def parse_client_metadata(self, message: A2AMessage) -> MetadataFromClientT | None:
+        """
+        Server should use this method to retrieve extension-associated metadata from a message.
+        """
+        return (
+            None
+            if not message.metadata or self.spec.URI not in message.metadata
+            else pydantic.TypeAdapter(self.MetadataFromClient).validate_python(message.metadata[self.spec.URI])
+        )
+
+    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)
+
+    def _fork(self) -> typing.Self:
+        """Creates a clone of this instance with the same arguments as the original"""
+        return type(self)(self.spec, *self._args, **self._kwargs)
+
+    def __call__(
+        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, run_context, request_context)
+        return instance
+
+    @asynccontextmanager
+    async def lifespan(self) -> AsyncIterator[None]:
+        """Called when entering the agent context after the first message was parsed (__call__ was already called)"""
+        yield
+
+
+class BaseExtensionClient(abc.ABC, typing.Generic[ExtensionSpecT, MetadataFromServerT]):
+    MetadataFromServer: type[MetadataFromServerT]
+    """
+    Type of the extension metadata, attached to messages.
+    """
+
+    def __init_subclass__(cls, **kwargs):
+        super().__init_subclass__(**kwargs)
+        cls.MetadataFromServer = _get_generic_args(cls, BaseExtensionClient)[1]
+
+    def __init__(self, spec: ExtensionSpecT) -> None:
+        self.spec = spec
+
+    def parse_server_metadata(self, message: A2AMessage) -> MetadataFromServerT | None:
+        """
+        Client should use this method to retrieve extension-associated metadata from a message.
+        """
+        return (
+            None
+            if not message.metadata or self.spec.URI not in message.metadata
+            else pydantic.TypeAdapter(self.MetadataFromServer).validate_python(message.metadata[self.spec.URI])
+        )

agentstack_sdk/a2a/extensions/common/__init__.py

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

agentstack_sdk/a2a/extensions/common/form.py

@@ -0,0 +1,149 @@
+# Copyright 2025 © BeeAI a Series of LF Projects, LLC
+# SPDX-License-Identifier: Apache-2.0
+
+from typing import Literal
+
+from pydantic import BaseModel, Field, model_validator
+
+
+class BaseField(BaseModel):
+    id: str
+    label: str
+    required: bool = False
+    col_span: int | None = Field(default=None, ge=1, le=4)
+
+
+class TextField(BaseField):
+    type: Literal["text"] = "text"
+    placeholder: str | None = None
+    default_value: str | None = None
+    auto_resize: bool | None = True
+
+
+class DateField(BaseField):
+    type: Literal["date"] = "date"
+    placeholder: str | None = None
+    default_value: str | None = None
+
+
+class FileItem(BaseModel):
+    uri: str
+    name: str | None = None
+    mime_type: str | None = None
+
+
+class FileField(BaseField):
+    type: Literal["file"] = "file"
+    accept: list[str]
+
+
+class OptionItem(BaseModel):
+    id: str
+    label: str
+
+
+class SingleSelectField(BaseField):
+    type: Literal["singleselect"] = "singleselect"
+    options: list[OptionItem]
+    default_value: str | None = None
+
+    @model_validator(mode="after")
+    def default_value_validator(self):
+        if self.default_value:
+            valid_values = {opt.id for opt in self.options}
+            if self.default_value not in valid_values:
+                raise ValueError(f"Invalid default_value: {self.default_value}. Must be one of {valid_values}")
+        return self
+
+
+class MultiSelectField(BaseField):
+    type: Literal["multiselect"] = "multiselect"
+    options: list[OptionItem]
+    default_value: list[str] | None = None
+
+    @model_validator(mode="after")
+    def default_values_validator(self):
+        if self.default_value:
+            valid_values = {opt.id for opt in self.options}
+            invalid_values = [v for v in self.default_value if v not in valid_values]
+            if invalid_values:
+                raise ValueError(f"Invalid default_value(s): {invalid_values}. Must be one of {valid_values}")
+        return self
+
+
+class CheckboxField(BaseField):
+    type: Literal["checkbox"] = "checkbox"
+    content: str
+    default_value: bool = False
+
+
+FormField = TextField | DateField | FileField | SingleSelectField | MultiSelectField | CheckboxField
+
+
+class FormRender(BaseModel):
+    title: str | None = None
+    description: str | None = None
+    columns: int | None = Field(default=None, ge=1, le=4)
+    submit_label: str | None = None
+    fields: list[FormField]
+
+
+class TextFieldValue(BaseModel):
+    type: Literal["text"] = "text"
+    value: str | None = None
+
+
+class DateFieldValue(BaseModel):
+    type: Literal["date"] = "date"
+    value: str | None = None
+
+
+class FileInfo(BaseModel):
+    uri: str
+    name: str | None = None
+    mime_type: str | None = None
+
+
+class FileFieldValue(BaseModel):
+    type: Literal["file"] = "file"
+    value: list[FileInfo] | None = None
+
+
+class SingleSelectFieldValue(BaseModel):
+    type: Literal["singleselect"] = "singleselect"
+    value: str | None = None
+
+
+class MultiSelectFieldValue(BaseModel):
+    type: Literal["multiselect"] = "multiselect"
+    value: list[str] | None = None
+
+
+class CheckboxFieldValue(BaseModel):
+    type: Literal["checkbox"] = "checkbox"
+    value: bool | None = None
+
+
+FormFieldValue = (
+    TextFieldValue
+    | DateFieldValue
+    | FileFieldValue
+    | SingleSelectFieldValue
+    | MultiSelectFieldValue
+    | CheckboxFieldValue
+)
+
+
+class FormResponse(BaseModel):
+    values: dict[str, FormFieldValue]
+
+    def __iter__(self):
+        for key, value in self.values.items():
+            match value:
+                case FileFieldValue():
+                    yield (
+                        key,
+                        [file.model_dump() for file in value.value] if value.value else None,
+                    )
+                case _:
+                    yield key, value.value