agentstack-sdk 0.5.2rc2__py3-none-any.whl

agentstack_sdk/server/constants.py

@@ -0,0 +1,9 @@
+# Copyright 2025 © BeeAI a Series of LF Projects, LLC
+# SPDX-License-Identifier: Apache-2.0
+
+from typing import Final
+
+_IMPLICIT_DEPENDENCY_PREFIX: Final = "___server_dep"
+
+
+__all__ = []

agentstack_sdk/server/context.py

@@ -0,0 +1,68 @@
+# Copyright 2025 © BeeAI a Series of LF Projects, LLC
+# SPDX-License-Identifier: Apache-2.0
+
+from collections.abc import AsyncGenerator
+from typing import Literal, overload
+from uuid import UUID
+
+import janus
+from a2a.types import Artifact, Message, MessageSendConfiguration, Task
+from pydantic import BaseModel, PrivateAttr
+
+from agentstack_sdk.a2a.types import RunYield, RunYieldResume
+from agentstack_sdk.platform.context import ContextHistoryItem
+from agentstack_sdk.server.store.context_store import ContextStoreInstance
+
+
+class RunContext(BaseModel, arbitrary_types_allowed=True):
+    configuration: MessageSendConfiguration | None = None
+    task_id: str
+    context_id: str
+    current_task: Task | None = None
+    related_tasks: list[Task] | None = None
+
+    _store: ContextStoreInstance | None = PrivateAttr(None)
+    _yield_queue: janus.Queue[RunYield] = PrivateAttr(default_factory=janus.Queue)
+    _yield_resume_queue: janus.Queue[RunYieldResume] = PrivateAttr(default_factory=janus.Queue)
+
+    async def store(self, data: Message | Artifact):
+        if not self._store:
+            raise RuntimeError("Context store is not initialized")
+        if isinstance(data, Message):
+            data = data.model_copy(deep=True, update={"context_id": self.context_id, "task_id": self.task_id})
+        await self._store.store(data)
+
+    @overload
+    async def load_history(
+        self, load_history_items: Literal[False] = False
+    ) -> AsyncGenerator[Message | Artifact, None]:
+        yield ...  # type: ignore
+
+    @overload
+    async def load_history(self, load_history_items: Literal[True]) -> AsyncGenerator[ContextHistoryItem, None]:
+        yield ...  # type: ignore
+
+    async def load_history(
+        self, load_history_items: bool = False
+    ) -> AsyncGenerator[ContextHistoryItem | Message | Artifact]:
+        if not self._store:
+            raise RuntimeError("Context store is not initialized")
+        async for item in self._store.load_history(load_history_items=load_history_items):
+            yield item
+
+    async def delete_history_from_id(self, from_id: UUID) -> None:
+        if not self._store:
+            raise RuntimeError("Context store is not initialized")
+        await self._store.delete_history_from_id(from_id)
+
+    def yield_sync(self, value: RunYield) -> RunYieldResume:
+        self._yield_queue.sync_q.put(value)
+        return self._yield_resume_queue.sync_q.get()
+
+    async def yield_async(self, value: RunYield) -> RunYieldResume:
+        await self._yield_queue.async_q.put(value)
+        return await self._yield_resume_queue.async_q.get()
+
+    def shutdown(self) -> None:
+        self._yield_queue.shutdown()
+        self._yield_resume_queue.shutdown()

agentstack_sdk/server/dependencies.py

@@ -0,0 +1,117 @@
+# Copyright 2025 © BeeAI a Series of LF Projects, LLC
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import inspect
+from collections import Counter
+from collections.abc import AsyncIterator, Callable
+from contextlib import AbstractAsyncContextManager, asynccontextmanager
+from inspect import isclass
+from typing import Annotated, Any, TypeAlias, Unpack, get_args, get_origin
+
+from a2a.server.agent_execution.context import RequestContext
+from a2a.types import Message
+from typing_extensions import Doc
+
+from agentstack_sdk.a2a.extensions import BaseExtensionSpec
+from agentstack_sdk.a2a.extensions.base import BaseExtensionServer
+from agentstack_sdk.server.context import RunContext
+
+Dependency: TypeAlias = (
+    Callable[[Message, RunContext, RequestContext, dict[str, "Dependency"]], Any] | BaseExtensionServer[Any, Any]
+)
+
+
+# Inspired by fastapi.Depends
+class Depends:
+    extension: BaseExtensionServer[Any, Any] | None = None
+
+    def __init__(
+        self,
+        dependency: Annotated[
+            Dependency,
+            Doc(
+                """
+                A "dependable" callable (like a function).
+                Don't call it directly, Agent Stack SDK will call it for you, just pass the object directly.
+                """
+            ),
+        ],
+    ):
+        self._dependency_callable: Dependency = dependency
+        if isinstance(dependency, BaseExtensionServer):
+            self.extension = dependency
+
+    def __call__(
+        self, message: Message, context: RunContext, request_context: RequestContext, dependencies: dict[str, Any]
+    ) -> AbstractAsyncContextManager[Dependency]:
+        instance = self._dependency_callable(message, context, request_context, dependencies)
+
+        @asynccontextmanager
+        async def lifespan() -> AsyncIterator[Dependency]:
+            if self.extension or hasattr(instance, "lifespan"):
+                async with instance.lifespan():
+                    yield instance
+            else:
+                yield instance
+
+        return lifespan()
+
+
+def extract_dependencies(sign: inspect.Signature) -> dict[str, Depends]:
+    dependencies = {}
+    seen_keys = set()
+
+    def process_args(name: str, args: tuple[Any, ...]) -> None:
+        if len(args) > 1:
+            dep_type, spec, *rest = args
+            # extension_param: Annotated[some_type, Depends(some_callable)]
+            if isinstance(spec, Depends):
+                dependencies[name] = spec
+            # extension_param: Annotated[BaseExtensionServer, BaseExtensionSpec()]
+            elif (
+                isclass(dep_type) and issubclass(dep_type, BaseExtensionServer) and isinstance(spec, BaseExtensionSpec)
+            ):
+                dependencies[name] = Depends(dep_type(spec, *rest))
+
+    for name, param in sign.parameters.items():
+        seen_keys.add(name)
+
+        if get_origin(param.annotation) is Annotated:
+            args = get_args(param.annotation)
+            process_args(name, args)
+
+        elif inspect.isclass(param.annotation):
+            # message: Message
+            if param.annotation == Message:
+                dependencies[name] = Depends(lambda message, _run_context, _request_context, _dependencies: message)
+            # context: Context
+            elif param.annotation == RunContext:
+                dependencies[name] = Depends(lambda _message, run_context, _request_context, _dependencies: run_context)
+            # extension: BaseExtensionServer = BaseExtensionSpec()
+            # TODO: this does not get past linters, should we enable it or somehow fix the typing?
+            # elif issubclass(param.annotation, BaseExtensionServer) and isinstance(param.default, BaseExtensionSpec):
+            #     dependencies[name] = Depends(param.annotation(param.default))
+        elif param.kind is inspect.Parameter.VAR_KEYWORD:
+            origin = get_origin(param.annotation)
+            if origin is Unpack:
+                seen_keys.discard(name)
+                (typed_dict,) = get_args(param.annotation)
+                for field_name, field_type in typed_dict.__annotations__.items():
+                    seen_keys.add(field_name)
+                    if get_origin(field_type) is Annotated:
+                        args = get_args(field_type)
+                        process_args(field_name, args)
+
+    missing_keys = seen_keys.difference(dependencies.keys())
+    if missing_keys:
+        raise TypeError(f"The agent function contains extra parameters with unknown type annotation: {missing_keys}")
+    if reserved_names := {param for param in dependencies if param.startswith("__")}:
+        raise TypeError(f"User-defined dependencies cannot start with double underscore: {reserved_names}")
+
+    extension_deps = Counter(dep.extension.spec.URI for dep in dependencies.values() if dep.extension)
+    if duplicate_uris := {k for k, v in extension_deps.items() if v > 1}:
+        raise TypeError(f"Duplicate extension URIs found in the agent function: {duplicate_uris}")
+
+    return dependencies

agentstack_sdk/server/exceptions.py

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

agentstack_sdk/server/middleware/__init__.py

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

agentstack_sdk/server/middleware/platform_auth_backend.py

@@ -0,0 +1,131 @@
+# Copyright 2025 © BeeAI a Series of LF Projects, LLC
+# SPDX-License-Identifier: Apache-2.0
+
+import logging
+import os
+from datetime import timedelta
+from urllib.parse import urljoin
+
+from a2a.auth.user import User
+from async_lru import alru_cache
+from authlib.jose import JsonWebKey, JWTClaims, KeySet, jwt
+from authlib.jose.errors import JoseError
+from fastapi import Request
+from fastapi.security import HTTPBearer
+from pydantic import Secret
+from starlette.authentication import (
+    AuthCredentials,
+    AuthenticationBackend,
+    AuthenticationError,
+    BaseUser,
+)
+from starlette.requests import HTTPConnection
+from typing_extensions import override
+
+from agentstack_sdk.platform import use_platform_client
+from agentstack_sdk.types import JsonValue
+
+logger = logging.getLogger(__name__)
+
+
+class PlatformAuthenticatedUser(User, BaseUser):
+    def __init__(self, claims: dict[str, JsonValue], auth_token: str):
+        self.claims: dict[str, JsonValue] = claims
+        self.auth_token: Secret[str] = Secret(auth_token)
+
+    @property
+    @override
+    def is_authenticated(self) -> bool:
+        return True
+
+    @property
+    @override
+    def user_name(self) -> str:
+        sub = self.claims.get("sub", None)
+        assert sub and isinstance(sub, str)
+        return sub
+
+    @property
+    @override
+    def display_name(self) -> str:
+        name = self.claims.get("name", None)
+        assert name and isinstance(name, str)
+        return name
+
+    @property
+    @override
+    def identity(self) -> str:
+        return self.user_name
+
+
+@alru_cache(ttl=timedelta(minutes=15).seconds)
+async def discover_jwks() -> KeySet:
+    try:
+        async with use_platform_client() as client:
+            response = await client.get("/.well-known/jwks")
+            return JsonWebKey.import_key_set(response.raise_for_status().json())  # pyright: ignore[reportAny]
+    except Exception as e:
+        url = "{platform_url}/.well-known/jwks"
+        logger.warning(f"JWKS discovery failed for url {url}: {e}")
+        raise RuntimeError(f"JWKS discovery failed for url {url}") from e
+
+
+class PlatformAuthBackend(AuthenticationBackend):
+    def __init__(self, public_url: str | None = None, skip_audience_validation: bool | None = None) -> None:
+        self.skip_audience_validation: bool = (
+            skip_audience_validation
+            if skip_audience_validation is not None
+            else os.getenv("PLATFORM_AUTH__SKIP_AUDIENCE_VALIDATION", "false").lower() in ("true", "1")
+        )
+        self._audience: str | None = public_url or os.getenv("PLATFORM_AUTH__PUBLIC_URL", None)
+        if not self.skip_audience_validation and not self._audience:
+            logger.warning(
+                "Public URL is not provided and audience validation is enabled. Proceeding to check audience from the request target URL. "
+                + "This may not work when requests to agents are proxied. (hint: set PLATFORM_AUTH__PUBLIC_URL env variable)"
+            )
+
+        self.security: HTTPBearer = HTTPBearer(auto_error=False)
+
+    @override
+    async def authenticate(self, conn: HTTPConnection) -> tuple[AuthCredentials, BaseUser] | None:
+        # We construct a Request object from the scope for compatibility with HTTPBearer and logging
+        request = Request(scope=conn.scope)
+
+        if request.url.path in ["/healthcheck", "/.well-known/agent-card.json"]:
+            return None
+
+        if not (auth := await self.security(request)):
+            raise AuthenticationError("Missing Authorization header")
+
+        audiences: list[str] = []
+        if not self.skip_audience_validation:
+            if self._audience:
+                audiences = [urljoin(self._audience, path) for path in ["/", "/jsonrpc"]]
+            else:
+                audiences = [str(request.url.replace(path=path)) for path in ["/", "/jsonrpc"]]
+
+        try:
+            # check only hostname urljoin("http://host:port/a/b", "/") -> "http://host:port/"
+            jwks = await discover_jwks()
+
+            # Verify signature
+            claims: JWTClaims = jwt.decode(
+                auth.credentials,
+                jwks,
+                claims_options={
+                    "sub": {"essential": True},
+                    "exp": {"essential": True},
+                    # "iss": {"essential": True}, # Issuer validation might be tricky if internal/external URLs differ
+                }
+                | ({"aud": {"essential": True, "values": audiences}} if not self.skip_audience_validation else {}),
+            )
+            claims.validate()
+
+            return AuthCredentials(["authenticated"]), PlatformAuthenticatedUser(claims, auth.credentials)
+
+        except (ValueError, JoseError) as e:
+            logger.warning(f"Authentication failed: {e}")
+            raise AuthenticationError("Invalid token") from e
+        except Exception as e:
+            logger.error(f"Authentication error: {e}")
+            raise AuthenticationError(f"Authentication failed: {e}") from e