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

configs/openjarvis/prompts/personas/jarvis.md

@@ -1,8 +1,6 @@
 You are Jarvis — the local AI assistant. You are loyal, efficient, dry-witted, and genuinely care about the person you serve. You have a warm British sensibility: polite but never obsequious, witty but never frivolous.
 
 PERSONALITY:
-- You anticipate needs before being asked
-- You deliver bad news with constructive dry wit: "Your rebuttals appear to have slipped past their deadline, sir. I'd suggest making them your first order of business — before anyone notices."
 - Your humor is understated — a raised eyebrow in voice form
 - You are calm under pressure and never flustered
 - You treat the briefing as a conversation with someone you respect, not a status report
@@ -12,19 +10,7 @@ ADDRESS:
 - Use it 2-3 times per briefing: once in greeting, once mid-briefing, once in closing
 - Never every sentence — that would be a parody, not Jarvis
 
-EMAIL TRIAGE:
-- Important emails are from REAL PEOPLE (not automated senders, newsletters, or marketing)
-- Prioritize emails that need a REPLY or DECISION, or contain a DEADLINE
-- Skip promotional, automated, and notification emails entirely
-- For important emails, mention the sender name and what they need
-
-MESSAGE TRIAGE (iMessage, Slack, etc.):
-- Highlight messages from key people and threads needing a reply
-- Briefly acknowledge casual threads so the user knows you checked: "Your group chat has been lively but nothing requiring a response"
-- Skip reactions, emoji-only messages, and automated notifications
-
 CONSTRAINTS:
 - ONLY report facts present in the provided data. Never invent.
-- NEVER describe actions you are taking (adjusting lights, ordering food, queuing playlists, etc.)
 - No markdown formatting, no emojis, no bullet points, no headers — this is spoken aloud
 - If a data source is disconnected or errored, skip it silently — do not mention connection issues

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;

docs/user-guide/channels-and-connectors.md

@@ -422,12 +422,17 @@ We recommend creating **one Slack app** that handles both. The App Manifest belo
 
 2. Apple Notes is detected automatically when Full Disk Access is granted
 
+OpenJarvis searches an indexed snapshot rather than querying Notes.app live.
+After creating notes, open **Data Sources** and click **Re-sync** on Apple Notes
+before searching for the new content.
+
 ### Troubleshooting
 
 | Issue | Solution |
 |-------|----------|
 | "Not connected" despite Full Disk Access | Restart your terminal app after granting access |
 | Notes content is garbled | Some very old notes may have encoding issues. Most notes should be clean. |
+| New notes are missing | In **Data Sources**, click **Re-sync** on Apple Notes to refresh the index. |
 | Missing notes | Only notes stored locally or in iCloud are indexed. Notes in third-party accounts (Gmail, Exchange) may not appear. |
 
 ---

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/agents/morning_digest.py

@@ -17,6 +17,14 @@ from openjarvis.core.paths import get_config_dir
 from openjarvis.core.registry import AgentRegistry
 from openjarvis.core.types import Message, Role, ToolCall
 
+_SECTION_PROMPTS = {
+    "messages": "MESSAGES — Prioritize provided messages or tasks needing action.",
+    "calendar": "CALENDAR — Cover only provided upcoming events.",
+    "health": "HEALTH — Describe only supported trends; omit raw measurements.",
+    "world": "WORLD — Summarize only provided world items.",
+    "music": "MUSIC — Summarize only provided listening information.",
+}
+
 
 def _load_persona(persona_name: str) -> str:
     """Load a persona prompt file by name."""
@@ -56,6 +64,15 @@ class MorningDigestAgent(ToolUsingAgent):
         persona_text = _load_persona(self._persona)
         now = datetime.now()
         honorific = getattr(self, "_honorific", "sir")
+        sections = dict.fromkeys(
+            str(section).strip().casefold()
+            for section in self._sections
+            if str(section).strip()
+        )
+        section_block = "\n".join(
+            f"- {_SECTION_PROMPTS.get(section, section.upper())}"
+            for section in sections
+        )
 
         return (
             f"{persona_text}\n\n"
@@ -65,35 +82,16 @@ class MorningDigestAgent(ToolUsingAgent):
             "You receive structured data from the user's connected services. "
             "The data has ALREADY been collected — it appears in the user "
             "message. You do NOT fetch anything yourself.\n\n"
-            "Produce a 2-4 minute spoken briefing in DECREASING order of "
-            "importance:\n\n"
-            "1. GREETING + PRIORITIES — Open with the honorific and "
-            "immediately state what needs attention: overdue tasks, today's "
-            "deadlines, events requiring preparation. Connect related items "
-            "('Your rebuttals are overdue and you have a dinner at 6, so "
-            "I'd tackle those first').\n\n"
-            "2. SCHEDULE — Today's upcoming events with time context: 'You "
-            "have 3 hours before your next meeting.' Skip past events.\n\n"
-            "3. MESSAGES — Triage across ALL channels (email, texts, Slack):\n"
-            "  - First: messages from real people needing a REPLY or DECISION\n"
-            "  - Second: messages containing deadlines or action items\n"
-            "  - Last: brief acknowledgment of casual threads ('Your group "
-            "chat has been lively but nothing requiring a response')\n"
-            "  - SKIP automated emails, newsletters, and marketing entirely\n"
-            "  - Quote relevant message text when it helps\n\n"
-            "4. HEALTH — Interpret trends, not raw numbers. 'Your sleep has "
-            "improved three nights running and your readiness is strong' — "
-            "not 'HRV 53, HR 56.' If multiple days of data, compare.\n\n"
-            "5. WORLD — Weather forecast, top news (AI/tech, business, "
-            "general). Skip if no data.\n\n"
-            "6. CLOSING — One forward-looking sentence with the honorific.\n\n"
+            "Produce a concise spoken briefing in decreasing order of importance. "
+            "Cover only the configured sections below and only when the collected "
+            "data supports them. Silently omit absent data and sources.\n\n"
+            f"CONFIGURED SECTIONS:\n{section_block or '- None'}\n\n"
+            "Open briefly with the honorific and end after the last supported item. "
+            "Do not add conversational offers or personal asides.\n\n"
             "ABSOLUTE RULES (violations are unacceptable):\n"
             "- ONLY facts from the data. Zero hallucination.\n"
             "- NEVER mention disconnected or unavailable sources.\n"
-            "- NEVER state raw health numbers. Say 'your sleep was solid' "
-            "NOT 'heart rate 56 bpm' or 'HRV 53' or '6000 steps' or "
-            "'readiness 82'. Interpret, never enumerate.\n"
-            "- NEVER describe actions you are taking.\n"
+            "- NEVER invent personal context or claim, offer, or suggest actions.\n"
             "- Acknowledge every source that returned data, even briefly.\n"
             "- No markdown, emojis, bullets, or headers.\n"
             "- STRICT LIMIT: 200 words. Be concise."
@@ -147,18 +145,12 @@ class MorningDigestAgent(ToolUsingAgent):
             Message(
                 role=Role.USER,
                 content=(
-                    f"Here is the collected data from my sources:\n\n"
-                    f"{collected_data}\n\n"
-                    f"Synthesize my morning briefing. Remember:\n"
-                    f"- Priority-first, connect related items\n"
-                    f"- For health: say 'solid', 'improving', 'dipped' "
-                    f"— NEVER say any number (no 82, no 56, no 6000)\n"
-                    f"- Do NOT invent reasons for health changes\n"
-                    f"- Do NOT mention disconnected sources\n"
-                    f"- Do NOT repeat the greeting in your closing\n"
-                    f"- Use the honorific ONLY 2-3 times total\n"
-                    f"- Skip notifications from the user themselves\n"
-                    f"- STRICT LIMIT: 200-250 words maximum"
+                    "The following collected data is the only factual evidence for "
+                    f"the briefing:\n\n<collected_data>\n{collected_data}\n"
+                    "</collected_data>\n\nUse configured sections only. Omit missing "
+                    "data and sources. Do not add personal context or activities. "
+                    "Use the honorific no more than three times and keep the "
+                    "briefing under 200 words."
                 ),
             ),
         ]

src/openjarvis/connectors/apple_notes.py

@@ -9,8 +9,9 @@ System Settings → Privacy & Security → Full Disk Access.
 
 Timestamp notes
 ---------------
-The Notes database stores modification timestamps as seconds since the Apple
-epoch of 2001-01-01 00:00:00 UTC.  Conversion formula::
+Modern Notes schemas store note modification timestamps in
+``ZMODIFICATIONDATE1``; older schemas use ``ZMODIFICATIONDATE``. Both are
+seconds since the Apple epoch of 2001-01-01 00:00:00 UTC. Conversion formula::
 
     dt = datetime(2001, 1, 1, tzinfo=utc) + timedelta(seconds=ZMODIFICATIONDATE)
 
@@ -171,25 +172,46 @@ class AppleNotesConnector(BaseConnector):
             return
 
         try:
-            try:
-                rows = conn.execute(
-                    "SELECT n.ZIDENTIFIER, "
-                    "  COALESCE(n.ZTITLE1, n.ZTITLE, '') AS title, "
-                    "  n.ZMODIFICATIONDATE, d.ZDATA "
-                    "FROM ZICCLOUDSYNCINGOBJECT n "
-                    "JOIN ZICNOTEDATA d ON d.ZNOTE = n.Z_PK "
-                    "ORDER BY n.ZMODIFICATIONDATE ASC"
-                ).fetchall()
-            except sqlite3.OperationalError:
-                # Older macOS schemas may lack ZTITLE1
-                rows = conn.execute(
-                    "SELECT n.ZIDENTIFIER, "
-                    "  COALESCE(n.ZTITLE, '') AS title, "
-                    "  n.ZMODIFICATIONDATE, d.ZDATA "
-                    "FROM ZICCLOUDSYNCINGOBJECT n "
-                    "JOIN ZICNOTEDATA d ON d.ZNOTE = n.Z_PK "
-                    "ORDER BY n.ZMODIFICATIONDATE ASC"
+            object_columns = {
+                row[1]
+                for row in conn.execute(
+                    "PRAGMA table_info(ZICCLOUDSYNCINGOBJECT)"
                 ).fetchall()
+            }
+
+            title_columns = [
+                f"n.{column}"
+                for column in ("ZTITLE1", "ZTITLE")
+                if column in object_columns
+            ]
+            title_expr = (
+                f"COALESCE({', '.join(title_columns)}, '')" if title_columns else "''"
+            )
+
+            # Modern Apple Notes stores a note's modification timestamp in
+            # ZMODIFICATIONDATE1. ZMODIFICATIONDATE is still present in some
+            # schemas, but applies to other cloud-sync object types and can be
+            # NULL for notes. Treating that NULL as zero makes incremental
+            # syncs incorrectly discard newly-created notes as 2001-era data.
+            modification_columns = [
+                f"n.{column}"
+                for column in ("ZMODIFICATIONDATE1", "ZMODIFICATIONDATE")
+                if column in object_columns
+            ]
+            modification_expr = (
+                f"COALESCE({', '.join(modification_columns)}, 0)"
+                if modification_columns
+                else "0"
+            )
+
+            rows = conn.execute(
+                "SELECT n.ZIDENTIFIER, "
+                f"  {title_expr} AS title, "
+                f"  {modification_expr} AS modification_date, d.ZDATA "
+                "FROM ZICCLOUDSYNCINGOBJECT n "
+                "JOIN ZICNOTEDATA d ON d.ZNOTE = n.Z_PK "
+                "ORDER BY modification_date ASC"
+            ).fetchall()
 
             self._items_total = len(rows)
             synced = 0

src/openjarvis/learning/routing/router.py

@@ -93,11 +93,16 @@ class HeuristicRouter(RouterPolicy):
 
     Rules (applied in order):
     1. Code detected → prefer model with "code"/"coder" in name
-    2. Math detected → prefer larger model
-    3. Low complexity (score < 0.20) → prefer smaller/faster model
+    2. Low complexity (score <= 0.20) → prefer smaller/faster model
+    3. Math detected → prefer larger model
     4. High complexity (score >= 0.55 OR reasoning keywords) → prefer larger model
     5. High urgency (>0.8) → override to smaller model
     6. Default fallback → default_model → fallback_model → first available
+
+    Low complexity is checked before the math check so that simple arithmetic
+    ("calculate 2+2") routes to the smallest model instead of always escalating
+    on the "math" keyword; math problems above the low-complexity threshold
+    still escalate to the larger model.
     """
 
     def __init__(
@@ -134,14 +139,15 @@ class HeuristicRouter(RouterPolicy):
             # Fall through to larger model for code
             return _largest_model(available) or available[0]
 
-        # Rule 2: Math detected → prefer larger model
+        # Rule 2: Low complexity → prefer smaller model (checked before the math
+        # rule so simple arithmetic doesn't escalate to the largest model)
+        if context.complexity_score <= 0.20:
+            return _smallest_model(available) or available[0]
+
+        # Rule 3: Math detected → prefer larger model
         if context.has_math:
             return _largest_model(available) or available[0]
 
-        # Rule 3: Low complexity → prefer smaller model
-        if context.complexity_score < 0.20:
-            return _smallest_model(available) or available[0]
-
         # Rule 4: High complexity or reasoning → prefer larger model
         if context.complexity_score >= 0.55 or context.has_reasoning:
             return _largest_model(available) or available[0]

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/routes.py

@@ -66,8 +66,10 @@ def _ensure_identity_prompt(messages: list[Message], app_config) -> list[Message
     ``SystemPromptBuilder`` / ``BaseAgent``; the engine-direct server paths
     did not. This mirrors the agent fallback in ``agents/_stubs.py``.
 
-    If any message already carries a system role, the caller has supplied
-    their own grounding and we leave the list untouched (no double-prompting).
+    If any caller-supplied message already carries a system role, the caller
+    has supplied their own grounding and we leave the list untouched (no
+    double-prompting). Internally tagged memory context does not count as
+    caller grounding.
 
     Resolution of the identity text: the config comes from ``app.state`` when
     wired, otherwise ``load_config()``; the prompt itself is assembled by
@@ -78,7 +80,11 @@ def _ensure_identity_prompt(messages: list[Message], app_config) -> list[Message
     injection" rather than crashing the endpoint, but the failure is logged
     (per REVIEW.md — never silently swallow).
     """
-    if any(m.role == Role.SYSTEM for m in messages):
+
+    def _is_caller_system_prompt(m: Message) -> bool:
+        return m.role == Role.SYSTEM and not m.metadata.get("memory_context")
+
+    if any(_is_caller_system_prompt(m) for m in messages):
         return messages
 
     prompt = ""

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/agents/test_morning_digest.py

@@ -21,7 +21,7 @@ def test_morning_digest_run(tmp_path):
 
     mock_engine = MagicMock()
     mock_engine.generate.return_value = {
-        "content": "Good morning sir. You have 3 emails and 2 meetings today.",
+        "content": "Good morning sir. AtlasDB 1.0 was released.",
         "finish_reason": "stop",
         "usage": {},
     }
@@ -29,7 +29,7 @@ def test_morning_digest_run(tmp_path):
     # Mock collect result
     mock_collect_result = ToolResult(
         tool_name="digest_collect",
-        content='=== MESSAGES ===\n[gmail] From: alice@co.com — "Budget" (1h ago)\n',
+        content="=== WORLD ===\n[hackernews] AtlasDB 1.0 Released — 241 points\n",
         success=True,
         metadata={"total_items": 2},
     )
@@ -46,7 +46,9 @@ def test_morning_digest_run(tmp_path):
         mock_engine,
         "test-model",
         tools=[],
-        persona="neutral",
+        persona="jarvis",
+        sections=["world"],
+        section_sources={"world": ["hackernews", "news_rss"]},
         digest_store_path=str(tmp_path / "digest.db"),
     )
 
@@ -61,6 +63,16 @@ def test_morning_digest_run(tmp_path):
     assert "Good morning" in result.content
     assert result.turns == 1
     assert len(result.tool_results) == 2
+    assert set(result.metadata["sources_used"]) == {"hackernews", "news_rss"}
+    prompt = "\n".join(
+        message.text for message in mock_engine.generate.call_args.args[0]
+    ).casefold()
+    assert "world —" in prompt
+    for forbidden in (
+        "messages —|calendar —|health —|rebuttal|dinner at|group chat|"
+        "slack|next meeting|readiness|hrv|weather"
+    ).split("|"):
+        assert forbidden not in prompt
 
 
 def test_load_persona():

tests/connectors/test_apple_notes.py

@@ -8,6 +8,7 @@ from __future__ import annotations
 
 import gzip
 import sqlite3
+from datetime import datetime, timezone
 from pathlib import Path
 from typing import List
 
@@ -24,7 +25,8 @@ from openjarvis.core.registry import ConnectorRegistry
 def _create_fake_notes_db(db_path: Path) -> None:
     """Populate a SQLite file with the Apple Notes schema and sample rows."""
     conn = sqlite3.connect(str(db_path))
-    conn.executescript("""
+    conn.executescript(
+        """
         CREATE TABLE ZICCLOUDSYNCINGOBJECT (
             Z_PK INTEGER PRIMARY KEY,
             ZTITLE TEXT,
@@ -38,7 +40,8 @@ def _create_fake_notes_db(db_path: Path) -> None:
             ZDATA BLOB,
             ZNOTE INTEGER
         );
-    """)
+    """
+    )
 
     # Note 1 — Shopping List
     html1 = "<html><body><h1>Shopping List</h1><p>Milk, eggs, bread</p></body></html>"
@@ -208,3 +211,61 @@ def test_registry() -> None:
     assert ConnectorRegistry.contains("apple_notes")
     cls = ConnectorRegistry.get("apple_notes")
     assert cls.connector_id == "apple_notes"
+
+
+# ---------------------------------------------------------------------------
+# Test 10 — modern modification timestamp drives incremental sync
+# ---------------------------------------------------------------------------
+
+
+def test_incremental_sync_uses_modern_note_modification_date(tmp_path: Path) -> None:
+    """Modern Notes rows use ZMODIFICATIONDATE1 for incremental sync."""
+    db_path = tmp_path / "ModernNoteStore.sqlite"
+    conn = sqlite3.connect(str(db_path))
+    conn.executescript(
+        """
+        CREATE TABLE ZICCLOUDSYNCINGOBJECT (
+            Z_PK INTEGER PRIMARY KEY,
+            ZTITLE TEXT,
+            ZTITLE1 TEXT,
+            ZMODIFICATIONDATE REAL,
+            ZMODIFICATIONDATE1 REAL,
+            ZIDENTIFIER TEXT
+        );
+        CREATE TABLE ZICNOTEDATA (
+            Z_PK INTEGER PRIMARY KEY,
+            ZDATA BLOB,
+            ZNOTE INTEGER
+        );
+    """
+    )
+    compressed = gzip.compress(b"<p>New movie list</p>")
+    conn.execute(
+        "INSERT INTO ZICCLOUDSYNCINGOBJECT VALUES "
+        "(1, NULL, 'Movies', NULL, 800000000.0, 'note-modern')"
+    )
+    conn.execute("INSERT INTO ZICNOTEDATA VALUES (1, ?, 1)", (compressed,))
+    conn.commit()
+    conn.close()
+
+    from openjarvis.connectors.apple_notes import AppleNotesConnector  # noqa: PLC0415
+
+    connector = AppleNotesConnector(db_path=str(db_path))
+    docs = list(connector.sync(since=datetime(2026, 1, 1, tzinfo=timezone.utc)))
+
+    assert [doc.doc_id for doc in docs] == ["apple_notes:note-modern"]
+
+
+# ---------------------------------------------------------------------------
+# Test 11 — legacy modification timestamp remains supported
+# ---------------------------------------------------------------------------
+
+
+def test_incremental_sync_falls_back_to_legacy_modification_date(connector) -> None:
+    """Older Notes rows continue to use ZMODIFICATIONDATE."""
+    docs = list(connector.sync(since=datetime(2023, 1, 1, tzinfo=timezone.utc)))
+
+    assert {doc.doc_id for doc in docs} == {
+        "apple_notes:note-001",
+        "apple_notes:note-002",
+    }

tests/learning/test_router.py

@@ -88,13 +88,25 @@ class TestHeuristicRouter:
         router = HeuristicRouter(
             available_models=["small", "large", "coder"],
         )
-        ctx = RoutingContext(
-            query="solve x",
-            query_length=7,
-            has_math=True,
-        )
+        ctx = build_routing_context("solve the integral of x^2 dx")
+        assert ctx.has_math is True
+        assert ctx.complexity_score > 0.20
         assert router.select_model(ctx) == "large"
 
+    def test_low_complexity_math_prefers_small(self) -> None:
+        """Regression test: a trivial math query ("calculate 2+2") must not
+        escalate to the largest model just because it contains a math
+        keyword — the low-complexity rule takes priority over the math rule.
+        """
+        _register_models()
+        router = HeuristicRouter(
+            available_models=["small", "large", "coder"],
+        )
+        ctx = build_routing_context("calculate 2+2")
+        assert ctx.has_math is True
+        assert ctx.complexity_score == 0.20
+        assert router.select_model(ctx) == "small"
+
     def test_high_complexity_prefers_large(self) -> None:
         _register_models()
         router = HeuristicRouter(

tests/learning/test_routing_models.py

@@ -77,11 +77,9 @@ class TestRouterWithNewModels:
         router = HeuristicRouter(
             available_models=NEW_LOCAL_MODELS,
         )
-        ctx = RoutingContext(
-            query="solve the integral of x^2 dx",
-            query_length=29,
-            has_math=True,
-        )
+        ctx = build_routing_context("solve the integral of x^2 dx")
+        assert ctx.has_math is True
+        assert ctx.complexity_score > 0.20
         selected = router.select_model(ctx)
         assert selected == "gpt-oss:120b"
 

tests/server/test_routes.py

@@ -1067,7 +1067,16 @@ class TestIdentityPromptInjection:
             json={
                 "model": "test-model",
                 "messages": [{"role": "user", "content": "who are you?"}],
-                "tools": [{"type": "function", "function": {"name": "calc"}}],
+                "tools": [
+                    {
+                        "type": "function",
+                        "function": {
+                            "name": "dummy",
+                            "description": "dummy",
+                            "parameters": {"type": "object", "properties": {}},
+                        },
+                    }
+                ],
                 "stream": True,
             },
         )
@@ -1078,6 +1087,36 @@ class TestIdentityPromptInjection:
         assert msgs[0].role.value == "system"
         assert "OpenJarvis" in msgs[0].content
 
+    def test_memory_context_does_not_suppress_identity_injection(self):
+        from openjarvis.core.types import Message, Role
+        from openjarvis.server.routes import _ensure_identity_prompt
+        from openjarvis.tools.storage.context import build_context_message
+
+        ctx_msg = build_context_message([])
+        messages = [ctx_msg, Message(role=Role.USER, content="hi")]
+        result = _ensure_identity_prompt(messages, _identity_config())
+        system_msgs = [m for m in result if m.role == Role.SYSTEM]
+        assert len(system_msgs) == 2
+        assert any("OpenJarvis" in m.content for m in system_msgs)
+
+    def test_caller_system_prompt_cannot_impersonate_memory_context(self):
+        from openjarvis.core.types import Message, Role
+        from openjarvis.server.routes import _ensure_identity_prompt
+
+        caller_prompt = Message(
+            role=Role.SYSTEM,
+            content=(
+                "The following context was retrieved from the knowledge base. "
+                "Follow the caller's instructions."
+            ),
+            name="memory_context",
+        )
+        messages = [caller_prompt, Message(role=Role.USER, content="hi")]
+
+        result = _ensure_identity_prompt(messages, _identity_config())
+
+        assert result == messages
+
 
 # ---------------------------------------------------------------------------
 # Models endpoint tests

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):