feat: improve sse

application/cache.py

@@ -1,6 +1,7 @@
 import hashlib
 import json
 import logging
+import socket as _socket
 import time
 from threading import Lock
 
@@ -11,6 +12,12 @@ from application.utils import get_hash
 
 logger = logging.getLogger(__name__)
 
+# Upper bound on any single blocking read by a pub/sub subscriber. Must stay
+# comfortably above Topic.subscribe's poll_timeout (1 s) — get_message's idle
+# wait polls with select() and never trips socket_timeout, but a half-open
+# connection's pending read (e.g. the health-check PONG) does.
+PUBSUB_SOCKET_TIMEOUT_SECONDS = 10
+
 
 def _cache_default(value):
     # Image attachments arrive inline as bytes (see GoogleLLM.prepare_messages_with_attachments);
@@ -51,6 +58,64 @@ def get_redis_instance():
     return _redis_instance
 
 
+_pubsub_redis_instance = None
+_pubsub_redis_creation_failed = False
+
+
+def _tcp_keepalive_options():
+    """Kernel keepalive knobs for long-lived, mostly-idle pub/sub sockets.
+
+    Probing well inside NAT/IPVS idle-expiry windows (Docker Swarm's IPVS
+    expires idle flows after ~15 min) keeps the flow-table entry alive and
+    lets the kernel surface a dead peer instead of leaving the socket
+    half-open. The constants are Linux-specific, so build the dict from
+    whatever this platform exposes.
+    """
+    options = {}
+    for name, value in (("TCP_KEEPIDLE", 300), ("TCP_KEEPINTVL", 60), ("TCP_KEEPCNT", 3)):
+        const = getattr(_socket, name, None)
+        if const is not None:
+            options[const] = value
+    return options
+
+
+def get_pubsub_redis_instance():
+    """Redis client dedicated to pub/sub subscribers.
+
+    Separate from ``get_redis_instance`` because subscribers hold a socket
+    open for the life of an SSE connection. Without ``socket_timeout``, a
+    connection silently dropped by NAT/IPVS blocks ``pubsub.get_message``
+    forever — including the ``health_check_interval`` PONG read — pinning
+    the subscriber's WSGI thread until the worker restarts. Bounding every
+    read lets a dead subscriber fail within seconds and release its thread.
+
+    Returns:
+        A shared ``redis.Redis`` client, or ``None`` if Redis is
+        unavailable or ``CACHE_REDIS_URL`` is invalid.
+    """
+    global _pubsub_redis_instance, _pubsub_redis_creation_failed
+    if _pubsub_redis_instance is None and not _pubsub_redis_creation_failed:
+        with _instance_lock:
+            if _pubsub_redis_instance is None and not _pubsub_redis_creation_failed:
+                try:
+                    _pubsub_redis_instance = redis.Redis.from_url(
+                        settings.CACHE_REDIS_URL,
+                        socket_connect_timeout=2,
+                        socket_timeout=PUBSUB_SOCKET_TIMEOUT_SECONDS,
+                        socket_keepalive=True,
+                        socket_keepalive_options=_tcp_keepalive_options(),
+                        health_check_interval=10,
+                    )
+                except ValueError as e:
+                    logger.error(f"Invalid Redis URL: {e}")
+                    _pubsub_redis_creation_failed = True  # Stop future attempts
+                    _pubsub_redis_instance = None
+                except redis.ConnectionError as e:
+                    logger.error(f"Redis connection error: {e}")
+                    _pubsub_redis_instance = None  # Keep trying for connection errors
+    return _pubsub_redis_instance
+
+
 def gen_cache_key(messages, model="docgpt", tools=None):
     if not all(isinstance(msg, dict) for msg in messages):
         raise ValueError("All messages must be dictionaries.")

application/streaming/broadcast_channel.py

@@ -19,7 +19,9 @@ from __future__ import annotations
 import logging
 from typing import Callable, Iterator, Optional
 
-from application.cache import get_redis_instance
+import redis as redis_lib
+
+from application.cache import get_pubsub_redis_instance, get_redis_instance
 
 logger = logging.getLogger(__name__)
 
@@ -65,8 +67,13 @@ class Topic:
 
         If Redis is unavailable, returns immediately without yielding.
         Cleanly unsubscribes on ``GeneratorExit`` (client disconnect).
+
+        Uses the pub/sub-dedicated client (bounded ``socket_timeout``): a
+        subscriber whose connection went half-open must fail within seconds
+        and release its WSGI thread, not block in ``get_message`` until the
+        worker restarts.
         """
-        redis = get_redis_instance()
+        redis = get_pubsub_redis_instance()
         if redis is None:
             logger.debug("Redis unavailable; subscribe to %s yielded nothing", self.name)
             return
@@ -86,6 +93,17 @@ class Topic:
             while True:
                 try:
                     msg = pubsub.get_message(timeout=poll_timeout)
+                except redis_lib.exceptions.TimeoutError:
+                    # A bounded read timed out mid-message/health-check: the
+                    # connection is half-open (NAT/IPVS dropped the flow).
+                    # End the subscription so the SSE client reconnects on a
+                    # fresh connection; expected during network churn, so no
+                    # stack trace.
+                    logger.info(
+                        "pubsub read timed out for %s; closing subscriber",
+                        self.name,
+                    )
+                    return
                 except Exception:
                     logger.exception("pubsub.get_message failed for %s", self.name)
                     return

tests/test_events_substrate.py

@@ -305,7 +305,7 @@ class TestTopic:
         # Must not raise.
         assert Topic("user:alice").publish("hi") == 0
 
-    @patch("application.streaming.broadcast_channel.get_redis_instance")
+    @patch("application.streaming.broadcast_channel.get_pubsub_redis_instance")
     def test_subscribe_returns_immediately_when_redis_unavailable(
         self, mock_redis
     ):
@@ -313,7 +313,7 @@ class TestTopic:
         # Generator should produce nothing, not raise.
         assert list(Topic("user:alice").subscribe(poll_timeout=0.01)) == []
 
-    @patch("application.streaming.broadcast_channel.get_redis_instance")
+    @patch("application.streaming.broadcast_channel.get_pubsub_redis_instance")
     def test_subscribe_yields_none_on_poll_timeout(self, mock_redis):
         client = MagicMock()
         pubsub = MagicMock()
@@ -331,7 +331,7 @@ class TestTopic:
         assert third == b"x"
         gen.close()
 
-    @patch("application.streaming.broadcast_channel.get_redis_instance")
+    @patch("application.streaming.broadcast_channel.get_pubsub_redis_instance")
     def test_subscribe_fires_on_subscribe_after_ack(self, mock_redis):
         client = MagicMock()
         pubsub = MagicMock()
@@ -355,7 +355,7 @@ class TestTopic:
         assert callback_calls == [1]
         gen.close()
 
-    @patch("application.streaming.broadcast_channel.get_redis_instance")
+    @patch("application.streaming.broadcast_channel.get_pubsub_redis_instance")
     def test_subscribe_cleans_up_on_generator_close(self, mock_redis):
         client = MagicMock()
         pubsub = MagicMock()
@@ -373,7 +373,7 @@ class TestTopic:
         pubsub.unsubscribe.assert_called_once_with("user:alice")
         pubsub.close.assert_called_once()
 
-    @patch("application.streaming.broadcast_channel.get_redis_instance")
+    @patch("application.streaming.broadcast_channel.get_pubsub_redis_instance")
     def test_subscribe_skips_unsubscribe_if_subscribe_never_acked(
         self, mock_redis
     ):

tests/test_pubsub_timeouts.py

@@ -0,0 +1,116 @@
+"""Tests for the pub/sub-dedicated Redis client and subscriber timeout handling.
+
+Regression tests for the 2026-07-05→07 prod incident: idle pub/sub
+connections silently dropped by NAT/IPVS never raised inside
+``pubsub.get_message`` (the shared cache client has no ``socket_timeout``),
+leaving each ``/api/events`` SSE generator blocked in ``_read_from_socket``
+forever. One WSGI thread-pool slot leaked per dead subscriber until the
+pool was exhausted and the worker stopped serving Flask requests.
+"""
+
+from __future__ import annotations
+
+from unittest.mock import MagicMock, patch
+
+import pytest
+import redis as redis_lib
+
+import application.cache as cache
+from application.streaming.broadcast_channel import Topic
+
+
+@pytest.fixture(autouse=True)
+def _reset_pubsub_singleton():
+    cache._pubsub_redis_instance = None
+    cache._pubsub_redis_creation_failed = False
+    yield
+    cache._pubsub_redis_instance = None
+    cache._pubsub_redis_creation_failed = False
+
+
+class TestGetPubsubRedisInstance:
+
+    @pytest.mark.unit
+    def test_bounds_blocking_reads_and_enables_keepalive(self):
+        with patch("application.cache.redis.Redis.from_url") as from_url:
+            client = cache.get_pubsub_redis_instance()
+        assert client is from_url.return_value
+        kwargs = from_url.call_args.kwargs
+        assert kwargs["socket_timeout"] == cache.PUBSUB_SOCKET_TIMEOUT_SECONDS
+        assert kwargs["socket_timeout"] > 0
+        assert kwargs["socket_keepalive"] is True
+        assert kwargs["health_check_interval"] == 10
+
+    @pytest.mark.unit
+    def test_socket_timeout_exceeds_subscribe_poll_interval(self):
+        """get_message(timeout=poll) waits ``poll`` seconds while healthy-idle;
+        socket_timeout must be comfortably larger or idle polling would be
+        misread as a dead connection."""
+        assert cache.PUBSUB_SOCKET_TIMEOUT_SECONDS >= 5
+
+    @pytest.mark.unit
+    def test_singleton(self):
+        with patch("application.cache.redis.Redis.from_url") as from_url:
+            first = cache.get_pubsub_redis_instance()
+            second = cache.get_pubsub_redis_instance()
+        assert first is second
+        assert from_url.call_count == 1
+
+    @pytest.mark.unit
+    def test_invalid_url_marks_failed_and_stops_retrying(self):
+        with patch(
+            "application.cache.redis.Redis.from_url", side_effect=ValueError("bad url")
+        ):
+            assert cache.get_pubsub_redis_instance() is None
+        with patch("application.cache.redis.Redis.from_url") as from_url:
+            assert cache.get_pubsub_redis_instance() is None
+        from_url.assert_not_called()
+
+    @pytest.mark.unit
+    def test_connection_error_allows_retry(self):
+        with patch(
+            "application.cache.redis.Redis.from_url",
+            side_effect=redis_lib.ConnectionError("down"),
+        ):
+            assert cache.get_pubsub_redis_instance() is None
+        with patch("application.cache.redis.Redis.from_url") as from_url:
+            assert cache.get_pubsub_redis_instance() is from_url.return_value
+
+
+class TestSubscribeTimeout:
+
+    @pytest.mark.unit
+    def test_socket_timeout_ends_subscription(self):
+        """A read timeout means the connection is half-open/dead: the generator
+        must exit — freeing its WSGI thread — rather than loop or raise."""
+        pubsub = MagicMock()
+        pubsub.get_message.side_effect = redis_lib.exceptions.TimeoutError(
+            "Timeout reading from socket"
+        )
+        client = MagicMock()
+        client.pubsub.return_value = pubsub
+        with patch(
+            "application.streaming.broadcast_channel.get_pubsub_redis_instance",
+            return_value=client,
+        ):
+            events = list(Topic("user:u1").subscribe())
+        assert events == []
+        pubsub.close.assert_called_once()
+
+    @pytest.mark.unit
+    def test_subscribe_uses_dedicated_pubsub_client(self):
+        """The shared cache client has no socket_timeout; subscribers must not
+        fall back to it."""
+        pubsub = MagicMock()
+        pubsub.get_message.side_effect = redis_lib.exceptions.TimeoutError()
+        client = MagicMock()
+        client.pubsub.return_value = pubsub
+        with patch(
+            "application.streaming.broadcast_channel.get_pubsub_redis_instance",
+            return_value=client,
+        ) as get_pubsub, patch(
+            "application.streaming.broadcast_channel.get_redis_instance"
+        ) as get_cache:
+            list(Topic("user:u1").subscribe())
+        get_pubsub.assert_called_once()
+        get_cache.assert_not_called()