deer-flow/backend/tests/test_sse_observer_disconnect.py
wutongyuonce 9ad79baf97
feat(gateway): support idempotent thread runs (#5258)
* feat(gateway): support idempotent thread runs

Accept Idempotency-Key on thread-scoped create, stream, and wait endpoints, scoped by owner and thread before durable admission.\n\nRefs #5257.

* fix(gateway): handle idempotent run reuse on wait and stream

Reused store-only records have no local task. /wait now waits on the
bridge when it can observe the stream, and otherwise returns durable
status instead of a stale checkpoint. A reused terminal stream that has
been evicted emits gap/reload_durable_state. Replay is bound to the
original input and assistant_id.

* fix(gateway): 409 reused in-flight streams on this worker

A store-only running record on a process-local bridge has no owner
stream. POST /runs/stream used to subscribe anyway, which created an
empty log and waited forever. Match join: 409 unless the run is already
terminal, so missing-stream retries can still emit gap.

* fix(gateway): keep observer joins off the idempotent stream-gap path

sse_consumer keyed missing-stream gap on the sticky
idempotency_reused flag, so a later join of a terminal run inherited
it. Gate that branch on apply_on_disconnect, which already separates
creating streams from joins. Document the retry outcomes clients have
to handle.

* fix(gateway): gate missing-stream gap on creating retry

Reuse apply_on_disconnect to pick gap vs end changed sse_consumer default path, so a missing stream started returning gap for default callers and for out-of-scope POST /api/runs/stream. Keep that branch behind emit_gap_on_missing_stream and pass it only from thread-scoped /runs/stream on this request reuse.

* fix(gateway): keep wait reuse off later checkpoints

Direct handler calls were crashing because FastAPI Header() leaked in as the Python default. Bind Idempotency-Key with Annotated so the default is None, and ignore non-str keys.

A reused completed /wait was still reading the latest thread checkpoint. After a later run on the same thread that is the later run's result. Return durable status instead of claiming the head as this run's output.

* fix(gateway): snapshot wait reuse and refresh store status

idempotency_reused lives on the shared cached record. Capture it before awaiting completion so an overlapping retry cannot suppress the original creating /wait checkpoint.

A store-only peer record still holds admission-time status after the owner publishes END. Refresh durable status/error before returning them.
2026-09-08 14:40:05 +08:00

116 lines
4.1 KiB
Python

"""Observer joins must not apply the creator's cancel-on-disconnect policy.
Review round 5 (PR #5041): every consumer of ``sse_consumer`` used to apply
the record's ``on_disconnect=cancel`` policy in its ``finally`` block, so a
read-only stream observer could cancel a locally-owned running run just by
closing the SSE connection. The fix separates creator streams
(``apply_on_disconnect=True``, the default) from join/observer streams
(``False``). These tests drive a real generator close — the same machinery
Starlette runs when a client drops the connection — against the production
consumer.
"""
import asyncio
import inspect
from types import SimpleNamespace
from app.gateway.services import sse_consumer
from deerflow.runtime import DisconnectMode, RunRecord, RunStatus
def _running_record() -> RunRecord:
return RunRecord(
run_id="run-1",
thread_id="t1",
assistant_id=None,
status=RunStatus.running,
on_disconnect=DisconnectMode.cancel,
)
class _StubBridge:
"""Yields one event, then parks until the consumer closes the generator."""
def subscribe(self, run_id, last_event_id=None):
async def _gen():
yield SimpleNamespace(event="message", data="{}", id="1")
await asyncio.Event().wait()
return _gen()
class _CancelRecorder:
"""Stands in for the RunManager: records cancel calls, mutates nothing."""
def __init__(self):
self.cancelled: list[str] = []
async def cancel(self, run_id, action="interrupt"):
self.cancelled.append(run_id)
class _StubRequest:
"""Minimal request: headers for Last-Event-ID, never-disconnected client
(the disconnect under test happens between events, via generator close)."""
def __init__(self):
self.headers = {}
async def is_disconnected(self) -> bool:
return False
def _request() -> _StubRequest:
return _StubRequest()
async def _drive_disconnect(consumer) -> None:
"""Start the generator (it yields one frame), then close it — a real
disconnect of the response stream, running the ``finally`` block."""
await consumer.__anext__()
await consumer.aclose()
def test_creator_stream_disconnect_applies_cancel_policy():
"""The stream returned by the creating endpoint keeps the creator's
cancel-on-disconnect semantics."""
async def scenario():
recorder = _CancelRecorder()
consumer = sse_consumer(_StubBridge(), _running_record(), _request(), recorder)
await _drive_disconnect(consumer)
return recorder.cancelled
assert asyncio.run(scenario()) == ["run-1"]
def test_observer_join_disconnect_does_not_cancel():
"""A join/observer stream closing must not cancel the run — including for
a read-only credential that never held runs:cancel."""
async def scenario():
recorder = _CancelRecorder()
consumer = sse_consumer(_StubBridge(), _running_record(), _request(), recorder, apply_on_disconnect=False)
await _drive_disconnect(consumer)
return recorder.cancelled
assert asyncio.run(scenario()) == []
def test_join_routes_wire_sse_consumer_as_observers():
"""Both join surfaces must be wired as observers, and the creator's
create-and-stream endpoints must keep the creator policy (default)."""
from app.gateway.routers import runs as runs_router
from app.gateway.routers import thread_runs
thread_runs_source = inspect.getsource(thread_runs)
# join_run + the shared existing-run stream implementation
assert thread_runs_source.count("sse_consumer(bridge, record, request, run_mgr, apply_on_disconnect=False)") == 2
# stream_run — creating retry opts into missing-stream gap; first create does not
assert "emit_gap_on_missing_stream=record.idempotency_reused" in thread_runs_source
runs_source = inspect.getsource(runs_router)
# stateless create-and-stream — creator on_disconnect policy, not the retry gap
assert "sse_consumer(bridge, record, request, run_mgr, apply_on_disconnect=False)" not in runs_source
assert "emit_gap_on_missing_stream" not in runs_source