mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-10-26 21:26:39 +00:00
fix(runtime): recover terminal Redis streams after producer loss (#7093)
Signed-off-by: CorgiBoyG <111257566+CorgiBoyG@users.noreply.github.com> Co-authored-by: Willem Jiang <willem.jiang@gmail.com>
This commit is contained in:
parent
46c0899e53
commit
bf2cb0e1df
@ -591,7 +591,9 @@ Administrators can suspend and restore accounts from **Settings → Users** (adm
|
||||
DeerFlow still uses `Forwarded` / `X-Forwarded-*` headers to recover the browser-facing scheme and origin behind a proxy. The bundled nginx sets `X-Forwarded-Proto`, but preserves an upstream HTTPS value and does not overwrite every forwarded header. Configure the outer trusted proxy to replace or strip client-supplied forwarding headers before traffic reaches DeerFlow.
|
||||
|
||||
> [!IMPORTANT]
|
||||
> The Gateway still owns active run tasks in process, so production defaults to a single Gateway worker (`GATEWAY_WORKERS=1`). Multi-worker deployments require Postgres, the Redis stream bridge (`stream_bridge.type: redis`), `run_ownership.heartbeat_enabled: true`, and `run_events.backend: db`; process-local memory/JSONL event stores cannot enforce singleton delivery receipts across workers. Kubernetes replicas run one worker per Pod, which the worker count cannot see: declare them with `deployment.multi_instance: true` (or `DEER_FLOW_MULTI_INSTANCE=1`, which deploy tooling such as a Helm chart can set from its replica count) so the startup gate enforces the same prerequisites instead of staying inert. Declared instances that store credentials (`channel_connections`) must also share one at-rest `DEER_FLOW_CREDENTIALS_KEY` (the Helm chart and `make up` generate it; back it up, see [backend/docs/CONFIGURATION.md](backend/docs/CONFIGURATION.md#credentials-encryption-key-deer_flow_credentials_key)). The bridge shares SSE delivery and bounded `Last-Event-ID` replay across workers. When a trimmed reconnect cursor cannot be verified for complete replay, or a subscriber that already established an empty-stream wait falls behind before its first delivery, Memory and Redis emit a machine-readable SSE `gap` event instead of silently returning a partial replay; the Web UI reloads durable thread/event state and resumes from the retained tail. The Memory bridge can replay without recovery when the cursor is its exact last evicted event: every successor is still retained. Older or unverifiable evicted cursors still trigger recovery. An already-observed Redis stream removed during replay also signals a gap, so slow subscribers reload durable state instead of heartbeating forever after cleanup. Lease reconciliation marks runs from dead workers as errors, persists their delivery receipts, publishes the terminal stream marker, schedules retained-stream cleanup, and updates the affected thread status. SSE, `/wait`, and internal stream consumers use `stream_bridge.heartbeat_interval_seconds` (default `15`) for idle liveness checks; changing it requires a Gateway restart. Malformed Redis reconnect IDs live-tail new events instead of replaying the retained buffer, and the rolling retained-buffer TTL (`stream_ttl_seconds`) remains a cleanup safety net rather than a run timeout. Failed-login counters and lockouts for `POST /api/v1/auth/login/local` (`auth.local.max_login_attempts` / `lockout_seconds`) are kept in the shared `login_throttle` table whenever the application database is SQLite or Postgres (`auth.local.throttle_storage: auto`, the default), so every replica enforces one lockout per client IP; `memory` keeps the historical per-process counter, which under N replicas hands an attacker N × `max_login_attempts` guesses and logs a startup warning. IM chat-to-thread bindings live in the shared `channel_thread_bindings` table whenever the application database is SQLite or Postgres (an existing `channels/store.json` is imported once at startup and renamed `store.json.migrated`); IM bot credentials entered in the web UI live in the shared `channel_runtime_configs` table, encrypted with `DEER_FLOW_CREDENTIALS_KEY` (an existing `channels/runtime-config.json` is imported once and then deleted), so a save on any instance reaches the provider's connection within one lease renewal; each provider's platform connection is held by one instance through the shared `channel_leases` table — `channels.leader_lease` (default `auto`, on wherever the multi-process signal is) hands it to another instance when the holder dies, and the startup gate refuses `enabled: false` while an IM channel other than GitHub is enabled (see [IM channel leader lease](backend/docs/CONFIGURATION.md#im-channel-leader-lease-channelsleader_lease)); follow-up buffers still need their own multi-worker coordination. To try this topology on one machine, `make dev-multi` starts two Gateways on shared Postgres, Redis and `DEER_FLOW_HOME`, and `make dev-multi-check` runs the cross-instance checks; see the local two-Gateway harness in [backend/docs/CONFIGURATION.md](backend/docs/CONFIGURATION.md#local-two-gateway-harness).
|
||||
> The Gateway still owns active run tasks in process, so production defaults to a single Gateway worker (`GATEWAY_WORKERS=1`). Multi-worker deployments require Postgres, the Redis stream bridge (`stream_bridge.type: redis`), `run_ownership.heartbeat_enabled: true`, and `run_events.backend: db`; process-local memory/JSONL event stores cannot enforce singleton delivery receipts across workers. Kubernetes replicas run one worker per Pod, which the worker count cannot see: declare them with `deployment.multi_instance: true` (or `DEER_FLOW_MULTI_INSTANCE=1`, which deploy tooling such as a Helm chart can set from its replica count) so the startup gate enforces the same prerequisites instead of staying inert. Declared instances that store credentials (`channel_connections`) must also share one at-rest `DEER_FLOW_CREDENTIALS_KEY` (the Helm chart and `make up` generate it; back it up, see [backend/docs/CONFIGURATION.md](backend/docs/CONFIGURATION.md#credentials-encryption-key-deer_flow_credentials_key)). The bridge shares SSE delivery and bounded `Last-Event-ID` replay across workers. When a trimmed reconnect cursor cannot be verified for complete replay, or a subscriber that already established an empty-stream wait falls behind before its first delivery, Memory and Redis emit a machine-readable SSE `gap` event instead of silently returning a partial replay; the Web UI reloads durable thread/event state and resumes from the retained tail. The Memory bridge can replay without recovery when the cursor is its exact last evicted event: every successor is still retained. Older or unverifiable evicted cursors still trigger recovery. An already-observed Redis stream removed during replay also signals a gap, so slow subscribers reload durable state instead of heartbeating forever after cleanup. Redis also tracks terminal stream producers with a short lease: if a worker dies after committing a terminal run status but before publishing END, SSE and `/wait` finish after the lease expires instead of waiting for the retained stream's TTL. Lease reconciliation marks runs from dead workers as errors, persists their delivery receipts, publishes the terminal stream marker, schedules retained-stream cleanup, and updates the affected thread status. SSE, `/wait`, and internal stream consumers use `stream_bridge.heartbeat_interval_seconds` (default `15`) for idle liveness checks; changing it requires a Gateway restart. Malformed Redis reconnect IDs live-tail new events instead of replaying the retained buffer, and the rolling retained-buffer TTL (`stream_ttl_seconds`) remains a cleanup safety net rather than a run timeout. Failed-login counters and lockouts for `POST /api/v1/auth/login/local` (`auth.local.max_login_attempts` / `lockout_seconds`) are kept in the shared `login_throttle` table whenever the application database is SQLite or Postgres (`auth.local.throttle_storage: auto`, the default), so every replica enforces one lockout per client IP; `memory` keeps the historical per-process counter, which under N replicas hands an attacker N × `max_login_attempts` guesses and logs a startup warning. IM chat-to-thread bindings live in the shared `channel_thread_bindings` table whenever the application database is SQLite or Postgres (an existing `channels/store.json` is imported once at startup and renamed `store.json.migrated`); IM bot credentials entered in the web UI live in the shared `channel_runtime_configs` table, encrypted with `DEER_FLOW_CREDENTIALS_KEY` (an existing `channels/runtime-config.json` is imported once and then deleted), so a save on any instance reaches the provider's connection within one lease renewal; each provider's platform connection is held by one instance through the shared `channel_leases` table — `channels.leader_lease` (default `auto`, on wherever the multi-process signal is) hands it to another instance when the holder dies, and the startup gate refuses `enabled: false` while an IM channel other than GitHub is enabled (see [IM channel leader lease](backend/docs/CONFIGURATION.md#im-channel-leader-lease-channelsleader_lease)); follow-up buffers still need their own multi-worker coordination. To try this topology on one machine, `make dev-multi` starts two Gateways on shared Postgres, Redis and `DEER_FLOW_HOME`, and `make dev-multi-check` runs the cross-instance checks; see the local two-Gateway harness in [backend/docs/CONFIGURATION.md](backend/docs/CONFIGURATION.md#local-two-gateway-harness).
|
||||
>
|
||||
> The producer-lease recovery behavior above is gated by `stream_bridge.producer_lease_recovery_enabled`. Keep it disabled while old Gateway binaries or their runs remain; enable it on every instance only after the writer rollout and drain complete.
|
||||
>
|
||||
> With run ownership heartbeats enabled, terminal success requires acknowledgment
|
||||
> from that worker's own durable write. A matching database row alone cannot
|
||||
|
||||
@ -489,6 +489,18 @@ async def _drain_inflight_runs(run_manager: RunManager) -> None:
|
||||
logger.exception("Failed to drain in-flight runs during shutdown")
|
||||
|
||||
|
||||
async def _retry_recovered_run_stream_end(
|
||||
bridge: StreamBridge,
|
||||
run_id: str,
|
||||
*,
|
||||
retry_delay: float,
|
||||
cleanup_delay: float,
|
||||
) -> None:
|
||||
await asyncio.sleep(max(0.0, retry_delay))
|
||||
await bridge.publish_end(run_id)
|
||||
await bridge.cleanup(run_id, delay=cleanup_delay)
|
||||
|
||||
|
||||
async def _publish_recovered_run_stream_end(
|
||||
bridge: StreamBridge,
|
||||
recovered_runs: list[RunRecord],
|
||||
@ -498,12 +510,57 @@ async def _publish_recovered_run_stream_end(
|
||||
) -> list[tuple[str, asyncio.Task[None]]]:
|
||||
"""Terminate retained streams for runs recovered as orphaned."""
|
||||
cleanup_tasks: list[tuple[str, asyncio.Task[None]]] = []
|
||||
|
||||
def schedule_retry(record: RunRecord) -> bool:
|
||||
retry_delay = getattr(
|
||||
bridge,
|
||||
"producer_lease_retry_delay_seconds",
|
||||
None,
|
||||
)
|
||||
if not isinstance(retry_delay, (int, float)) or isinstance(retry_delay, bool) or retry_delay < 0:
|
||||
return False
|
||||
task = asyncio.create_task(
|
||||
_retry_recovered_run_stream_end(
|
||||
bridge,
|
||||
record.run_id,
|
||||
retry_delay=float(retry_delay),
|
||||
cleanup_delay=cleanup_delay,
|
||||
)
|
||||
)
|
||||
task.add_done_callback(
|
||||
lambda task, run_id=record.run_id: _log_recovered_stream_cleanup_result(
|
||||
task,
|
||||
run_id,
|
||||
)
|
||||
)
|
||||
cleanup_tasks.append((record.run_id, task))
|
||||
if on_cleanup_scheduled is not None:
|
||||
on_cleanup_scheduled(record.run_id, task)
|
||||
return True
|
||||
|
||||
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)
|
||||
cleanup_expired_producer = getattr(
|
||||
bridge,
|
||||
"cleanup_expired_producer",
|
||||
None,
|
||||
)
|
||||
if callable(cleanup_expired_producer):
|
||||
try:
|
||||
cleaned = bool(await cleanup_expired_producer(record.run_id))
|
||||
except Exception:
|
||||
logger.warning(
|
||||
"Failed to clean expired stream producer marker for %s",
|
||||
record.run_id,
|
||||
exc_info=True,
|
||||
)
|
||||
cleaned = False
|
||||
if not cleaned:
|
||||
schedule_retry(record)
|
||||
continue
|
||||
except Exception:
|
||||
logger.debug("Failed to check recovered stream existence for %s", record.run_id, exc_info=True)
|
||||
@ -515,6 +572,7 @@ async def _publish_recovered_run_stream_end(
|
||||
record.run_id,
|
||||
exc_info=True,
|
||||
)
|
||||
schedule_retry(record)
|
||||
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))
|
||||
|
||||
@ -9,6 +9,7 @@ from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import hashlib
|
||||
import inspect
|
||||
import json
|
||||
import logging
|
||||
import math
|
||||
@ -74,6 +75,7 @@ from deerflow.runtime import (
|
||||
RunStatus,
|
||||
StreamBridge,
|
||||
StreamGap,
|
||||
StreamProducerLeaseState,
|
||||
ThreadOperationKind,
|
||||
UnsupportedStrategyError,
|
||||
build_state_mutation_graph,
|
||||
@ -115,6 +117,16 @@ from deerflow.utils.thread_id import validate_thread_id
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_MISSING_BRIDGE_METHOD = object()
|
||||
|
||||
|
||||
def _declared_bridge_method(bridge: Any, name: str) -> Any | None:
|
||||
"""Return a real bridge method without triggering dynamic mock attributes."""
|
||||
if inspect.getattr_static(bridge, name, _MISSING_BRIDGE_METHOD) is _MISSING_BRIDGE_METHOD:
|
||||
return None
|
||||
method = getattr(bridge, name, None)
|
||||
return method if callable(method) else None
|
||||
|
||||
|
||||
class BusyThreadConflict(HTTPException):
|
||||
"""A retryable run-manager admission conflict exposed as HTTP 409."""
|
||||
@ -147,6 +159,8 @@ _TERMINAL_RUN_STATUSES = {
|
||||
RunStatus.timeout,
|
||||
RunStatus.interrupted,
|
||||
}
|
||||
_TERMINAL_STREAM_RECOVERED = "recovered"
|
||||
_TERMINAL_STREAM_PRODUCER_EXPIRED = "producer_expired"
|
||||
|
||||
_THREAD_METADATA_SETUP_TIMEOUT_SECONDS = 5.0
|
||||
|
||||
@ -282,15 +296,59 @@ async def _ensure_thread_metadata(
|
||||
return existing
|
||||
|
||||
|
||||
async def _stream_producer_lease_state(
|
||||
bridge: StreamBridge,
|
||||
run_id: str,
|
||||
) -> StreamProducerLeaseState | None:
|
||||
producer_lease_state = getattr(bridge, "producer_lease_state", None)
|
||||
if producer_lease_state is None:
|
||||
return StreamProducerLeaseState.untracked
|
||||
try:
|
||||
state = await producer_lease_state(run_id)
|
||||
return StreamProducerLeaseState(state)
|
||||
except (TypeError, ValueError):
|
||||
return StreamProducerLeaseState.untracked
|
||||
except Exception:
|
||||
logger.debug(
|
||||
"Failed to inspect stream producer lease for run %s",
|
||||
sanitize_log_param(run_id),
|
||||
exc_info=True,
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
async def _cleanup_expired_stream_producer(
|
||||
bridge: StreamBridge,
|
||||
run_id: str,
|
||||
) -> bool:
|
||||
cleanup = _declared_bridge_method(
|
||||
bridge,
|
||||
"cleanup_expired_producer",
|
||||
)
|
||||
if cleanup is None:
|
||||
return True
|
||||
try:
|
||||
return bool(await cleanup(run_id))
|
||||
except Exception:
|
||||
logger.warning(
|
||||
"Failed to clean expired stream producer marker for run %s",
|
||||
sanitize_log_param(run_id),
|
||||
exc_info=True,
|
||||
)
|
||||
return False
|
||||
|
||||
|
||||
async def _terminal_record_stream_missing(bridge: StreamBridge, record: RunRecord) -> bool:
|
||||
"""True when a terminal run has no retained stream on bridges that can tell."""
|
||||
if not _run_is_terminal(record):
|
||||
return False
|
||||
if not getattr(record, "terminal_status_persisted", True):
|
||||
return False
|
||||
stream_exists = getattr(bridge, "stream_exists", None)
|
||||
if stream_exists is None:
|
||||
return False
|
||||
try:
|
||||
return not bool(await stream_exists(record.run_id))
|
||||
missing = not bool(await stream_exists(record.run_id))
|
||||
except Exception:
|
||||
logger.debug(
|
||||
"Failed to probe stream existence for terminal run %s",
|
||||
@ -298,25 +356,81 @@ async def _terminal_record_stream_missing(bridge: StreamBridge, record: RunRecor
|
||||
exc_info=True,
|
||||
)
|
||||
return False
|
||||
if not missing:
|
||||
return False
|
||||
if getattr(bridge, "producer_lease_recovery_enabled", False) is not True:
|
||||
return True
|
||||
state = await _stream_producer_lease_state(bridge, record.run_id)
|
||||
if state == StreamProducerLeaseState.untracked:
|
||||
return True
|
||||
if state != StreamProducerLeaseState.expired:
|
||||
return False
|
||||
try:
|
||||
# The producer may have appended its tail and released the lease after
|
||||
# the first EXISTS. Recheck after observing expiry before synthesizing.
|
||||
still_missing = not bool(await stream_exists(record.run_id))
|
||||
except Exception:
|
||||
logger.debug(
|
||||
"Failed to recheck stream existence for terminal run %s",
|
||||
sanitize_log_param(record.run_id),
|
||||
exc_info=True,
|
||||
)
|
||||
return False
|
||||
if not still_missing:
|
||||
return False
|
||||
return await _cleanup_expired_stream_producer(
|
||||
bridge,
|
||||
record.run_id,
|
||||
)
|
||||
|
||||
|
||||
async def _orphan_recovery_observed_after_heartbeat(
|
||||
record: RunRecord,
|
||||
run_mgr: RunManager,
|
||||
) -> bool:
|
||||
"""Return whether durable orphan recovery is the consumer's liveness edge.
|
||||
bridge: StreamBridge,
|
||||
) -> str | None:
|
||||
"""Return why durable state proves that the stream producer is gone.
|
||||
|
||||
A normal terminal status is not sufficient: the producer persists status
|
||||
before publishing its final error/data frames and END. Orphan recovery is
|
||||
different because the producer is known to be gone and the durable
|
||||
``stop_reason`` is written atomically with the terminal status by run or
|
||||
scheduled-task reconciliation. Only those explicit recovery signals may
|
||||
synthesize END after a heartbeat.
|
||||
scheduled-task reconciliation.
|
||||
|
||||
Redis producers additionally hold a short cross-process lease from before
|
||||
durable admission until END publication. A durable terminal row
|
||||
plus an explicitly expired producer lease closes the crash window where the
|
||||
row committed but the process died before publishing END. Untracked streams
|
||||
retain the historical conservative behavior.
|
||||
"""
|
||||
if not record.store_only:
|
||||
return False
|
||||
refreshed = await run_mgr.get(record.run_id, user_id=record.user_id)
|
||||
return refreshed is not None and _run_is_terminal(refreshed) and refreshed.stop_reason in (ORPHAN_RECOVERY_STOP_REASON, "scheduled_task_orphan_recovered")
|
||||
if refreshed is None or not _run_is_terminal(refreshed):
|
||||
return None
|
||||
if refreshed.stop_reason in (
|
||||
ORPHAN_RECOVERY_STOP_REASON,
|
||||
"scheduled_task_orphan_recovered",
|
||||
):
|
||||
if getattr(bridge, "producer_lease_recovery_enabled", False) is True:
|
||||
state = await _stream_producer_lease_state(
|
||||
bridge,
|
||||
record.run_id,
|
||||
)
|
||||
if state in (
|
||||
None,
|
||||
StreamProducerLeaseState.active,
|
||||
):
|
||||
return None
|
||||
if state == StreamProducerLeaseState.expired:
|
||||
return _TERMINAL_STREAM_PRODUCER_EXPIRED
|
||||
return _TERMINAL_STREAM_RECOVERED
|
||||
if not getattr(refreshed, "terminal_status_persisted", True):
|
||||
return None
|
||||
if getattr(bridge, "producer_lease_recovery_enabled", False) is not True:
|
||||
return None
|
||||
state = await _stream_producer_lease_state(bridge, record.run_id)
|
||||
if state == StreamProducerLeaseState.expired:
|
||||
return _TERMINAL_STREAM_PRODUCER_EXPIRED
|
||||
return None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
@ -1948,6 +2062,9 @@ async def start_run(
|
||||
bridge = get_stream_bridge(request)
|
||||
run_mgr = get_run_manager(request)
|
||||
run_ctx = get_run_context(request)
|
||||
admission_run_id: str | None = None
|
||||
producer_lease_started = False
|
||||
producer_lease_handed_off = False
|
||||
|
||||
disconnect = DisconnectMode.cancel if body.on_disconnect == "cancel" else DisconnectMode.continue_
|
||||
|
||||
@ -2378,6 +2495,15 @@ async def start_run(
|
||||
# cannot both succeed across Gateway workers.
|
||||
if require_existing_thread and not await thread_access_allowed():
|
||||
raise HTTPException(status_code=404, detail=f"Thread {thread_id} not found")
|
||||
admission_run_id = str(uuid.uuid4())
|
||||
start_producer_lease = _declared_bridge_method(
|
||||
bridge,
|
||||
"start_producer_lease",
|
||||
)
|
||||
if start_producer_lease is not None:
|
||||
producer_lease_started = await start_producer_lease(admission_run_id)
|
||||
if getattr(bridge, "supports_cross_process", False) is True and not producer_lease_started:
|
||||
raise RuntimeError("Cross-process stream producer tracking is unavailable")
|
||||
try:
|
||||
record = await run_mgr.create_or_reject(
|
||||
thread_id,
|
||||
@ -2403,6 +2529,7 @@ async def start_run(
|
||||
# Canonical agent keys, not routing aliases or display names.
|
||||
# Bootstrap and unresolved legacy identities stay read-only.
|
||||
evidence_agent_id=(None if scope_runtime_config.get("is_bootstrap") else (_DEFAULT_ASSISTANT_ID if scope_assistant_id == _DEFAULT_ASSISTANT_ID else getattr(agent_config, "name", None))),
|
||||
run_id=admission_run_id,
|
||||
)
|
||||
except RunIdempotencyUnsupported as exc:
|
||||
raise HTTPException(status_code=503, detail=str(exc)) from exc
|
||||
@ -2437,12 +2564,21 @@ async def start_run(
|
||||
|
||||
worker = run_after_metadata(record)
|
||||
try:
|
||||
# No await is allowed between durable admission and task
|
||||
# attachment. Metadata setup runs inside the attached
|
||||
# worker so a pending cancellation can bypass stalled
|
||||
# thread-store IO and still reach run_agent's startup
|
||||
# barrier / stream finalization.
|
||||
# The stream lease was established before durable admission.
|
||||
# Keep the admission-to-task-attachment window await-free so
|
||||
# cancellation cannot strand a pending row.
|
||||
record.task = asyncio.create_task(worker)
|
||||
if producer_lease_started:
|
||||
transfer_producer_lease = _declared_bridge_method(
|
||||
bridge,
|
||||
"transfer_producer_lease",
|
||||
)
|
||||
if transfer_producer_lease is not None:
|
||||
transfer_producer_lease(
|
||||
record.run_id,
|
||||
record.task,
|
||||
)
|
||||
producer_lease_handed_off = True
|
||||
except Exception as exc:
|
||||
worker.close()
|
||||
await run_mgr.fail_start_if_pending(
|
||||
@ -2461,8 +2597,19 @@ async def start_run(
|
||||
|
||||
return record
|
||||
finally:
|
||||
if owner_context_token is not None:
|
||||
reset_current_user(owner_context_token)
|
||||
try:
|
||||
if producer_lease_started and not producer_lease_handed_off and admission_run_id is not None:
|
||||
try:
|
||||
await await_drained(bridge.abort_producer_lease(admission_run_id))
|
||||
except Exception:
|
||||
logger.warning(
|
||||
"Failed to release unused stream producer lease for run %s",
|
||||
sanitize_log_param(admission_run_id),
|
||||
exc_info=True,
|
||||
)
|
||||
finally:
|
||||
if owner_context_token is not None:
|
||||
reset_current_user(owner_context_token)
|
||||
|
||||
|
||||
async def launch_scheduled_thread_run(
|
||||
@ -2702,6 +2849,7 @@ async def sse_consumer(
|
||||
gap_emitted = False
|
||||
terminal_emitted = False
|
||||
disconnect_observed = False
|
||||
expired_heartbeat_seen = False
|
||||
try:
|
||||
async for entry in bridge.subscribe(record.run_id, last_event_id=last_event_id):
|
||||
if await request.is_disconnected():
|
||||
@ -2724,13 +2872,27 @@ async def sse_consumer(
|
||||
return
|
||||
|
||||
if entry is HEARTBEAT_SENTINEL:
|
||||
if await _orphan_recovery_observed_after_heartbeat(record, run_mgr):
|
||||
completion_reason = await _orphan_recovery_observed_after_heartbeat(
|
||||
record,
|
||||
run_mgr,
|
||||
bridge,
|
||||
)
|
||||
if completion_reason == _TERMINAL_STREAM_RECOVERED or (completion_reason == _TERMINAL_STREAM_PRODUCER_EXPIRED and expired_heartbeat_seen):
|
||||
if completion_reason == _TERMINAL_STREAM_PRODUCER_EXPIRED and not await _cleanup_expired_stream_producer(
|
||||
bridge,
|
||||
record.run_id,
|
||||
):
|
||||
expired_heartbeat_seen = False
|
||||
yield ": heartbeat\n\n"
|
||||
continue
|
||||
terminal_emitted = True
|
||||
yield format_sse("end", None)
|
||||
return
|
||||
expired_heartbeat_seen = completion_reason == _TERMINAL_STREAM_PRODUCER_EXPIRED
|
||||
yield ": heartbeat\n\n"
|
||||
continue
|
||||
|
||||
expired_heartbeat_seen = False
|
||||
if entry is END_SENTINEL:
|
||||
terminal_emitted = True
|
||||
yield format_sse("end", None, event_id=entry.id or None)
|
||||
@ -2810,6 +2972,7 @@ async def wait_for_run_completion(
|
||||
return True
|
||||
|
||||
resume_from_event_id: str | None = None
|
||||
expired_heartbeat_seen = False
|
||||
try:
|
||||
while True:
|
||||
gap_seen = False
|
||||
@ -2832,9 +2995,24 @@ async def wait_for_run_completion(
|
||||
resume_from_event_id = entry.latest_available_event_id
|
||||
gap_seen = True
|
||||
break
|
||||
if entry is HEARTBEAT_SENTINEL and await _orphan_recovery_observed_after_heartbeat(record, run_mgr):
|
||||
completed = True
|
||||
return True
|
||||
if entry is HEARTBEAT_SENTINEL:
|
||||
completion_reason = await _orphan_recovery_observed_after_heartbeat(
|
||||
record,
|
||||
run_mgr,
|
||||
bridge,
|
||||
)
|
||||
if completion_reason == _TERMINAL_STREAM_RECOVERED or (completion_reason == _TERMINAL_STREAM_PRODUCER_EXPIRED and expired_heartbeat_seen):
|
||||
if completion_reason == _TERMINAL_STREAM_PRODUCER_EXPIRED and not await _cleanup_expired_stream_producer(
|
||||
bridge,
|
||||
record.run_id,
|
||||
):
|
||||
expired_heartbeat_seen = False
|
||||
continue
|
||||
completed = True
|
||||
return True
|
||||
expired_heartbeat_seen = completion_reason == _TERMINAL_STREAM_PRODUCER_EXPIRED
|
||||
else:
|
||||
expired_heartbeat_seen = False
|
||||
if await request.is_disconnected():
|
||||
disconnect_observed = True
|
||||
return False
|
||||
|
||||
@ -760,6 +760,7 @@ Notes:
|
||||
- Set `multi_instance: true` (or export `DEER_FLOW_MULTI_INSTANCE=1`, which lets deploy tooling such as a Helm chart set it from its replica count) on every instance that shares one database. Startup then enforces the same prerequisites as `GATEWAY_WORKERS > 1`: `database.backend: postgres`, `run_events.backend: db`, `run_ownership.heartbeat_enabled: true`, and a Redis stream bridge (`stream_bridge.type: redis` or `DEER_FLOW_STREAM_BRIDGE_REDIS_URL`). It refuses an explicit `sandbox.ownership.type: memory`, process-local browser tools, and `scheduler.enabled: true` without `scheduler.multi_instance: true`.
|
||||
- `DEER_FLOW_MULTI_INSTANCE` treats blank, `0`, `false`, `no` and `off` as "not declared"; any other value declares a multi-instance deployment, so a typo fails closed.
|
||||
- `GET /health/ready` pings the Redis stream bridge on every probe and answers 503 (`stream_bridge: unreachable`) while Redis is down, so the orchestrator drains that instance instead of routing it runs it cannot publish or stream; the memory bridge reports `not_configured`. When `sandbox.provisioner_url` is set, the body also carries the provisioner's `/health` verdict (`provisioner: ok|unreachable`), which never changes the status code because every instance shares one provisioner. That probe follows the sandbox clients' proxy policy: loopback, private, link-local and cluster-local provisioner addresses bypass `HTTP_PROXY`, external hosts keep the environment's proxy settings.
|
||||
- `stream_bridge.producer_lease_recovery_enabled` is a restart-required activation gate for recovering a terminal Redis stream whose producer died before publishing END. Keep it `false` during the binary rollout: upgraded writers already publish producer markers, while old writers do not. After every old Gateway has stopped and its runs have drained, set it to `true` on every instance and restart them together. Enabling it while an old worker may still finalize a run can make a new consumer mistake that untracked legacy producer for completed history.
|
||||
- The declaration also drives the `agent_storage.backend: file` divergence warning, the inbound webhook dedupe warning, the `auth.local.throttle_storage: memory` warning, and the WeChat QR-login guard, which otherwise only look at the worker count.
|
||||
- Login lockouts are shared through the database: `auth.local.throttle_storage` (startup-only) defaults to `auto`, which keeps the per-IP failed-login counters for `POST /api/v1/auth/login/local` in the `login_throttle` table whenever `database.backend` is `sqlite` or `postgres`, so every instance enforces one `max_login_attempts` limit per client IP and a lockout on one instance holds on all of them. `memory` keeps the historical per-process counter (N instances give an attacker N × `max_login_attempts` guesses; a declared multi-instance deployment logs a warning), and `db` forces the table (it falls back to memory with a warning when the database backend is `memory`, and refuses to start when the configured database's engine is unavailable; `auto` falls back with a warning in that case). `max_login_attempts` and `lockout_seconds` stay live-read.
|
||||
- IM chat-to-thread bindings are shared through the database too: with `database.backend: sqlite` or `postgres` the `ChannelManager` keeps the `channel_name:chat_id[:topic_id]` → thread mapping of unbound IM conversations in the `channel_thread_bindings` table (migration `0038_channel_thread_bindings`), so a conversation created on one instance continues on the same thread when its next message lands on another. On the first start after upgrading, an existing `{base_dir}/channels/store.json` is imported once — only into an empty table, with `INSERT … ON CONFLICT DO NOTHING` so concurrently starting instances cannot duplicate a binding — and renamed `store.json.migrated`; a populated table leaves the file untouched. `memory` keeps the per-process JSON file.
|
||||
|
||||
@ -143,7 +143,8 @@ 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。`stream_bridge.heartbeat_interval_seconds`(默认 15 秒)是 bridge 实例的默认心跳周期,统一作用于 Gateway SSE、`/wait` 和内部 channel watcher;显式传给 `subscribe()` 的值仍可覆盖单次订阅。Memory 和 Redis 都只保留 `queue_maxsize` 条数据事件;游标早于保留水位线时返回 `StreamGap`,不会从当前最早事件静默部分重放。Redis backend 会在每次 `publish()` / `publish_end()` 刷新 retained stream key TTL;启动恢复与基于 worker lease 的周期恢复共用 Gateway stream terminalization 路径:`RunManager` 先将 orphan run 持久化为 `error` 并写入显式的 `stop_reason=orphan_recovered`,随后 Gateway 发布 `END_SENTINEL` 并安排 stream cleanup。周期扫描、逐行状态写入和 Gateway callback 作为一个受监督的 single-flight 后台 task 执行;慢任务不会堆积,也不会阻塞唯一的 lease heartbeat。shutdown 优先收敛活跃 run,再处理恢复 task;尚未执行的延迟 stream cleanup 会改为立即删除。只有 runtime `yield` 前、无并发请求的启动恢复会把最新受影响 thread 标记为 error;周期恢复不做非原子的 thread 投影。store-only SSE 与 `/wait` consumer 不能把普通 durable terminal status 当成流已完成,否则可能跳过延迟发布的 error 等尾部事件;只有 `orphan_recovered` 信号能在 heartbeat 时触发 END fallback,因为此时 producer 已被确认失联。TTL 仍是 Redis 内存和故障安全网,不是正常的 subscriber 终止机制。
|
||||
- `runtime/stream_bridge` — 抽象 Queue。`publish/subscribe` 解耦生产者和消费者,支持 `Last-Event-ID` 重连、心跳、多订阅者 fan-out。`stream_bridge.heartbeat_interval_seconds`(默认 15 秒)是 bridge 实例的默认心跳周期,统一作用于 Gateway SSE、`/wait` 和内部 channel watcher;显式传给 `subscribe()` 的值仍可覆盖单次订阅。Memory 和 Redis 都只保留 `queue_maxsize` 条数据事件;游标早于保留水位线时返回 `StreamGap`,不会从当前最早事件静默部分重放。Redis backend 会在每次 `publish()` / `publish_end()` 刷新 retained stream key TTL;启动恢复与基于 worker lease 的周期恢复共用 Gateway stream terminalization 路径:`RunManager` 先将 orphan run 持久化为 `error` 并写入显式的 `stop_reason=orphan_recovered`,随后 Gateway 发布 `END_SENTINEL` 并安排 stream cleanup。周期扫描、逐行状态写入和 Gateway callback 作为一个受监督的 single-flight 后台 task 执行;慢任务不会堆积,也不会阻塞唯一的 lease heartbeat。shutdown 优先收敛活跃 run,再处理恢复 task;尚未执行的延迟 stream cleanup 会改为立即删除。只有 runtime `yield` 前、无并发请求的启动恢复会把最新受影响 thread 标记为 error;周期恢复不做非原子的 thread 投影。普通 durable terminal status 本身仍不足以证明尾帧已发布;Redis producer 在 durable admission 前取得短租约,经 request task、worker、finalizer 转移并保留到 END。只有显式 `orphan_recovered` 或 durable terminal 加已失效的 tracked producer lease 才能在 heartbeat 上触发 END fallback;lease loss 会阻止后续 stream 写入,active 或 untracked producer 继续等待,避免截断尾帧。TTL 仍是 Redis 内存和故障安全网,不是正常的 subscriber 终止机制。
|
||||
- `stream_bridge.producer_lease_recovery_enabled` 是上述 producer-lease fallback 的 restart-required activate gate。升级时先保持 `false` 完成所有 writer 的滚动发布并排空旧 run,再在所有 Gateway 上同时设为 `true` 并重启;mixed-version 阶段不得启用。
|
||||
- `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]`。
|
||||
|
||||
|
||||
@ -54,6 +54,10 @@ class StreamBridgeConfig(BaseModel):
|
||||
"even if cleanup never runs. Set to 0 to disable. Only applies to the redis bridge."
|
||||
),
|
||||
)
|
||||
producer_lease_recovery_enabled: bool = Field(
|
||||
default=False,
|
||||
description=("Allow terminal Redis consumers to recover a missing END from an expired producer lease. Enable only after every Gateway writer has been upgraded and pre-upgrade runs have drained."),
|
||||
)
|
||||
recovered_stream_cleanup_delay_seconds: float = Field(
|
||||
default=60.0,
|
||||
ge=0,
|
||||
|
||||
@ -28,7 +28,7 @@ from .store import get_store, make_store, reset_store, store_context
|
||||
# NOTE: ``RedisStreamBridge`` is intentionally not re-exported — ``redis`` is an
|
||||
# optional extra and importing it here would load ``redis.asyncio`` in every
|
||||
# process. Import it from ``deerflow.runtime.stream_bridge.redis`` when needed.
|
||||
from .stream_bridge import END_SENTINEL, HEARTBEAT_SENTINEL, MemoryStreamBridge, StreamBridge, StreamEvent, StreamGap, StreamItem, make_stream_bridge
|
||||
from .stream_bridge import END_SENTINEL, HEARTBEAT_SENTINEL, MemoryStreamBridge, StreamBridge, StreamEvent, StreamGap, StreamItem, StreamProducerLeaseState, make_stream_bridge
|
||||
|
||||
__all__ = [
|
||||
# checkpoint state
|
||||
@ -73,5 +73,6 @@ __all__ = [
|
||||
"StreamEvent",
|
||||
"StreamGap",
|
||||
"StreamItem",
|
||||
"StreamProducerLeaseState",
|
||||
"make_stream_bridge",
|
||||
]
|
||||
|
||||
@ -422,6 +422,24 @@ class RunManager:
|
||||
payload["stop_reason"] = record.stop_reason
|
||||
return payload
|
||||
|
||||
@staticmethod
|
||||
def _terminal_row_matches_payload(
|
||||
row: dict[str, Any],
|
||||
payload: dict[str, Any],
|
||||
) -> bool:
|
||||
return all(
|
||||
row.get(key) == payload.get(key)
|
||||
for key in (
|
||||
"thread_id",
|
||||
"user_id",
|
||||
"owner_worker_id",
|
||||
"status",
|
||||
"error",
|
||||
"stop_reason",
|
||||
"goal_verdict",
|
||||
)
|
||||
)
|
||||
|
||||
async def _call_store_with_retry(
|
||||
self,
|
||||
operation_name: str,
|
||||
@ -646,6 +664,7 @@ class RunManager:
|
||||
await self._fence_unconfirmed_terminal(record)
|
||||
return False
|
||||
row_recovery_payload = self._store_put_payload(record, error=error, stop_reason=stop_reason)
|
||||
row_recovery_payload["status"] = status.value
|
||||
try:
|
||||
updated = await self._call_store_with_retry(
|
||||
"update_status",
|
||||
@ -662,7 +681,10 @@ class RunManager:
|
||||
existing = await self._store.get(record.run_id, user_id=record.user_id)
|
||||
if existing is not None:
|
||||
existing_status = existing.get("status")
|
||||
if existing_status == status.value:
|
||||
if existing_status == status.value and self._terminal_row_matches_payload(
|
||||
existing,
|
||||
row_recovery_payload,
|
||||
):
|
||||
logger.info(
|
||||
"Run %s status update to %s was already persisted",
|
||||
record.run_id,
|
||||
@ -1226,6 +1248,7 @@ class RunManager:
|
||||
async with self._lock:
|
||||
if record.status == RunStatus.pending:
|
||||
record.status = RunStatus.interrupted
|
||||
record.terminal_status_persisted = False
|
||||
record.abort_event.set()
|
||||
record.updated_at = _now_iso()
|
||||
return RunStartOutcome.cancelled
|
||||
@ -1257,11 +1280,19 @@ class RunManager:
|
||||
if record is None or record.status != RunStatus.pending:
|
||||
return False
|
||||
record.status = RunStatus.error
|
||||
record.terminal_status_persisted = self._store is None
|
||||
record.error = error
|
||||
record.abort_event.set()
|
||||
record.updated_at = _now_iso()
|
||||
|
||||
await self._persist_status(record, RunStatus.error, error=error)
|
||||
persisted = await self._persist_status(
|
||||
record,
|
||||
RunStatus.error,
|
||||
error=error,
|
||||
)
|
||||
async with self._lock:
|
||||
if self._runs.get(run_id) is record:
|
||||
record.terminal_status_persisted = persisted
|
||||
return True
|
||||
|
||||
async def get_many_by_thread(
|
||||
@ -1735,6 +1766,7 @@ class RunManager:
|
||||
if task_active and record.status == RunStatus.running:
|
||||
record.task.cancel()
|
||||
record.status = RunStatus.interrupted
|
||||
record.terminal_status_persisted = self._store is None
|
||||
record.updated_at = _now_iso()
|
||||
|
||||
# Persist outside the lock so store calls don't block other mutations.
|
||||
@ -1745,6 +1777,9 @@ class RunManager:
|
||||
# cancellation; ordinary runs retain immediate persistence.
|
||||
return CancelOutcome.cancelled
|
||||
persisted = await self._persist_status(record, RunStatus.interrupted)
|
||||
async with self._lock:
|
||||
if self._runs.get(run_id) is record:
|
||||
record.terminal_status_persisted = self._has_final_terminal_ack(record) if self.heartbeat_enabled and self._store is not None else persisted
|
||||
if not persisted and self._store is not None:
|
||||
# ``_persist_status`` already fetched ``existing`` internally;
|
||||
# re-check the store to see if a peer takeover flipped the
|
||||
@ -1878,6 +1913,7 @@ class RunManager:
|
||||
evidence_origin: str = "unknown",
|
||||
evidence_agent_id: str | None = None,
|
||||
idempotency_request: dict[str, Any] | None = None,
|
||||
run_id: str | None = None,
|
||||
) -> RunRecord:
|
||||
"""Atomically admit a normal agent run for a thread."""
|
||||
return await self._admit_thread_operation(
|
||||
@ -1894,6 +1930,7 @@ class RunManager:
|
||||
evidence_origin=evidence_origin,
|
||||
evidence_agent_id=evidence_agent_id,
|
||||
idempotency_request=idempotency_request,
|
||||
run_id=run_id,
|
||||
)
|
||||
|
||||
async def _close_cancelled_admission(self, record: RunRecord) -> None:
|
||||
@ -1942,6 +1979,10 @@ class RunManager:
|
||||
async with self._lock:
|
||||
if self._runs.get(record.run_id) is record:
|
||||
record.status = stored_status
|
||||
record.terminal_status_persisted = stored_status not in (
|
||||
RunStatus.pending,
|
||||
RunStatus.running,
|
||||
)
|
||||
record.error = stored.get("error")
|
||||
record.stop_reason = stored.get("stop_reason")
|
||||
record.updated_at = _now_iso()
|
||||
@ -1962,6 +2003,7 @@ class RunManager:
|
||||
evidence_origin: str = "unknown",
|
||||
evidence_agent_id: str | None = None,
|
||||
idempotency_request: dict[str, Any] | None = None,
|
||||
run_id: str | None = None,
|
||||
) -> RunRecord:
|
||||
"""Atomically check for inflight runs and create a new one.
|
||||
|
||||
@ -1980,7 +2022,7 @@ class RunManager:
|
||||
from deerflow.runtime.assembly_evidence import EVIDENCE_KEY
|
||||
|
||||
self._validate_evidence_origin(evidence_origin)
|
||||
run_id = str(uuid.uuid4())
|
||||
run_id = run_id or str(uuid.uuid4())
|
||||
now = _now_iso()
|
||||
# Resolve before the idempotency checks below compare it with stored rows.
|
||||
user_id = _resolve_record_user_id(user_id)
|
||||
@ -2180,6 +2222,7 @@ class RunManager:
|
||||
if task_active:
|
||||
r.task.cancel()
|
||||
r.status = RunStatus.interrupted
|
||||
r.terminal_status_persisted = self._store is None
|
||||
r.updated_at = now
|
||||
interrupted_records.append(r)
|
||||
|
||||
@ -2192,7 +2235,13 @@ class RunManager:
|
||||
# a terminal write already acknowledged by the atomic transaction.
|
||||
await asyncio.sleep(0)
|
||||
for interrupted_record in interrupted_records:
|
||||
await self._persist_status(interrupted_record, RunStatus.interrupted)
|
||||
persisted = await self._persist_status(
|
||||
interrupted_record,
|
||||
RunStatus.interrupted,
|
||||
)
|
||||
async with self._lock:
|
||||
if self._runs.get(interrupted_record.run_id) is interrupted_record:
|
||||
interrupted_record.terminal_status_persisted = persisted
|
||||
except asyncio.CancelledError:
|
||||
cleanup = asyncio.create_task(self._close_cancelled_admission(record))
|
||||
cleanup.set_name(f"deerflow-close-cancelled-admission-{record.run_id}")
|
||||
@ -2361,6 +2410,7 @@ class RunManager:
|
||||
)
|
||||
continue
|
||||
record.status = RunStatus.error
|
||||
record.terminal_status_persisted = True
|
||||
record.error = error
|
||||
record.stop_reason = stop_reason
|
||||
record.updated_at = now
|
||||
@ -3000,6 +3050,7 @@ class RunManager:
|
||||
continue
|
||||
if record.status in (RunStatus.pending, RunStatus.running):
|
||||
record.status = RunStatus.interrupted
|
||||
record.terminal_status_persisted = self._store is None
|
||||
record.updated_at = _now_iso()
|
||||
to_persist.append(record)
|
||||
|
||||
@ -3028,6 +3079,7 @@ class RunManager:
|
||||
logger.warning("Unexpected error persisting interrupted status for run %s during shutdown: %r", record.run_id, result)
|
||||
elif result is False:
|
||||
logger.warning("Could not persist interrupted status for run %s during shutdown", record.run_id)
|
||||
record.terminal_status_persisted = result is True
|
||||
|
||||
if pending:
|
||||
logger.warning("Run drain exceeded %.1fs on shutdown; %d run task(s) still active and may race checkpointer teardown", timeout, len(pending))
|
||||
|
||||
@ -111,6 +111,15 @@ from .store.base import canonical_run_created_at
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
_THREAD_INCARNATION_UNSET = object()
|
||||
_MISSING_BRIDGE_METHOD = object()
|
||||
|
||||
|
||||
def _declared_bridge_method(bridge: Any, name: str) -> Any | None:
|
||||
"""Return a real bridge method without triggering dynamic mock attributes."""
|
||||
if inspect.getattr_static(bridge, name, _MISSING_BRIDGE_METHOD) is _MISSING_BRIDGE_METHOD:
|
||||
return None
|
||||
method = getattr(bridge, name, None)
|
||||
return method if callable(method) else None
|
||||
|
||||
|
||||
def _log_cancelled_stream_close_failure(
|
||||
@ -999,6 +1008,7 @@ async def run_agent(
|
||||
scheduled_goal: GoalState | None = None
|
||||
final_goal_verdict: dict[str, Any] | None = None
|
||||
produced_output_paths: list[str] | None = None
|
||||
terminal_error_frame: dict[str, str] | None = None
|
||||
# Journal construction moved ahead of preflight so every terminal run can
|
||||
# emit a receipt. Completion persistence keeps its prior boundary: before
|
||||
# #4272 the journal did not exist until preflight had succeeded, so early
|
||||
@ -1012,6 +1022,24 @@ async def run_agent(
|
||||
# finally is safe even if an exception fires before streaming begins.
|
||||
subagent_events: _SubagentEventBuffer | None = None
|
||||
started = False
|
||||
stream_producer_lease_started = False
|
||||
|
||||
async def _ensure_stream_producer_lease() -> None:
|
||||
nonlocal stream_producer_lease_started
|
||||
if stream_producer_lease_started:
|
||||
return
|
||||
start_producer_lease = _declared_bridge_method(
|
||||
bridge,
|
||||
"start_producer_lease",
|
||||
)
|
||||
if start_producer_lease is None:
|
||||
return
|
||||
tracked = await start_producer_lease(run_id)
|
||||
if not tracked:
|
||||
# A registration failure whose marker rollback was confirmed keeps
|
||||
# the historical conservative recovery behavior.
|
||||
return
|
||||
stream_producer_lease_started = True
|
||||
|
||||
async def _finish_cancellation(
|
||||
action: str,
|
||||
@ -1021,6 +1049,7 @@ async def run_agent(
|
||||
nonlocal checkpoint_rollback_completed
|
||||
if record.ownership_lost:
|
||||
return
|
||||
await _ensure_stream_producer_lease()
|
||||
await run_manager.set_finalizing(run_id, True)
|
||||
if record.ownership_lost:
|
||||
return
|
||||
@ -1610,6 +1639,7 @@ async def run_agent(
|
||||
record.goal_verdict = final_goal_verdict
|
||||
|
||||
# 8. Final status
|
||||
await _ensure_stream_producer_lease()
|
||||
if record.abort_event.is_set():
|
||||
await _finish_cancellation(record.abort_action)
|
||||
elif llm_error_fallback_message or (journal is not None and journal.had_llm_error_fallback):
|
||||
@ -1671,6 +1701,7 @@ async def run_agent(
|
||||
except Exception as exc:
|
||||
error_msg = f"{exc}"
|
||||
logger.exception("Run %s failed: %s", run_id, error_msg)
|
||||
await _ensure_stream_producer_lease()
|
||||
await _ensure_finalizing_before_edit_failure(run_manager, record)
|
||||
cancel_action = await run_manager.set_status_if_not_cancelled(
|
||||
run_id,
|
||||
@ -1680,15 +1711,19 @@ async def run_agent(
|
||||
)
|
||||
if cancel_action is not None:
|
||||
await _finish_cancellation(cancel_action)
|
||||
else:
|
||||
await bridge.publish(
|
||||
run_id,
|
||||
"error",
|
||||
{
|
||||
"message": error_msg,
|
||||
"name": type(exc).__name__,
|
||||
},
|
||||
)
|
||||
elif not record.ownership_lost:
|
||||
error_frame = {
|
||||
"message": error_msg,
|
||||
"name": type(exc).__name__,
|
||||
}
|
||||
if record.status not in (RunStatus.pending, RunStatus.running) and getattr(record, "terminal_status_persisted", True):
|
||||
await bridge.publish(
|
||||
run_id,
|
||||
"error",
|
||||
error_frame,
|
||||
)
|
||||
else:
|
||||
terminal_error_frame = error_frame
|
||||
|
||||
finally:
|
||||
completion_hook_task: asyncio.Task[None] | None = None
|
||||
@ -1706,7 +1741,7 @@ async def run_agent(
|
||||
nonlocal accessor, agent, agent_result, cancel_action, checkpoint_rollback_completed
|
||||
nonlocal completion_data, deferred_finalization_interrupt, delivery_content, goal_completion, goal_evaluator_model
|
||||
nonlocal graph_input, pre_run_workspace_snapshot, produced_output_paths, rollback_point, runtime
|
||||
nonlocal runtime_ctx, subagent_events, task_info, task_store
|
||||
nonlocal runtime_ctx, subagent_events, task_info, task_store, terminal_error_frame
|
||||
|
||||
try:
|
||||
if record.ownership_lost:
|
||||
@ -2016,7 +2051,20 @@ async def run_agent(
|
||||
if record.finalizing:
|
||||
await run_manager.set_finalizing(run_id, False)
|
||||
|
||||
await bridge.publish_end(run_id)
|
||||
terminal_stream_authorized = not record.ownership_lost and record.status not in (RunStatus.pending, RunStatus.running) and getattr(record, "terminal_status_persisted", True)
|
||||
if terminal_stream_authorized:
|
||||
if terminal_error_frame is not None:
|
||||
await bridge.publish(
|
||||
run_id,
|
||||
"error",
|
||||
terminal_error_frame,
|
||||
)
|
||||
await bridge.publish_end(run_id)
|
||||
else:
|
||||
logger.warning(
|
||||
"Skipping stream end for run %s because its durable terminal status is unconfirmed",
|
||||
run_id,
|
||||
)
|
||||
|
||||
if deferred_finalization_interrupt is not None:
|
||||
raise deferred_finalization_interrupt
|
||||
@ -2028,6 +2076,15 @@ async def run_agent(
|
||||
except Exception:
|
||||
logger.warning("Failed to close journal for run %s", run_id, exc_info=True)
|
||||
finally:
|
||||
if stream_producer_lease_started:
|
||||
try:
|
||||
await bridge.finish_producer_lease(run_id)
|
||||
except Exception:
|
||||
logger.warning(
|
||||
"Failed to finish stream producer lease for run %s",
|
||||
run_id,
|
||||
exc_info=True,
|
||||
)
|
||||
lease_cleanup_interrupt: BaseException | None = None
|
||||
try:
|
||||
from deerflow.sandbox.lease import release_sandbox_execution_lease_async
|
||||
@ -2089,6 +2146,13 @@ async def run_agent(
|
||||
# The registered worker remains the owner until every terminal and
|
||||
# resource-cleanup await settles. Never cancel or detach this task.
|
||||
finalizer = asyncio.create_task(finalize(), name=f"run-finalization-{run_id}")
|
||||
if stream_producer_lease_started:
|
||||
transfer_producer_lease = _declared_bridge_method(
|
||||
bridge,
|
||||
"transfer_producer_lease",
|
||||
)
|
||||
if transfer_producer_lease is not None:
|
||||
transfer_producer_lease(run_id, finalizer)
|
||||
await _drain_run_finalization(finalizer, on_host_cancel=interrupt_completion_hook)
|
||||
|
||||
|
||||
|
||||
@ -8,7 +8,7 @@ by :mod:`asyncio.Queue`.
|
||||
"""
|
||||
|
||||
from .async_provider import make_stream_bridge
|
||||
from .base import END_SENTINEL, HEARTBEAT_SENTINEL, StreamBridge, StreamEvent, StreamGap, StreamItem
|
||||
from .base import END_SENTINEL, HEARTBEAT_SENTINEL, StreamBridge, StreamEvent, StreamGap, StreamItem, StreamProducerLeaseState
|
||||
from .memory import MemoryStreamBridge
|
||||
|
||||
# NOTE: ``RedisStreamBridge`` is intentionally NOT imported here. ``redis`` is an
|
||||
@ -27,5 +27,6 @@ __all__ = [
|
||||
"StreamEvent",
|
||||
"StreamGap",
|
||||
"StreamItem",
|
||||
"StreamProducerLeaseState",
|
||||
"make_stream_bridge",
|
||||
]
|
||||
|
||||
@ -85,6 +85,11 @@ async def make_stream_bridge(app_config: AppConfig | None = None) -> AsyncIterat
|
||||
heartbeat_interval=config.heartbeat_interval_seconds,
|
||||
max_connections=config.max_connections,
|
||||
stream_ttl_seconds=config.stream_ttl_seconds,
|
||||
producer_lease_recovery_enabled=getattr(
|
||||
config,
|
||||
"producer_lease_recovery_enabled",
|
||||
False,
|
||||
),
|
||||
)
|
||||
logger.info(
|
||||
"Stream bridge initialised: redis (queue_maxsize=%d, heartbeat_interval_seconds=%.1f, max_connections=%s, stream_ttl_seconds=%d)",
|
||||
|
||||
@ -11,6 +11,7 @@ import abc
|
||||
import math
|
||||
from collections.abc import AsyncIterator
|
||||
from dataclasses import dataclass
|
||||
from enum import StrEnum
|
||||
from typing import Any
|
||||
|
||||
from deerflow.config.stream_bridge_config import DEFAULT_HEARTBEAT_INTERVAL_SECONDS, MAX_HEARTBEAT_INTERVAL_SECONDS
|
||||
@ -54,10 +55,19 @@ END_SENTINEL = StreamEvent(id="", event="__end__", data=None)
|
||||
type StreamItem = StreamEvent | StreamGap
|
||||
|
||||
|
||||
class StreamProducerLeaseState(StrEnum):
|
||||
"""Cross-process liveness state for a run's stream producer."""
|
||||
|
||||
untracked = "untracked"
|
||||
active = "active"
|
||||
expired = "expired"
|
||||
|
||||
|
||||
class StreamBridge(abc.ABC):
|
||||
"""Abstract base for stream bridges."""
|
||||
|
||||
supports_cross_process: bool = False
|
||||
producer_lease_recovery_enabled: bool = False
|
||||
|
||||
def __init__(self, *, heartbeat_interval: float = DEFAULT_HEARTBEAT_INTERVAL_SECONDS) -> None:
|
||||
self._heartbeat_interval = self._validate_heartbeat_interval(heartbeat_interval)
|
||||
@ -123,5 +133,31 @@ class StreamBridge(abc.ABC):
|
||||
"""
|
||||
return True
|
||||
|
||||
async def start_producer_lease(self, run_id: str) -> bool:
|
||||
"""Start producer liveness tracking; return whether it is active."""
|
||||
return False
|
||||
|
||||
async def finish_producer_lease(self, run_id: str) -> None:
|
||||
"""Stop producer liveness tracking when supported."""
|
||||
|
||||
async def abort_producer_lease(self, run_id: str) -> None:
|
||||
"""Discard tracking for a producer whose run was never admitted."""
|
||||
await self.finish_producer_lease(run_id)
|
||||
|
||||
def transfer_producer_lease(
|
||||
self,
|
||||
run_id: str,
|
||||
owner_task: Any,
|
||||
) -> None:
|
||||
"""Transfer cancellation ownership to the task publishing the tail."""
|
||||
|
||||
async def producer_lease_state(self, run_id: str) -> StreamProducerLeaseState:
|
||||
"""Return producer liveness without inferring it from stream contents."""
|
||||
return StreamProducerLeaseState.untracked
|
||||
|
||||
async def cleanup_expired_producer(self, run_id: str) -> bool:
|
||||
"""Delete an expired producer marker without touching an active owner."""
|
||||
return False
|
||||
|
||||
async def close(self) -> None:
|
||||
"""Release backend resources. Default is a no-op."""
|
||||
|
||||
@ -8,7 +8,10 @@ import json
|
||||
import logging
|
||||
import re
|
||||
from collections.abc import AsyncIterator, Mapping
|
||||
from contextvars import Context
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
from uuid import uuid4
|
||||
|
||||
try:
|
||||
from redis.asyncio import Redis
|
||||
@ -27,7 +30,9 @@ except ImportError: # pragma: no cover - only hit when the optional extra is mi
|
||||
"Or switch to stream_bridge.type: memory in config.yaml for single-process deployment."
|
||||
) from None
|
||||
|
||||
from .base import DEFAULT_HEARTBEAT_INTERVAL_SECONDS, END_SENTINEL, HEARTBEAT_SENTINEL, StreamBridge, StreamEvent, StreamGap, StreamItem
|
||||
from ...utils.file_io import await_drained
|
||||
from ..keyed_lock import AsyncKeyedLockTable
|
||||
from .base import DEFAULT_HEARTBEAT_INTERVAL_SECONDS, END_SENTINEL, HEARTBEAT_SENTINEL, StreamBridge, StreamEvent, StreamGap, StreamItem, StreamProducerLeaseState
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@ -47,6 +52,70 @@ _XREAD_COUNT = 64
|
||||
# propagates to the caller. Brief blips are retried with exponential backoff
|
||||
# capped at ``heartbeat_interval``.
|
||||
_MAX_SUBSCRIBE_RETRIES = 3
|
||||
_PRODUCER_LEASE_SECONDS = 30
|
||||
_PRODUCER_LEASE_RENEW_INTERVAL_SECONDS = 10
|
||||
_PRODUCER_LEASE_START_TIMEOUT_SECONDS = 5.0
|
||||
_PRODUCER_LEASE_RELEASE_TIMEOUT_SECONDS = 5.0
|
||||
_PRODUCER_LEASE_RELEASE_RETRY_SECONDS = 0.05
|
||||
_PRODUCER_TRACKING_MIN_TTL_SECONDS = _PRODUCER_LEASE_SECONDS * 4
|
||||
_COMPARE_DELETE_SCRIPT = """
|
||||
if redis.call("get", KEYS[1]) == ARGV[1] then
|
||||
return redis.call("del", KEYS[1])
|
||||
end
|
||||
return 0
|
||||
"""
|
||||
_COMPARE_RENEW_SCRIPT = """
|
||||
if redis.call("get", KEYS[1]) ~= ARGV[1] then return 0 end
|
||||
if redis.call("get", KEYS[2]) ~= ARGV[1] then return 0 end
|
||||
redis.call("expire", KEYS[1], ARGV[2])
|
||||
if tonumber(ARGV[3]) > 0 then redis.call("expire", KEYS[2], ARGV[3]) end
|
||||
return 1
|
||||
"""
|
||||
_COMPARE_XADD_SCRIPT = """
|
||||
if redis.call("get", KEYS[1]) ~= ARGV[1] then return false end
|
||||
if redis.call("get", KEYS[2]) ~= ARGV[1] then return false end
|
||||
local event_id = redis.call(
|
||||
"xadd", KEYS[3], "MAXLEN", "=", ARGV[2], "*",
|
||||
"kind", ARGV[4], "event", ARGV[5], "data", ARGV[6]
|
||||
)
|
||||
if tonumber(ARGV[3]) > 0 then redis.call("expire", KEYS[3], ARGV[3]) end
|
||||
return event_id
|
||||
"""
|
||||
_XADD_IF_NO_ACTIVE_PRODUCER_SCRIPT = """
|
||||
if redis.call("exists", KEYS[1]) == 1 then return false end
|
||||
local event_id = redis.call(
|
||||
"xadd", KEYS[2], "MAXLEN", "=", ARGV[1], "*",
|
||||
"kind", ARGV[3], "event", ARGV[4], "data", ARGV[5]
|
||||
)
|
||||
if tonumber(ARGV[2]) > 0 then redis.call("expire", KEYS[2], ARGV[2]) end
|
||||
return event_id
|
||||
"""
|
||||
_CLEANUP_IF_TOKEN_SCRIPT = """
|
||||
local tracked = redis.call("get", KEYS[2])
|
||||
local active = redis.call("get", KEYS[3])
|
||||
if tracked and tracked ~= ARGV[1] then return 0 end
|
||||
if active and active ~= ARGV[1] then return 0 end
|
||||
return redis.call("del", KEYS[1], KEYS[2], KEYS[3])
|
||||
"""
|
||||
_CLEANUP_IF_NO_ACTIVE_PRODUCER_SCRIPT = """
|
||||
if redis.call("exists", KEYS[3]) == 1 then return 0 end
|
||||
return redis.call("del", KEYS[1], KEYS[2])
|
||||
"""
|
||||
_DELETE_EXPIRED_TRACKING_SCRIPT = """
|
||||
if redis.call("exists", KEYS[2]) == 1 then return 0 end
|
||||
if not redis.call("get", KEYS[1]) then return 2 end
|
||||
return redis.call("del", KEYS[1])
|
||||
"""
|
||||
|
||||
|
||||
@dataclass
|
||||
class _ProducerLease:
|
||||
token: str
|
||||
owner_task: asyncio.Task[Any]
|
||||
renewal_task: asyncio.Task[None] | None = None
|
||||
lost: bool = False
|
||||
stopped: bool = False
|
||||
discard_tracking_requested: bool = False
|
||||
|
||||
|
||||
class RedisStreamBridge(StreamBridge):
|
||||
@ -58,6 +127,7 @@ class RedisStreamBridge(StreamBridge):
|
||||
"""
|
||||
|
||||
supports_cross_process = True
|
||||
producer_lease_retry_delay_seconds = _PRODUCER_LEASE_SECONDS + 1
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
@ -68,6 +138,7 @@ class RedisStreamBridge(StreamBridge):
|
||||
heartbeat_interval: float = DEFAULT_HEARTBEAT_INTERVAL_SECONDS,
|
||||
max_connections: int | None = None,
|
||||
stream_ttl_seconds: int | None = 86400,
|
||||
producer_lease_recovery_enabled: bool = False,
|
||||
client: Redis | None = None,
|
||||
) -> None:
|
||||
super().__init__(heartbeat_interval=heartbeat_interval)
|
||||
@ -83,10 +154,438 @@ class RedisStreamBridge(StreamBridge):
|
||||
# caps that pool; ``None`` keeps redis-py's effectively-unbounded default.
|
||||
self._redis = client if client is not None else Redis.from_url(redis_url, decode_responses=True, max_connections=max_connections)
|
||||
self._owns_client = client is None
|
||||
self.producer_lease_recovery_enabled = producer_lease_recovery_enabled
|
||||
self._producer_leases: dict[str, _ProducerLease] = {}
|
||||
self._producer_lease_locks = AsyncKeyedLockTable[str]()
|
||||
self._closing = False
|
||||
self._producer_setups = 0
|
||||
self._producer_setups_drained = asyncio.Event()
|
||||
self._producer_setups_drained.set()
|
||||
|
||||
def _stream_key(self, run_id: str) -> str:
|
||||
return f"{self._key_prefix}:{run_id}"
|
||||
|
||||
def _producer_tracking_key(self, run_id: str) -> str:
|
||||
return f"{self._stream_key(run_id)}:producer-tracked"
|
||||
|
||||
def _producer_active_key(self, run_id: str) -> str:
|
||||
return f"{self._stream_key(run_id)}:producer-active"
|
||||
|
||||
def _producer_tracking_ttl_seconds(self) -> int:
|
||||
if self._stream_ttl_seconds is None:
|
||||
return 0
|
||||
return max(
|
||||
self._stream_ttl_seconds,
|
||||
_PRODUCER_TRACKING_MIN_TTL_SECONDS,
|
||||
)
|
||||
|
||||
async def _delete_if_token(self, key: str, token: str) -> bool:
|
||||
return bool(
|
||||
await self._redis.eval(
|
||||
_COMPARE_DELETE_SCRIPT,
|
||||
1,
|
||||
key,
|
||||
token,
|
||||
)
|
||||
)
|
||||
|
||||
async def _rollback_producer_lease(
|
||||
self,
|
||||
run_id: str,
|
||||
token: str,
|
||||
*,
|
||||
timeout: float = _PRODUCER_LEASE_RELEASE_TIMEOUT_SECONDS,
|
||||
) -> bool:
|
||||
keys = (
|
||||
self._producer_tracking_key(run_id),
|
||||
self._producer_active_key(run_id),
|
||||
)
|
||||
loop = asyncio.get_running_loop()
|
||||
deadline = loop.time() + max(0.0, timeout)
|
||||
last_error: BaseException | None = None
|
||||
while True:
|
||||
remaining_seconds = deadline - loop.time()
|
||||
if remaining_seconds <= 0:
|
||||
break
|
||||
try:
|
||||
async with asyncio.timeout(remaining_seconds):
|
||||
for key in keys:
|
||||
await self._delete_if_token(key, token)
|
||||
remaining = await self._redis.mget(*keys)
|
||||
return all(value is None or self._decode(value) != token for value in remaining)
|
||||
except RedisError as exc:
|
||||
last_error = exc
|
||||
await asyncio.sleep(
|
||||
min(
|
||||
_PRODUCER_LEASE_RELEASE_RETRY_SECONDS,
|
||||
max(0.0, deadline - loop.time()),
|
||||
)
|
||||
)
|
||||
except TimeoutError as exc:
|
||||
last_error = exc
|
||||
break
|
||||
logger.warning(
|
||||
"Failed to roll back Redis stream producer lease for run %s",
|
||||
run_id,
|
||||
exc_info=((type(last_error), last_error, last_error.__traceback__) if last_error is not None else False),
|
||||
)
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
def _cancel_lost_producer(run_id: str, lease: _ProducerLease) -> None:
|
||||
lease.lost = True
|
||||
if lease.owner_task.done():
|
||||
return
|
||||
logger.error(
|
||||
"Redis stream producer lease expired for run %s; cancelling its owner",
|
||||
run_id,
|
||||
)
|
||||
lease.owner_task.cancel(f"Redis stream producer lease expired for run {run_id}")
|
||||
|
||||
def _watch_producer_owner(
|
||||
self,
|
||||
run_id: str,
|
||||
lease: _ProducerLease,
|
||||
owner_task: asyncio.Task[Any],
|
||||
) -> None:
|
||||
def owner_done(done_task: asyncio.Task[Any]) -> None:
|
||||
current = self._producer_leases.get(run_id)
|
||||
if current is not lease or current.owner_task is not done_task:
|
||||
return
|
||||
cleanup = asyncio.create_task(
|
||||
self._stop_producer_lease(
|
||||
run_id,
|
||||
discard_tracking=False,
|
||||
expected_owner_task=done_task,
|
||||
),
|
||||
name=f"redis-stream-producer-owner-cleanup:{run_id}",
|
||||
context=Context(),
|
||||
)
|
||||
cleanup.add_done_callback(
|
||||
lambda task: self._log_owner_cleanup_result(
|
||||
run_id,
|
||||
task,
|
||||
)
|
||||
)
|
||||
|
||||
owner_task.add_done_callback(owner_done)
|
||||
|
||||
@staticmethod
|
||||
def _log_owner_cleanup_result(
|
||||
run_id: str,
|
||||
task: asyncio.Task[None],
|
||||
) -> None:
|
||||
if task.cancelled():
|
||||
return
|
||||
try:
|
||||
task.result()
|
||||
except Exception:
|
||||
logger.warning(
|
||||
"Failed to release Redis stream producer lease after owner exit for run %s",
|
||||
run_id,
|
||||
exc_info=True,
|
||||
)
|
||||
|
||||
def _ensure_producer_can_publish(self, run_id: str) -> None:
|
||||
lease = self._producer_leases.get(run_id)
|
||||
if lease is None:
|
||||
return
|
||||
if lease.lost:
|
||||
raise RuntimeError(f"Redis stream producer lease was lost for run {run_id}")
|
||||
if lease.stopped:
|
||||
raise RuntimeError(f"Redis stream producer lease was already stopped for run {run_id}")
|
||||
|
||||
async def _renew_producer_lease(self, run_id: str, lease: _ProducerLease) -> None:
|
||||
loop = asyncio.get_running_loop()
|
||||
deadline = loop.time() + _PRODUCER_LEASE_SECONDS
|
||||
while True:
|
||||
remaining = deadline - loop.time()
|
||||
if remaining <= 0:
|
||||
self._cancel_lost_producer(run_id, lease)
|
||||
return
|
||||
try:
|
||||
await asyncio.sleep(min(_PRODUCER_LEASE_RENEW_INTERVAL_SECONDS, remaining))
|
||||
remaining = deadline - loop.time()
|
||||
if remaining <= 0:
|
||||
self._cancel_lost_producer(run_id, lease)
|
||||
return
|
||||
async with asyncio.timeout(remaining):
|
||||
if not await self._renew_owned_producer_lease(run_id, lease):
|
||||
self._cancel_lost_producer(run_id, lease)
|
||||
return
|
||||
deadline = loop.time() + _PRODUCER_LEASE_SECONDS
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
except (RedisError, TimeoutError):
|
||||
if loop.time() >= deadline:
|
||||
self._cancel_lost_producer(run_id, lease)
|
||||
return
|
||||
logger.warning(
|
||||
"Failed to renew Redis stream producer lease for run %s; will retry before the confirmed deadline",
|
||||
run_id,
|
||||
exc_info=True,
|
||||
)
|
||||
|
||||
async def start_producer_lease(self, run_id: str) -> bool:
|
||||
owner_task = asyncio.current_task()
|
||||
if owner_task is None:
|
||||
raise RuntimeError("Stream producer lease requires an owning asyncio task")
|
||||
if self._closing:
|
||||
raise RuntimeError("Stream bridge is closing")
|
||||
|
||||
self._producer_setups += 1
|
||||
self._producer_setups_drained.clear()
|
||||
try:
|
||||
return await self._start_producer_lease(run_id, owner_task)
|
||||
finally:
|
||||
self._producer_setups -= 1
|
||||
if self._producer_setups == 0:
|
||||
self._producer_setups_drained.set()
|
||||
|
||||
async def _start_producer_lease(
|
||||
self,
|
||||
run_id: str,
|
||||
owner_task: asyncio.Task[Any],
|
||||
) -> bool:
|
||||
async with self._producer_lease_locks.hold(run_id):
|
||||
if self._closing:
|
||||
raise RuntimeError("Stream bridge is closing")
|
||||
existing = self._producer_leases.get(run_id)
|
||||
if existing is not None:
|
||||
if existing.owner_task is owner_task and not existing.stopped:
|
||||
return True
|
||||
raise RuntimeError(f"Run {run_id} already has a local stream producer")
|
||||
|
||||
token = uuid4().hex
|
||||
active_key = self._producer_active_key(run_id)
|
||||
try:
|
||||
async with asyncio.timeout(_PRODUCER_LEASE_START_TIMEOUT_SECONDS):
|
||||
claimed = await self._redis.set(
|
||||
active_key,
|
||||
token,
|
||||
ex=_PRODUCER_LEASE_SECONDS,
|
||||
nx=True,
|
||||
)
|
||||
if not claimed:
|
||||
raise RuntimeError(f"Run {run_id} already has an active stream producer")
|
||||
tracking_ttl = self._producer_tracking_ttl_seconds()
|
||||
tracking_created = await self._redis.set(
|
||||
self._producer_tracking_key(run_id),
|
||||
token,
|
||||
**({"ex": tracking_ttl} if tracking_ttl > 0 else {}),
|
||||
)
|
||||
if not tracking_created:
|
||||
raise RuntimeError(f"Run {run_id} stream producer tracking marker disappeared")
|
||||
except asyncio.CancelledError:
|
||||
await await_drained(
|
||||
self._rollback_producer_lease(run_id, token),
|
||||
)
|
||||
raise
|
||||
except (RedisError, TimeoutError) as exc:
|
||||
if await self._rollback_producer_lease(run_id, token):
|
||||
logger.warning(
|
||||
"Could not start Redis stream producer lease for run %s; the partial registration was rolled back",
|
||||
run_id,
|
||||
exc_info=True,
|
||||
)
|
||||
return False
|
||||
raise RuntimeError(f"Could not safely establish stream producer lease for run {run_id}") from exc
|
||||
except BaseException:
|
||||
await await_drained(
|
||||
self._rollback_producer_lease(run_id, token),
|
||||
)
|
||||
raise
|
||||
|
||||
lease = _ProducerLease(token=token, owner_task=owner_task)
|
||||
lease.renewal_task = asyncio.create_task(
|
||||
self._renew_producer_lease(run_id, lease),
|
||||
name=f"redis-stream-producer-lease:{run_id}",
|
||||
)
|
||||
self._producer_leases[run_id] = lease
|
||||
self._watch_producer_owner(run_id, lease, owner_task)
|
||||
return True
|
||||
|
||||
async def _renew_owned_producer_lease(
|
||||
self,
|
||||
run_id: str,
|
||||
lease: _ProducerLease,
|
||||
) -> bool:
|
||||
return bool(
|
||||
await self._redis.eval(
|
||||
_COMPARE_RENEW_SCRIPT,
|
||||
2,
|
||||
self._producer_active_key(run_id),
|
||||
self._producer_tracking_key(run_id),
|
||||
lease.token,
|
||||
_PRODUCER_LEASE_SECONDS,
|
||||
self._producer_tracking_ttl_seconds(),
|
||||
)
|
||||
)
|
||||
|
||||
async def _xadd_owned(
|
||||
self,
|
||||
run_id: str,
|
||||
lease: _ProducerLease,
|
||||
fields: dict[str, str],
|
||||
*,
|
||||
maxlen: int,
|
||||
) -> None:
|
||||
result = await self._redis.eval(
|
||||
_COMPARE_XADD_SCRIPT,
|
||||
3,
|
||||
self._producer_active_key(run_id),
|
||||
self._producer_tracking_key(run_id),
|
||||
self._stream_key(run_id),
|
||||
lease.token,
|
||||
maxlen,
|
||||
self._stream_ttl_seconds or 0,
|
||||
fields.get("kind", _KIND_EVENT),
|
||||
fields.get("event", ""),
|
||||
fields.get("data", ""),
|
||||
)
|
||||
if result:
|
||||
return
|
||||
self._cancel_lost_producer(run_id, lease)
|
||||
raise RuntimeError(f"Redis stream producer lease was lost for run {run_id}")
|
||||
|
||||
async def _xadd_without_active_producer(
|
||||
self,
|
||||
run_id: str,
|
||||
fields: dict[str, str],
|
||||
*,
|
||||
maxlen: int,
|
||||
) -> None:
|
||||
result = await self._redis.eval(
|
||||
_XADD_IF_NO_ACTIVE_PRODUCER_SCRIPT,
|
||||
2,
|
||||
self._producer_active_key(run_id),
|
||||
self._stream_key(run_id),
|
||||
maxlen,
|
||||
self._stream_ttl_seconds or 0,
|
||||
fields.get("kind", _KIND_EVENT),
|
||||
fields.get("event", ""),
|
||||
fields.get("data", ""),
|
||||
)
|
||||
if result:
|
||||
return
|
||||
raise RuntimeError(f"Cannot publish stream data for run {run_id} because another active producer exists")
|
||||
|
||||
async def _stop_producer_lease(
|
||||
self,
|
||||
run_id: str,
|
||||
*,
|
||||
discard_tracking: bool,
|
||||
expected_owner_task: asyncio.Task[Any] | None = None,
|
||||
closing: bool = False,
|
||||
) -> None:
|
||||
async with self._producer_lease_locks.hold(run_id):
|
||||
lease = self._producer_leases.get(run_id)
|
||||
if lease is not None and expected_owner_task is not None and lease.owner_task is not expected_owner_task:
|
||||
return
|
||||
if lease is None:
|
||||
return
|
||||
discard_tracking = discard_tracking or lease.discard_tracking_requested
|
||||
|
||||
renewal_task = lease.renewal_task
|
||||
if renewal_task is not None:
|
||||
renewal_task.cancel()
|
||||
await asyncio.gather(renewal_task, return_exceptions=True)
|
||||
lease.renewal_task = None
|
||||
|
||||
if discard_tracking:
|
||||
lease.discard_tracking_requested = True
|
||||
discarded = await self._rollback_producer_lease(
|
||||
run_id,
|
||||
lease.token,
|
||||
)
|
||||
lease.stopped = True
|
||||
if discarded:
|
||||
self._producer_leases.pop(run_id, None)
|
||||
elif closing:
|
||||
raise RuntimeError(f"Could not discard Redis stream producer tracking for run {run_id}")
|
||||
return
|
||||
|
||||
if lease.stopped:
|
||||
if closing:
|
||||
self._producer_leases.pop(run_id, None)
|
||||
return
|
||||
release_deadline = asyncio.get_running_loop().time() + _PRODUCER_LEASE_RELEASE_TIMEOUT_SECONDS
|
||||
try:
|
||||
async with asyncio.timeout(
|
||||
max(
|
||||
0.0,
|
||||
release_deadline - asyncio.get_running_loop().time(),
|
||||
)
|
||||
):
|
||||
await self._delete_if_token(
|
||||
self._producer_active_key(run_id),
|
||||
lease.token,
|
||||
)
|
||||
except (RedisError, TimeoutError):
|
||||
# The key still has a short TTL; consumers remain conservative until it expires.
|
||||
logger.warning(
|
||||
"Failed to stop Redis stream producer lease for run %s",
|
||||
run_id,
|
||||
exc_info=True,
|
||||
)
|
||||
finally:
|
||||
lease.stopped = True
|
||||
if closing:
|
||||
if lease.discard_tracking_requested:
|
||||
discarded = await self._rollback_producer_lease(
|
||||
run_id,
|
||||
lease.token,
|
||||
)
|
||||
if not discarded:
|
||||
raise RuntimeError(f"Could not discard Redis stream producer tracking for run {run_id}")
|
||||
self._producer_leases.pop(run_id, None)
|
||||
|
||||
async def finish_producer_lease(self, run_id: str) -> None:
|
||||
await self._stop_producer_lease(
|
||||
run_id,
|
||||
discard_tracking=False,
|
||||
)
|
||||
|
||||
async def abort_producer_lease(self, run_id: str) -> None:
|
||||
lease = self._producer_leases.get(run_id)
|
||||
if lease is not None:
|
||||
lease.discard_tracking_requested = True
|
||||
await self._stop_producer_lease(
|
||||
run_id,
|
||||
discard_tracking=True,
|
||||
)
|
||||
|
||||
def transfer_producer_lease(
|
||||
self,
|
||||
run_id: str,
|
||||
owner_task: asyncio.Task[Any],
|
||||
) -> None:
|
||||
lease = self._producer_leases.get(run_id)
|
||||
if lease is not None and not lease.stopped:
|
||||
lease.owner_task = owner_task
|
||||
self._watch_producer_owner(run_id, lease, owner_task)
|
||||
|
||||
async def producer_lease_state(self, run_id: str) -> StreamProducerLeaseState:
|
||||
tracked, active = await self._redis.mget(
|
||||
self._producer_tracking_key(run_id),
|
||||
self._producer_active_key(run_id),
|
||||
)
|
||||
if tracked is None:
|
||||
return StreamProducerLeaseState.untracked
|
||||
if active is not None and self._decode(active) == self._decode(tracked):
|
||||
return StreamProducerLeaseState.active
|
||||
return StreamProducerLeaseState.expired
|
||||
|
||||
async def cleanup_expired_producer(self, run_id: str) -> bool:
|
||||
return bool(
|
||||
await self._redis.eval(
|
||||
_DELETE_EXPIRED_TRACKING_SCRIPT,
|
||||
2,
|
||||
self._producer_tracking_key(run_id),
|
||||
self._producer_active_key(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(
|
||||
@ -172,23 +671,43 @@ class RedisStreamBridge(StreamBridge):
|
||||
return None
|
||||
|
||||
async def publish(self, run_id: str, event: str, data: Any) -> None:
|
||||
key = self._stream_key(run_id)
|
||||
await self._xadd_retained(
|
||||
key,
|
||||
{
|
||||
"kind": _KIND_EVENT,
|
||||
"event": event,
|
||||
"data": self._encode_data(data),
|
||||
},
|
||||
self._ensure_producer_can_publish(run_id)
|
||||
lease = self._producer_leases.get(run_id)
|
||||
fields = {
|
||||
"kind": _KIND_EVENT,
|
||||
"event": event,
|
||||
"data": self._encode_data(data),
|
||||
}
|
||||
if lease is not None:
|
||||
await self._xadd_owned(
|
||||
run_id,
|
||||
lease,
|
||||
fields,
|
||||
maxlen=self._maxsize,
|
||||
)
|
||||
return
|
||||
await self._xadd_without_active_producer(
|
||||
run_id,
|
||||
fields,
|
||||
maxlen=self._maxsize,
|
||||
)
|
||||
|
||||
async def publish_end(self, run_id: str) -> None:
|
||||
self._ensure_producer_can_publish(run_id)
|
||||
lease = self._producer_leases.get(run_id)
|
||||
fields = {"kind": _KIND_END}
|
||||
if lease is not None:
|
||||
await self._xadd_owned(
|
||||
run_id,
|
||||
lease,
|
||||
fields,
|
||||
maxlen=self._maxsize + 1,
|
||||
)
|
||||
return
|
||||
# Keep the configured number of data events plus the internal end marker.
|
||||
key = self._stream_key(run_id)
|
||||
await self._xadd_retained(
|
||||
key,
|
||||
{"kind": _KIND_END},
|
||||
await self._xadd_without_active_producer(
|
||||
run_id,
|
||||
fields,
|
||||
maxlen=self._maxsize + 1,
|
||||
)
|
||||
|
||||
@ -425,14 +944,57 @@ class RedisStreamBridge(StreamBridge):
|
||||
async def cleanup(self, run_id: str, *, delay: float = 0) -> None:
|
||||
if delay > 0:
|
||||
await asyncio.sleep(delay)
|
||||
await self._redis.delete(self._stream_key(run_id))
|
||||
lease = self._producer_leases.get(run_id)
|
||||
token = lease.token if lease is not None else None
|
||||
await self.finish_producer_lease(run_id)
|
||||
if token is None:
|
||||
await self._redis.eval(
|
||||
_CLEANUP_IF_NO_ACTIVE_PRODUCER_SCRIPT,
|
||||
3,
|
||||
self._stream_key(run_id),
|
||||
self._producer_tracking_key(run_id),
|
||||
self._producer_active_key(run_id),
|
||||
)
|
||||
else:
|
||||
await self._redis.eval(
|
||||
_CLEANUP_IF_TOKEN_SCRIPT,
|
||||
3,
|
||||
self._stream_key(run_id),
|
||||
self._producer_tracking_key(run_id),
|
||||
self._producer_active_key(run_id),
|
||||
token,
|
||||
)
|
||||
async with self._producer_lease_locks.hold(run_id):
|
||||
current = self._producer_leases.get(run_id)
|
||||
if current is lease:
|
||||
self._producer_leases.pop(run_id, None)
|
||||
|
||||
async def close(self) -> None:
|
||||
self._closing = True
|
||||
await self._producer_setups_drained.wait()
|
||||
run_ids = tuple(self._producer_leases)
|
||||
failures: list[BaseException] = []
|
||||
if run_ids:
|
||||
results = await asyncio.gather(
|
||||
*(
|
||||
self._stop_producer_lease(
|
||||
run_id,
|
||||
discard_tracking=False,
|
||||
closing=True,
|
||||
)
|
||||
for run_id in run_ids
|
||||
),
|
||||
return_exceptions=True,
|
||||
)
|
||||
failures = [result for result in results if isinstance(result, BaseException)]
|
||||
if not self._owns_client:
|
||||
if failures:
|
||||
raise failures[0]
|
||||
return
|
||||
close = getattr(self._redis, "aclose", None) or getattr(self._redis, "close", None)
|
||||
if close is None:
|
||||
return
|
||||
result = close()
|
||||
if inspect.isawaitable(result):
|
||||
await result
|
||||
if close is not None:
|
||||
result = close()
|
||||
if inspect.isawaitable(result):
|
||||
await result
|
||||
if failures:
|
||||
raise failures[0]
|
||||
|
||||
@ -385,3 +385,30 @@ async def test_delayed_atomic_interrupted_ack_cannot_restart_acknowledged_error(
|
||||
release_ack.set()
|
||||
await asyncio.gather(admission, return_exceptions=True)
|
||||
await manager.stop_heartbeat(timeout=0.1)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
@pytest.mark.parametrize("action", ["interrupt", "rollback"])
|
||||
async def test_direct_cancel_keeps_terminal_confirmation_behind_rollback(store, action):
|
||||
from deerflow.runtime.runs.manager import CancelOutcome
|
||||
|
||||
manager = manager_for(store)
|
||||
record = await manager.create_or_reject("thread", user_id="review")
|
||||
await manager.try_start(record.run_id)
|
||||
record.task = asyncio.create_task(asyncio.Event().wait())
|
||||
try:
|
||||
assert await manager.cancel(record.run_id, action=action) == CancelOutcome.cancelled
|
||||
assert (await store.get(record.run_id, user_id="review"))["status"] == "interrupted"
|
||||
assert record.terminal_commit_proof.status == RunStatus.interrupted
|
||||
assert record.rollback_finalization_pending is (action == "rollback")
|
||||
assert record.terminal_status_persisted is (action == "interrupt")
|
||||
|
||||
if action == "rollback":
|
||||
await manager.set_status(record.run_id, RunStatus.error, error="Rolled back by user")
|
||||
assert record.terminal_status_persisted
|
||||
assert not record.rollback_finalization_pending
|
||||
assert record.terminal_commit_proof.status == RunStatus.error
|
||||
finally:
|
||||
record.task.cancel()
|
||||
await asyncio.gather(record.task, return_exceptions=True)
|
||||
await manager.stop_heartbeat(timeout=0.1)
|
||||
|
||||
@ -89,6 +89,7 @@ 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.expired_producer_cleanup_calls: list[str] = []
|
||||
self._existing_streams: set[str] = existing_streams if existing_streams is not None else set()
|
||||
|
||||
async def stream_exists(self, run_id: str) -> bool:
|
||||
@ -100,6 +101,10 @@ class _FakeStreamBridge:
|
||||
async def cleanup(self, run_id: str, *, delay: float = 0) -> None:
|
||||
self.cleanup_calls.append((run_id, delay))
|
||||
|
||||
async def cleanup_expired_producer(self, run_id: str) -> bool:
|
||||
self.expired_producer_cleanup_calls.append(run_id)
|
||||
return True
|
||||
|
||||
|
||||
class _RetainedMemoryStreamBridge(MemoryStreamBridge):
|
||||
"""Memory bridge that records cleanup without deleting test history."""
|
||||
@ -130,6 +135,34 @@ class _DelayedCleanupStreamBridge(_FakeStreamBridge):
|
||||
raise
|
||||
|
||||
|
||||
class _RetryingRecoveredStreamBridge(_FakeStreamBridge):
|
||||
producer_lease_retry_delay_seconds = 0
|
||||
|
||||
def __init__(self) -> None:
|
||||
super().__init__(existing_streams={"run-1"})
|
||||
self.publish_attempts = 0
|
||||
|
||||
async def publish_end(self, run_id: str) -> None:
|
||||
self.publish_attempts += 1
|
||||
if self.publish_attempts == 1:
|
||||
raise RuntimeError("producer still active")
|
||||
await super().publish_end(run_id)
|
||||
|
||||
|
||||
class _MissingActiveRecoveredStreamBridge(_RetryingRecoveredStreamBridge):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self._existing_streams.clear()
|
||||
|
||||
async def cleanup_expired_producer(self, run_id: str) -> bool:
|
||||
self.expired_producer_cleanup_calls.append(run_id)
|
||||
return False
|
||||
|
||||
async def publish_end(self, run_id: str) -> None:
|
||||
self.publish_attempts += 1
|
||||
await _FakeStreamBridge.publish_end(self, run_id)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_recovered_run_stream_end_skips_expired_stream():
|
||||
"""Startup recovery should not recreate an already-expired retained stream."""
|
||||
@ -142,6 +175,7 @@ async def test_recovered_run_stream_end_skips_expired_stream():
|
||||
|
||||
assert stream_bridge.publish_end_calls == []
|
||||
assert stream_bridge.cleanup_calls == []
|
||||
assert stream_bridge.expired_producer_cleanup_calls == ["expired-run"]
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
@ -165,6 +199,39 @@ async def test_shutdown_flushes_delayed_recovered_stream_cleanup_immediately():
|
||||
assert stream_bridge.cleanup_calls == [("run-1", 60.0), ("run-1", 0)]
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_recovered_stream_end_retries_after_producer_lease_delay():
|
||||
stream_bridge = _RetryingRecoveredStreamBridge()
|
||||
|
||||
cleanups = await gateway_deps._publish_recovered_run_stream_end(
|
||||
stream_bridge,
|
||||
[SimpleNamespace(run_id="run-1", thread_id="thread-1")],
|
||||
cleanup_delay=0,
|
||||
)
|
||||
await asyncio.gather(*(task for _run_id, task in cleanups))
|
||||
|
||||
assert stream_bridge.publish_attempts == 2
|
||||
assert stream_bridge.publish_end_calls == ["run-1"]
|
||||
assert stream_bridge.cleanup_calls == [("run-1", 0)]
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_missing_recovered_stream_retries_while_producer_is_active():
|
||||
stream_bridge = _MissingActiveRecoveredStreamBridge()
|
||||
|
||||
cleanups = await gateway_deps._publish_recovered_run_stream_end(
|
||||
stream_bridge,
|
||||
[SimpleNamespace(run_id="run-1", thread_id="thread-1")],
|
||||
cleanup_delay=0,
|
||||
)
|
||||
await asyncio.gather(*(task for _run_id, task in cleanups))
|
||||
|
||||
assert stream_bridge.expired_producer_cleanup_calls == ["run-1"]
|
||||
assert stream_bridge.publish_attempts == 1
|
||||
assert stream_bridge.publish_end_calls == ["run-1"]
|
||||
assert stream_bridge.cleanup_calls == [("run-1", 0)]
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_periodic_recovery_terminalizes_stream_without_thread_projection():
|
||||
"""Periodic recovery must close streams without racing thread projection."""
|
||||
|
||||
@ -24,7 +24,7 @@ def _stub_app_config():
|
||||
reset_app_config()
|
||||
|
||||
|
||||
def _make_start_run_request(run_manager, *, thread_store=None, auth_source=None, state_user=None, headers=None):
|
||||
def _make_start_run_request(run_manager, *, thread_store=None, auth_source=None, state_user=None, headers=None, stream_bridge=None):
|
||||
from langgraph.checkpoint.memory import InMemorySaver
|
||||
from langgraph.store.memory import InMemoryStore
|
||||
|
||||
@ -39,7 +39,7 @@ def _make_start_run_request(run_manager, *, thread_store=None, auth_source=None,
|
||||
state=state,
|
||||
app=SimpleNamespace(
|
||||
state=SimpleNamespace(
|
||||
stream_bridge=SimpleNamespace(),
|
||||
stream_bridge=stream_bridge or SimpleNamespace(),
|
||||
run_manager=run_manager,
|
||||
checkpointer=InMemorySaver(),
|
||||
store=store,
|
||||
@ -2044,6 +2044,111 @@ async def test_start_run_checkpoint_validation_failure_does_not_admit_run(_stub_
|
||||
assert await run_store.list_by_thread(thread_id, user_id=None) == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_start_run_establishes_stream_producer_before_durable_admission(
|
||||
_stub_app_config,
|
||||
):
|
||||
from unittest.mock import patch
|
||||
|
||||
from app.gateway.services import start_run
|
||||
from deerflow.runtime import RunManager
|
||||
from deerflow.runtime.runs.store.memory import MemoryRunStore
|
||||
|
||||
run_store = MemoryRunStore()
|
||||
run_manager = RunManager(store=run_store)
|
||||
events = []
|
||||
agent_started = asyncio.Event()
|
||||
release = asyncio.Event()
|
||||
|
||||
class TrackingBridge:
|
||||
supports_cross_process = True
|
||||
|
||||
async def start_producer_lease(self, run_id):
|
||||
assert await run_store.get(run_id) is None
|
||||
events.append(("start", run_id))
|
||||
return True
|
||||
|
||||
def transfer_producer_lease(self, run_id, owner_task):
|
||||
events.append(("transfer", run_id, owner_task))
|
||||
|
||||
async def finish_producer_lease(self, run_id):
|
||||
events.append(("finish", run_id))
|
||||
|
||||
bridge = TrackingBridge()
|
||||
|
||||
async def fake_run_agent(_bridge, _manager, record, **_kwargs):
|
||||
agent_started.set()
|
||||
try:
|
||||
await release.wait()
|
||||
finally:
|
||||
await _bridge.finish_producer_lease(record.run_id)
|
||||
|
||||
request = _make_start_run_request(
|
||||
run_manager,
|
||||
stream_bridge=bridge,
|
||||
)
|
||||
body = _run_create_request()
|
||||
with (
|
||||
patch("app.gateway.services.resolve_agent_factory", return_value=object()),
|
||||
patch("app.gateway.services.run_agent", side_effect=fake_run_agent),
|
||||
):
|
||||
record = await start_run(body, "thread-stream-producer-order", request)
|
||||
try:
|
||||
await asyncio.wait_for(agent_started.wait(), timeout=1)
|
||||
assert events[:2] == [
|
||||
("start", record.run_id),
|
||||
("transfer", record.run_id, record.task),
|
||||
]
|
||||
assert (await run_store.get(record.run_id))["status"] == "pending"
|
||||
finally:
|
||||
release.set()
|
||||
await asyncio.wait_for(record.task, timeout=1)
|
||||
assert events[-1] == ("finish", record.run_id)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_start_run_releases_preallocated_producer_after_admission_conflict(
|
||||
_stub_app_config,
|
||||
):
|
||||
from unittest.mock import patch
|
||||
|
||||
from app.gateway.services import BusyThreadConflict, start_run
|
||||
from deerflow.runtime import RunManager
|
||||
from deerflow.runtime.runs.store.memory import MemoryRunStore
|
||||
|
||||
run_manager = RunManager(store=MemoryRunStore())
|
||||
await run_manager.create_or_reject("thread-stream-producer-conflict")
|
||||
events = []
|
||||
|
||||
class TrackingBridge:
|
||||
supports_cross_process = True
|
||||
|
||||
async def start_producer_lease(self, run_id):
|
||||
events.append(("start", run_id))
|
||||
return True
|
||||
|
||||
async def abort_producer_lease(self, run_id):
|
||||
events.append(("abort", run_id))
|
||||
|
||||
request = _make_start_run_request(
|
||||
run_manager,
|
||||
stream_bridge=TrackingBridge(),
|
||||
)
|
||||
with (
|
||||
patch("app.gateway.services.resolve_agent_factory", return_value=object()),
|
||||
pytest.raises(BusyThreadConflict),
|
||||
):
|
||||
await start_run(
|
||||
_run_create_request(),
|
||||
"thread-stream-producer-conflict",
|
||||
request,
|
||||
)
|
||||
|
||||
assert len(events) == 2
|
||||
assert events[0][0] == "start"
|
||||
assert events[1] == ("abort", events[0][1])
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pending_cancel_bypasses_thread_metadata_and_logs_failure(_stub_app_config, caplog):
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
@ -2128,6 +2128,37 @@ async def test_local_owner_cancel_falls_back_when_durable_request_fails():
|
||||
await record.task
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_local_cancel_does_not_claim_unpersisted_terminal_status():
|
||||
class FailingStatusStore(MemoryRunStore):
|
||||
async def update_status(self, run_id, status, **kwargs):
|
||||
if status == RunStatus.interrupted.value:
|
||||
raise OSError("status store unavailable")
|
||||
return await super().update_status(run_id, status, **kwargs)
|
||||
|
||||
store = FailingStatusStore()
|
||||
manager = _make_manager(
|
||||
store=store,
|
||||
worker_id="worker-a",
|
||||
run_ownership_config=_lease_config(heartbeat_enabled=True, lease_seconds=5),
|
||||
)
|
||||
record = await manager.create_or_reject("thread-1")
|
||||
await manager.set_status(record.run_id, RunStatus.running)
|
||||
record.task = asyncio.create_task(asyncio.sleep(3600))
|
||||
|
||||
try:
|
||||
assert await manager.cancel(record.run_id) == CancelOutcome.cancelled
|
||||
assert record.ownership_lost
|
||||
assert record.status == RunStatus.error
|
||||
assert record.terminal_commit_proof is None
|
||||
assert record.terminal_status_persisted is False
|
||||
assert (await store.get(record.run_id))["status"] == RunStatus.running.value
|
||||
finally:
|
||||
if not record.task.done():
|
||||
record.task.cancel()
|
||||
await asyncio.gather(record.task, return_exceptions=True)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_owner_cancel_retry_on_peer_is_accepted():
|
||||
"""A peer must treat the owner's interrupted status as an accepted cancel."""
|
||||
@ -2537,6 +2568,7 @@ async def test_heartbeat_cancels_task_on_lease_loss():
|
||||
# doesn't await).
|
||||
await asyncio.sleep(0)
|
||||
assert record.ownership_lost is True
|
||||
assert record.terminal_status_persisted is False
|
||||
assert record.abort_event.is_set() is True
|
||||
assert record.task.cancelled()
|
||||
|
||||
@ -2637,6 +2669,41 @@ async def test_peer_reconciliation_fences_late_success_and_completion():
|
||||
assert row["error"] == "peer takeover"
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_peer_error_takeover_is_not_mistaken_for_local_error_commit():
|
||||
store = MemoryRunStore()
|
||||
config = _lease_config(heartbeat_enabled=True, grace_seconds=0)
|
||||
owner = _make_manager(
|
||||
store=store,
|
||||
worker_id="worker-a",
|
||||
run_ownership_config=config,
|
||||
)
|
||||
peer = _make_manager(
|
||||
store=store,
|
||||
worker_id="worker-b",
|
||||
run_ownership_config=config,
|
||||
)
|
||||
record = await owner.create("thread-1")
|
||||
await owner.set_status(record.run_id, RunStatus.running)
|
||||
expired = (datetime.now(UTC) - timedelta(seconds=1)).isoformat()
|
||||
record.lease_expires_at = expired
|
||||
store._runs[record.run_id]["lease_expires_at"] = expired
|
||||
|
||||
recovered = await peer.reconcile_orphaned_inflight_runs(error="peer takeover")
|
||||
assert [item.run_id for item in recovered] == [record.run_id]
|
||||
|
||||
await owner.set_status(
|
||||
record.run_id,
|
||||
RunStatus.error,
|
||||
error="local failure",
|
||||
)
|
||||
|
||||
row = await store.get(record.run_id)
|
||||
assert row["error"] == "peer takeover"
|
||||
assert record.ownership_lost is True
|
||||
assert record.terminal_status_persisted is False
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_unconfirmed_success_is_fenced_when_heartbeat_is_enabled():
|
||||
"""A store outage cannot turn an unconfirmed terminal write into local success."""
|
||||
|
||||
@ -7,7 +7,13 @@ import anyio
|
||||
import pytest
|
||||
from test_stream_bridge import REDIS_TEST_URL, _DelayedBlockingReadRedis, _FakeRedis, requires_redis
|
||||
|
||||
from deerflow.runtime import RunManager, RunStatus, StreamGap
|
||||
from deerflow.runtime import (
|
||||
ORPHAN_RECOVERY_STOP_REASON,
|
||||
RunManager,
|
||||
RunStatus,
|
||||
StreamEvent,
|
||||
StreamGap,
|
||||
)
|
||||
from deerflow.runtime.stream_bridge.base import END_SENTINEL, HEARTBEAT_SENTINEL
|
||||
from deerflow.runtime.stream_bridge.redis import RedisStreamBridge
|
||||
|
||||
@ -19,6 +25,79 @@ class _ConnectedRequest:
|
||||
return False
|
||||
|
||||
|
||||
class _ProducerLeaseBridge:
|
||||
"""Retained stream with a controlled cross-process producer state."""
|
||||
|
||||
producer_lease_recovery_enabled = True
|
||||
|
||||
def __init__(self, state="expired"):
|
||||
self.state = state
|
||||
self.expired_producer_cleanup_calls: list[str] = []
|
||||
|
||||
async def stream_exists(self, run_id):
|
||||
return True
|
||||
|
||||
async def producer_lease_state(self, run_id):
|
||||
return self.state
|
||||
|
||||
async def cleanup_expired_producer(self, run_id):
|
||||
self.expired_producer_cleanup_calls.append(run_id)
|
||||
return self.state == "expired"
|
||||
|
||||
async def subscribe(self, run_id, *, last_event_id=None):
|
||||
yield HEARTBEAT_SENTINEL
|
||||
yield HEARTBEAT_SENTINEL
|
||||
await asyncio.Event().wait()
|
||||
|
||||
|
||||
class _TailAfterExpiredHeartbeatBridge(_ProducerLeaseBridge):
|
||||
def __init__(self):
|
||||
super().__init__("expired")
|
||||
self.tail_consumed = asyncio.Event()
|
||||
|
||||
async def subscribe(self, run_id, *, last_event_id=None):
|
||||
yield HEARTBEAT_SENTINEL
|
||||
yield StreamEvent(id="2-0", event="values", data={"tail": True})
|
||||
self.tail_consumed.set()
|
||||
yield END_SENTINEL
|
||||
|
||||
|
||||
class _ReactivatedProducerBridge(_ProducerLeaseBridge):
|
||||
async def cleanup_expired_producer(self, run_id):
|
||||
self.expired_producer_cleanup_calls.append(run_id)
|
||||
self.state = "active"
|
||||
return False
|
||||
|
||||
|
||||
class _FailingProducerCleanupBridge(_ProducerLeaseBridge):
|
||||
async def cleanup_expired_producer(self, run_id):
|
||||
self.expired_producer_cleanup_calls.append(run_id)
|
||||
raise TimeoutError("producer cleanup unavailable")
|
||||
|
||||
|
||||
class _MissingActiveProducerBridge(_ProducerLeaseBridge):
|
||||
def __init__(self):
|
||||
super().__init__("active")
|
||||
|
||||
async def stream_exists(self, run_id):
|
||||
return False
|
||||
|
||||
|
||||
class _MissingThenTailBridge(_TailAfterExpiredHeartbeatBridge):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.exists_calls = 0
|
||||
|
||||
async def stream_exists(self, run_id):
|
||||
self.exists_calls += 1
|
||||
return self.exists_calls > 1
|
||||
|
||||
async def subscribe(self, run_id, *, last_event_id=None):
|
||||
yield StreamEvent(id="2-0", event="values", data={"tail": True})
|
||||
self.tail_consumed.set()
|
||||
yield END_SENTINEL
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def cleanup_bridge(request):
|
||||
real = request.param == "redis"
|
||||
@ -216,6 +295,321 @@ async def test_wait_survives_cleanup_with_real_bridge(cleanup_bridge, store_only
|
||||
await manager.shutdown()
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_sse_terminalizes_retained_stream_after_producer_lease_expires():
|
||||
from app.gateway.services import sse_consumer
|
||||
from deerflow.runtime.runs.store.memory import MemoryRunStore
|
||||
|
||||
store = MemoryRunStore()
|
||||
await store.put("expired-producer", thread_id="cleanup-thread", status="success")
|
||||
manager = RunManager(store=store)
|
||||
record = await manager.get("expired-producer")
|
||||
bridge = _ProducerLeaseBridge()
|
||||
stream = sse_consumer(
|
||||
bridge,
|
||||
record,
|
||||
_ConnectedRequest(),
|
||||
manager,
|
||||
apply_on_disconnect=False,
|
||||
)
|
||||
try:
|
||||
with anyio.fail_after(0.5):
|
||||
frames = [frame async for frame in stream]
|
||||
assert frames[0] == ": heartbeat\n\n"
|
||||
assert frames[1].startswith("event: end\n")
|
||||
assert bridge.expired_producer_cleanup_calls == ["expired-producer"]
|
||||
finally:
|
||||
await stream.aclose()
|
||||
await manager.shutdown()
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_wait_terminalizes_retained_stream_after_producer_lease_expires():
|
||||
from app.gateway.services import wait_for_run_completion
|
||||
from deerflow.runtime.runs.store.memory import MemoryRunStore
|
||||
|
||||
store = MemoryRunStore()
|
||||
await store.put("expired-producer", thread_id="cleanup-thread", status="success")
|
||||
manager = RunManager(store=store)
|
||||
record = await manager.get("expired-producer")
|
||||
bridge = _ProducerLeaseBridge()
|
||||
try:
|
||||
with anyio.fail_after(0.5):
|
||||
assert await wait_for_run_completion(
|
||||
bridge,
|
||||
record,
|
||||
_ConnectedRequest(),
|
||||
manager,
|
||||
)
|
||||
assert bridge.expired_producer_cleanup_calls == ["expired-producer"]
|
||||
finally:
|
||||
await manager.shutdown()
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_sse_does_not_terminalize_when_producer_reactivates_before_cleanup():
|
||||
from app.gateway.services import sse_consumer
|
||||
from deerflow.runtime.runs.store.memory import MemoryRunStore
|
||||
|
||||
store = MemoryRunStore()
|
||||
await store.put("reactivated-producer", thread_id="cleanup-thread", status="success")
|
||||
manager = RunManager(store=store)
|
||||
record = await manager.get("reactivated-producer")
|
||||
bridge = _ReactivatedProducerBridge()
|
||||
stream = sse_consumer(
|
||||
bridge,
|
||||
record,
|
||||
_ConnectedRequest(),
|
||||
manager,
|
||||
apply_on_disconnect=False,
|
||||
)
|
||||
try:
|
||||
with anyio.fail_after(0.5):
|
||||
assert await anext(stream) == ": heartbeat\n\n"
|
||||
assert await anext(stream) == ": heartbeat\n\n"
|
||||
assert bridge.expired_producer_cleanup_calls == ["reactivated-producer"]
|
||||
finally:
|
||||
await stream.aclose()
|
||||
await manager.shutdown()
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_sse_does_not_terminalize_when_producer_cleanup_fails():
|
||||
from app.gateway.services import sse_consumer
|
||||
from deerflow.runtime.runs.store.memory import MemoryRunStore
|
||||
|
||||
store = MemoryRunStore()
|
||||
await store.put("cleanup-failed", thread_id="cleanup-thread", status="success")
|
||||
manager = RunManager(store=store)
|
||||
record = await manager.get("cleanup-failed")
|
||||
bridge = _FailingProducerCleanupBridge()
|
||||
stream = sse_consumer(
|
||||
bridge,
|
||||
record,
|
||||
_ConnectedRequest(),
|
||||
manager,
|
||||
apply_on_disconnect=False,
|
||||
)
|
||||
try:
|
||||
with anyio.fail_after(0.5):
|
||||
assert await anext(stream) == ": heartbeat\n\n"
|
||||
assert await anext(stream) == ": heartbeat\n\n"
|
||||
assert bridge.expired_producer_cleanup_calls == ["cleanup-failed"]
|
||||
finally:
|
||||
await stream.aclose()
|
||||
await manager.shutdown()
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_sse_orphan_recovery_waits_for_active_stream_producer():
|
||||
from app.gateway.services import sse_consumer
|
||||
from deerflow.runtime.runs.store.memory import MemoryRunStore
|
||||
|
||||
store = MemoryRunStore()
|
||||
await store.put(
|
||||
"orphan-active-producer",
|
||||
thread_id="cleanup-thread",
|
||||
status="error",
|
||||
stop_reason=ORPHAN_RECOVERY_STOP_REASON,
|
||||
)
|
||||
manager = RunManager(store=store)
|
||||
record = await manager.get("orphan-active-producer")
|
||||
bridge = _ProducerLeaseBridge("active")
|
||||
stream = sse_consumer(
|
||||
bridge,
|
||||
record,
|
||||
_ConnectedRequest(),
|
||||
manager,
|
||||
apply_on_disconnect=False,
|
||||
)
|
||||
try:
|
||||
with anyio.fail_after(0.5):
|
||||
assert await anext(stream) == ": heartbeat\n\n"
|
||||
finally:
|
||||
await stream.aclose()
|
||||
await manager.shutdown()
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_sse_drains_tail_published_after_expired_heartbeat():
|
||||
from app.gateway.services import sse_consumer
|
||||
from deerflow.runtime.runs.store.memory import MemoryRunStore
|
||||
|
||||
store = MemoryRunStore()
|
||||
await store.put("expired-tail", thread_id="cleanup-thread", status="success")
|
||||
manager = RunManager(store=store)
|
||||
record = await manager.get("expired-tail")
|
||||
bridge = _TailAfterExpiredHeartbeatBridge()
|
||||
stream = sse_consumer(
|
||||
bridge,
|
||||
record,
|
||||
_ConnectedRequest(),
|
||||
manager,
|
||||
apply_on_disconnect=False,
|
||||
)
|
||||
try:
|
||||
with anyio.fail_after(0.5):
|
||||
frames = [frame async for frame in stream]
|
||||
assert frames[0] == ": heartbeat\n\n"
|
||||
assert frames[1].startswith("event: values\n")
|
||||
assert frames[2].startswith("event: end\n")
|
||||
assert bridge.tail_consumed.is_set()
|
||||
finally:
|
||||
await stream.aclose()
|
||||
await manager.shutdown()
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_wait_drains_tail_published_after_expired_heartbeat():
|
||||
from app.gateway.services import wait_for_run_completion
|
||||
from deerflow.runtime.runs.store.memory import MemoryRunStore
|
||||
|
||||
store = MemoryRunStore()
|
||||
await store.put("expired-tail", thread_id="cleanup-thread", status="success")
|
||||
manager = RunManager(store=store)
|
||||
record = await manager.get("expired-tail")
|
||||
bridge = _TailAfterExpiredHeartbeatBridge()
|
||||
try:
|
||||
with anyio.fail_after(0.5):
|
||||
assert await wait_for_run_completion(
|
||||
bridge,
|
||||
record,
|
||||
_ConnectedRequest(),
|
||||
manager,
|
||||
)
|
||||
assert bridge.tail_consumed.is_set()
|
||||
finally:
|
||||
await manager.shutdown()
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_missing_stream_does_not_terminalize_an_active_producer():
|
||||
from app.gateway.services import sse_consumer
|
||||
from deerflow.runtime.runs.store.memory import MemoryRunStore
|
||||
|
||||
store = MemoryRunStore()
|
||||
await store.put("active-missing", thread_id="cleanup-thread", status="success")
|
||||
manager = RunManager(store=store)
|
||||
record = await manager.get("active-missing")
|
||||
stream = sse_consumer(
|
||||
_MissingActiveProducerBridge(),
|
||||
record,
|
||||
_ConnectedRequest(),
|
||||
manager,
|
||||
apply_on_disconnect=False,
|
||||
)
|
||||
try:
|
||||
with anyio.fail_after(0.5):
|
||||
assert await anext(stream) == ": heartbeat\n\n"
|
||||
finally:
|
||||
await stream.aclose()
|
||||
await manager.shutdown()
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_missing_stream_rechecks_after_expiry_and_drains_new_tail():
|
||||
from app.gateway.services import sse_consumer
|
||||
from deerflow.runtime.runs.store.memory import MemoryRunStore
|
||||
|
||||
store = MemoryRunStore()
|
||||
await store.put("missing-then-tail", thread_id="cleanup-thread", status="success")
|
||||
manager = RunManager(store=store)
|
||||
record = await manager.get("missing-then-tail")
|
||||
bridge = _MissingThenTailBridge()
|
||||
stream = sse_consumer(
|
||||
bridge,
|
||||
record,
|
||||
_ConnectedRequest(),
|
||||
manager,
|
||||
apply_on_disconnect=False,
|
||||
)
|
||||
try:
|
||||
with anyio.fail_after(0.5):
|
||||
frames = [frame async for frame in stream]
|
||||
assert bridge.exists_calls == 2
|
||||
assert frames[0].startswith("event: values\n")
|
||||
assert frames[1].startswith("event: end\n")
|
||||
assert bridge.tail_consumed.is_set()
|
||||
finally:
|
||||
await stream.aclose()
|
||||
await manager.shutdown()
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_wait_rechecks_missing_stream_after_expiry_and_drains_new_tail():
|
||||
from app.gateway.services import wait_for_run_completion
|
||||
from deerflow.runtime.runs.store.memory import MemoryRunStore
|
||||
|
||||
store = MemoryRunStore()
|
||||
await store.put("missing-then-tail", thread_id="cleanup-thread", status="success")
|
||||
manager = RunManager(store=store)
|
||||
record = await manager.get("missing-then-tail")
|
||||
bridge = _MissingThenTailBridge()
|
||||
try:
|
||||
with anyio.fail_after(0.5):
|
||||
assert await wait_for_run_completion(
|
||||
bridge,
|
||||
record,
|
||||
_ConnectedRequest(),
|
||||
manager,
|
||||
)
|
||||
assert bridge.exists_calls == 2
|
||||
assert bridge.tail_consumed.is_set()
|
||||
finally:
|
||||
await manager.shutdown()
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
@pytest.mark.parametrize("producer_state", ["active", "untracked"])
|
||||
async def test_sse_keeps_waiting_for_nonexpired_terminal_producer(producer_state):
|
||||
from app.gateway.services import sse_consumer
|
||||
from deerflow.runtime.runs.store.memory import MemoryRunStore
|
||||
|
||||
store = MemoryRunStore()
|
||||
await store.put("live-producer", thread_id="cleanup-thread", status="success")
|
||||
manager = RunManager(store=store)
|
||||
record = await manager.get("live-producer")
|
||||
stream = sse_consumer(
|
||||
_ProducerLeaseBridge(producer_state),
|
||||
record,
|
||||
_ConnectedRequest(),
|
||||
manager,
|
||||
apply_on_disconnect=False,
|
||||
)
|
||||
try:
|
||||
with anyio.fail_after(0.5):
|
||||
assert await anext(stream) == ": heartbeat\n\n"
|
||||
finally:
|
||||
await stream.aclose()
|
||||
await manager.shutdown()
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_sse_keeps_waiting_for_uncommitted_local_terminal_status():
|
||||
from app.gateway.services import sse_consumer
|
||||
from deerflow.runtime.runs.store.memory import MemoryRunStore
|
||||
|
||||
store = MemoryRunStore()
|
||||
manager = RunManager(store=store)
|
||||
record = await manager.create_or_reject("cleanup-thread")
|
||||
await manager.try_start(record.run_id)
|
||||
await manager.set_status(record.run_id, RunStatus.success, persist=False)
|
||||
assert record.terminal_status_persisted is False
|
||||
stream = sse_consumer(
|
||||
_ProducerLeaseBridge(),
|
||||
record,
|
||||
_ConnectedRequest(),
|
||||
manager,
|
||||
apply_on_disconnect=False,
|
||||
)
|
||||
try:
|
||||
with anyio.fail_after(0.5):
|
||||
assert await anext(stream) == ": heartbeat\n\n"
|
||||
finally:
|
||||
await stream.aclose()
|
||||
await manager.shutdown(timeout=0)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
@pytest.mark.parametrize("terminal_with_retained_tail", [False, True])
|
||||
async def test_empty_gap_does_not_complete_running_or_retained_run(terminal_with_retained_tail):
|
||||
|
||||
@ -298,6 +298,20 @@ async def test_create_and_get(manager: RunManager):
|
||||
assert fetched is record
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_create_or_reject_uses_preallocated_run_id():
|
||||
store = MemoryRunStore()
|
||||
manager = RunManager(store=store)
|
||||
|
||||
record = await manager.create_or_reject(
|
||||
"thread-preallocated",
|
||||
run_id="preallocated-run",
|
||||
)
|
||||
|
||||
assert record.run_id == "preallocated-run"
|
||||
assert (await store.get("preallocated-run"))["thread_id"] == "thread-preallocated"
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_status_transitions(manager: RunManager):
|
||||
"""Status should transition pending -> running -> success."""
|
||||
@ -1245,6 +1259,7 @@ async def test_create_or_reject_retries_replacement_when_cancel_status_cannot_pe
|
||||
assert stored_replacement is not None
|
||||
assert stored_replacement["status"] == RunStatus.interrupted.value
|
||||
assert replacement.status == RunStatus.interrupted
|
||||
assert replacement.terminal_status_persisted is True
|
||||
assert not await manager.has_inflight("thread-1")
|
||||
|
||||
|
||||
|
||||
@ -613,6 +613,57 @@ async def test_delivery_is_durable_before_terminal_run_status():
|
||||
assert (await run_store.get(record.run_id))["status"] == "success"
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_deferred_error_stream_frames_follow_durable_terminal_status():
|
||||
events = MemoryRunEventStore()
|
||||
order: list[str] = []
|
||||
|
||||
class OrderingRunStore(MemoryRunStore):
|
||||
async def update_status(self, run_id, status, *, error=None, stop_reason=None):
|
||||
result = await super().update_status(
|
||||
run_id,
|
||||
status,
|
||||
error=error,
|
||||
stop_reason=stop_reason,
|
||||
)
|
||||
if status == "error":
|
||||
order.append("durable-error")
|
||||
return result
|
||||
|
||||
class OrderingBridge:
|
||||
async def publish(self, _run_id, event, _data):
|
||||
if event == "error":
|
||||
order.append("stream-error")
|
||||
|
||||
async def publish_end(self, _run_id):
|
||||
order.append("stream-end")
|
||||
|
||||
async def cleanup(self, _run_id, *, delay=0):
|
||||
del delay
|
||||
|
||||
class FailingAgent:
|
||||
async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False):
|
||||
del graph_input, config, stream_mode, subgraphs
|
||||
raise RuntimeError("boom")
|
||||
yield # pragma: no cover - make this an async generator
|
||||
|
||||
run_store = OrderingRunStore()
|
||||
run_manager = RunManager(store=run_store)
|
||||
record = await run_manager.create("thread-1")
|
||||
|
||||
await run_agent(
|
||||
OrderingBridge(),
|
||||
run_manager,
|
||||
record,
|
||||
ctx=RunContext(checkpointer=None, event_store=events),
|
||||
agent_factory=lambda *, config: FailingAgent(),
|
||||
graph_input={},
|
||||
config={},
|
||||
)
|
||||
|
||||
assert order == ["durable-error", "stream-error", "stream-end"]
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_delivery_write_retries_before_persisting_success():
|
||||
class FlakyReceiptStore(MemoryRunEventStore):
|
||||
|
||||
@ -192,6 +192,44 @@ async def test_cancel_is_preserved_when_guarded_flush_fails(monkeypatch):
|
||||
await reap(run)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_producer_lease_brackets_terminal_commit_and_end(monkeypatch):
|
||||
run_store = MemoryRunStore()
|
||||
order = []
|
||||
|
||||
class TrackingBridge(MemoryStreamBridge):
|
||||
async def start_producer_lease(self, run_id):
|
||||
order.append(("start", (await run_store.get(run_id))["status"]))
|
||||
return True
|
||||
|
||||
async def publish_end(self, run_id):
|
||||
order.append(("end", (await run_store.get(run_id))["status"]))
|
||||
await super().publish_end(run_id)
|
||||
|
||||
def transfer_producer_lease(self, run_id, owner_task):
|
||||
order.append(("transfer", owner_task.get_name()))
|
||||
|
||||
async def finish_producer_lease(self, run_id):
|
||||
order.append(("finish", (await run_store.get(run_id))["status"]))
|
||||
|
||||
run = await start_run(
|
||||
monkeypatch,
|
||||
event_store=MemoryRunEventStore(),
|
||||
run_store=run_store,
|
||||
bridge=TrackingBridge(),
|
||||
)
|
||||
try:
|
||||
await asyncio.wait_for(run.worker, timeout=5)
|
||||
assert order == [
|
||||
("start", "running"),
|
||||
("transfer", f"run-finalization-{run.record.run_id}"),
|
||||
("end", "success"),
|
||||
("finish", "success"),
|
||||
]
|
||||
finally:
|
||||
await reap(run)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_failed_finalizer_preserves_host_cancellation_as_primary():
|
||||
entered, release = asyncio.Event(), asyncio.Event()
|
||||
|
||||
@ -4182,7 +4182,7 @@ async def test_worker_finally_block_swallows_helper_exceptions(monkeypatch):
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_worker_skips_execution_and_finalization_after_ownership_loss():
|
||||
"""A fenced worker closes its stream without starting or finalizing work."""
|
||||
"""A fenced worker leaves terminal publication to durable recovery."""
|
||||
run_manager = RunManager()
|
||||
record = await run_manager.create("thread-lease-lost")
|
||||
record.ownership_lost = True
|
||||
@ -4219,7 +4219,7 @@ async def test_worker_skips_execution_and_finalization_after_ownership_loss():
|
||||
thread_store.update_display_name.assert_not_awaited()
|
||||
thread_store.update_status.assert_not_awaited()
|
||||
on_run_completed.assert_not_awaited()
|
||||
bridge.publish_end.assert_awaited_once_with(record.run_id)
|
||||
bridge.publish_end.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
|
||||
@ -278,8 +278,8 @@ def test_creator_stream_close_after_orphan_recovery_end_does_not_cancel(monkeypa
|
||||
async def scenario():
|
||||
recorder = _CancelRecorder()
|
||||
|
||||
async def orphan_recovered(record, run_mgr):
|
||||
return True
|
||||
async def orphan_recovered(record, run_mgr, bridge):
|
||||
return "recovered"
|
||||
|
||||
monkeypatch.setattr(
|
||||
"app.gateway.services._orphan_recovery_observed_after_heartbeat",
|
||||
|
||||
@ -5,6 +5,7 @@ import os
|
||||
import re
|
||||
import uuid
|
||||
from collections import defaultdict
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import anyio
|
||||
import pytest
|
||||
@ -16,7 +17,7 @@ from deerflow.config.stream_bridge_config import (
|
||||
StreamBridgeConfig,
|
||||
set_stream_bridge_config,
|
||||
)
|
||||
from deerflow.runtime import END_SENTINEL, HEARTBEAT_SENTINEL, MemoryStreamBridge, StreamGap, make_stream_bridge
|
||||
from deerflow.runtime import END_SENTINEL, HEARTBEAT_SENTINEL, MemoryStreamBridge, StreamGap, StreamProducerLeaseState, make_stream_bridge
|
||||
|
||||
# RedisStreamBridge is no longer re-exported from deerflow.runtime (redis is an
|
||||
# optional extra; see the NOTE in runtime/stream_bridge/__init__.py). Import it
|
||||
@ -45,6 +46,7 @@ class _FakeRedis:
|
||||
self.counters = defaultdict(int)
|
||||
self.deleted = []
|
||||
self.expirations = []
|
||||
self.values = {}
|
||||
self.closed = False
|
||||
self.ping_calls = 0
|
||||
self.ping_error: Exception | None = None
|
||||
@ -88,17 +90,111 @@ class _FakeRedis:
|
||||
entries = list(self.streams.get(name, []))
|
||||
return entries[:count] if count is not None else entries
|
||||
|
||||
async def delete(self, name):
|
||||
self.deleted.append(name)
|
||||
self.streams.pop(name, None)
|
||||
return 1
|
||||
async def set(self, name, value, *, ex=None, nx=False, xx=False):
|
||||
exists = name in self.values
|
||||
if nx and exists:
|
||||
return False
|
||||
if xx and not exists:
|
||||
return False
|
||||
self.values[name] = value
|
||||
if ex is not None:
|
||||
self.expirations.append((name, ex))
|
||||
return True
|
||||
|
||||
async def get(self, name):
|
||||
return self.values.get(name)
|
||||
|
||||
async def mget(self, *names):
|
||||
return [self.values.get(name) for name in names]
|
||||
|
||||
async def eval(self, script, numkeys, *args):
|
||||
keys = args[:numkeys]
|
||||
values = args[numkeys:]
|
||||
if numkeys == 2 and not values:
|
||||
tracking_key, active_key = keys
|
||||
if active_key in self.values:
|
||||
return 0
|
||||
if tracking_key not in self.values:
|
||||
return 2
|
||||
self.values.pop(tracking_key, None)
|
||||
self.deleted.append(tracking_key)
|
||||
return 1
|
||||
if numkeys == 3 and not values:
|
||||
stream_key, tracking_key, active_key = keys
|
||||
if active_key in self.values:
|
||||
return 0
|
||||
return await self.delete(stream_key, tracking_key)
|
||||
if numkeys == 3 and len(values) == 1:
|
||||
stream_key, tracking_key, active_key = keys
|
||||
(token,) = values
|
||||
tracked = self.values.get(tracking_key)
|
||||
active = self.values.get(active_key)
|
||||
if tracked is not None and tracked != token or active is not None and active != token:
|
||||
return 0
|
||||
await self.delete(stream_key, tracking_key, active_key)
|
||||
return 1
|
||||
if numkeys == 2 and len(values) == 5:
|
||||
active_key, stream_key = keys
|
||||
maxlen, stream_ttl, kind, event, data = values
|
||||
if active_key in self.values:
|
||||
return None
|
||||
event_id = await self.xadd(
|
||||
stream_key,
|
||||
{"kind": kind, "event": event, "data": data},
|
||||
maxlen=int(maxlen),
|
||||
approximate=False,
|
||||
)
|
||||
if int(stream_ttl) > 0:
|
||||
await self.expire(stream_key, int(stream_ttl))
|
||||
return event_id
|
||||
if numkeys == 1:
|
||||
(name,) = keys
|
||||
(token,) = values
|
||||
if self.values.get(name) != token:
|
||||
return 0
|
||||
self.values.pop(name, None)
|
||||
self.deleted.append(name)
|
||||
return 1
|
||||
if numkeys == 2:
|
||||
active_key, tracking_key = keys
|
||||
token, active_ttl, tracking_ttl = values
|
||||
if self.values.get(active_key) != token or self.values.get(tracking_key) != token:
|
||||
return 0
|
||||
await self.expire(active_key, int(active_ttl))
|
||||
if int(tracking_ttl) > 0:
|
||||
await self.expire(tracking_key, int(tracking_ttl))
|
||||
return 1
|
||||
if numkeys == 3:
|
||||
active_key, tracking_key, stream_key = keys
|
||||
token, maxlen, stream_ttl, kind, event, data = values
|
||||
if self.values.get(active_key) != token or self.values.get(tracking_key) != token:
|
||||
return None
|
||||
event_id = await self.xadd(
|
||||
stream_key,
|
||||
{"kind": kind, "event": event, "data": data},
|
||||
maxlen=int(maxlen),
|
||||
approximate=False,
|
||||
)
|
||||
if int(stream_ttl) > 0:
|
||||
await self.expire(stream_key, int(stream_ttl))
|
||||
return event_id
|
||||
raise AssertionError(f"unexpected eval key count: {numkeys}")
|
||||
|
||||
async def delete(self, *names):
|
||||
deleted = 0
|
||||
for name in names:
|
||||
self.deleted.append(name)
|
||||
deleted += int(name in self.streams or name in self.values)
|
||||
self.streams.pop(name, None)
|
||||
self.values.pop(name, None)
|
||||
return deleted
|
||||
|
||||
async def exists(self, name):
|
||||
return 1 if name in self.streams else 0
|
||||
return 1 if name in self.streams or name in self.values else 0
|
||||
|
||||
async def expire(self, name, seconds):
|
||||
self.expirations.append((name, seconds))
|
||||
return True
|
||||
return name in self.streams or name in self.values
|
||||
|
||||
def pipeline(self, *, transaction=True):
|
||||
return _FakeRedisPipeline(self)
|
||||
@ -864,7 +960,633 @@ async def test_redis_cleanup_deletes_stream(redis_bridge: RedisStreamBridge):
|
||||
await redis_bridge.publish(run_id, "event", {})
|
||||
await redis_bridge.cleanup(run_id)
|
||||
|
||||
assert fake.deleted == ["deerflow:stream_bridge:redis-run-cleanup"]
|
||||
assert fake.deleted == [
|
||||
"deerflow:stream_bridge:redis-run-cleanup",
|
||||
"deerflow:stream_bridge:redis-run-cleanup:producer-tracked",
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_redis_producer_lease_distinguishes_active_expired_and_untracked():
|
||||
fake = _FakeRedis()
|
||||
bridge = RedisStreamBridge(redis_url="redis://fake", client=fake)
|
||||
run_id = "redis-run-producer-state"
|
||||
|
||||
assert await bridge.producer_lease_state(run_id) == StreamProducerLeaseState.untracked
|
||||
|
||||
await bridge.start_producer_lease(run_id)
|
||||
assert await bridge.producer_lease_state(run_id) == StreamProducerLeaseState.active
|
||||
|
||||
await bridge.finish_producer_lease(run_id)
|
||||
assert await bridge.producer_lease_state(run_id) == StreamProducerLeaseState.expired
|
||||
|
||||
await bridge.cleanup(run_id)
|
||||
assert await bridge.producer_lease_state(run_id) == StreamProducerLeaseState.untracked
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_redis_recovery_cleanup_deletes_only_expired_tracking_marker():
|
||||
fake = _FakeRedis()
|
||||
bridge = RedisStreamBridge(
|
||||
redis_url="redis://fake",
|
||||
stream_ttl_seconds=0,
|
||||
client=fake,
|
||||
)
|
||||
run_id = "redis-run-expired-tracking-cleanup"
|
||||
|
||||
await bridge.start_producer_lease(run_id)
|
||||
assert await bridge.cleanup_expired_producer(run_id) is False
|
||||
assert await bridge.producer_lease_state(run_id) == StreamProducerLeaseState.active
|
||||
|
||||
await bridge.finish_producer_lease(run_id)
|
||||
assert await bridge.cleanup_expired_producer(run_id) is True
|
||||
assert await bridge.producer_lease_state(run_id) == StreamProducerLeaseState.untracked
|
||||
await bridge.cleanup(run_id)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_redis_recovery_cleanup_rejects_active_registration_window():
|
||||
fake = _FakeRedis()
|
||||
bridge = RedisStreamBridge(redis_url="redis://fake", client=fake)
|
||||
run_id = "redis-run-active-registration-window"
|
||||
|
||||
await bridge.start_producer_lease(run_id)
|
||||
fake.values.pop(bridge._producer_tracking_key(run_id))
|
||||
try:
|
||||
assert await bridge.cleanup_expired_producer(run_id) is False
|
||||
assert bridge._producer_active_key(run_id) in fake.values
|
||||
finally:
|
||||
await bridge.abort_producer_lease(run_id)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_redis_lost_producer_lease_cancels_the_owning_task(monkeypatch):
|
||||
import deerflow.runtime.stream_bridge.redis as redis_bridge_module
|
||||
|
||||
monkeypatch.setattr(
|
||||
redis_bridge_module,
|
||||
"_PRODUCER_LEASE_RENEW_INTERVAL_SECONDS",
|
||||
0.01,
|
||||
)
|
||||
fake = _FakeRedis()
|
||||
bridge = RedisStreamBridge(redis_url="redis://fake", client=fake)
|
||||
run_id = "redis-run-producer-lost"
|
||||
started = asyncio.Event()
|
||||
cancelled = asyncio.Event()
|
||||
release = asyncio.Event()
|
||||
|
||||
async def producer():
|
||||
await bridge.start_producer_lease(run_id)
|
||||
started.set()
|
||||
try:
|
||||
await asyncio.Event().wait()
|
||||
except asyncio.CancelledError:
|
||||
cancelled.set()
|
||||
await release.wait()
|
||||
raise
|
||||
finally:
|
||||
await bridge.finish_producer_lease(run_id)
|
||||
|
||||
owner = asyncio.create_task(producer())
|
||||
await asyncio.wait_for(started.wait(), timeout=1)
|
||||
fake.values.pop(bridge._producer_active_key(run_id))
|
||||
|
||||
try:
|
||||
await asyncio.wait_for(cancelled.wait(), timeout=1)
|
||||
with pytest.raises(RuntimeError, match="producer lease was lost"):
|
||||
await bridge.publish(run_id, "late", {})
|
||||
with pytest.raises(RuntimeError, match="producer lease was lost"):
|
||||
await bridge.publish_end(run_id)
|
||||
finally:
|
||||
release.set()
|
||||
with pytest.raises(asyncio.CancelledError, match="producer lease expired"):
|
||||
await asyncio.wait_for(owner, timeout=1)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_redis_tracking_marker_outlives_short_stream_ttl():
|
||||
fake = _FakeRedis()
|
||||
bridge = RedisStreamBridge(
|
||||
redis_url="redis://fake",
|
||||
stream_ttl_seconds=1,
|
||||
client=fake,
|
||||
)
|
||||
run_id = "redis-run-short-stream-ttl"
|
||||
|
||||
await bridge.start_producer_lease(run_id)
|
||||
try:
|
||||
assert (
|
||||
bridge._producer_tracking_key(run_id),
|
||||
120,
|
||||
) in fake.expirations
|
||||
assert await bridge.producer_lease_state(run_id) == StreamProducerLeaseState.active
|
||||
finally:
|
||||
await bridge.finish_producer_lease(run_id)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_redis_finish_does_not_delete_a_successor_token():
|
||||
fake = _FakeRedis()
|
||||
bridge = RedisStreamBridge(redis_url="redis://fake", client=fake)
|
||||
run_id = "redis-run-successor-token"
|
||||
active_key = bridge._producer_active_key(run_id)
|
||||
|
||||
await bridge.start_producer_lease(run_id)
|
||||
fake.values[active_key] = "successor-token"
|
||||
await bridge.finish_producer_lease(run_id)
|
||||
|
||||
assert fake.values[active_key] == "successor-token"
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_redis_publish_is_atomically_fenced_by_producer_token():
|
||||
fake = _FakeRedis()
|
||||
bridge = RedisStreamBridge(redis_url="redis://fake", client=fake)
|
||||
run_id = "redis-run-publish-token"
|
||||
active_key = bridge._producer_active_key(run_id)
|
||||
replacement_owner = asyncio.create_task(asyncio.Event().wait())
|
||||
|
||||
await bridge.start_producer_lease(run_id)
|
||||
bridge.transfer_producer_lease(run_id, replacement_owner)
|
||||
fake.values[active_key] = "successor-token"
|
||||
try:
|
||||
with pytest.raises(RuntimeError, match="producer lease was lost"):
|
||||
await bridge.publish(run_id, "late", {"value": 1})
|
||||
assert bridge._stream_key(run_id) not in fake.streams
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await replacement_owner
|
||||
finally:
|
||||
await bridge.finish_producer_lease(run_id)
|
||||
replacement_owner.cancel()
|
||||
await asyncio.gather(replacement_owner, return_exceptions=True)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_redis_producer_lease_releases_when_owner_never_starts():
|
||||
fake = _FakeRedis()
|
||||
bridge = RedisStreamBridge(redis_url="redis://fake", client=fake)
|
||||
run_id = "redis-run-owner-never-started"
|
||||
|
||||
await bridge.start_producer_lease(run_id)
|
||||
owner = asyncio.create_task(asyncio.sleep(3600))
|
||||
bridge.transfer_producer_lease(run_id, owner)
|
||||
owner.cancel()
|
||||
await asyncio.gather(owner, return_exceptions=True)
|
||||
for _ in range(10):
|
||||
lease = bridge._producer_leases.get(run_id)
|
||||
if lease is not None and lease.stopped and lease.renewal_task is None and await bridge.producer_lease_state(run_id) == StreamProducerLeaseState.expired:
|
||||
break
|
||||
await asyncio.sleep(0)
|
||||
|
||||
lease = bridge._producer_leases[run_id]
|
||||
assert lease.stopped is True
|
||||
assert lease.renewal_task is None
|
||||
assert await bridge.producer_lease_state(run_id) == StreamProducerLeaseState.expired
|
||||
await bridge.cleanup(run_id)
|
||||
assert run_id not in bridge._producer_leases
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_redis_old_owner_exit_cannot_release_transferred_lease():
|
||||
fake = _FakeRedis()
|
||||
bridge = RedisStreamBridge(redis_url="redis://fake", client=fake)
|
||||
run_id = "redis-run-transferred-owner"
|
||||
old_owner = asyncio.create_task(asyncio.sleep(3600))
|
||||
current_owner = asyncio.create_task(asyncio.sleep(3600))
|
||||
|
||||
await bridge.start_producer_lease(run_id)
|
||||
bridge.transfer_producer_lease(run_id, old_owner)
|
||||
bridge.transfer_producer_lease(run_id, current_owner)
|
||||
old_owner.cancel()
|
||||
await asyncio.gather(old_owner, return_exceptions=True)
|
||||
for _ in range(5):
|
||||
await asyncio.sleep(0)
|
||||
|
||||
try:
|
||||
assert bridge._producer_leases[run_id].owner_task is current_owner
|
||||
assert await bridge.producer_lease_state(run_id) == StreamProducerLeaseState.active
|
||||
finally:
|
||||
current_owner.cancel()
|
||||
await asyncio.gather(current_owner, return_exceptions=True)
|
||||
await bridge.finish_producer_lease(run_id)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
@pytest.mark.parametrize("operation", ["abort", "close"])
|
||||
async def test_redis_explicit_stop_waits_for_owner_cleanup(operation):
|
||||
fake = _FakeRedis()
|
||||
bridge = RedisStreamBridge(redis_url="redis://fake", client=fake)
|
||||
run_id = f"redis-run-owner-cleanup-{operation}"
|
||||
owner = asyncio.create_task(asyncio.sleep(3600))
|
||||
release_started = asyncio.Event()
|
||||
allow_release = asyncio.Event()
|
||||
original_delete = bridge._delete_if_token
|
||||
|
||||
async def block_active_release(name, token):
|
||||
if name == bridge._producer_active_key(run_id):
|
||||
release_started.set()
|
||||
await allow_release.wait()
|
||||
return await original_delete(name, token)
|
||||
|
||||
await bridge.start_producer_lease(run_id)
|
||||
bridge.transfer_producer_lease(run_id, owner)
|
||||
bridge._delete_if_token = block_active_release
|
||||
owner.cancel()
|
||||
await asyncio.gather(owner, return_exceptions=True)
|
||||
await asyncio.wait_for(release_started.wait(), timeout=1)
|
||||
|
||||
if operation == "abort":
|
||||
explicit_stop = asyncio.create_task(bridge.abort_producer_lease(run_id))
|
||||
else:
|
||||
explicit_stop = asyncio.create_task(bridge.close())
|
||||
await asyncio.sleep(0)
|
||||
assert not explicit_stop.done()
|
||||
|
||||
allow_release.set()
|
||||
await asyncio.wait_for(explicit_stop, timeout=1)
|
||||
expected = StreamProducerLeaseState.untracked if operation == "abort" else StreamProducerLeaseState.expired
|
||||
assert await bridge.producer_lease_state(run_id) == expected
|
||||
assert run_id not in bridge._producer_leases
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_redis_unowned_end_cannot_bypass_an_active_remote_producer():
|
||||
fake = _FakeRedis()
|
||||
owner = RedisStreamBridge(
|
||||
redis_url="redis://fake",
|
||||
producer_lease_recovery_enabled=True,
|
||||
client=fake,
|
||||
)
|
||||
recovering_peer = RedisStreamBridge(
|
||||
redis_url="redis://fake",
|
||||
producer_lease_recovery_enabled=True,
|
||||
client=fake,
|
||||
)
|
||||
run_id = "redis-run-active-remote-producer"
|
||||
|
||||
await owner.start_producer_lease(run_id)
|
||||
try:
|
||||
with pytest.raises(RuntimeError, match="active producer"):
|
||||
await recovering_peer.publish_end(run_id)
|
||||
assert owner._stream_key(run_id) not in fake.streams
|
||||
finally:
|
||||
await owner.abort_producer_lease(run_id)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_redis_abort_discards_unadmitted_tracking_marker():
|
||||
fake = _FakeRedis()
|
||||
bridge = RedisStreamBridge(redis_url="redis://fake", client=fake)
|
||||
run_id = "redis-run-unadmitted"
|
||||
|
||||
await bridge.start_producer_lease(run_id)
|
||||
await bridge.abort_producer_lease(run_id)
|
||||
|
||||
assert await bridge.producer_lease_state(run_id) == StreamProducerLeaseState.untracked
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_redis_abort_retries_transient_release_failure():
|
||||
from redis.exceptions import RedisError
|
||||
|
||||
fake = _FakeRedis()
|
||||
bridge = RedisStreamBridge(
|
||||
redis_url="redis://fake",
|
||||
stream_ttl_seconds=0,
|
||||
client=fake,
|
||||
)
|
||||
run_id = "redis-run-abort-failure"
|
||||
|
||||
await bridge.start_producer_lease(run_id)
|
||||
assert all(name != bridge._producer_tracking_key(run_id) for name, _seconds in fake.expirations)
|
||||
|
||||
original_delete = bridge._delete_if_token
|
||||
failed_once = False
|
||||
|
||||
async def fail_delete_once(name, token):
|
||||
nonlocal failed_once
|
||||
if not failed_once:
|
||||
failed_once = True
|
||||
raise RedisError("delete temporarily unavailable")
|
||||
return await original_delete(name, token)
|
||||
|
||||
bridge._delete_if_token = fail_delete_once
|
||||
await bridge.abort_producer_lease(run_id)
|
||||
|
||||
assert run_id not in bridge._producer_leases
|
||||
assert await bridge.producer_lease_state(run_id) == StreamProducerLeaseState.untracked
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_redis_close_retries_failed_abort_cleanup():
|
||||
fake = _FakeRedis()
|
||||
bridge = RedisStreamBridge(redis_url="redis://fake", client=fake)
|
||||
run_id = "redis-run-abort-close-retry"
|
||||
await bridge.start_producer_lease(run_id)
|
||||
bridge._rollback_producer_lease = AsyncMock(
|
||||
side_effect=[False, True],
|
||||
)
|
||||
|
||||
await bridge.abort_producer_lease(run_id)
|
||||
assert bridge._producer_leases[run_id].discard_tracking_requested is True
|
||||
|
||||
await bridge.close()
|
||||
|
||||
assert bridge._rollback_producer_lease.await_count == 2
|
||||
assert run_id not in bridge._producer_leases
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_redis_delayed_cleanup_preserves_successor_stream_and_lease():
|
||||
fake = _FakeRedis()
|
||||
bridge = RedisStreamBridge(redis_url="redis://fake", client=fake)
|
||||
run_id = "redis-run-cleanup-successor"
|
||||
stream_key = bridge._stream_key(run_id)
|
||||
tracking_key = bridge._producer_tracking_key(run_id)
|
||||
active_key = bridge._producer_active_key(run_id)
|
||||
|
||||
await bridge.start_producer_lease(run_id)
|
||||
await bridge.publish(run_id, "values", {"owner": "old"})
|
||||
await bridge.finish_producer_lease(run_id)
|
||||
fake.values[tracking_key] = "successor-token"
|
||||
fake.values[active_key] = "successor-token"
|
||||
await fake.xadd(
|
||||
stream_key,
|
||||
{"kind": "event", "event": "values", "data": '{"owner":"successor"}'},
|
||||
)
|
||||
|
||||
await bridge.cleanup(run_id)
|
||||
|
||||
assert fake.values[tracking_key] == "successor-token"
|
||||
assert fake.values[active_key] == "successor-token"
|
||||
assert fake.streams[stream_key]
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_redis_unowned_cleanup_preserves_active_successor():
|
||||
fake = _FakeRedis()
|
||||
owner = RedisStreamBridge(redis_url="redis://fake", client=fake)
|
||||
recovering_peer = RedisStreamBridge(redis_url="redis://fake", client=fake)
|
||||
run_id = "redis-run-unowned-cleanup-successor"
|
||||
stream_key = owner._stream_key(run_id)
|
||||
tracking_key = owner._producer_tracking_key(run_id)
|
||||
active_key = owner._producer_active_key(run_id)
|
||||
|
||||
await owner.start_producer_lease(run_id)
|
||||
await owner.publish(run_id, "values", {"owner": "successor"})
|
||||
|
||||
try:
|
||||
await recovering_peer.cleanup(run_id)
|
||||
assert fake.values[tracking_key] == fake.values[active_key]
|
||||
assert fake.streams[stream_key]
|
||||
finally:
|
||||
await owner.abort_producer_lease(run_id)
|
||||
await owner.cleanup(run_id)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_redis_producer_starts_for_unrelated_runs_do_not_serialize():
|
||||
fake = _FakeRedis()
|
||||
bridge = RedisStreamBridge(redis_url="redis://fake", client=fake)
|
||||
first_started = asyncio.Event()
|
||||
release_first = asyncio.Event()
|
||||
original_set = fake.set
|
||||
|
||||
async def delayed_set(name, value, **kwargs):
|
||||
if name == bridge._producer_active_key("run-a"):
|
||||
first_started.set()
|
||||
await release_first.wait()
|
||||
return await original_set(name, value, **kwargs)
|
||||
|
||||
fake.set = delayed_set
|
||||
first = asyncio.create_task(bridge.start_producer_lease("run-a"))
|
||||
second = None
|
||||
try:
|
||||
await asyncio.wait_for(first_started.wait(), timeout=1)
|
||||
second = asyncio.create_task(bridge.start_producer_lease("run-b"))
|
||||
assert await asyncio.wait_for(asyncio.shield(second), timeout=0.2)
|
||||
assert not first.done()
|
||||
finally:
|
||||
release_first.set()
|
||||
await asyncio.wait_for(first, timeout=1)
|
||||
if second is not None:
|
||||
await asyncio.gather(second, return_exceptions=True)
|
||||
await bridge.finish_producer_lease("run-a")
|
||||
await bridge.finish_producer_lease("run-b")
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_redis_producer_start_failure_rolls_back_before_falling_back():
|
||||
from redis.exceptions import RedisError
|
||||
|
||||
fake = _FakeRedis()
|
||||
bridge = RedisStreamBridge(redis_url="redis://fake", client=fake)
|
||||
run_id = "redis-run-producer-start-failure"
|
||||
original_set = fake.set
|
||||
|
||||
async def fail_tracking_write(name, value, **kwargs):
|
||||
created = await original_set(name, value, **kwargs)
|
||||
if name == bridge._producer_tracking_key(run_id):
|
||||
raise RedisError("tracking write acknowledgement unavailable")
|
||||
return created
|
||||
|
||||
fake.set = fail_tracking_write
|
||||
|
||||
assert await bridge.start_producer_lease(run_id) is False
|
||||
assert await bridge.producer_lease_state(run_id) == StreamProducerLeaseState.untracked
|
||||
assert run_id not in bridge._producer_leases
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_redis_ambiguous_producer_start_fails_closed():
|
||||
from redis.exceptions import RedisError
|
||||
|
||||
fake = _FakeRedis()
|
||||
bridge = RedisStreamBridge(redis_url="redis://fake", client=fake)
|
||||
run_id = "redis-run-producer-start-ambiguous"
|
||||
original_set = fake.set
|
||||
|
||||
async def fail_tracking_write(name, value, **kwargs):
|
||||
created = await original_set(name, value, **kwargs)
|
||||
if name == bridge._producer_tracking_key(run_id):
|
||||
raise RedisError("tracking write acknowledgement unavailable")
|
||||
return created
|
||||
|
||||
async def fail_readback(*names):
|
||||
raise RedisError("rollback read unavailable")
|
||||
|
||||
fake.set = fail_tracking_write
|
||||
fake.mget = fail_readback
|
||||
|
||||
with pytest.raises(
|
||||
RuntimeError,
|
||||
match="Could not safely establish stream producer lease",
|
||||
):
|
||||
await bridge.start_producer_lease(run_id)
|
||||
assert run_id not in bridge._producer_leases
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_redis_producer_start_rollback_drains_repeated_cancellation():
|
||||
fake = _FakeRedis()
|
||||
bridge = RedisStreamBridge(
|
||||
redis_url="redis://fake",
|
||||
stream_ttl_seconds=0,
|
||||
client=fake,
|
||||
)
|
||||
run_id = "redis-run-repeated-start-cancellation"
|
||||
tracking_written = asyncio.Event()
|
||||
rollback_started = asyncio.Event()
|
||||
allow_rollback = asyncio.Event()
|
||||
original_set = fake.set
|
||||
original_delete = bridge._delete_if_token
|
||||
|
||||
async def block_after_tracking_write(name, value, **kwargs):
|
||||
created = await original_set(name, value, **kwargs)
|
||||
if name == bridge._producer_tracking_key(run_id):
|
||||
tracking_written.set()
|
||||
await asyncio.Event().wait()
|
||||
return created
|
||||
|
||||
async def block_rollback(name, token):
|
||||
rollback_started.set()
|
||||
await allow_rollback.wait()
|
||||
return await original_delete(name, token)
|
||||
|
||||
fake.set = block_after_tracking_write
|
||||
bridge._delete_if_token = block_rollback
|
||||
task = asyncio.create_task(bridge.start_producer_lease(run_id))
|
||||
try:
|
||||
await asyncio.wait_for(tracking_written.wait(), timeout=1)
|
||||
task.cancel()
|
||||
await asyncio.wait_for(rollback_started.wait(), timeout=1)
|
||||
task.cancel()
|
||||
await asyncio.sleep(0)
|
||||
assert not task.done()
|
||||
|
||||
allow_rollback.set()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await asyncio.wait_for(task, timeout=1)
|
||||
assert await bridge.producer_lease_state(run_id) == StreamProducerLeaseState.untracked
|
||||
finally:
|
||||
allow_rollback.set()
|
||||
if not task.done():
|
||||
task.cancel()
|
||||
await asyncio.gather(task, return_exceptions=True)
|
||||
await fake.delete(
|
||||
bridge._producer_active_key(run_id),
|
||||
bridge._producer_tracking_key(run_id),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_redis_close_releases_producer_leases_concurrently():
|
||||
fake = _FakeRedis()
|
||||
bridge = RedisStreamBridge(redis_url="redis://fake", client=fake)
|
||||
run_ids = ("close-run-a", "close-run-b", "close-run-c")
|
||||
for run_id in run_ids:
|
||||
await bridge.start_producer_lease(run_id)
|
||||
|
||||
releases_started = 0
|
||||
all_releases_started = asyncio.Event()
|
||||
allow_releases = asyncio.Event()
|
||||
original_delete = bridge._delete_if_token
|
||||
|
||||
async def block_active_release(name, token):
|
||||
nonlocal releases_started
|
||||
if name.endswith(":producer-active"):
|
||||
releases_started += 1
|
||||
if releases_started == len(run_ids):
|
||||
all_releases_started.set()
|
||||
await allow_releases.wait()
|
||||
return await original_delete(name, token)
|
||||
|
||||
bridge._delete_if_token = block_active_release
|
||||
close_task = asyncio.create_task(bridge.close())
|
||||
try:
|
||||
await asyncio.wait_for(all_releases_started.wait(), timeout=0.2)
|
||||
finally:
|
||||
allow_releases.set()
|
||||
await asyncio.wait_for(close_task, timeout=1)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_redis_close_waits_for_inflight_producer_start():
|
||||
fake = _FakeRedis()
|
||||
bridge = RedisStreamBridge(redis_url="redis://fake", client=fake)
|
||||
run_id = "redis-run-close-during-start"
|
||||
tracking_write_started = asyncio.Event()
|
||||
allow_tracking_write = asyncio.Event()
|
||||
original_set = fake.set
|
||||
|
||||
async def block_tracking_write(name, value, **kwargs):
|
||||
if name == bridge._producer_tracking_key(run_id):
|
||||
tracking_write_started.set()
|
||||
await allow_tracking_write.wait()
|
||||
return await original_set(name, value, **kwargs)
|
||||
|
||||
fake.set = block_tracking_write
|
||||
start_task = asyncio.create_task(bridge.start_producer_lease(run_id))
|
||||
close_task = None
|
||||
try:
|
||||
await asyncio.wait_for(tracking_write_started.wait(), timeout=1)
|
||||
close_task = asyncio.create_task(bridge.close())
|
||||
await asyncio.sleep(0)
|
||||
assert not close_task.done()
|
||||
|
||||
allow_tracking_write.set()
|
||||
assert await asyncio.wait_for(start_task, timeout=1)
|
||||
await asyncio.wait_for(close_task, timeout=1)
|
||||
assert run_id not in bridge._producer_leases
|
||||
assert await bridge.producer_lease_state(run_id) == StreamProducerLeaseState.expired
|
||||
finally:
|
||||
allow_tracking_write.set()
|
||||
if not start_task.done():
|
||||
start_task.cancel()
|
||||
await asyncio.gather(start_task, return_exceptions=True)
|
||||
if close_task is not None and not close_task.done():
|
||||
await bridge.close()
|
||||
await asyncio.gather(close_task, return_exceptions=True)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_redis_close_honors_abort_intent_declared_while_stopping():
|
||||
fake = _FakeRedis()
|
||||
bridge = RedisStreamBridge(
|
||||
redis_url="redis://fake",
|
||||
stream_ttl_seconds=0,
|
||||
client=fake,
|
||||
)
|
||||
run_id = "redis-run-close-abort-intent"
|
||||
active_release_started = asyncio.Event()
|
||||
allow_active_release = asyncio.Event()
|
||||
original_delete = bridge._delete_if_token
|
||||
|
||||
async def block_active_release(name, token):
|
||||
if name == bridge._producer_active_key(run_id):
|
||||
active_release_started.set()
|
||||
await allow_active_release.wait()
|
||||
return await original_delete(name, token)
|
||||
|
||||
await bridge.start_producer_lease(run_id)
|
||||
bridge._delete_if_token = block_active_release
|
||||
close_task = asyncio.create_task(bridge.close())
|
||||
abort_task = None
|
||||
try:
|
||||
await asyncio.wait_for(active_release_started.wait(), timeout=1)
|
||||
abort_task = asyncio.create_task(bridge.abort_producer_lease(run_id))
|
||||
await asyncio.sleep(0)
|
||||
assert bridge._producer_leases[run_id].discard_tracking_requested is True
|
||||
|
||||
allow_active_release.set()
|
||||
await asyncio.wait_for(close_task, timeout=1)
|
||||
await asyncio.wait_for(abort_task, timeout=1)
|
||||
assert await bridge.producer_lease_state(run_id) == StreamProducerLeaseState.untracked
|
||||
finally:
|
||||
allow_active_release.set()
|
||||
await asyncio.gather(close_task, return_exceptions=True)
|
||||
if abort_task is not None:
|
||||
await asyncio.gather(abort_task, return_exceptions=True)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
@ -1044,6 +1766,11 @@ def test_stream_bridge_config_accepts_numeric_heartbeat_string():
|
||||
assert config.heartbeat_interval_seconds == 2.5
|
||||
|
||||
|
||||
def test_stream_bridge_producer_lease_recovery_requires_explicit_activation():
|
||||
assert StreamBridgeConfig().producer_lease_recovery_enabled is False
|
||||
assert StreamBridgeConfig(producer_lease_recovery_enabled=True).producer_lease_recovery_enabled is True
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"cleanup_delay",
|
||||
[
|
||||
@ -1375,6 +2102,7 @@ async def test_make_stream_bridge_passes_redis_options(monkeypatch):
|
||||
heartbeat_interval_seconds=2.5,
|
||||
max_connections=50,
|
||||
stream_ttl_seconds=42,
|
||||
producer_lease_recovery_enabled=True,
|
||||
)
|
||||
)
|
||||
try:
|
||||
@ -1382,6 +2110,7 @@ async def test_make_stream_bridge_passes_redis_options(monkeypatch):
|
||||
assert isinstance(bridge, RedisStreamBridge)
|
||||
assert bridge.heartbeat_interval == 2.5
|
||||
assert bridge._stream_ttl_seconds == 42
|
||||
assert bridge.producer_lease_recovery_enabled is True
|
||||
assert captured["max_connections"] == 50
|
||||
assert captured["decode_responses"] is True
|
||||
finally:
|
||||
|
||||
@ -2960,6 +2960,11 @@ run_ownership:
|
||||
# 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.
|
||||
# producer_lease_recovery_enabled: false # restart-required activation gate.
|
||||
# # First roll this binary to every Gateway and
|
||||
# # drain pre-upgrade runs; then enable on every
|
||||
# # instance together. Never enable during a
|
||||
# # mixed-version rollout.
|
||||
# recovered_stream_cleanup_delay_seconds: 60 # seconds to wait after
|
||||
# # publishing END for a recovered orphan run
|
||||
# # before deleting the stream key (max: 86400).
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user