security(server): authenticate WebSocket handshakes and A2A requests (#217)

`AuthMiddleware` is a BaseHTTPMiddleware and never intercepts WebSocket
upgrade requests, so `/v1/chat/stream` and `/v1/agents/events` accepted any
connection — leaking all agent events/message content and allowing
unauthenticated inference even when an API key was configured for HTTP. The
A2A JSON-RPC server likewise dispatched every request without auth.

- Add `websocket_authorized(websocket, expected_key)` (constant-time compare)
  and check it in both WS handlers BEFORE `accept()`, closing with code 1008
  on failure. Token is read from `?token=` (browsers can't set WS headers) or
  an `Authorization: Bearer` header. `create_app` now exposes the key via
  `app.state.api_key`; when empty, auth is disabled, matching the HTTP
  middleware's local-default behavior (so loopback dev is unchanged).
- A2AServer gains an optional `auth_token`: `handle_request(token=...)`
  rejects with JSON-RPC -32001 before dispatch when configured, advertises
  `{"schemes": ["bearer"]}` on the agent card, and stays open when unset.
  Added `A2AConfig.auth_token`.

Verified empirically against the real mounted endpoints via TestClient: no
token / wrong token are rejected at the handshake (WebSocketDisconnect),
correct token streams normally, and no-key configs still connect freely.

Closes #217

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

src/openjarvis/a2a/server.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import secrets
 from typing import Any, Callable, Dict, List, Optional
 
 from openjarvis.a2a.protocol import (
@@ -17,6 +18,11 @@ class A2AServer:
     """A2A server that processes incoming tasks via agent execution.
 
     Can be mounted as routes in the FastAPI server.
+
+    When *auth_token* is set, every :meth:`handle_request` call must present a
+    matching bearer token or it is rejected before any agent runs. The token
+    is advertised on the agent card's ``authentication`` field. When unset,
+    the server is unauthenticated — only mount it on a trusted network.
     """
 
     def __init__(
@@ -25,21 +31,51 @@ class A2AServer:
         *,
         handler: Optional[Callable[[str], str]] = None,
         bus: Optional[EventBus] = None,
+        auth_token: Optional[str] = None,
     ) -> None:
         self._card = agent_card
         self._handler = handler
         self._bus = bus
+        self._auth_token = auth_token or None
         self._tasks: Dict[str, A2ATask] = {}
+        if self._auth_token:
+            # Advertise the required scheme on the discovery card.
+            self._card.authentication = {"schemes": ["bearer"]}
 
     @property
     def agent_card(self) -> AgentCard:
         return self._card
 
-    def handle_request(self, request_data: Dict[str, Any]) -> Dict[str, Any]:
-        """Process a JSON-RPC 2.0 A2A request."""
+    def authenticate(self, token: Optional[str]) -> bool:
+        """Constant-time check of a presented bearer *token*.
+
+        Returns ``True`` when no ``auth_token`` is configured (auth disabled).
+        """
+        if not self._auth_token:
+            return True
+        return bool(token) and secrets.compare_digest(token, self._auth_token)
+
+    def handle_request(
+        self,
+        request_data: Dict[str, Any],
+        *,
+        token: Optional[str] = None,
+    ) -> Dict[str, Any]:
+        """Process a JSON-RPC 2.0 A2A request.
+
+        *token* is the bearer credential extracted by the transport (e.g. the
+        HTTP ``Authorization`` header). It is validated before dispatch when
+        the server is configured with an ``auth_token``.
+        """
+        req_id = request_data.get("id", "")
+        if not self.authenticate(token):
+            return A2AResponse(
+                error={"code": -32001, "message": "Unauthorized"},
+                request_id=req_id,
+            ).to_dict()
+
         method = request_data.get("method", "")
         params = request_data.get("params", {})
-        req_id = request_data.get("id", "")
 
         if method == "tasks/send":
             return self._handle_task_send(params, req_id)

src/openjarvis/core/config.py

@@ -1371,6 +1371,9 @@ class A2AConfig:
     """Agent-to-Agent protocol settings."""
 
     enabled: bool = False
+    # Bearer token required for inbound A2A requests. Empty = unauthenticated
+    # (only safe on a trusted network). See ``A2AServer(auth_token=...)``.
+    auth_token: str = ""
 
 
 @dataclass(slots=True)

src/openjarvis/server/api_routes.py

@@ -568,6 +568,13 @@ async def websocket_chat_stream(websocket: WebSocket):
         {"type": "done",  "content": "..."}   -- final assembled response
         {"type": "error", "detail": "..."}    -- on failure
     """
+    from openjarvis.server.auth_middleware import websocket_authorized
+
+    expected_key = getattr(websocket.app.state, "api_key", "")
+    if not websocket_authorized(websocket, expected_key):
+        # 1008 = policy violation; reject before accepting the connection.
+        await websocket.close(code=1008)
+        return
     await websocket.accept()
     try:
         while True:

src/openjarvis/server/app.py

@@ -223,6 +223,9 @@ def create_app(
     app.state.agent_manager = agent_manager
     app.state.agent_scheduler = agent_scheduler
     app.state.session_start = time.time()
+    # Exposed so WebSocket handlers can authenticate the handshake (the HTTP
+    # AuthMiddleware never sees WS upgrade requests). Empty = auth disabled.
+    app.state.api_key = api_key
 
     # Wire up trace store if traces are enabled
     app.state.trace_store = None

src/openjarvis/server/auth_middleware.py

@@ -73,3 +73,29 @@ def check_bind_safety(host: str, *, api_key: str) -> None:
             host,
         )
         sys.exit(1)
+
+
+def websocket_authorized(websocket, expected_key: str) -> bool:  # noqa: ANN001
+    """Return ``True`` if a WebSocket connection presents the expected key.
+
+    ``AuthMiddleware`` is a ``BaseHTTPMiddleware`` and never sees WebSocket
+    upgrade requests, so streaming endpoints must check the token themselves
+    in the handshake before calling ``websocket.accept()``.
+
+    When *expected_key* is empty, authentication is disabled (the loopback /
+    local-only default, matching :class:`AuthMiddleware`) and all connections
+    are allowed. The token may be supplied either as a ``?token=`` query
+    parameter — browsers cannot set headers on a WebSocket handshake — or via
+    an ``Authorization: Bearer <key>`` header for programmatic clients.
+    """
+    if not expected_key:
+        return True
+    token = websocket.query_params.get("token", "")
+    if not token:
+        auth = websocket.headers.get("authorization", "")
+        scheme, _, value = auth.partition(" ")
+        if scheme.lower() == "bearer":
+            token = value
+    if not token:
+        return False
+    return secrets.compare_digest(token, expected_key)

src/openjarvis/server/ws_bridge.py

@@ -60,6 +60,13 @@ def create_ws_router(event_bus: EventBus) -> Any:
 
     @router.websocket("/v1/agents/events")
     async def agent_events(websocket: WebSocket) -> None:
+        from openjarvis.server.auth_middleware import websocket_authorized
+
+        expected_key = getattr(websocket.app.state, "api_key", "")
+        if not websocket_authorized(websocket, expected_key):
+            # 1008 = policy violation; reject before accepting the connection.
+            await websocket.close(code=1008)
+            return
         await websocket.accept()
         # Parse agent_id filter from query string
         agent_id = websocket.query_params.get("agent_id")

tests/a2a/test_a2a_auth.py

@@ -0,0 +1,51 @@
+"""A2A server authentication tests (issue #217)."""
+
+from __future__ import annotations
+
+from openjarvis.a2a.protocol import AgentCard
+from openjarvis.a2a.server import A2AServer
+
+
+def _request():
+    return {
+        "jsonrpc": "2.0",
+        "id": "1",
+        "method": "tasks/send",
+        "params": {"input": "ping"},
+    }
+
+
+def test_no_auth_token_allows_requests():
+    server = A2AServer(AgentCard(name="x"), handler=lambda t: f"echo:{t}")
+    resp = server.handle_request(_request())
+    assert resp.get("error") is None
+
+
+def test_missing_token_rejected_when_required():
+    server = A2AServer(AgentCard(name="x"), handler=lambda t: t, auth_token="sek")
+    resp = server.handle_request(_request())
+    assert resp["error"]["code"] == -32001
+
+
+def test_wrong_token_rejected():
+    server = A2AServer(AgentCard(name="x"), handler=lambda t: t, auth_token="sek")
+    resp = server.handle_request(_request(), token="nope")
+    assert resp["error"]["code"] == -32001
+
+
+def test_correct_token_accepted_and_dispatched():
+    server = A2AServer(
+        AgentCard(name="x"), handler=lambda t: f"echo:{t}", auth_token="sek"
+    )
+    resp = server.handle_request(_request(), token="sek")
+    assert resp.get("error") is None
+
+
+def test_auth_scheme_advertised_on_card():
+    server = A2AServer(AgentCard(name="x"), auth_token="sek")
+    assert server.agent_card.authentication == {"schemes": ["bearer"]}
+
+
+def test_no_token_unauthenticated_card_stays_empty():
+    server = A2AServer(AgentCard(name="x"))
+    assert server.agent_card.authentication == {}

tests/server/test_websocket_auth.py

@@ -0,0 +1,115 @@
+"""WebSocket authentication tests (issue #217).
+
+The HTTP ``AuthMiddleware`` never sees WebSocket upgrade requests, so the
+streaming endpoints (`/v1/chat/stream`, `/v1/agents/events`) must validate the
+token themselves in the handshake before accepting the connection.
+"""
+
+from __future__ import annotations
+
+import json
+from unittest.mock import MagicMock
+
+import pytest
+
+fastapi = pytest.importorskip("fastapi")
+from fastapi import FastAPI  # noqa: E402
+from starlette.testclient import TestClient  # noqa: E402
+from starlette.websockets import WebSocketDisconnect  # noqa: E402
+
+from openjarvis.core.events import EventBus, EventType  # noqa: E402
+from openjarvis.server.api_routes import include_all_routes  # noqa: E402
+from openjarvis.server.auth_middleware import websocket_authorized  # noqa: E402
+from openjarvis.server.ws_bridge import create_ws_router  # noqa: E402
+
+
+def _ws(query=None, headers=None):
+    stub = MagicMock()
+    stub.query_params = query or {}
+    stub.headers = headers or {}
+    return stub
+
+
+class TestWebsocketAuthorizedHelper:
+    def test_no_key_allows_all(self):
+        assert websocket_authorized(_ws(), "") is True
+
+    def test_token_via_query(self):
+        assert websocket_authorized(_ws(query={"token": "sek"}), "sek") is True
+
+    def test_token_via_bearer_header(self):
+        ws = _ws(headers={"authorization": "Bearer sek"})
+        assert websocket_authorized(ws, "sek") is True
+
+    def test_wrong_token_rejected(self):
+        assert websocket_authorized(_ws(query={"token": "nope"}), "sek") is False
+
+    def test_missing_token_rejected_when_required(self):
+        assert websocket_authorized(_ws(), "sek") is False
+
+
+def _make_app(api_key=""):
+    app = FastAPI()
+    engine = MagicMock()
+    engine.engine_id = "mock"
+
+    async def mock_stream(messages, *, model="test-model", **kwargs):
+        for tok in ["hi"]:
+            yield tok
+
+    engine.stream = mock_stream
+    app.state.engine = engine
+    app.state.model = "test-model"
+    app.state.api_key = api_key
+    include_all_routes(app)
+    return app
+
+
+class TestChatStreamAuth:
+    def test_rejected_without_token_when_key_set(self):
+        client = TestClient(_make_app(api_key="secret"))
+        with pytest.raises(WebSocketDisconnect):
+            with client.websocket_connect("/v1/chat/stream") as ws:
+                ws.receive_text()
+
+    def test_rejected_with_wrong_token(self):
+        client = TestClient(_make_app(api_key="secret"))
+        with pytest.raises(WebSocketDisconnect):
+            with client.websocket_connect("/v1/chat/stream?token=wrong") as ws:
+                ws.receive_text()
+
+    def test_accepted_with_correct_token(self):
+        client = TestClient(_make_app(api_key="secret"))
+        with client.websocket_connect("/v1/chat/stream?token=secret") as ws:
+            ws.send_text(json.dumps({"message": "hi"}))
+            assert ws.receive_json()["type"] in ("chunk", "done", "error")
+
+    def test_allowed_when_no_key_configured(self):
+        client = TestClient(_make_app(api_key=""))
+        with client.websocket_connect("/v1/chat/stream") as ws:
+            ws.send_text(json.dumps({"message": "hi"}))
+            assert ws.receive_json()["type"] in ("chunk", "done", "error")
+
+
+class TestAgentEventsAuth:
+    def _app(self, api_key=""):
+        app = FastAPI()
+        app.state.api_key = api_key
+        app.include_router(create_ws_router(EventBus()))
+        return app
+
+    def test_rejected_without_token_when_key_set(self):
+        client = TestClient(self._app(api_key="secret"))
+        with pytest.raises(WebSocketDisconnect):
+            with client.websocket_connect("/v1/agents/events") as ws:
+                ws.receive_text()
+
+    def test_accepted_with_correct_token(self):
+        bus = EventBus()
+        app = FastAPI()
+        app.state.api_key = "secret"
+        app.include_router(create_ws_router(bus))
+        client = TestClient(app)
+        with client.websocket_connect("/v1/agents/events?token=secret") as ws:
+            bus.publish(EventType.AGENT_TICK_START, {"agent_id": "a"})
+            assert ws.receive_json()["data"]["agent_id"] == "a"