fix(redis): stream retention recovery (#3933)

* fix(gateway): retain redis streams safely

* Potential fix for pull request finding

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

---------

Co-authored-by: Willem Jiang <willem.jiang@gmail.com>
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
This commit is contained in:
Janlay
2026-07-04 16:07:42 +08:00
committed by GitHub
co-authored by Willem Jiang Copilot Autofix powered by AI
parent 53a80d3ad1
commit 38342b15a3
10 changed files with 259 additions and 14 deletions
+1 -1
View File
@@ -262,7 +262,7 @@ section, when present, overrides the first two for backward compatibility.
The unified nginx endpoint is same-origin by default and does not emit browser CORS headers. If you run a split-origin or port-forwarded browser client, set `GATEWAY_CORS_ORIGINS` to comma-separated exact origins such as `http://localhost:3000`; the Gateway then applies the CORS allowlist and matching CSRF origin checks.
> [!IMPORTANT]
> The Gateway holds run state (RunManager and the stream bridge) in process, so production defaults to a single Gateway worker (`GATEWAY_WORKERS=1`). Raising the worker count without a shared cross-worker stream bridge — which is not yet available — breaks run cancellation, SSE reconnects, request de-duplication, and IM channels, because nginx uses no sticky sessions and each worker keeps its own run state. Scale a single worker up with more CPU/RAM (or move the database and sandbox onto dedicated tiers) instead of raising `GATEWAY_WORKERS`.
> The Gateway still owns active run tasks in process, so production defaults to a single Gateway worker (`GATEWAY_WORKERS=1`). The Redis stream bridge (`stream_bridge.type: redis`) shares SSE delivery and `Last-Event-ID` replay across workers, with a rolling retained-buffer TTL (`stream_ttl_seconds`) as a cleanup safety net. It does not make run cancellation, request de-duplication, or IM channel state fully cross-worker by itself; use single-worker Gateway or explicit sticky routing/ownership before raising `GATEWAY_WORKERS`.
See [CONTRIBUTING.md](CONTRIBUTING.md) for detailed Docker development guide.
+1
View File
@@ -311,6 +311,7 @@ CORS is same-origin by default when requests enter through nginx on port 2026. S
- `cancel()` and `create_or_reject(..., multitask_strategy="interrupt"|"rollback")` persist interrupted status through `RunStore.update_status()`, matching normal `set_status()` transitions.
- Store-only hydrated runs are readable history. If the current worker has no in-memory task/control state for that run, cancellation APIs can return 409 because this worker cannot stop the task.
- `POST /wait` (both thread-scoped and `/api/runs/wait`) drains the stream bridge via `wait_for_run_completion()` instead of bare `await record.task`, so it honours the run's `on_disconnect` setting and cancels the background run on real client disconnect rather than returning a stale checkpoint (issue #3265).
- Redis `StreamBridge` keys use a rolling retained-buffer TTL (`stream_bridge.stream_ttl_seconds`, refreshed on `publish()` / `publish_end()`) as a leak safety net, not as a run timeout. Startup orphan recovery publishes `END_SENTINEL` and schedules stream cleanup for recovered runs; do not broaden this into a shared-database multi-pod reaper without adding worker ownership/liveness first.
- Thread-scoped run creation accepts `checkpoint` / `checkpoint_id`; Gateway validates the checkpoint belongs to the request thread before writing `checkpoint_id` / `checkpoint_ns` into `config.configurable` for LangGraph branching.
- Thread-scoped Gateway runs evaluate an active `ThreadState.goal` after the visible turn completes. `runtime/goal.py` asks a non-thinking evaluator model to judge only visible conversation evidence and return a typed blocker; the evaluator model is created once per run and reused across hidden continuation checks. Satisfied goals are cleared; every non-satisfied evaluation — continuable or stand-down — is persisted with `last_evaluation` (the blocker, reason, and evidence summary; outcomes that stop the loop additionally record a `stand_down_reason` for observability), but only `goal_not_met_yet` evaluations are streamed as hidden `HumanMessage` continuations, and only when a durable assistant end-of-turn checkpoint exists, the run has not been aborted, the thread did not change during evaluation, and the no-progress breaker has not fired. The continuation cap is 8 — a hard maximum in the `0``8` range; callers requesting more are clamped (`set_goal`/TUI) or rejected with 422 (`PUT /goal`). The no-progress breaker keys on the latest visible assistant evidence (not the evaluator's free-text reason, which an LLM rewords every turn), so two consecutive continuations that add no new visible assistant output stop the loop after 2 attempts. Model-response cleanup helpers such as think-block stripping and code-fence stripping live in `deerflow.utils.llm_text` so `runtime/goal.py` and Gateway suggestion parsing share the same JSON-prep behavior.
+41
View File
@@ -71,6 +71,44 @@ async def _drain_inflight_runs(run_manager: RunManager) -> None:
logger.exception("Failed to drain in-flight runs during shutdown")
async def _publish_recovered_run_stream_end(
bridge: StreamBridge,
recovered_runs: list[RunRecord],
*,
cleanup_delay: float = 60.0,
) -> None:
"""Terminate retained streams for runs recovered as orphaned at startup."""
for record in recovered_runs:
stream_exists = getattr(bridge, "stream_exists", None)
if stream_exists is not None:
try:
if not await stream_exists(record.run_id):
logger.debug("Skipping recovered stream end for %s: stream already expired", record.run_id)
continue
except Exception:
logger.debug("Failed to check recovered stream existence for %s", record.run_id, exc_info=True)
try:
await bridge.publish_end(record.run_id)
except Exception:
logger.warning(
"Failed to publish recovered run stream end for %s",
record.run_id,
exc_info=True,
)
continue
task = asyncio.create_task(bridge.cleanup(record.run_id, delay=cleanup_delay))
task.add_done_callback(lambda task, run_id=record.run_id: _log_recovered_stream_cleanup_result(task, run_id))
def _log_recovered_stream_cleanup_result(task: asyncio.Task[None], run_id: str) -> None:
if task.cancelled():
return
try:
task.result()
except Exception:
logger.warning("Failed to clean up recovered run stream for %s", run_id, exc_info=True)
if TYPE_CHECKING:
from app.gateway.auth.local_provider import LocalAuthProvider
from app.gateway.auth.repositories.sqlite import SQLiteUserRepository
@@ -220,6 +258,9 @@ async def langgraph_runtime(app: FastAPI, startup_config: AppConfig) -> AsyncGen
error="Gateway restarted before this run reached a durable final state.",
before=now_iso(),
)
sb_config = getattr(config, "stream_bridge", None)
cleanup_delay = getattr(sb_config, "recovered_stream_cleanup_delay_seconds", 60.0) if sb_config else 60.0
await _publish_recovered_run_stream_end(app.state.stream_bridge, recovered_runs, cleanup_delay=cleanup_delay)
await _mark_latest_recovered_threads_error(app.state.run_manager, app.state.thread_store, recovered_runs)
try:
+1 -1
View File
@@ -137,7 +137,7 @@ sequenceDiagram
关键组件:
- `runtime/runs/worker.py::run_agent` — 在 `asyncio.Task` 里跑 `agent.astream()`,把每个 chunk 通过 `serialize(chunk, mode=mode)` 转成 JSON,再 `bridge.publish()`
- `runtime/stream_bridge` — 抽象 Queue。`publish/subscribe` 解耦生产者和消费者,支持 `Last-Event-ID` 重连、心跳、多订阅者 fan-out。
- `runtime/stream_bridge` — 抽象 Queue。`publish/subscribe` 解耦生产者和消费者,支持 `Last-Event-ID` 重连、心跳、多订阅者 fan-out。Redis backend 会在每次 `publish()` / `publish_end()` 刷新 retained stream key TTL;启动恢复将 orphan run 标记为 error 后也会发布 `END_SENTINEL`,避免重连的 SSE 客户端只收到心跳。注意:TTL 是 Redis 内存安全网(防止 key 泄漏),不是 subscriber 终止机制——如果 worker 和 gateway 同时挂掉,已连接的 SSE 客户端在 TTL 过期后仍无法收到 END 信号,需要依赖客户端侧超时。完整的跨 pod subscriber 终止需要 worker 存活检测(liveness),当前版本不包含此功能。
- `app/gateway/services.py::sse_consumer` — 从 bridge 订阅,格式化为 SSE wire 帧。
- `runtime/serialization.py::serialize` — mode-aware 序列化;`messages` mode 下 `serialize_messages_tuple``(chunk, metadata)` 转成 `[chunk.model_dump(), metadata]`
@@ -32,6 +32,20 @@ class StreamBridgeConfig(BaseModel):
"concurrent SSE clients. Only applies to the redis bridge."
),
)
stream_ttl_seconds: int = Field(
default=86400,
ge=0,
description=(
"Rolling Redis stream key TTL in seconds. The redis bridge refreshes this TTL after "
"each publish and publish_end so retained SSE replay buffers are eventually reclaimed "
"even if cleanup never runs. Set to 0 to disable. Only applies to the redis bridge."
),
)
recovered_stream_cleanup_delay_seconds: float = Field(
default=60.0,
ge=0,
description=("Seconds to wait after publishing an END marker for a recovered orphaned run before deleting the stream key. Gives reconnecting SSE clients time to drain the end signal. Only applies to the redis bridge."),
)
# Global configuration instance — None means no stream bridge is configured
@@ -74,11 +74,13 @@ async def make_stream_bridge(app_config: AppConfig | None = None) -> AsyncIterat
redis_url=redis_url,
queue_maxsize=config.queue_maxsize,
max_connections=config.max_connections,
stream_ttl_seconds=config.stream_ttl_seconds,
)
logger.info(
"Stream bridge initialised: redis (queue_maxsize=%d, max_connections=%s)",
"Stream bridge initialised: redis (queue_maxsize=%d, max_connections=%s, stream_ttl_seconds=%d)",
config.queue_maxsize,
config.max_connections,
config.stream_ttl_seconds,
)
try:
yield bridge
@@ -63,11 +63,16 @@ class RedisStreamBridge(StreamBridge):
queue_maxsize: int = 256,
key_prefix: str = "deerflow:stream_bridge",
max_connections: int | None = None,
stream_ttl_seconds: int | None = 86400,
client: Redis | None = None,
) -> None:
self._redis_url = redis_url
self._maxsize = max(1, queue_maxsize)
self._key_prefix = key_prefix.rstrip(":")
if stream_ttl_seconds is not None and stream_ttl_seconds > 0:
self._stream_ttl_seconds = stream_ttl_seconds
else:
self._stream_ttl_seconds = None
# Each live SSE subscriber holds one pooled connection blocked in
# ``XREAD ... BLOCK`` for up to ``heartbeat_interval``. ``max_connections``
# caps that pool; ``None`` keeps redis-py's effectively-unbounded default.
@@ -77,6 +82,26 @@ class RedisStreamBridge(StreamBridge):
def _stream_key(self, run_id: str) -> str:
return f"{self._key_prefix}:{run_id}"
async def _xadd_retained(self, key: str, fields: dict[str, str], *, maxlen: int) -> None:
if self._stream_ttl_seconds is None:
await self._redis.xadd(
key,
fields,
maxlen=maxlen,
approximate=False,
)
return
async with self._redis.pipeline(transaction=True) as pipe:
pipe.xadd(
key,
fields,
maxlen=maxlen,
approximate=False,
)
pipe.expire(key, self._stream_ttl_seconds)
await pipe.execute()
@staticmethod
def _decode(value: Any) -> str:
if isinstance(value, bytes):
@@ -113,24 +138,24 @@ class RedisStreamBridge(StreamBridge):
)
async def publish(self, run_id: str, event: str, data: Any) -> None:
await self._redis.xadd(
self._stream_key(run_id),
key = self._stream_key(run_id)
await self._xadd_retained(
key,
{
"kind": _KIND_EVENT,
"event": event,
"data": self._encode_data(data),
},
maxlen=self._maxsize,
approximate=False,
)
async def publish_end(self, run_id: str) -> None:
# Keep the configured number of data events plus the internal end marker.
await self._redis.xadd(
self._stream_key(run_id),
key = self._stream_key(run_id)
await self._xadd_retained(
key,
{"kind": _KIND_END},
maxlen=self._maxsize + 1,
approximate=False,
)
async def stream_exists(self, run_id: str) -> bool:
+43 -2
View File
@@ -5,6 +5,7 @@ from __future__ import annotations
from contextlib import asynccontextmanager
from types import SimpleNamespace
import anyio
import pytest
from fastapi import FastAPI
@@ -57,6 +58,36 @@ class _FakeThreadStore:
self.status_updates.append((thread_id, status, user_id))
class _FakeStreamBridge:
def __init__(self, *, existing_streams: set[str] | None = None) -> None:
self.publish_end_calls: list[str] = []
self.cleanup_calls: list[tuple[str, float]] = []
self._existing_streams: set[str] = existing_streams if existing_streams is not None else set()
async def stream_exists(self, run_id: str) -> bool:
return run_id in self._existing_streams
async def publish_end(self, run_id: str) -> None:
self.publish_end_calls.append(run_id)
async def cleanup(self, run_id: str, *, delay: float = 0) -> None:
self.cleanup_calls.append((run_id, delay))
@pytest.mark.anyio
async def test_recovered_run_stream_end_skips_expired_stream():
"""Startup recovery should not recreate an already-expired retained stream."""
stream_bridge = _FakeStreamBridge(existing_streams=set())
await gateway_deps._publish_recovered_run_stream_end(
stream_bridge,
[SimpleNamespace(run_id="expired-run", thread_id="thread-1")],
)
assert stream_bridge.publish_end_calls == []
assert stream_bridge.cleanup_calls == []
@pytest.mark.anyio
async def test_sqlite_runtime_reconciles_orphaned_runs_on_startup(monkeypatch):
"""SQLite startup should recover stale active runs before serving requests."""
@@ -64,8 +95,10 @@ async def test_sqlite_runtime_reconciles_orphaned_runs_on_startup(monkeypatch):
config = SimpleNamespace(
database=SimpleNamespace(backend="sqlite"),
run_events=SimpleNamespace(backend="memory"),
stream_bridge=SimpleNamespace(recovered_stream_cleanup_delay_seconds=60.0),
)
thread_store = _FakeThreadStore()
stream_bridge = _FakeStreamBridge(existing_streams={"run-1"})
_FakeRunManager.instances.clear()
_FakeRunManager.recovered_runs = [SimpleNamespace(run_id="run-1", thread_id="thread-1")]
_FakeRunManager.latest_by_thread = {}
@@ -79,7 +112,7 @@ async def test_sqlite_runtime_reconciles_orphaned_runs_on_startup(monkeypatch):
monkeypatch.setattr(engine_module, "init_engine_from_config", fake_init_engine_from_config)
monkeypatch.setattr(engine_module, "get_session_factory", lambda: None)
monkeypatch.setattr(engine_module, "close_engine", fake_close_engine)
monkeypatch.setattr(runtime_module, "make_stream_bridge", lambda _config: _fake_context(object()))
monkeypatch.setattr(runtime_module, "make_stream_bridge", lambda _config: _fake_context(stream_bridge))
monkeypatch.setattr(checkpointer_module, "make_checkpointer", lambda _config: _fake_context(object()))
monkeypatch.setattr(runtime_module, "make_store", lambda _config: _fake_context(object()))
monkeypatch.setattr(thread_meta_module, "make_thread_store", lambda _sf, _store: thread_store)
@@ -88,12 +121,15 @@ async def test_sqlite_runtime_reconciles_orphaned_runs_on_startup(monkeypatch):
async with gateway_deps.langgraph_runtime(app, config):
pass
await anyio.sleep(0)
assert len(_FakeRunManager.instances) == 1
assert _FakeRunManager.instances[0].reconcile_calls
assert _FakeRunManager.instances[0].reconcile_calls[0]["error"]
assert _FakeRunManager.instances[0].list_by_thread_calls == [{"thread_id": "thread-1", "user_id": None, "limit": 1}]
assert thread_store.status_updates == [("thread-1", "error", None)]
assert stream_bridge.publish_end_calls == ["run-1"]
assert stream_bridge.cleanup_calls == [("run-1", 60.0)]
@pytest.mark.anyio
@@ -103,8 +139,10 @@ async def test_sqlite_runtime_does_not_mark_thread_error_when_newer_run_is_succe
config = SimpleNamespace(
database=SimpleNamespace(backend="sqlite"),
run_events=SimpleNamespace(backend="memory"),
stream_bridge=SimpleNamespace(recovered_stream_cleanup_delay_seconds=60.0),
)
thread_store = _FakeThreadStore()
stream_bridge = _FakeStreamBridge(existing_streams={"old-running"})
_FakeRunManager.instances.clear()
_FakeRunManager.recovered_runs = [SimpleNamespace(run_id="old-running", thread_id="thread-1")]
_FakeRunManager.latest_by_thread = {"thread-1": [SimpleNamespace(run_id="newer-success", thread_id="thread-1", status="success")]}
@@ -118,7 +156,7 @@ async def test_sqlite_runtime_does_not_mark_thread_error_when_newer_run_is_succe
monkeypatch.setattr(engine_module, "init_engine_from_config", fake_init_engine_from_config)
monkeypatch.setattr(engine_module, "get_session_factory", lambda: None)
monkeypatch.setattr(engine_module, "close_engine", fake_close_engine)
monkeypatch.setattr(runtime_module, "make_stream_bridge", lambda _config: _fake_context(object()))
monkeypatch.setattr(runtime_module, "make_stream_bridge", lambda _config: _fake_context(stream_bridge))
monkeypatch.setattr(checkpointer_module, "make_checkpointer", lambda _config: _fake_context(object()))
monkeypatch.setattr(runtime_module, "make_store", lambda _config: _fake_context(object()))
monkeypatch.setattr(thread_meta_module, "make_thread_store", lambda _sf, _store: thread_store)
@@ -127,7 +165,10 @@ async def test_sqlite_runtime_does_not_mark_thread_error_when_newer_run_is_succe
async with gateway_deps.langgraph_runtime(app, config):
pass
await anyio.sleep(0)
assert len(_FakeRunManager.instances) == 1
assert _FakeRunManager.instances[0].list_by_thread_calls == [{"thread_id": "thread-1", "user_id": None, "limit": 1}]
assert thread_store.status_updates == []
assert stream_bridge.publish_end_calls == ["old-running"]
assert stream_bridge.cleanup_calls == [("old-running", 60.0)]
+118 -3
View File
@@ -30,6 +30,7 @@ class _FakeRedis:
self.conditions = defaultdict(asyncio.Condition)
self.counters = defaultdict(int)
self.deleted = []
self.expirations = []
self.closed = False
async def xadd(self, name, fields, maxlen=None, approximate=True):
@@ -65,10 +66,48 @@ class _FakeRedis:
async def exists(self, name):
return 1 if name in self.streams else 0
async def expire(self, name, seconds):
self.expirations.append((name, seconds))
return True
def pipeline(self, *, transaction=True):
return _FakeRedisPipeline(self)
async def aclose(self):
self.closed = True
class _FakeRedisPipeline:
def __init__(self, redis: _FakeRedis) -> None:
self.redis = redis
self.ops = []
async def __aenter__(self):
return self
async def __aexit__(self, exc_type, exc, tb):
return None
def xadd(self, name, fields, maxlen=None, approximate=True):
self.ops.append(("xadd", name, fields, maxlen, approximate))
return self
def expire(self, name, seconds):
self.ops.append(("expire", name, seconds))
return self
async def execute(self):
results = []
for op in self.ops:
if op[0] == "xadd":
_, name, fields, maxlen, approximate = op
results.append(await self.redis.xadd(name, fields, maxlen=maxlen, approximate=approximate))
elif op[0] == "expire":
_, name, seconds = op
results.append(await self.redis.expire(name, seconds))
return results
# ---------------------------------------------------------------------------
# Unit tests for MemoryStreamBridge
# ---------------------------------------------------------------------------
@@ -487,6 +526,43 @@ async def test_redis_cleanup_deletes_stream(redis_bridge: RedisStreamBridge):
assert fake.deleted == ["deerflow:stream_bridge:redis-run-cleanup"]
@pytest.mark.anyio
async def test_redis_publish_refreshes_stream_ttl():
"""Redis stream TTL should be rolling on publish and publish_end."""
fake = _FakeRedis()
bridge = RedisStreamBridge(
redis_url="redis://fake",
queue_maxsize=2,
stream_ttl_seconds=42,
client=fake,
)
run_id = "redis-run-ttl"
key = "deerflow:stream_bridge:redis-run-ttl"
await bridge.publish(run_id, "event-1", {"n": 1})
await bridge.publish(run_id, "event-2", {"n": 2})
await bridge.publish_end(run_id)
assert fake.expirations == [(key, 42), (key, 42), (key, 42)]
@pytest.mark.anyio
async def test_redis_stream_ttl_can_be_disabled():
"""A zero TTL disables the Redis leak safety net for installations that need it."""
fake = _FakeRedis()
bridge = RedisStreamBridge(
redis_url="redis://fake",
queue_maxsize=2,
stream_ttl_seconds=0,
client=fake,
)
await bridge.publish("redis-run-no-ttl", "event", {})
await bridge.publish_end("redis-run-no-ttl")
assert fake.expirations == []
@pytest.mark.anyio
async def test_redis_subscribe_waits_for_first_publish(redis_bridge: RedisStreamBridge):
"""A subscriber that starts before the first XADD must not receive END."""
@@ -673,8 +749,8 @@ async def test_make_stream_bridge_uses_docker_redis_env(monkeypatch):
@pytest.mark.anyio
async def test_make_stream_bridge_passes_max_connections(monkeypatch):
"""max_connections from config should be forwarded to Redis.from_url."""
async def test_make_stream_bridge_passes_redis_options(monkeypatch):
"""Redis options from config should be forwarded to Redis bridge setup."""
import deerflow.runtime.stream_bridge.redis as redis_module
captured: dict = {}
@@ -685,10 +761,18 @@ async def test_make_stream_bridge_passes_max_connections(monkeypatch):
return _FakeRedis()
monkeypatch.setattr(redis_module.Redis, "from_url", staticmethod(fake_from_url))
set_stream_bridge_config(StreamBridgeConfig(type="redis", redis_url="redis://fake:6379/0", max_connections=50))
set_stream_bridge_config(
StreamBridgeConfig(
type="redis",
redis_url="redis://fake:6379/0",
max_connections=50,
stream_ttl_seconds=42,
)
)
try:
async with make_stream_bridge() as bridge:
assert isinstance(bridge, RedisStreamBridge)
assert bridge._stream_ttl_seconds == 42
assert captured["max_connections"] == 50
assert captured["decode_responses"] is True
finally:
@@ -829,3 +913,34 @@ async def test_redis_integration_maxlen_trims_history(real_redis_bridge):
key = real_redis_bridge._stream_key(run_id)
length = await real_redis_bridge._redis.xlen(key)
assert length == 2
@pytest.mark.integration
@requires_redis
@pytest.mark.anyio
async def test_redis_integration_stream_ttl_reclaims_key():
"""Redis should reclaim retained stream data when cleanup never runs."""
from redis.asyncio import Redis
client = Redis.from_url(REDIS_TEST_URL, decode_responses=True)
key_prefix = f"deerflow:test:{uuid.uuid4().hex}"
bridge = RedisStreamBridge(
redis_url=REDIS_TEST_URL,
queue_maxsize=2,
key_prefix=key_prefix,
stream_ttl_seconds=1,
client=client,
)
run_id = "integ-ttl"
key = bridge._stream_key(run_id)
try:
await bridge.publish(run_id, "metadata", {"run_id": run_id})
assert await client.exists(key) == 1
assert await client.ttl(key) >= 0
await anyio.sleep(2.0)
assert await client.exists(key) == 0
finally:
async for existing_key in client.scan_iter(f"{key_prefix}:*"):
await client.delete(existing_key)
await client.aclose()
+6
View File
@@ -1423,6 +1423,12 @@ run_events:
# type: redis # recommended for Docker / multi-worker gateway
# redis_url: redis://redis:6379/0
# queue_maxsize: 256 # events retained per run (redis stream MAXLEN)
# stream_ttl_seconds: 86400 # rolling TTL for retained stream buffers.
# # Refreshed on each publish/publish_end; set 0
# # to disable. This is not a run timeout.
# recovered_stream_cleanup_delay_seconds: 60 # seconds to wait after
# # publishing END for a recovered orphan run
# # before deleting the stream key.
# max_connections: 100 # optional pool ceiling. Each live SSE client
# # holds one connection blocked in XREAD ... BLOCK
# # for up to heartbeat_interval (15s), so hundreds