fix: address review comments on first-message auth

- Add comment explaining 10-second timeout rationale
- Add comment explaining why accept() must precede first-message read
- Log specific failure reasons (timeout, malformed JSON, disconnect,
  wrong type, invalid key) instead of a single generic message
- Log info on successful first-message auth for adoption tracking
- Add test for query-param-takes-precedence-over-first-message case

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: Debug Agent

openhands-agent-server/openhands/agent_server/sockets.py

@@ -85,6 +85,9 @@ def _resolve_websocket_session_api_key(
     return None
 
 
+# Give clients 10 seconds to send auth frame after connection opens.
+# This balances security (don't hold connections indefinitely) with
+# accommodating slow networks and client startup time.
 _FIRST_MESSAGE_AUTH_TIMEOUT_SECONDS = 10
 
 
@@ -124,31 +127,61 @@ async def _accept_authenticated_websocket(
         await websocket.close(code=4001, reason="Authentication failed")
         return False
 
-    # First-message auth: accept the connection, then read the first frame.
+    # First-message auth: we must accept() before reading frames because the
+    # WebSocket protocol requires the handshake to complete first.  The legacy
+    # path above can reject *before* accepting (close on an un-accepted socket
+    # sends an HTTP 403-style response), but here we need to read a frame.
     await websocket.accept()
     try:
         raw = await asyncio.wait_for(
             websocket.receive_text(),
             timeout=_FIRST_MESSAGE_AUTH_TIMEOUT_SECONDS,
         )
         data = json.loads(raw)
-    except (TimeoutError, json.JSONDecodeError, WebSocketDisconnect):
-        logger.warning("WebSocket first-message auth failed: bad or missing payload")
+    except TimeoutError:
+        logger.warning(
+            "WebSocket first-message auth failed: timeout waiting for auth frame"
+        )
+        await _safe_close_websocket(
+            websocket, code=4001, reason="Authentication failed"
+        )
+        return False
+    except json.JSONDecodeError:
+        logger.warning("WebSocket first-message auth failed: malformed JSON")
+        await _safe_close_websocket(
+            websocket, code=4001, reason="Authentication failed"
+        )
+        return False
+    except WebSocketDisconnect:
+        logger.warning("WebSocket first-message auth failed: client disconnected")
         await _safe_close_websocket(
             websocket, code=4001, reason="Authentication failed"
         )
         return False
 
-    if (
-        isinstance(data, dict)
-        and data.get("type") == "auth"
-        and data.get("session_api_key") in config.session_api_keys
-    ):
-        return True
+    if not isinstance(data, dict):
+        logger.warning(
+            "WebSocket first-message auth failed: payload is not a JSON object"
+        )
+        await _safe_close_websocket(
+            websocket, code=4001, reason="Authentication failed"
+        )
+        return False
+    if data.get("type") != "auth":
+        logger.warning("WebSocket first-message auth failed: wrong message type")
+        await _safe_close_websocket(
+            websocket, code=4001, reason="Authentication failed"
+        )
+        return False
+    if data.get("session_api_key") not in config.session_api_keys:
+        logger.warning("WebSocket first-message auth failed: invalid API key")
+        await _safe_close_websocket(
+            websocket, code=4001, reason="Authentication failed"
+        )
+        return False
 
-    logger.warning("WebSocket first-message auth failed: invalid key or payload")
-    await _safe_close_websocket(websocket, code=4001, reason="Authentication failed")
-    return False
+    logger.info("WebSocket authenticated via first-message auth")
+    return True
 
 
 @sockets_router.websocket("/events/{conversation_id}")

tests/agent_server/test_websocket_first_message_auth.py

@@ -69,6 +69,25 @@ async def test_legacy_query_param_invalid_key():
     ws.accept.assert_not_called()
 
 
+@pytest.mark.asyncio
+async def test_legacy_query_param_takes_precedence_over_first_message():
+    """When both query param and first-message auth could apply, query param wins."""
+    ws = _make_mock_websocket()
+    ws.receive_text.return_value = json.dumps(
+        {"type": "auth", "session_api_key": "sk-oh-different"}
+    )
+    with patch("openhands.agent_server.sockets.get_default_config") as mock_config:
+        mock_config.return_value.session_api_keys = ["sk-oh-valid"]
+        result = await _accept_authenticated_websocket(
+            ws, session_api_key="sk-oh-valid"
+        )
+
+    assert result is True
+    ws.accept.assert_called_once()
+    # Should NOT read first message because query param already authenticated.
+    ws.receive_text.assert_not_called()
+
+
 # -- Legacy header auth (deprecated) --