wip

Author: jezekra1

agents/chat/src/chat/agent.py

@@ -9,6 +9,7 @@
     AgentSkill,
     Message,
 )
+from agentstack_sdk.server.middleware.platform_auth_backend import PlatformAuthBackend
 from beeai_framework.agents.requirement.utils._tool import FinalAnswerTool
 from beeai_framework.errors import FrameworkError
 from pydantic import BaseModel
@@ -300,6 +301,7 @@ def serve():
             port=int(os.getenv("PORT", 8000)),
             configure_telemetry=True,
             context_store=PlatformContextStore(),
+            auth_backend=PlatformAuthBackend(skip_audience_validation=True),
         )
     except KeyboardInterrupt:
         pass

apps/agentstack-cli/src/agentstack_cli/api.py

@@ -14,6 +14,7 @@
 import openai
 from a2a.client import A2AClientHTTPError, Client, ClientConfig, ClientFactory
 from a2a.types import AgentCard
+from agentstack_sdk.platform.context import ContextToken
 from httpx import HTTPStatusError
 from httpx._types import RequestFiles
 
@@ -103,14 +104,10 @@ async def api_stream(
 
 
 @asynccontextmanager
-async def a2a_client(agent_card: AgentCard, use_auth: bool = True) -> AsyncIterator[Client]:
+async def a2a_client(agent_card: AgentCard, context_token: ContextToken) -> AsyncIterator[Client]:
     try:
         async with httpx.AsyncClient(
-            headers=(
-                {"Authorization": f"Bearer {token}"}
-                if use_auth and (token := await config.auth_manager.load_auth_token())
-                else {}
-            ),
+            headers={"Authorization": f"Bearer {context_token.token.get_secret_value()}"},
             follow_redirects=True,
             timeout=timedelta(hours=1).total_seconds(),
         ) as httpx_client:

apps/agentstack-cli/src/agentstack_cli/commands/agent.py

@@ -556,7 +556,7 @@ async def _run_agent(
                     console.print()  # Add newline after completion
                     return
                 case Task(id=task_id), TaskStatusUpdateEvent(
-                    status=TaskStatus(state=TaskState.working, message=message)
+                    status=TaskStatus(state=TaskState.working | TaskState.submitted, message=message)
                 ):
                     # Handle streaming content during working state
                     if message:
@@ -944,7 +944,7 @@ async def run_agent(
         if interaction_mode == InteractionMode.MULTI_TURN:
             console.print(f"{user_greeting}\n")
             turn_input = await _ask_form_questions(initial_form_render) if initial_form_render else handle_input()
-            async with a2a_client(provider.agent_card) as client:
+            async with a2a_client(provider.agent_card, context_token=context_token) as client:
                 while True:
                     console.print()
                     await _run_agent(
@@ -961,7 +961,7 @@ async def run_agent(
             user_greeting = ui_annotations.get("user_greeting", None) or "Enter your instructions."
             console.print(f"{user_greeting}\n")
             console.print()
-            async with a2a_client(provider.agent_card) as client:
+            async with a2a_client(provider.agent_card, context_token=context_token) as client:
                 await _run_agent(
                     client,
                     input=await _ask_form_questions(initial_form_render) if initial_form_render else handle_input(),
@@ -972,7 +972,7 @@ async def run_agent(
                 )
 
     else:
-        async with a2a_client(provider.agent_card) as client:
+        async with a2a_client(provider.agent_card, context_token=context_token) as client:
             await _run_agent(
                 client,
                 input,

apps/agentstack-cli/uv.lock

@@ -73,7 +73,7 @@ def _get_header_token(self, request_context: RequestContext) -> pydantic.Secret[
         assert call_context
         if isinstance(call_context.user, PlatformAuthenticatedUser):
             header_token = call_context.user.auth_token.get_secret_value()
-        elif (headers := call_context.state.get("headers")) and (auth_header := headers.get("authorization")):
+        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
 

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

@@ -76,7 +76,7 @@ def __init__(self, public_url: str | None = None, skip_audience_validation: bool
             if skip_audience_validation is not None
             else os.getenv("PLATFORM_AUTH__SKIP_AUDIENCE_VALIDATION", "false").lower() in ("true", "1")
         )
-        _audience = public_url or os.getenv("PLATFORM_AUTH__PUBLIC_URL")
+        _audience = public_url or os.getenv("PLATFORM_AUTH__PUBLIC_URL", "http://host.docker.internal:8333")
         if not self.skip_audience_validation and not _audience:
             raise ValueError(
                 "Public URL must be provided if audience validation is enabled (hint: set PLATFORM_AUTH__PUBLIC_URL env variable)"
@@ -89,7 +89,7 @@ async def authenticate(self, conn: HTTPConnection) -> tuple[AuthCredentials, Bas
         # 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"]:
+        if request.url.path in ["/healthcheck", "/.well-known/agent-card.json"]:
             return None
 
         if not (auth := await self.security(request)):

apps/agentstack-sdk-py/src/agentstack_sdk/server/middleware/platform_auth_backend.py

@@ -21,8 +21,12 @@
 from a2a.types import AgentExtension
 from fastapi import FastAPI
 from fastapi.applications import AppType
+from fastapi.responses import PlainTextResponse
 from httpx import HTTPError, HTTPStatusError
 from pydantic import AnyUrl
+from starlette.authentication import AuthenticationBackend, AuthenticationError
+from starlette.middleware.authentication import AuthenticationMiddleware
+from starlette.requests import HTTPConnection
 from starlette.types import Lifespan
 from tenacity import AsyncRetrying, retry_if_exception_type, stop_after_attempt, wait_exponential
 
@@ -127,7 +131,7 @@ async def serve(
         factory: bool = False,
         h11_max_incomplete_event_size: int | None = None,
         self_registration_client_factory: Callable[[], PlatformClient] | None = None,
-        auth_middleware: Any | None = None,
+        auth_backend: AuthenticationBackend | None = None,
     ) -> None:
         if self.server:
             raise RuntimeError("The server is already running")
@@ -199,6 +203,13 @@ async def _lifespan_fn(app: FastAPI) -> AsyncGenerator[None, None]:
             request_context_builder=request_context_builder,
         )
 
+        if auth_backend:
+
+            def on_error(connection: HTTPConnection, error: AuthenticationError) -> PlainTextResponse:
+                return PlainTextResponse("Unauthorized", status_code=401)
+
+            app.add_middleware(AuthenticationMiddleware, backend=auth_backend, on_error=on_error)
+
         if configure_logger:
             configure_logger_func(log_level)
 

apps/agentstack-sdk-py/src/agentstack_sdk/server/server.py

@@ -17,6 +17,7 @@
 from authlib.oidc.discovery import OpenIDProviderMetadata
 from authlib.oidc.discovery import get_well_known_url as oidc_get_well_known_url
 from fastapi.security import HTTPAuthorizationCredentials
+from kink import inject
 from pydantic import AwareDatetime, BaseModel
 
 from agentstack_server.api.auth.errors import (

apps/agentstack-server/src/agentstack_server/api/auth/auth.py

@@ -26,8 +26,5 @@ def protected_resource_metadata(
 @well_known_router.get("/jwks")
 def jwks():
     config = get_configuration()
-    if not config.auth.jwt_public_key or config.auth.jwt_public_key.get_secret_value() == "dummy":
-        return {"keys": []}
-
     key = JsonWebKey.import_key(config.auth.jwt_public_key.get_secret_value(), {"use": "sig", "alg": "RS256"})
     return {"keys": [key.as_dict()]}

apps/agentstack-server/src/agentstack_server/api/routes/auth.py

@@ -134,6 +134,17 @@ def validate_auth(self):
             raise ValueError("JWT private and public keys must be provided if authentication is enabled")
         return self
 
+    @model_validator(mode="after")
+    def set_default_jwt_keys(self):
+        if self.jwt_private_key.get_secret_value() == "dummy" or self.jwt_public_key.get_secret_value() == "dummy":
+            logger.warning("JWT private and public keys are not set. Generating default keys.")
+            from authlib.jose import JsonWebKey
+
+            key = JsonWebKey.generate_key("RSA", 1024, is_private=True)
+            self.jwt_private_key = Secret(key.as_pem(is_private=True).decode("utf-8"))
+            self.jwt_public_key = Secret(key.as_pem(is_private=False).decode("utf-8"))
+        return self
+
 
 class McpConfiguration(BaseModel):
     gateway_endpoint_url: AnyUrl = AnyUrl("http://forge-svc:4444")