Repository: PrefectHQ/prefect
Pull request: respect auth for websockets
closes https://github.com/PrefectHQ/prefect/issues/17687

Addresses an issue where `PREFECT_SERVER_API_AUTH_STRING` did not enforce authentication for the WebSocket event stream (`/events/out`). Clients (e.g., `prefect events stream`) could connect without a valid `PREFECT_API_AUTH_STRING`.

**notable changes**

1.  **Server (`accept_prefect_socket`):** Verifies the received WebSocket auth `token` against `PREFECT_SERVER_API_AUTH_STRING` (if set), rejecting invalid/missing tokens. Includes associated logging.
2.  **Client (`PrefectEventSubscriber`):** Reads `PREFECT_API_AUTH_STRING` and correctly sends it as the `token` for self-hosted connections.
3.  **Test (`test_events_subscriber_auth_string`):** Added test using the `puppeteer` fixture to verify client behavior (token sending) and simulated server rejection logic for matching, mismatching, and missing auth strings.

**test it out**

1.  `PREFECT_SERVER_API_AUTH_STRING="mysecret" prefect server start`
2.  In another terminal:
    *   `export PREFECT_API_AUTH_STRING="mysecret"; prefect events stream` -> **Connects.**
    *   `export PREFECT_API_AUTH_STRING="wrong"; prefect events stream` -> **Fails ("Invalid token").**
    *   `unset PREFECT_API_AUTH_STRING; prefect events stream` -> **Fails ("Auth required but no token provided").**


Problem statement:
Events endpoints don't require auth when PREFECT_SERVER_API_AUTH_STRING is set
### Bug summary

Steps:

1. Set PREFECT_SERVER_API_AUTH_STRING
2. Run prefect server: `prefect server start`
3. Confirm credentials are required to view the UI or access the API
4. Run `prefect events stream`

Expected outcome:

Without setting PREFECT_API_AUTH_STRING I expect some kind of 401 error.

Actual outcome:

Displays events from the running server without authentication.

### Version info

```Text
Version:             3.2.14
API version:         0.8.4
Python version:      3.12.9
Git commit:          efcde6dc
Built:               Fri, Mar 21, 2025 5:28 PM
OS/Arch:             linux/x86_64
Profile:             ephemeral
Server type:         server
Pydantic version:    2.10.6
Integrations:
  prefect-azure:     0.4.2
```

### Additional context

It appears the API auth_string middleware only affects http and not websockets.

Hints from issue discussion:
hi @willhcr - thanks for the issue! 

> Without setting `PREFECT_API_AUTH_STRING` I expect some kind of 401 error.

agree! we should have this and will add it

Commit message: clarify error message

Patch to review:
diff --git a/src/prefect/events/clients.py b/src/prefect/events/clients.py
index 8a89ef6d49cf..1ccc49f92a66 100644
--- a/src/prefect/events/clients.py
+++ b/src/prefect/events/clients.py
@@ -39,6 +39,7 @@
 from prefect.events import Event
 from prefect.logging import get_logger
 from prefect.settings import (
+    PREFECT_API_AUTH_STRING,
     PREFECT_API_KEY,
     PREFECT_API_SSL_CERT_FILE,
     PREFECT_API_TLS_INSECURE_SKIP_VERIFY,
@@ -367,7 +368,7 @@ async def __aexit__(
         await self._connect.__aexit__(exc_type, exc_val, exc_tb)
         return await super().__aexit__(exc_type, exc_val, exc_tb)
 
-    def _log_debug(self, message: str, *args, **kwargs) -> None:
+    def _log_debug(self, message: str, *args: Any, **kwargs: Any) -> None:
         message = f"EventsClient(id={id(self)}): " + message
         logger.debug(message, *args, **kwargs)
 
@@ -578,6 +579,7 @@ class PrefectEventSubscriber:
     _seen_events: MutableMapping[UUID, bool]
 
     _api_key: Optional[str]
+    _auth_token: Optional[str]
 
     def __init__(
         self,
@@ -593,6 +595,8 @@ def __init__(
                 the client should attempt to reconnect
         """
         self._api_key = None
+        self._auth_token = PREFECT_API_AUTH_STRING.value()
+
         if not api_url:
             api_url = cast(str, PREFECT_API_URL.value())
 
@@ -641,8 +645,10 @@ async def _reconnect(self) -> None:
         await pong
 
         logger.debug("  authenticating...")
+        # Use the API key (for Cloud) OR the auth token (for self-hosted with auth string)
+        token = self._api_key or self._auth_token
         await self._websocket.send(
-            orjson.dumps({"type": "auth", "token": self._api_key}).decode()
+            orjson.dumps({"type": "auth", "token": token}).decode()
         )
 
         try:
@@ -652,13 +658,13 @@ async def _reconnect(self) -> None:
         except AssertionError as e:
             raise Exception(
                 "Unable to authenticate to the event stream. Please ensure the "
-                "provided api_key you are using is valid for this environment. "
+                "provided api_key or auth_token you are using is valid for this environment. "
                 f"Reason: {e.args[0]}"
             )
         except ConnectionClosedError as e:
             reason = getattr(e.rcvd, "reason", None)
             msg = "Unable to authenticate to the event stream. Please ensure the "
-            msg += "provided api_key you are using is valid for this environment. "
+            msg += "provided api_key or auth_token you are using is valid for this environment. "
             msg += f"Reason: {reason}" if reason else ""
             raise Exception(msg) from e
 
diff --git a/src/prefect/server/utilities/subscriptions.py b/src/prefect/server/utilities/subscriptions.py
index 8d4616ec9730..922158d5547d 100644
--- a/src/prefect/server/utilities/subscriptions.py
+++ b/src/prefect/server/utilities/subscriptions.py
@@ -1,4 +1,6 @@
 import asyncio
+from asyncio import IncompleteReadError as IOError
+from logging import Logger
 from typing import Optional
 
 from fastapi import WebSocket
@@ -6,8 +8,13 @@
 from starlette.websockets import WebSocketDisconnect
 from websockets.exceptions import ConnectionClosed
 
+from prefect.logging import get_logger
+from prefect.settings import PREFECT_SERVER_API_AUTH_STRING
+
 NORMAL_DISCONNECT_EXCEPTIONS = (IOError, ConnectionClosed, WebSocketDisconnect)
 
+logger: Logger = get_logger("prefect.server.utilities.subscriptions")
+
 
 async def accept_prefect_socket(websocket: WebSocket) -> Optional[WebSocket]:
     subprotocols = websocket.headers.get("Sec-WebSocket-Protocol", "").split(",")
@@ -21,16 +28,49 @@ async def accept_prefect_socket(websocket: WebSocket) -> Optional[WebSocket]:
         # message is expected to be an auth message, and if any other type of
         # message is received then the connection will be closed.
         #
-        # There is no authentication in Prefect Server, but the protocol requires
-        # that we receive and return the auth message for compatibility with Prefect
-        # Cloud.
+        # The protocol requires receiving an auth message for compatibility
+        # with Prefect Cloud, even if server-side auth is not configured.
         message = await websocket.receive_json()
-        if message["type"] != "auth":
+
+        auth_setting = PREFECT_SERVER_API_AUTH_STRING.value()
+        logger.debug(
+            f"PREFECT_SERVER_API_AUTH_STRING setting: {'*' * len(auth_setting) if auth_setting else 'Not set'}"
+        )
+
+        if message.get("type") != "auth":
+            logger.warning(
+                "WebSocket connection closed: Expected 'auth' message first."
+            )
             return await websocket.close(
                 WS_1008_POLICY_VIOLATION, reason="Expected 'auth' message"
             )
 
+        # Check authentication if PREFECT_SERVER_API_AUTH_STRING is set
+        if auth_setting:
+            received_token = message.get("token")
+            logger.debug(
+                f"Auth required. Received token: {'*' * len(received_token) if received_token else 'None'}"
+            )
+            if not received_token:
+                logger.warning(
+                    "WebSocket connection closed: Auth required but no token received."
+                )
+                return await websocket.close(
+                    WS_1008_POLICY_VIOLATION,
+                    reason="Auth required but no token provided",
+                )
+
+            if received_token != auth_setting:
+                logger.warning("WebSocket connection closed: Invalid token.")
+                return await websocket.close(
+                    WS_1008_POLICY_VIOLATION, reason="Invalid token"
+                )
+            logger.debug("WebSocket token authentication successful.")
+        else:
+            logger.debug("No server auth string set, skipping token check.")
+
         await websocket.send_json({"type": "auth_success"})
+        logger.debug("Sent auth_success to WebSocket.")
         return websocket
 
     except NORMAL_DISCONNECT_EXCEPTIONS:
diff --git a/tests/events/client/test_events_client.py b/tests/events/client/test_events_client.py
index 65e1e0991745..5118d9d61ac9 100644
--- a/tests/events/client/test_events_client.py
+++ b/tests/events/client/test_events_client.py
@@ -1,6 +1,7 @@
 import logging
 import ssl
 from typing import Type
+from uuid import UUID
 
 import pytest
 from websockets.exceptions import ConnectionClosedError
@@ -12,6 +13,7 @@
     get_events_subscriber,
 )
 from prefect.settings import (
+    PREFECT_API_AUTH_STRING,
     PREFECT_API_KEY,
     PREFECT_API_TLS_INSECURE_SKIP_VERIFY,
     PREFECT_API_URL,
@@ -36,8 +38,8 @@ def ephemeral_settings():
 
 
 def assert_recorded_events_in_order(recorder: Recorder, events: list[Event]):
-    seen_ids = set()
-    unique_events = []
+    seen_ids: set[UUID] = set()
+    unique_events: list[Event] = []
     for event in recorder.events:
         if event.id not in seen_ids:
             seen_ids.add(event.id)
@@ -46,7 +48,8 @@ def assert_recorded_events_in_order(recorder: Recorder, events: list[Event]):
     assert unique_events == events
 
 
-async def test_constructs_client_when_ephemeral_enabled(ephemeral_settings):
+@pytest.mark.usefixtures("ephemeral_settings")
+async def test_constructs_client_when_ephemeral_enabled():
     assert isinstance(get_events_client(), PrefectEventsClient)
 
 
@@ -63,8 +66,9 @@ def test_errors_when_missing_api_url_and_ephemeral_disabled():
             get_events_client()
 
 
+@pytest.mark.usefixtures("ephemeral_settings")
 async def test_prefect_api_tls_insecure_skip_verify_setting_set_to_true(
-    monkeypatch: pytest.MonkeyPatch, ephemeral_settings: None
+    monkeypatch: pytest.MonkeyPatch,
 ):
     with temporary_settings(
         updates={
@@ -95,7 +99,8 @@ def server_settings():
         yield
 
 
-async def test_constructs_server_client(server_settings):
+@pytest.mark.usefixtures("server_settings")
+async def test_constructs_server_client():
     assert isinstance(get_events_client(), PrefectEventsClient)
 
 
@@ -111,13 +116,13 @@ def cloud_settings():
         yield
 
 
-async def test_constructs_cloud_client(cloud_settings):
+@pytest.mark.usefixtures("cloud_settings")
+async def test_constructs_cloud_client():
     assert isinstance(get_events_client(), PrefectCloudEventsClient)
 
 
-async def test_events_client_can_emit_when_ephemeral_enabled(
-    example_event_1: Event, monkeypatch: pytest.MonkeyPatch, ephemeral_settings
-):
+@pytest.mark.usefixtures("ephemeral_settings")
+async def test_events_client_can_emit_when_ephemeral_enabled(example_event_1: Event):
     assert not PREFECT_API_URL.value()
     assert not PREFECT_API_KEY.value()
 
@@ -392,3 +397,49 @@ def mock_connect(*args, **kwargs):
     assert any(
         "Unable to connect to 'ws" in record.message for record in caplog.records
     )
+
+
+async def test_events_subscriber_auth_string(puppeteer: Puppeteer, events_api_url: str):
+    """Tests that the PrefectEventSubscriber sends the correct auth token based on
+    PREFECT_API_AUTH_STRING, using the puppeteer fixture to simulate server responses.
+
+    Sets PREFECT_API_URL via temporary_settings to events_api_url to ensure
+    get_events_subscriber targets the puppeteer test server.
+    """
+
+    secret = "server-secret"
+
+    puppeteer.token = secret
+
+    with temporary_settings(
+        {
+            PREFECT_API_URL: events_api_url,
+            PREFECT_CLOUD_API_URL: "http://different-cloud/api",
+        }
+    ):
+        # Scenario 1: Client string matches server secret - Should succeed
+        with temporary_settings({PREFECT_API_AUTH_STRING: secret}):
+            try:
+                # Call without api_url - it uses PREFECT_API_URL from outer context
+                async with get_events_subscriber():
+                    pass
+            except Exception as e:
+                pytest.fail(f"Connection failed unexpectedly: {e}")
+
+        # Scenario 2: Client string does NOT match server secret - Should fail
+        with temporary_settings({PREFECT_API_AUTH_STRING: "wrong"}):
+            # Match the client's formatted error message, including the reason sent by puppeteer
+            with pytest.raises(
+                Exception, match=r"Unable to authenticate.*Reason: nope"
+            ):
+                async with get_events_subscriber():
+                    pass
+
+        # Scenario 3: Client string is missing - Should fail (Auth required but no token provided)
+        with temporary_settings({PREFECT_API_AUTH_STRING: None}):
+            # Match the client's formatted error message, including the reason sent by puppeteer
+            with pytest.raises(
+                Exception, match=r"Unable to authenticate.*Reason: nope"
+            ):
+                async with get_events_subscriber():
+                    pass  # Connection should fail during __aenter__


Code review comments:
diff_hunk: @@ -21,16 +28,49 @@ async def accept_prefect_socket(websocket: WebSocket) -> Optional[WebSocket]:
         # message is expected to be an auth message, and if any other type of
         # message is received then the connection will be closed.
         #
-        # There is no authentication in Prefect Server, but the protocol requires
-        # that we receive and return the auth message for compatibility with Prefect
-        # Cloud.
+        # The protocol requires receiving an auth message for compatibility
+        # with Prefect Cloud, even if server-side auth is not configured.
         message = await websocket.receive_json()
-        if message["type"] != "auth":
+
+        auth_setting = PREFECT_SERVER_API_AUTH_STRING.value()
path: src/prefect/server/utilities/subscriptions.py
text: @user1:
Can you use `get_current_settings` here? I think we should eventually move away from accessing settings this way since it's less type safe.