fix(security): stop leaking API key via WebSocket URL query param (#747)

* fix(security): stop leaking API key via WebSocket URL query param

Why: automated security review flagged useAgentEvents.ts sending the API
key as a ?token= query param; browser WS URLs can end up in server
access logs. Browsers can't set custom headers on a WS handshake, so
switch to Sec-WebSocket-Protocol (bearer, <key>) instead.

- auth_middleware.py: websocket_authorized() drops the ?token= query
  path, adds Sec-WebSocket-Protocol extraction; new
  websocket_subprotocol() helper to echo the negotiated protocol back
- ws_bridge.py / api_routes.py: accept() now passes subprotocol=...
- useAgentEvents.ts: buildWsUrl() no longer embeds the key; new
  buildWsProtocols() sends it via WS subprotocols instead
- tests: updated auth assertions for both endpoints; fixed 3
  FakeWebSocket stubs in test_ws_bridge.py whose no-arg accept()
  turned the new subprotocol kwarg into a silent hang instead of a
  clean failure in test_cancelling_handler_cleans_up_child_tasks

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

* fix(security): encode WebSocket auth protocols safely

Use a versioned subprotocol marker and base64url credential token so custom API keys always satisfy browser WebSocket protocol syntax. Parse the ASGI protocol offer unambiguously, accept either valid auth channel, document the query-token migration, and run frontend tests in CI.

---------

Co-authored-by: Claude Sonnet 5 <noreply@anthropic.com>
Co-authored-by: Elliot Slusky <elliot@slusky.com>

.github/clone-stats/badge.json

@@ -1,7 +1,7 @@
 {
   "schemaVersion": 1,
   "label": "Git Clones",
-  "message": "190,252",
+  "message": "191,195",
   "color": "green",
   "namedLogo": "git"
 }

.github/clone-stats/clone-data.json

@@ -1,6 +1,6 @@
 {
-  "total_clones": 190252,
-  "last_updated": "2026-08-14T07:19:51Z",
+  "total_clones": 191195,
+  "last_updated": "2026-08-15T06:32:36Z",
   "daily": {
     "2026-03-27": 2189,
     "2026-03-28": 1874,
@@ -141,6 +141,7 @@
     "2026-08-10": 1060,
     "2026-08-11": 2182,
     "2026-08-12": 641,
-    "2026-08-13": 770
+    "2026-08-13": 770,
+    "2026-08-14": 943
   }
 }

.github/workflows/frontend.yml

@@ -33,6 +33,7 @@ jobs:
           cache-dependency-path: frontend/package-lock.json
 
       - run: npm ci
+      - run: npm test
       - run: npx tsc --noEmit
       - run: npm run build
         env:

CHANGELOG.md

@@ -21,6 +21,19 @@ stack with `mss`/`Pillow` fallbacks on other platforms. Adds the
 `JARVIS_NUM_CTX` environment variable to tune the Ollama context window
 (default `16384`).
 
+### Security
+
+**WebSocket API keys no longer appear in request URLs.** Browser clients now
+send a marked, base64url-encoded credential through
+`Sec-WebSocket-Protocol`; programmatic clients can continue to use an
+`Authorization: Bearer <key>` handshake header. The encoding only makes the
+credential safe for WebSocket protocol syntax and does not encrypt it, so use
+`wss://` for remote connections.
+
+The former `?token=<key>` WebSocket authentication path is no longer accepted.
+Custom browser clients must migrate to the `openjarvis.auth.v1` subprotocol
+format documented in the API server guide.
+
 ## [1.0.2] - 2026-05-24
 
 A patch release that fixes a packaging bug which broke the v1.0.1

docs/deployment/api-server.md

@@ -248,6 +248,33 @@ Show connection status for all configured channels.
 !!! note "Channel endpoints"
     Channel endpoints require `[channel] enabled = true` in your config and platform-specific credentials configured in `[channel.<platform>]` sub-sections. When not configured, `GET /v1/channels` returns an empty list and other channel endpoints return 503.
 
+### WebSocket endpoints
+
+- `WS /v1/chat/stream` streams interactive chat messages.
+- `WS /v1/agents/events` streams agent lifecycle events and accepts an optional
+  `agent_id` query parameter as a filter.
+
+When `OPENJARVIS_API_KEY` or `[server.auth].api_key` is configured,
+programmatic WebSocket clients should send the same
+`Authorization: Bearer <key>` header used by HTTP requests. Browsers cannot set
+that header on a WebSocket handshake, so browser clients must offer exactly
+these two subprotocol values:
+
+1. `openjarvis.auth.v1`
+2. `openjarvis.key.b64url.<encoded-key>`, where `<encoded-key>` is the unpadded
+   base64url encoding of the API key's UTF-8 bytes
+
+The server selects `openjarvis.auth.v1` in its handshake response. The built-in
+frontend handles this encoding automatically. Base64url is only a transport
+encoding, not encryption; use `wss://` for remote connections and treat the
+`Sec-WebSocket-Protocol` request header as credential-bearing.
+
+!!! warning "WebSocket authentication migration"
+    The former `?token=<key>` query parameter is not accepted because request
+    URLs commonly appear in access logs and browser history. Existing custom
+    WebSocket clients must migrate to `Authorization` or the browser
+    subprotocol format above.
+
 ## Streaming via SSE
 
 When `"stream": true` is set in the request, the server returns a `text/event-stream` response using Server-Sent Events (SSE). The response follows the same format as the OpenAI streaming API.
@@ -456,6 +483,11 @@ For production deployments, run OpenJarvis behind a reverse proxy like Nginx or
 ### Nginx
 
 ```nginx
+map $http_upgrade $connection_upgrade {
+    default upgrade;
+    '' close;
+}
+
 server {
     listen 443 ssl;
     server_name jarvis.example.com;
@@ -470,6 +502,11 @@ server {
         proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for;
         proxy_set_header X-Forwarded-Proto $scheme;
 
+        # WebSocket upgrade support
+        proxy_http_version 1.1;
+        proxy_set_header Upgrade $http_upgrade;
+        proxy_set_header Connection $connection_upgrade;
+
         # SSE streaming support
         proxy_buffering off;
         proxy_cache off;

frontend/src/lib/useAgentEvents.test.ts

@@ -1,6 +1,6 @@
 import { afterEach, beforeEach, describe, expect, it } from 'vitest';
 
-import { buildWsUrl } from './useAgentEvents';
+import { buildWsProtocols, buildWsUrl } from './useAgentEvents';
 
 const SETTINGS_KEY = 'openjarvis-settings';
 
@@ -27,7 +27,7 @@ afterEach(() => {
 });
 
 describe('buildWsUrl', () => {
-  it('authenticates agent events with the configured API key', () => {
+  it('builds the agent-events URL without leaking the API key into it', () => {
     localStorage.setItem(
       SETTINGS_KEY,
       JSON.stringify({
@@ -41,7 +41,8 @@ describe('buildWsUrl', () => {
     expect(url.origin).toBe('wss://jarvis.example.com:8443');
     expect(url.pathname).toBe('/v1/agents/events');
     expect(url.searchParams.get('agent_id')).toBe('agent/one');
-    expect(url.searchParams.get('token')).toBe('secret+/=');
+    expect(url.searchParams.has('token')).toBe(false);
+    expect(url.toString()).not.toContain('secret');
   });
 
   it('normalizes a versioned API base without duplicating /v1', () => {
@@ -52,15 +53,33 @@ describe('buildWsUrl', () => {
 
     expect(buildWsUrl()).toBe('ws://192.0.2.10:8000/v1/agents/events');
   });
+});
 
-  it('omits the token for a keyless server', () => {
+describe('buildWsProtocols', () => {
+  it.each([
+    ['secret+/=', 'c2VjcmV0Ky89'],
+    ['bearer', 'YmVhcmVy'],
+    ['sëcret🔑', 'c8OrY3JldPCflJE'],
+  ])('offers a browser-safe encoding of API key %j', (apiKey, encoded) => {
+    localStorage.setItem(SETTINGS_KEY, JSON.stringify({ apiKey }));
+
+    const protocols = buildWsProtocols();
+
+    expect(protocols).toEqual([
+      'openjarvis.auth.v1',
+      `openjarvis.key.b64url.${encoded}`,
+    ]);
+    expect(new Set(protocols).size).toBe(protocols?.length);
+    expect(protocols?.every((value) => /^[!#$%&'*+\-.^_`|~0-9A-Za-z]+$/.test(value)))
+      .toBe(true);
+  });
+
+  it('omits protocols for a keyless server', () => {
     localStorage.setItem(
       SETTINGS_KEY,
       JSON.stringify({ apiUrl: 'http://localhost:8000' }),
     );
 
-    const url = new URL(buildWsUrl('agent-one'));
-
-    expect(url.searchParams.has('token')).toBe(false);
+    expect(buildWsProtocols()).toBeUndefined();
   });
 });

frontend/src/lib/useAgentEvents.ts

@@ -7,18 +7,45 @@ export interface AgentEvent {
   data: Record<string, unknown>;
 }
 
+const WS_AUTH_PROTOCOL = 'openjarvis.auth.v1';
+const WS_KEY_PROTOCOL_PREFIX = 'openjarvis.key.b64url.';
+
+function utf8ToBase64Url(value: string): string {
+  const bytes = new TextEncoder().encode(value);
+  let binary = '';
+  for (const byte of bytes) binary += String.fromCharCode(byte);
+  return btoa(binary)
+    .replace(/\+/g, '-')
+    .replace(/\//g, '_')
+    .replace(/=+$/, '');
+}
+
 export function buildWsUrl(agentId?: string): string {
   const base = getBase();
   const url = new URL('/v1/agents/events', base || window.location.origin);
   url.protocol = url.protocol === 'https:' ? 'wss:' : 'ws:';
 
   if (agentId) url.searchParams.set('agent_id', agentId);
-  const apiKey = getApiKey();
-  if (apiKey) url.searchParams.set('token', apiKey);
 
   return url.toString();
 }
 
+/**
+ * WebSocket auth protocols carrying the API key, if configured. The key is
+ * UTF-8/base64url encoded so every value satisfies browser subprotocol syntax.
+ * This keeps the key out of the request URL and request-line access logs; the
+ * encoding is transport-safe, not encryption.
+ */
+export function buildWsProtocols(): string[] | undefined {
+  const apiKey = getApiKey();
+  return apiKey
+    ? [
+        WS_AUTH_PROTOCOL,
+        `${WS_KEY_PROTOCOL_PREFIX}${utf8ToBase64Url(apiKey)}`,
+      ]
+    : undefined;
+}
+
 /**
  * Subscribe to agent events over WebSocket.
  * Auto-reconnects with backoff when the socket drops.
@@ -43,7 +70,7 @@ export function useAgentEvents(
     const connect = () => {
       if (closed) return;
       try {
-        ws = new WebSocket(buildWsUrl(agentId));
+        ws = new WebSocket(buildWsUrl(agentId), buildWsProtocols());
       } catch {
         schedule();
         return;

src/openjarvis/server/api_routes.py

@@ -667,14 +667,15 @@ 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
+    from openjarvis.server.auth_middleware import authenticate_websocket
 
     expected_key = getattr(websocket.app.state, "api_key", "")
-    if not websocket_authorized(websocket, expected_key):
-        # 1008 = policy violation; reject before accepting the connection.
+    authorized, subprotocol = authenticate_websocket(websocket, expected_key)
+    if not authorized:
+        # Closing before accept rejects the HTTP upgrade request.
         await websocket.close(code=1008)
         return
-    await websocket.accept()
+    await websocket.accept(subprotocol=subprotocol)
     try:
         while True:
             raw = await websocket.receive_text()

src/openjarvis/server/auth_middleware.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import base64
 import logging
 import os
 import secrets
@@ -12,6 +13,19 @@ from starlette.responses import JSONResponse
 
 logger = logging.getLogger(__name__)
 
+_WS_AUTH_PROTOCOL = "openjarvis.auth.v1"
+_WS_KEY_PROTOCOL_PREFIX = "openjarvis.key.b64url."
+
+
+def _api_keys_match(presented: str, expected: str) -> bool:
+    """Compare API keys as bytes so non-ASCII values do not raise ``TypeError``."""
+    try:
+        presented_bytes = presented.encode("utf-8")
+        expected_bytes = expected.encode("utf-8")
+    except UnicodeEncodeError:
+        return False
+    return secrets.compare_digest(presented_bytes, expected_bytes)
+
 
 class AuthMiddleware(BaseHTTPMiddleware):
     """Validates ``Authorization: Bearer <key>`` on ``/v1/*`` and ``/api/*`` routes.
@@ -34,9 +48,7 @@ class AuthMiddleware(BaseHTTPMiddleware):
                 )
             scheme, _, token = auth.partition(" ")
             # Constant-time comparison to avoid leaking the key via timing.
-            if scheme.lower() != "bearer" or not secrets.compare_digest(
-                token, self._api_key
-            ):
+            if scheme.lower() != "bearer" or not _api_keys_match(token, self._api_key):
                 return JSONResponse(
                     {"detail": "Invalid API key"},
                     status_code=401,
@@ -87,6 +99,71 @@ def check_bind_safety(host: str, *, api_key: str) -> None:
         sys.exit(1)
 
 
+def _websocket_key_protocol(api_key: str) -> str:
+    """Encode an API key as a browser-safe WebSocket protocol token."""
+    try:
+        encoded = base64.urlsafe_b64encode(api_key.encode("utf-8")).decode("ascii")
+    except UnicodeEncodeError:
+        return ""
+    return f"{_WS_KEY_PROTOCOL_PREFIX}{encoded.rstrip('=')}"
+
+
+def _offered_websocket_auth(websocket) -> tuple[str, str | None]:  # noqa: ANN001
+    """Return the credential protocol and protocol to negotiate, if well formed.
+
+    ASGI exposes the complete, flattened protocol offer in ``scope``. Reading
+    it avoids ambiguity from repeated ``Sec-WebSocket-Protocol`` header fields.
+    Exactly one stable protocol marker and one encoded credential are required.
+    """
+    offered = getattr(websocket, "scope", {}).get("subprotocols", [])
+    if not isinstance(offered, (list, tuple)) or len(offered) != 2:
+        return "", None
+    if offered.count(_WS_AUTH_PROTOCOL) != 1:
+        return "", None
+    credentials = [
+        protocol
+        for protocol in offered
+        if isinstance(protocol, str) and protocol != _WS_AUTH_PROTOCOL
+    ]
+    if len(credentials) != 1 or not credentials[0].startswith(_WS_KEY_PROTOCOL_PREFIX):
+        return "", None
+    if credentials[0] == _WS_KEY_PROTOCOL_PREFIX:
+        return "", None
+    return credentials[0], _WS_AUTH_PROTOCOL
+
+
+def authenticate_websocket(
+    websocket,
+    expected_key: str,  # noqa: ANN001
+) -> tuple[bool, str | None]:
+    """Authenticate a WebSocket and return its negotiated auth subprotocol.
+
+    Programmatic clients can send ``Authorization: Bearer <key>``. Browser
+    clients, which cannot set that header, offer ``openjarvis.auth.v1`` plus a
+    marked, unpadded base64url encoding of the UTF-8 key. The encoding only
+    makes the credential valid subprotocol syntax; it does not make it secret.
+    """
+    credential_protocol, selected_protocol = _offered_websocket_auth(websocket)
+
+    # Match AuthMiddleware's local, keyless behavior. If a stale client still
+    # offers a well-formed auth protocol, negotiate it so the browser does not
+    # fail an otherwise allowed handshake.
+    if not expected_key:
+        return True, selected_protocol
+
+    auth = websocket.headers.get("authorization", "")
+    scheme, _, header_token = auth.partition(" ")
+    header_valid = scheme.lower() == "bearer" and _api_keys_match(
+        header_token, expected_key
+    )
+
+    expected_protocol = _websocket_key_protocol(expected_key)
+    protocol_valid = bool(credential_protocol and expected_protocol) and (
+        secrets.compare_digest(credential_protocol, expected_protocol)
+    )
+    return header_valid or protocol_valid, selected_protocol
+
+
 def websocket_authorized(websocket, expected_key: str) -> bool:  # noqa: ANN001
     """Return ``True`` if a WebSocket connection presents the expected key.
 
@@ -96,18 +173,8 @@ def websocket_authorized(websocket, expected_key: str) -> bool:  # noqa: ANN001
 
     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.
+    are allowed. See :func:`authenticate_websocket` for the supported
+    credential transports. URL query parameters are deliberately not accepted
+    because request targets commonly appear in access logs and browser history.
     """
-    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)
+    return authenticate_websocket(websocket, expected_key)[0]

src/openjarvis/server/ws_bridge.py

@@ -65,14 +65,15 @@ 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
+        from openjarvis.server.auth_middleware import authenticate_websocket
 
         expected_key = getattr(websocket.app.state, "api_key", "")
-        if not websocket_authorized(websocket, expected_key):
-            # 1008 = policy violation; reject before accepting the connection.
+        authorized, subprotocol = authenticate_websocket(websocket, expected_key)
+        if not authorized:
+            # Closing before accept rejects the HTTP upgrade request.
             await websocket.close(code=1008)
             return
-        await websocket.accept()
+        await websocket.accept(subprotocol=subprotocol)
         # Parse agent_id filter from query string
         agent_id = websocket.query_params.get("agent_id")
         websocket._agent_filter = agent_id  # type: ignore[attr-defined]

tests/server/test_websocket_auth.py

@@ -7,6 +7,7 @@ token themselves in the handshake before accepting the connection.
 
 from __future__ import annotations
 
+import base64
 import json
 from unittest.mock import MagicMock
 
@@ -22,11 +23,20 @@ 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
 
+AUTH_PROTOCOL = "openjarvis.auth.v1"
+KEY_PROTOCOL_PREFIX = "openjarvis.key.b64url."
 
-def _ws(query=None, headers=None):
+
+def _auth_subprotocols(api_key: str) -> list[str]:
+    encoded = base64.urlsafe_b64encode(api_key.encode()).decode().rstrip("=")
+    return [AUTH_PROTOCOL, f"{KEY_PROTOCOL_PREFIX}{encoded}"]
+
+
+def _ws(query=None, headers=None, subprotocols=None):
     stub = MagicMock()
     stub.query_params = query or {}
     stub.headers = headers or {}
+    stub.scope = {"subprotocols": subprotocols or []}
     return stub
 
 
@@ -34,15 +44,57 @@ 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
 
+    @pytest.mark.parametrize(
+        "api_key,credential_protocol",
+        [
+            ("secret+/=", "openjarvis.key.b64url.c2VjcmV0Ky89"),
+            ("bearer", "openjarvis.key.b64url.YmVhcmVy"),
+            ("sëcret🔑", "openjarvis.key.b64url.c8OrY3JldPCflJE"),
+        ],
+    )
+    def test_token_via_subprotocol(self, api_key, credential_protocol):
+        ws = _ws(subprotocols=[AUTH_PROTOCOL, credential_protocol])
+        assert websocket_authorized(ws, api_key) is True
+
+    def test_valid_subprotocol_survives_conflicting_authorization(self):
+        ws = _ws(
+            headers={"authorization": "Bearer proxy-token"},
+            subprotocols=_auth_subprotocols("secret+/="),
+        )
+        assert websocket_authorized(ws, "secret+/=") is True
+
+    def test_valid_authorization_survives_conflicting_subprotocol(self):
+        ws = _ws(
+            headers={"authorization": "Bearer secret"},
+            subprotocols=_auth_subprotocols("wrong"),
+        )
+        assert websocket_authorized(ws, "secret") is True
+
+    def test_query_token_no_longer_accepted(self):
+        # A ?token= query param would leak the key into server access logs;
+        # only headers and Sec-WebSocket-Protocol are honored.
+        assert websocket_authorized(_ws(query={"token": "sek"}), "sek") is False
+
     def test_wrong_token_rejected(self):
-        assert websocket_authorized(_ws(query={"token": "nope"}), "sek") is False
+        ws = _ws(subprotocols=_auth_subprotocols("nope"))
+        assert websocket_authorized(ws, "sek") is False
+
+    @pytest.mark.parametrize(
+        "subprotocols",
+        [
+            [AUTH_PROTOCOL],
+            [AUTH_PROTOCOL, KEY_PROTOCOL_PREFIX],
+            [AUTH_PROTOCOL, AUTH_PROTOCOL],
+            [AUTH_PROTOCOL, f"{KEY_PROTOCOL_PREFIX}c2Vr", "extra"],
+            ["bearer", "sek"],
+        ],
+    )
+    def test_malformed_subprotocol_offer_rejected(self, subprotocols):
+        assert websocket_authorized(_ws(subprotocols=subprotocols), "sek") is False
 
     def test_missing_token_rejected_when_required(self):
         assert websocket_authorized(_ws(), "sek") is False
@@ -75,12 +127,34 @@ class TestChatStreamAuth:
     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:
+            with client.websocket_connect(
+                "/v1/chat/stream", subprotocols=_auth_subprotocols("wrong")
+            ) as ws:
+                ws.receive_text()
+
+    def test_rejected_with_query_token(self):
+        # ?token= is no longer an accepted auth channel (would leak into logs).
+        client = TestClient(_make_app(api_key="secret"))
+        with pytest.raises(WebSocketDisconnect):
+            with client.websocket_connect("/v1/chat/stream?token=secret") as ws:
                 ws.receive_text()
 
     def test_accepted_with_correct_token(self):
+        api_key = "secret+/="
+        client = TestClient(_make_app(api_key=api_key))
+        with client.websocket_connect(
+            "/v1/chat/stream", subprotocols=_auth_subprotocols(api_key)
+        ) as ws:
+            assert ws.accepted_subprotocol == AUTH_PROTOCOL
+            ws.send_text(json.dumps({"message": "hi"}))
+            assert ws.receive_json()["type"] in ("chunk", "done", "error")
+
+    def test_accepted_with_authorization_header(self):
         client = TestClient(_make_app(api_key="secret"))
-        with client.websocket_connect("/v1/chat/stream?token=secret") as ws:
+        with client.websocket_connect(
+            "/v1/chat/stream", headers={"Authorization": "Bearer secret"}
+        ) as ws:
+            assert ws.accepted_subprotocol is None
             ws.send_text(json.dumps({"message": "hi"}))
             assert ws.receive_json()["type"] in ("chunk", "done", "error")
 
@@ -90,6 +164,15 @@ class TestChatStreamAuth:
             ws.send_text(json.dumps({"message": "hi"}))
             assert ws.receive_json()["type"] in ("chunk", "done", "error")
 
+    def test_keyless_server_negotiates_stale_client_auth_protocol(self):
+        client = TestClient(_make_app(api_key=""))
+        with client.websocket_connect(
+            "/v1/chat/stream", subprotocols=_auth_subprotocols("stale")
+        ) as ws:
+            assert ws.accepted_subprotocol == AUTH_PROTOCOL
+            ws.send_text(json.dumps({"message": "hi"}))
+            assert ws.receive_json()["type"] in ("chunk", "done", "error")
+
 
 class TestAgentEventsAuth:
     def _app(self, api_key=""):
@@ -104,12 +187,35 @@ class TestAgentEventsAuth:
             with client.websocket_connect("/v1/agents/events") as ws:
                 ws.receive_text()
 
+    def test_rejected_with_query_token(self):
+        client = TestClient(self._app(api_key="secret"))
+        with pytest.raises(WebSocketDisconnect):
+            with client.websocket_connect("/v1/agents/events?token=secret") as ws:
+                ws.receive_text()
+
     def test_accepted_with_correct_token(self):
+        bus = EventBus()
+        app = FastAPI()
+        api_key = "secret+/="
+        app.state.api_key = api_key
+        app.include_router(create_ws_router(bus))
+        client = TestClient(app)
+        with client.websocket_connect(
+            "/v1/agents/events", subprotocols=_auth_subprotocols(api_key)
+        ) as ws:
+            assert ws.accepted_subprotocol == AUTH_PROTOCOL
+            bus.publish(EventType.AGENT_TICK_START, {"agent_id": "a"})
+            assert ws.receive_json()["data"]["agent_id"] == "a"
+
+    def test_accepted_with_authorization_header(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:
+        with client.websocket_connect(
+            "/v1/agents/events", headers={"Authorization": "Bearer secret"}
+        ) as ws:
+            assert ws.accepted_subprotocol is None
             bus.publish(EventType.AGENT_TICK_START, {"agent_id": "a"})
             assert ws.receive_json()["data"]["agent_id"] == "a"

tests/server/test_ws_bridge.py

@@ -74,7 +74,7 @@ class TestWSBridge:
                 query_params = {}
                 headers = {}
 
-                async def accept(self):
+                async def accept(self, subprotocol=None):
                     pass
 
                 async def receive(self):
@@ -98,7 +98,7 @@ class TestWSBridge:
                     self.receive_count = 0
                     self.disconnect = asyncio.Event()
 
-                async def accept(self):
+                async def accept(self, subprotocol=None):
                     pass
 
                 async def receive(self):
@@ -136,7 +136,7 @@ class TestWSBridge:
                     self.receiving = asyncio.Event()
                     self.receive_cancelled = asyncio.Event()
 
-                async def accept(self):
+                async def accept(self, subprotocol=None):
                     pass
 
                 async def receive(self):