fix: stability improvements on event loop

tests/api/test_async_sse_routes.py

@@ -184,6 +184,24 @@ class TestOwnershipGate:
         assert ": connected" in r.text
 
 
+    def test_stream_logs_carry_request_context(self):
+        from docsgpt.core import log_context
+
+        seen: dict = {}
+
+        async def _gen(message_id, last_event_id=None, **kwargs):
+            seen.update(log_context.snapshot())
+            yield ": connected\n\n"
+
+        with patch(_AUTH, return_value={"sub": "alice"}), patch(
+            _OWNS, return_value=True
+        ), patch(_STREAM, _gen):
+            _client().get(f"/api/messages/{VALID_UUID}/events")
+        assert seen["endpoint"] == "message_events"
+        assert seen["user_id"] == "alice"
+        assert seen["activity_id"]
+
+
 # ── Last-Event-ID parsing ───────────────────────────────────────────────────
 
 

tests/api/test_events_routes.py

@@ -197,6 +197,18 @@ class TestStreamShape:
         assert record["poll_timeout"] == events_module.SUBSCRIBE_POLL_INTERVAL_SECONDS
         assert record["liveness_timeout"] == PUBSUB_SOCKET_TIMEOUT_SECONDS
 
+    def test_stream_logs_carry_the_flask_endpoint_name(self):
+        from docsgpt.core import log_context
+
+        seen: dict = {}
+        with patch(_AUTH, return_value=ALICE), patch(
+            _AREDIS, AsyncMock(return_value=_redis())
+        ), patch(_SUBSCRIBE, _subscribe(probe=lambda: seen.update(log_context.snapshot()))):
+            _get()
+        # Same value the Flask route logged, so saved log queries keep matching.
+        assert seen["endpoint"] == "event_stream.stream_events"
+        assert seen["user_id"] == "alice"
+
     def test_keepalive_emitted_on_idle_ticks(self, monkeypatch):
         monkeypatch.setattr(events_module.settings, "SSE_KEEPALIVE_SECONDS", 0)
         with patch(_AUTH, return_value=ALICE), patch(
@@ -474,7 +486,7 @@ class TestReplayBudget:
         assert await events_module._allow_replay(redis, "alice", OLD_CURSOR) is True
 
     async def test_fail_open_when_redis_stalls(self, monkeypatch):
-        monkeypatch.setattr(events_module, "_REDIS_OP_TIMEOUT_SECONDS", 0.05)
+        monkeypatch.setattr(events_module, "ASYNC_REDIS_OP_TIMEOUT_SECONDS", 0.05)
         redis = _redis()
 
         async def _stalled(_key):

tests/asgi_stream.py

@@ -19,6 +19,7 @@ async def stream_then_disconnect(
     *,
     headers: Optional[Mapping[str, str]] = None,
     query_string: bytes = b"",
+    root_path: str = "",
     chunks_before_disconnect: int = 1,
     timeout: float = 5.0,
 ) -> tuple[Optional[int], dict[str, str], list[bytes]]:
@@ -26,8 +27,9 @@ async def stream_then_disconnect(
 
     A response that finishes on its own (a 401, a stream that ends) returns
     without any disconnect. ``chunks_before_disconnect=0`` disconnects as soon
-    as the request is read, before any body. ``timeout`` fails the test
-    instead of hanging it when a stream never ends.
+    as the request is read, before any body. ``root_path`` mounts the app under
+    a prefix, which ASGI servers include in ``path``. ``timeout`` fails the
+    test instead of hanging it when a stream never ends.
 
     Returns:
         tuple: ``(status, headers, body_chunks)`` with lower-cased header names.
@@ -46,9 +48,9 @@ async def stream_then_disconnect(
         "http_version": "1.1",
         "method": "GET",
         "scheme": "http",
-        "path": path,
-        "raw_path": path.encode("utf-8"),
-        "root_path": "",
+        "path": root_path + path,
+        "raw_path": (root_path + path).encode("utf-8"),
+        "root_path": root_path,
         "query_string": query_string,
         "headers": [
             (name.lower().encode("latin-1"), value.encode("latin-1"))

tests/devices/test_session_events_async.py

@@ -16,13 +16,16 @@ import json
 import time
 from types import SimpleNamespace
 
+import anyio
 import pytest
 from flask import Flask
 from starlette.applications import Starlette
+from starlette.requests import Request
 from starlette.testclient import TestClient
 
 from docsgpt.api.devices import auth as auth_module
 from docsgpt.api.devices import session_events as events_module
+from docsgpt.core import log_context
 from docsgpt.core.shutdown import begin_shutdown, reset_shutdown
 from tests.asgi_stream import stream_then_disconnect
 
@@ -95,6 +98,27 @@ def _path(session_id: str) -> str:
     return f"/api/devices/sessions/{session_id}/events"
 
 
+def _scope(session_id: str, headers: dict | None = None) -> dict:
+    """A raw ASGI GET scope for the stream route."""
+    return {
+        "type": "http",
+        "asgi": {"version": "3.0", "spec_version": "2.3"},
+        "http_version": "1.1",
+        "method": "GET",
+        "scheme": "http",
+        "path": _path(session_id),
+        "raw_path": _path(session_id).encode(),
+        "root_path": "",
+        "query_string": b"",
+        "headers": [
+            (name.lower().encode("latin-1"), value.encode("latin-1"))
+            for name, value in (AUTH if headers is None else headers).items()
+        ],
+        "client": ("testclient", 50000),
+        "server": ("testserver", 80),
+    }
+
+
 def _open(session_id: str, headers: dict | None = None):
     return TestClient(_app()).get(_path(session_id), headers=headers if headers is not None else AUTH)
 
@@ -150,6 +174,42 @@ class TestAuth:
         assert _open(ticket).status_code == 200
         assert device.touched == [DEVICE["id"]]
 
+    @pytest.mark.asyncio
+    async def test_unauthenticated_request_body_is_never_read(self, device, broker):
+        # An endless body must not be buffered before the token check rejects it.
+        reads = 0
+        sent: list = []
+
+        async def receive():
+            nonlocal reads
+            reads += 1
+            await anyio.sleep(0)
+            return {"type": "http.request", "body": b"x" * 65536, "more_body": True}
+
+        async def send(message):
+            sent.append(message)
+
+        with anyio.fail_after(2):
+            await _app()(_scope("st_x", headers={}), receive, send)
+        assert sent[0]["status"] == 401
+        assert reads == 0
+
+    def test_stream_logs_carry_request_context(self, device, broker, monkeypatch):
+        b, _ = broker
+        ticket = _ticket(b)
+        seen: dict = {}
+        original = b.next_command_async
+
+        async def _capture(sess, timeout=1.0):
+            seen.update(log_context.snapshot())
+            return await original(sess, timeout=timeout)
+
+        monkeypatch.setattr(b, "next_command_async", _capture)
+        assert _open(ticket).status_code == 200
+        assert seen["endpoint"] == "devices.session_events"
+        assert seen["user_id"] == DEVICE["user_id"]
+        assert seen["activity_id"]
+
 
 # ── stream ──────────────────────────────────────────────────────────────
 
@@ -271,6 +331,27 @@ class TestStreamLifecycle:
         assert b.get_session(ticket) is None
         assert DEVICE["id"] not in b._sessions_by_device
 
+    async def test_ticket_is_claimed_before_the_body_starts(self, device, broker):
+        b, fake = broker
+        ticket = _ticket(b)
+        scope = {**_scope(ticket), "path_params": {"session_id": ticket}}
+
+        async def receive():
+            return {"type": "http.disconnect"}
+
+        response = await events_module.device_session_events(Request(scope, receive))
+        assert response.status_code == 200
+        # A response whose body never runs has still used the ticket up.
+        assert fake.get(f"dev:ticket:{DEVICE['id']}") is None
+        assert b.get_session(ticket) is not None
+
+        async def send(_message):
+            return None
+
+        with anyio.fail_after(2):
+            await response(scope, receive, send)
+        assert b.get_session(ticket) is None
+
 
 # ── machine-key signatures ──────────────────────────────────────────────
 
@@ -320,6 +401,26 @@ class TestSignature:
         # Rejected before the session opened: the ticket is still unclaimed.
         assert fake.get(f"dev:ticket:{DEVICE['id']}") == ticket.encode()
 
+    @pytest.mark.asyncio
+    async def test_signed_path_excludes_root_path(self, device, broker, monkeypatch):
+        # The Flask device routes verify against PATH_INFO, which has the mount
+        # prefix stripped; the stream must check the same path.
+        monkeypatch.setattr(auth_module.settings, "REMOTE_DEVICE_REQUIRE_SIGNATURE", True)
+        b, _ = broker
+        ticket = _ticket(b)
+        fingerprint, headers = _signed_headers(_path(ticket))
+        device.row["machine_pubkey_fingerprint"] = fingerprint
+        status, _, chunks = await stream_then_disconnect(
+            _app(),
+            _path(ticket),
+            headers=headers,
+            root_path="/docsgpt",
+            chunks_before_disconnect=10**6,
+            timeout=3,
+        )
+        assert status == 200
+        assert b"event: invocation" in b"".join(chunks)
+
     def test_missing_signature_rejected(self, device, broker, monkeypatch):
         monkeypatch.setattr(auth_module.settings, "REMOTE_DEVICE_REQUIRE_SIGNATURE", True)
         b, _ = broker

tests/fake_async_redis.py

@@ -3,8 +3,9 @@
 Sorted sets and transactional pipelines behave like Redis, enough for the
 per-connection leases. ``incr``, ``expire``, ``xrange`` and ``xinfo_stream``
 are ``AsyncMock`` attributes that tests configure directly. Set ``fail`` to
-make sorted-set commands and pipelines raise, or ``hang`` to make them never
-return (a dead connection).
+make sorted-set commands and pipelines raise, ``hang`` to make them never
+return (a dead connection), or ``hang_after_execute`` to make a pipeline apply
+its commands and then never reply (a reply lost after the server ran EXEC).
 """
 
 from __future__ import annotations
@@ -37,6 +38,7 @@ class FakeAsyncRedis:
         self.key_ttls: dict[str, int] = {}
         self.fail: Optional[BaseException] = None
         self.hang = False
+        self.hang_after_execute = False
         self.incr = AsyncMock(return_value=1)
         self.expire = AsyncMock(return_value=True)
         self.xrange = AsyncMock(return_value=[])
@@ -146,4 +148,6 @@ class FakePipeline:
         await self._redis._gate()
         results = [fn(*args, **kwargs) for fn, args, kwargs in self._ops]
         self._ops.clear()
+        if self._redis.hang_after_execute:
+            await anyio.sleep_forever()
         return results

tests/test_async_redis.py

@@ -40,6 +40,17 @@ class TestGetAsyncRedisInstance:
         assert kwargs["socket_connect_timeout"] == 2
         assert kwargs["health_check_interval"] == 10
 
+    async def test_reads_bounded_like_the_sync_pubsub_client(self):
+        # A silently dropped connection must fail a BLPOP or XRANGE, not hang it.
+        from docsgpt.cache import PUBSUB_SOCKET_TIMEOUT_SECONDS
+
+        with patch(_FROM_URL) as from_url:
+            await async_redis.get_async_redis_instance()
+        kwargs = from_url.call_args.kwargs
+        assert kwargs["socket_timeout"] == PUBSUB_SOCKET_TIMEOUT_SECONDS
+        assert kwargs["socket_keepalive"] is True
+        assert "socket_keepalive_options" in kwargs
+
     async def test_real_client_pool_honours_setting(self, monkeypatch):
         # from_url connects lazily, so this never opens a socket.
         monkeypatch.setattr(async_redis.settings, "ASYNC_REDIS_MAX_CONNECTIONS", 777)

tests/test_sse_leases.py

@@ -89,12 +89,22 @@ class TestAcquire:
         assert await acquire_stream_lease(redis, "alice") is None
 
     async def test_stalled_redis_fails_open_within_timeout(self, monkeypatch):
-        monkeypatch.setattr(sse_leases, "_REDIS_OP_TIMEOUT_SECONDS", 0.05)
+        monkeypatch.setattr(sse_leases, "ASYNC_REDIS_OP_TIMEOUT_SECONDS", 0.05)
         redis = FakeAsyncRedis()
         redis.hang = True
         with anyio.fail_after(2):
             assert await acquire_stream_lease(redis, "alice") is None
 
+    async def test_timed_out_acquire_leaves_no_orphan_lease(self, monkeypatch):
+        # EXEC ran on the server but its reply never arrived: the stream opens
+        # without a lease, so the member it added must not keep counting.
+        monkeypatch.setattr(sse_leases, "ASYNC_REDIS_OP_TIMEOUT_SECONDS", 0.05)
+        redis = FakeAsyncRedis()
+        redis.hang_after_execute = True
+        with anyio.fail_after(2):
+            assert await acquire_stream_lease(redis, "alice") is None
+        assert redis.zsets == {}
+
 
 @pytest.mark.unit
 @pytest.mark.asyncio
@@ -110,7 +120,7 @@ class TestLeaseLifecycle:
         assert set(redis.zsets[KEY]) == {second.lease_id}
 
     async def test_release_is_bounded_when_redis_stalls(self, monkeypatch):
-        monkeypatch.setattr(sse_leases, "_REDIS_OP_TIMEOUT_SECONDS", 0.05)
+        monkeypatch.setattr(sse_leases, "ASYNC_REDIS_OP_TIMEOUT_SECONDS", 0.05)
         redis = FakeAsyncRedis()
         lease = await acquire_stream_lease(redis, "alice")
         redis.hang = True
@@ -144,6 +154,26 @@ class TestLeaseLifecycle:
         assert long_lived.lease_id not in redis.zsets[KEY]
         assert len(redis.zsets[KEY]) == 7
 
+    async def test_refresh_does_not_revive_a_pruned_lease(self, clock):
+        # Refreshes failed past the TTL, a newer connect pruned the lease and
+        # took the slot. Re-adding it would put the user over the cap.
+        redis = FakeAsyncRedis()
+        lease = await acquire_stream_lease(redis, "alice")
+        redis.zsets[KEY].pop(lease.lease_id)
+        _fill(redis, 8, clock["t"])
+        clock["t"] += 30
+        await lease.refresh_if_due()
+        assert lease.lease_id not in redis.zsets[KEY]
+        assert len(redis.zsets[KEY]) == 8
+
+    async def test_refresh_does_not_undo_a_manual_clear(self, clock):
+        redis = FakeAsyncRedis()
+        lease = await acquire_stream_lease(redis, "alice")
+        redis.zsets.clear()  # redis-cli DEL user:<id>:sse_leases
+        clock["t"] += 30
+        await lease.refresh_if_due()
+        assert redis.zsets == {}
+
     async def test_refresh_failure_is_swallowed(self, clock):
         redis = FakeAsyncRedis()
         lease = await acquire_stream_lease(redis, "alice")