mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-09-09 21:49:37 +00:00
feat(gateway): add /health/ready readiness probe backed by the database (#5166)
* feat(gateway): add /health/ready readiness probe backed by the database
## Why
GET /health only proves the process is up: it returns 200 even when the persistence engine cannot reach the database. Orchestrators already treat it as a readiness gate (docker-compose.yaml marks the gateway service healthy and nginx depends_on service_healthy), so a DB outage or a still-migrating Postgres leaves the stack 'healthy' while every request fails.
## What changed
- New GET /health/ready endpoint: bounded SELECT 1 against the existing persistence engine (deerflow.persistence.engine.get_engine) with a 2s timeout.
- Response is 200 {'status': 'ready', 'database': 'ok'} when reachable, 503 {'status': 'degraded', 'database': 'unreachable'} when the probe fails, and 200 ready with database=not_configured for backend=memory (nothing to probe).
- GET /health is unchanged (pure liveness), and /health/ready is public through the existing /health auth whitelist.
- docker-compose.yaml gateway healthcheck now polls /health/ready so service_healthy reflects database reachability.
- Documented both endpoints in backend/app/gateway/AGENTS.md.
## Surface area
- [x] Backend API - new GET /health/ready endpoint under backend/app/gateway
- [x] Sandbox / Docker - gateway healthcheck in docker/docker-compose.yaml now gates on readiness
- [ ] Frontend UI / Agents / Skills / Dependencies
- [x] Default behavior change - existing /health unchanged; the prod compose healthcheck is stricter (503 while the database is unreachable)
## Validation
- New unit tests in backend/tests/test_gateway_health.py cover ok / unreachable / not_configured probe results and the 200/503 payload mapping (6 passed).
- app.gateway.app imports cleanly and registers both /health and /health/ready.
- ruff check + ruff format clean.
## AI assistance
**Tool(s) used:** Codex (coding agent)
**How you used it:** design, implementation, and unit tests produced with AI assistance; reviewed before commit.
- [ ] I've read and understand every line of this change and take responsibility for it — it's not unreviewed AI output.
* fix(helm): point the gateway readiness probe at /health/ready
## Why
Review on #5166 (willem-bd, P1): the chart still probed /health for readiness,
so Kubernetes marked the pod ready and routed traffic while the database was
unreachable - exactly the failure mode /health/ready was added to catch.
## What changed
- deploy/helm/deer-flow/templates/gateway-deployment.yaml: readinessProbe
httpGet.path now hits /health/ready (DB-backed, 503 while the database is
unreachable). The liveness probe stays on /health.
## Verification
- One-line path change inside the existing readinessProbe block; git diff
confirms only the readiness path changed (liveness untouched).
## AI assistance
**Tool(s) used:** Codex (coding agent)
**How you used it:** implemented the reviewer-requested probe path change; reviewed before commit.
- [ ] I've read and understand every line of this change and take responsibility for it — it's not unreviewed AI output.
* fix(gateway): readiness probe also checks the effective checkpointer/Store backend
## Why
Follow-up review on #5166 (willem-bd, P1): get_engine() only represents the
ORM backend selected by `database:`. The legacy `checkpointer:` section takes
precedence for the LangGraph checkpointer and Store, so a split configuration
(a local SQLite/memory `database:` with `checkpointer.type: postgres`) could
report 200 while the PostgreSQL backend agent runs depend on was down.
## What changed
- GET /health/ready now probes both persistence halves: the ORM engine behind
`database:` (unchanged) and the effective LangGraph checkpointer/Store
backend resolved with the runtime's own rule (legacy `checkpointer:` config,
otherwise derived from `database:`), for memory/sqlite/postgres backends.
- The payload gains a `checkpointer` field with the same
ok / not_configured / unreachable vocabulary as `database`; 503 degraded is
returned when either probe is unreachable.
- Probes are bounded by the existing 2s timeout: sqlite via aiosqlite SELECT 1
on the resolved path, postgres via a bounded psycopg AsyncConnection SELECT 1
on the DSN with the configured search_path. A missing driver for a configured
backend degrades readiness (the runtime could not run either).
- Documented the two-probe semantics in the endpoint docstring and
backend/app/gateway/AGENTS.md.
## Verification
- New tests: healthy ORM engine + unreachable legacy checkpointer backend ->
503 degraded with database: ok / checkpointer: unreachable; checkpointer
probe mapping for memory/sqlite(postgres missing-driver) backends; existing
payload tests now pin the checkpointer field.
- cd backend && python -m pytest tests/test_gateway_health.py: 11 passed.
- app.gateway.app imports cleanly; ruff check + ruff format clean.
## AI assistance
**Tool(s) used:** Codex (coding agent)
**How you used it:** design, implementation, and unit tests produced with AI assistance; reviewed before commit.
- [ ] I've read and understand every line of this change and take responsibility for it — it's not unreviewed AI output.
* fix(gateway): bound /health/ready to one deadline and probe the startup checkpointer snapshot
## Why
Second round of review on #5166 (zhfeng P1/P2, willem-bd P1/P1/P2). Three
correctness issues remained in the readiness endpoint:
- The database and checkpointer probes ran sequentially, each allowed 2s, so a
healthy response could take almost 4s - past Kubernetes' 1s default
readinessProbe timeout and inside Docker's 3s client timeout. A slow but
healthy backend could make every replica unready.
- The checkpointer probe re-resolved process-wide, hot-reloaded configuration
per request, while app.state.checkpointer/store are built once from the
startup_config snapshot in langgraph_runtime(). After a live config edit the
endpoint could probe a backend the running gateway does not use, and a
resolution failure was swallowed into None -> not_configured -> 200.
- The SQLite probe opened the path with aiosqlite.connect(), which creates the
file when missing: a deleted checkpoint database was silently resurrected as
an empty file and reported ok instead of surfacing the outage.
## What changed
- backend/app/gateway/health.py: the two probes now run concurrently beneath a
single endpoint-wide deadline (_READINESS_DEADLINE_SECONDS=3.0) so a healthy
response completes within one probe window (~2s), never the sum of both.
A probe that overruns the deadline degrades the endpoint instead of hanging.
- langgraph_runtime() now records the checkpointer/Store config resolved from
the same startup_config snapshot its checkpointer/store singletons are built
from (app.state.checkpointer_config); /health/ready probes that snapshot and
never re-resolves hot-reloaded config. resolve_checkpointer_config() returns
None on resolution failure and the endpoint fails closed (503, checkpointer:
unreachable) instead of reporting not_configured.
- The SQLite probe opens disk-backed paths with the non-creating mode=rw URI
flag, so a missing database file stays missing and yields unreachable;
in-memory forms (:memory:, file:...mode=memory) have nothing external to
probe and report not_configured like the memory backend.
- Orchestrator timeouts now sit above the endpoint bound: Helm readinessProbe
gains timeoutSeconds: 5 (Kubernetes default is 1s) and the docker-compose
gateway healthcheck client timeout moves from 3s to 5s.
## Verification
- New regression tests: concurrent probes keep total elapsed time within one
probe window; a probe ignoring its budget trips the endpoint deadline to 503;
missing SQLite file stays absent and yields unreachable; in-memory SQLite
forms map to not_configured; missing startup snapshot / config resolution
failure fail closed to 503; resolve_checkpointer_config() raising is covered.
- cd backend && python -m pytest tests/test_gateway_health.py: 21 passed;
tests/test_gateway_docs_toggle.py and lifespan/shutdown gateway suites pass.
- ruff check + ruff format clean on all changed files.
## AI assistance
**Tool(s) used:** Codex (coding agent)
**How you used it:** implemented the reviewer-requested concurrency/deadline, startup-snapshot probing, fail-closed resolution, and non-creating SQLite probe; reviewed before commit.
- [ ] I've read and understand every line of this change and take responsibility for it — it's not unreviewed AI output.
* fix(gateway): serialize connection-opening readiness probes behind a strict gate
## Why
Review on #5166 (willem-bd, P1): every request to /health/ready opened a new
PostgreSQL connection in _probe_postgres_backend, outside both the ORM pool and
the runtime checkpointer pool. The route is public through the /health auth
prefix and nginx proxies /health/*, so concurrent unauthenticated requests
could create an unbounded number of connections (each held for up to two
seconds), exhaust PostgreSQL max_connections, and take down both normal
traffic and the readiness probe itself.
## What changed
- backend/app/gateway/health.py: connection-opening checkpointer probes
(sqlite connect, postgres AsyncConnection.connect) now run inside a strict
per-process gate - an asyncio.Lock cached per running event loop - so at
most one probe connection can be in flight per worker process. Requests
that queue behind the gate are still shed by the existing endpoint-wide
deadline, so a flood cannot pile up new connections or open files.
- Memory and unknown-backend decisions stay outside the gate; payload and
probe semantics are unchanged. The serialization is documented in the
module docstring and backend/app/gateway/AGENTS.md.
## Verification
- New regression test: 8 concurrent readiness_payload() requests against an
instrumented sqlite probe assert the maximum number of in-flight probe
connections is 1 while every request still returns 200.
- cd backend && python -m pytest tests/test_gateway_health.py: 22 passed;
tests/test_gateway_docs_toggle.py and tests/test_gateway_lifespan_shutdown.py
also pass on the merged main head.
- ruff check + ruff format clean on all changed files.
## AI assistance
**Tool(s) used:** Codex (coding agent)
**How you used it:** implemented the reviewer-requested strict concurrency bound for the public readiness probe; reviewed before commit.
- [ ] I've read and understand every line of this change and take responsibility for it — it's not unreviewed AI output.
This commit is contained in:
Michael
GitHub
parent
e21245fd5b
commit
4791e94a73
@ -1,6 +1,6 @@
|
||||
### Gateway API (`app/gateway/`)
|
||||
|
||||
FastAPI listens on port 8001; health: `GET /health`. Set `GATEWAY_ENABLE_DOCS=false` to disable the default `/docs`, `/redoc`, and `/openapi.json` endpoints.
|
||||
FastAPI listens on port 8001; health: `GET /health` (liveness) and `GET /health/ready` (readiness; concurrently probes the ORM engine behind `database:` plus the effective LangGraph checkpointer/Store backend - the legacy `checkpointer:` section, otherwise derived from `database:`, resolved from the startup config snapshot recorded on `app.state` - beneath a single bounded deadline, with connection-opening probes serialized behind a strict per-process gate, 503 while either is unreachable or the startup backend cannot be resolved, `not_configured` for process-local backends such as `backend=memory`). Set `GATEWAY_ENABLE_DOCS=false` to disable the default `/docs`, `/redoc`, and `/openapi.json` endpoints.
|
||||
|
||||
Durable MCP notifications use internal Agent runs. Keep their trusted delivery instruction outside the user-input boundary, and frame serialized remote events as untrusted before model invocation. Strict thread existence/ownership admission dead-letters events whose task outlives its deleted chat instead of recreating the thread.
|
||||
|
||||
|
||||
@ -4,7 +4,7 @@ from collections.abc import AsyncGenerator
|
||||
from contextlib import asynccontextmanager
|
||||
|
||||
from deerflow_extension_api import EXTENSION_PRINCIPAL_RESOLVER_KEY, ExtensionPrincipal
|
||||
from fastapi import FastAPI
|
||||
from fastapi import FastAPI, Request, Response
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
|
||||
from app.gateway.auth_disabled import AUTH_SOURCE_INTERNAL, AUTH_SOURCE_PAT, warn_if_auth_disabled_enabled
|
||||
@ -13,6 +13,7 @@ from app.gateway.browser_capability import ensure_browser_runtime_available
|
||||
from app.gateway.config import get_gateway_config
|
||||
from app.gateway.csrf_middleware import CORS_EXPOSED_HEADERS, CSRFMiddleware, get_configured_cors_origins
|
||||
from app.gateway.deps import langgraph_runtime
|
||||
from app.gateway.health import READINESS_CHECKPOINTER_CONFIG_ATTR, readiness_payload
|
||||
from app.gateway.routers import (
|
||||
agents,
|
||||
artifacts,
|
||||
@ -882,6 +883,24 @@ This gateway provides runtime endpoints for agent runs plus custom endpoints for
|
||||
"""
|
||||
return {"status": "healthy", "service": "deer-flow-gateway"}
|
||||
|
||||
@app.get("/health/ready", tags=["health"])
|
||||
async def readiness_check(request: Request, response: Response) -> dict[str, str]:
|
||||
"""Readiness endpoint: 200 when the persistence backends are reachable.
|
||||
|
||||
Probes the ORM engine behind ``database:`` and the effective LangGraph
|
||||
checkpointer/Store backend (legacy ``checkpointer:`` section, otherwise
|
||||
derived from ``database:``) concurrently beneath one bounded deadline.
|
||||
The checkpointer config comes from the startup snapshot recorded by
|
||||
``langgraph_runtime`` (never hot-reloaded config), so orchestrators can
|
||||
gate on the gateway actually being ready rather than merely alive.
|
||||
Returns 503 with ``status: degraded`` when either probe fails or the
|
||||
startup backend cannot be resolved.
|
||||
"""
|
||||
checkpointer_config = getattr(request.app.state, READINESS_CHECKPOINTER_CONFIG_ATTR, None)
|
||||
status_code, payload = await readiness_payload(checkpointer_config)
|
||||
response.status_code = status_code
|
||||
return payload
|
||||
|
||||
# Extension routes are deliberately last: FastAPI/Starlette dispatches in
|
||||
# registration order, so every host route (including conditional routes
|
||||
# and /health) keeps precedence. Definite shadows are rejected with an
|
||||
|
||||
@ -449,6 +449,15 @@ async def langgraph_runtime(app: FastAPI, startup_config: AppConfig) -> AsyncGen
|
||||
app.state.checkpointer = await stack.enter_async_context(make_checkpointer(config))
|
||||
app.state.store = await stack.enter_async_context(make_store(config))
|
||||
|
||||
# Record the checkpointer/Store backend selected from this startup
|
||||
# snapshot so GET /health/ready probes what the running process
|
||||
# actually uses. These singletons are restart-required by design and
|
||||
# are never rebuilt on config.yaml hot reload, so the probe must not
|
||||
# re-resolve process-wide configuration per request.
|
||||
from app.gateway.health import READINESS_CHECKPOINTER_CONFIG_ATTR, resolve_checkpointer_config
|
||||
|
||||
setattr(app.state, READINESS_CHECKPOINTER_CONFIG_ATTR, resolve_checkpointer_config(config))
|
||||
|
||||
# Initialize repositories — one get_session_factory() call for all.
|
||||
sf = get_session_factory()
|
||||
if sf is not None:
|
||||
|
||||
282
backend/app/gateway/health.py
Normal file
282
backend/app/gateway/health.py
Normal file
@ -0,0 +1,282 @@
|
||||
"""Readiness probe helpers for the gateway health endpoints.
|
||||
|
||||
``GET /health`` stays a pure liveness signal: 200 whenever the process is up.
|
||||
``GET /health/ready`` additionally probes the persistence the gateway actually
|
||||
depends on, so orchestrators (Docker healthchecks, Kubernetes probes) treat the
|
||||
gateway as ready only when the databases behind agent runs are reachable. Two
|
||||
backends can be configured independently:
|
||||
|
||||
* the ORM engine behind ``database:`` (application repositories), and
|
||||
* the effective LangGraph checkpointer/Store backend - the legacy
|
||||
``checkpointer:`` section when present, otherwise derived from ``database:``
|
||||
(memory/sqlite/postgres).
|
||||
|
||||
Both probes run concurrently beneath a single endpoint-wide deadline
|
||||
(:data:`_READINESS_DEADLINE_SECONDS`), so a healthy response completes within
|
||||
one probe window rather than the sum of both budgets. The checkpointer config
|
||||
is resolved once at startup from the same snapshot ``langgraph_runtime`` builds
|
||||
its resources from and is stored on ``app.state``; probing a hot-reloaded
|
||||
config instead could check a backend the running process is not using. A
|
||||
``backend=memory`` deployment has nothing to probe and is always considered
|
||||
ready; a startup config that cannot be resolved fails closed as unreachable.
|
||||
Connection-opening probes are serialized behind a strict per-process gate: the
|
||||
route is public through the ``/health`` auth prefix, so unlimited concurrent
|
||||
requests must never translate into unlimited new PostgreSQL connections.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import pathlib
|
||||
import urllib.parse
|
||||
import weakref
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from sqlalchemy import text
|
||||
|
||||
from deerflow.persistence.engine import get_engine
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from deerflow.config.app_config import AppConfig
|
||||
from deerflow.config.checkpointer_config import CheckpointerConfig
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Upper bound for a single probe attempt. The endpoint must never hang behind
|
||||
# a dead database (for example a TCP connect timeout to Postgres).
|
||||
_PROBE_TIMEOUT_SECONDS = 2.0
|
||||
|
||||
# Whole-endpoint deadline covering both probes. They run concurrently, so a
|
||||
# healthy response completes within a single probe window; the extra margin
|
||||
# absorbs scheduling/cancellation overhead without letting the request
|
||||
# approach the sum of both probe budgets. Orchestrator timeouts (Helm
|
||||
# readinessProbe ``timeoutSeconds``, the docker-compose healthcheck client
|
||||
# timeout) must be configured above this bound.
|
||||
_READINESS_DEADLINE_SECONDS = 3.0
|
||||
|
||||
# ``app.state`` attribute under which :func:`app.gateway.deps.langgraph_runtime`
|
||||
# records the startup-bound checkpointer/Store config the probe targets.
|
||||
READINESS_CHECKPOINTER_CONFIG_ATTR = "checkpointer_config"
|
||||
|
||||
# One gate per running event loop (one per worker process in production; one
|
||||
# per test loop in the suite). ``/health/ready`` is public and unauthenticated,
|
||||
# so a thundering herd of probes - or an attacker - must never be able to open
|
||||
# an unbounded number of new connections: every connection-opening probe below
|
||||
# is serialized through this gate, bounding in-flight probe connections to one
|
||||
# per process. Waiting requests are still shed by the endpoint-wide deadline.
|
||||
_PROBE_GATES: weakref.WeakKeyDictionary[asyncio.AbstractEventLoop, asyncio.Lock] = weakref.WeakKeyDictionary()
|
||||
|
||||
|
||||
def _probe_gate() -> asyncio.Lock:
|
||||
"""Return the serialization gate bound to the currently running loop."""
|
||||
loop = asyncio.get_running_loop()
|
||||
gate = _PROBE_GATES.get(loop)
|
||||
if gate is None:
|
||||
gate = asyncio.Lock()
|
||||
_PROBE_GATES[loop] = gate
|
||||
return gate
|
||||
|
||||
|
||||
# Result vocabulary for the database probe.
|
||||
DATABASE_OK = "ok"
|
||||
DATABASE_NOT_CONFIGURED = "not_configured"
|
||||
DATABASE_UNREACHABLE = "unreachable"
|
||||
|
||||
|
||||
async def check_database_health() -> str:
|
||||
"""Probe the persistence engine; return one of the DATABASE_* values."""
|
||||
engine = get_engine()
|
||||
if engine is None:
|
||||
# backend=memory, or the engine has not been initialized yet: there is
|
||||
# no database to probe.
|
||||
return DATABASE_NOT_CONFIGURED
|
||||
try:
|
||||
async with asyncio.timeout(_PROBE_TIMEOUT_SECONDS):
|
||||
async with engine.connect() as connection:
|
||||
await connection.execute(text("SELECT 1"))
|
||||
except Exception:
|
||||
logger.warning("Readiness database probe failed", exc_info=True)
|
||||
return DATABASE_UNREACHABLE
|
||||
return DATABASE_OK
|
||||
|
||||
|
||||
def resolve_checkpointer_config(startup_config: AppConfig) -> CheckpointerConfig | None:
|
||||
"""Resolve the checkpointer/Store backend bound to a startup config snapshot.
|
||||
|
||||
Mirrors the runtime's own selection (the legacy ``checkpointer`` section
|
||||
first, otherwise derived from the unified ``database`` section), so the
|
||||
probe targets the exact backend ``langgraph_runtime`` built at startup -
|
||||
which can differ from the ORM ``database:`` backend and from a later,
|
||||
hot-reloaded config. Returns None when the config cannot be resolved;
|
||||
callers must treat that as a failure (unreachable), never as
|
||||
``not_configured``.
|
||||
"""
|
||||
from deerflow.runtime.checkpointer.provider import _resolve_checkpointer_config
|
||||
|
||||
try:
|
||||
return _resolve_checkpointer_config(startup_config)
|
||||
except Exception:
|
||||
logger.warning(
|
||||
"Readiness probe: unable to resolve the startup checkpointer config; failing closed",
|
||||
exc_info=True,
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
def _sqlite_is_in_memory(conn_str: str) -> bool:
|
||||
"""Return True when *conn_str* refers to a purely in-memory SQLite database."""
|
||||
if conn_str == ":memory:":
|
||||
return True
|
||||
if not conn_str.startswith("file:"):
|
||||
return False
|
||||
parts = urllib.parse.urlsplit(conn_str)
|
||||
if parts.path in (":memory:", ""):
|
||||
return True
|
||||
return any(key == "mode" and value == "memory" for key, value in urllib.parse.parse_qsl(parts.query))
|
||||
|
||||
|
||||
def _sqlite_disk_uri(conn_str: str) -> str:
|
||||
"""Return a non-creating (``mode=rw``) SQLite URI for a disk-backed database.
|
||||
|
||||
Opening with ``mode=rw`` refuses to create a missing database file, so a
|
||||
readiness probe can never resurrect a checkpointer/Store file that was
|
||||
deleted or lost after startup - absence must surface as unreachable. Plain
|
||||
filesystem paths (already absolute after
|
||||
``deerflow.runtime.store._sqlite_utils.resolve_sqlite_conn_str``) are
|
||||
converted with ``Path.as_uri`` for correct percent-encoding; existing
|
||||
``file:`` URIs keep their path bytes and get ``mode=rw`` merged into the
|
||||
query, replacing any pinned mode.
|
||||
"""
|
||||
if not conn_str.startswith("file:"):
|
||||
return f"{pathlib.Path(conn_str).as_uri()}?mode=rw"
|
||||
parts = urllib.parse.urlsplit(conn_str)
|
||||
query_pairs = urllib.parse.parse_qsl(parts.query, keep_blank_values=True)
|
||||
if not any(key == "mode" for key, _ in query_pairs):
|
||||
separator = "&" if parts.query else "?"
|
||||
return f"{conn_str}{separator}mode=rw"
|
||||
replaced = urllib.parse.urlencode([(key, "rw") if key == "mode" else (key, value) for key, value in query_pairs])
|
||||
return urllib.parse.urlunsplit((parts.scheme, parts.netloc, parts.path, replaced, parts.fragment))
|
||||
|
||||
|
||||
async def _probe_sqlite_backend(conn_string: str | None) -> str:
|
||||
"""Probe a SQLite checkpointer/Store database with a bounded SELECT 1.
|
||||
|
||||
Disk-backed databases are opened non-creating (``mode=rw``): a missing
|
||||
file stays missing and fails the probe instead of being recreated empty.
|
||||
In-memory forms (``:memory:`` and ``file:`` URIs with ``mode=memory``)
|
||||
only exist inside the running process, so there is nothing external to
|
||||
probe and they report ``not_configured`` like the memory backend.
|
||||
"""
|
||||
try:
|
||||
import aiosqlite
|
||||
except ImportError:
|
||||
logger.error("Readiness probe: aiosqlite is not installed for the sqlite checkpointer backend")
|
||||
return DATABASE_UNREACHABLE
|
||||
from deerflow.runtime.store._sqlite_utils import resolve_sqlite_conn_str
|
||||
|
||||
conn_str = resolve_sqlite_conn_str(conn_string or "store.db")
|
||||
if _sqlite_is_in_memory(conn_str):
|
||||
return DATABASE_NOT_CONFIGURED
|
||||
try:
|
||||
async with asyncio.timeout(_PROBE_TIMEOUT_SECONDS):
|
||||
connection = await aiosqlite.connect(_sqlite_disk_uri(conn_str), uri=True)
|
||||
try:
|
||||
await connection.execute("SELECT 1")
|
||||
finally:
|
||||
await connection.close()
|
||||
except Exception:
|
||||
logger.warning("Readiness sqlite checkpointer probe failed", exc_info=True)
|
||||
return DATABASE_UNREACHABLE
|
||||
return DATABASE_OK
|
||||
|
||||
|
||||
async def _probe_postgres_backend(conn_string: str, schema: str) -> str:
|
||||
"""Probe a PostgreSQL checkpointer/Store database with a bounded SELECT 1."""
|
||||
try:
|
||||
from psycopg import AsyncConnection
|
||||
except ImportError:
|
||||
logger.error("Readiness probe: psycopg is not installed for the postgres checkpointer backend")
|
||||
return DATABASE_UNREACHABLE
|
||||
try:
|
||||
from deerflow.persistence.postgres_schema import dsn_with_search_path, normalize_libpq_dsn
|
||||
|
||||
dsn = dsn_with_search_path(normalize_libpq_dsn(conn_string), schema)
|
||||
async with asyncio.timeout(_PROBE_TIMEOUT_SECONDS):
|
||||
connection = await AsyncConnection.connect(dsn, connect_timeout=int(_PROBE_TIMEOUT_SECONDS))
|
||||
try:
|
||||
async with connection.cursor() as cursor:
|
||||
await cursor.execute("SELECT 1")
|
||||
finally:
|
||||
await connection.close()
|
||||
except Exception:
|
||||
logger.warning("Readiness postgres checkpointer probe failed", exc_info=True)
|
||||
return DATABASE_UNREACHABLE
|
||||
return DATABASE_OK
|
||||
|
||||
|
||||
async def _probe_checkpointer_backend(config: CheckpointerConfig) -> str:
|
||||
"""Probe the LangGraph checkpointer/Store backend described by *config*.
|
||||
|
||||
*config* is the startup-bound snapshot (see :func:`resolve_checkpointer_config`);
|
||||
an in-process memory backend has nothing external to probe. Probes that
|
||||
open a connection (sqlite file, postgres) are serialized so concurrent
|
||||
unauthenticated requests cannot exhaust the database's connections.
|
||||
"""
|
||||
if config.type == "memory":
|
||||
# In-process backend: there is nothing external to probe.
|
||||
return DATABASE_NOT_CONFIGURED
|
||||
if config.type not in ("sqlite", "postgres"):
|
||||
logger.warning("Readiness probe: unknown checkpointer backend %r", config.type)
|
||||
return DATABASE_UNREACHABLE
|
||||
async with _probe_gate():
|
||||
if config.type == "sqlite":
|
||||
return await _probe_sqlite_backend(config.connection_string)
|
||||
if not config.connection_string:
|
||||
return DATABASE_UNREACHABLE
|
||||
return await _probe_postgres_backend(config.connection_string, config.postgres_schema)
|
||||
|
||||
|
||||
async def readiness_payload(checkpointer_config: CheckpointerConfig | None = None) -> tuple[int, dict[str, str]]:
|
||||
"""Return the (status_code, body) pair served by ``GET /health/ready``.
|
||||
|
||||
Probes both persistence halves the gateway depends on: the ORM engine
|
||||
behind ``database:`` (repositories) and the effective LangGraph
|
||||
checkpointer/Store backend (the legacy ``checkpointer:`` section, otherwise
|
||||
derived from ``database:``). The probes run concurrently beneath one
|
||||
endpoint-wide deadline so the request duration is bounded by the slowest
|
||||
single probe, not their sum. ``checkpointer_config`` is the startup
|
||||
snapshot recorded by ``langgraph_runtime``; None means no snapshot could be
|
||||
resolved, which fails closed as an unreachable backend rather than
|
||||
reporting ready. Either backend can be configured independently of the
|
||||
other, so an unreachable probe on either degrades the endpoint.
|
||||
"""
|
||||
|
||||
async def _probe_engine() -> str:
|
||||
return await check_database_health()
|
||||
|
||||
async def _probe_checkpointer() -> str:
|
||||
if checkpointer_config is None:
|
||||
# Fail closed: without the startup-bound config we cannot know what
|
||||
# backend agent runs use, so readiness must not be claimed.
|
||||
logger.error("Readiness probe: no startup checkpointer config snapshot recorded; failing closed")
|
||||
return DATABASE_UNREACHABLE
|
||||
return await _probe_checkpointer_backend(checkpointer_config)
|
||||
|
||||
try:
|
||||
async with asyncio.timeout(_READINESS_DEADLINE_SECONDS):
|
||||
database, checkpointer = await asyncio.gather(_probe_engine(), _probe_checkpointer())
|
||||
except TimeoutError:
|
||||
logger.error(
|
||||
"Readiness probes exceeded the %.1fs endpoint deadline",
|
||||
_READINESS_DEADLINE_SECONDS,
|
||||
)
|
||||
database = checkpointer = DATABASE_UNREACHABLE
|
||||
degraded = DATABASE_UNREACHABLE in (database, checkpointer)
|
||||
payload = {
|
||||
"status": "degraded" if degraded else "ready",
|
||||
"service": "deer-flow-gateway",
|
||||
"database": database,
|
||||
"checkpointer": checkpointer,
|
||||
}
|
||||
return (503 if degraded else 200, payload)
|
||||
303
backend/tests/test_gateway_health.py
Normal file
303
backend/tests/test_gateway_health.py
Normal file
@ -0,0 +1,303 @@
|
||||
"""Unit tests for the gateway readiness probe (app.gateway.health)."""
|
||||
|
||||
import asyncio
|
||||
import pathlib
|
||||
import sqlite3
|
||||
import sys
|
||||
import time
|
||||
from contextlib import asynccontextmanager
|
||||
|
||||
import pytest
|
||||
|
||||
import app.gateway.health as health_module
|
||||
from app.gateway.health import (
|
||||
DATABASE_NOT_CONFIGURED,
|
||||
DATABASE_OK,
|
||||
DATABASE_UNREACHABLE,
|
||||
_probe_checkpointer_backend,
|
||||
check_database_health,
|
||||
readiness_payload,
|
||||
resolve_checkpointer_config,
|
||||
)
|
||||
from deerflow.config.checkpointer_config import CheckpointerConfig
|
||||
|
||||
|
||||
class _FakeConnection:
|
||||
async def execute(self, *args, **kwargs):
|
||||
return None
|
||||
|
||||
|
||||
class _FakeEngine:
|
||||
def __init__(self, *, unreachable: bool = False):
|
||||
self._unreachable = unreachable
|
||||
|
||||
def connect(self):
|
||||
@asynccontextmanager
|
||||
async def _connect():
|
||||
if self._unreachable:
|
||||
raise RuntimeError("database is down")
|
||||
yield _FakeConnection()
|
||||
|
||||
return _connect()
|
||||
|
||||
|
||||
def _create_sqlite_file(path: pathlib.Path) -> None:
|
||||
"""Create a valid (empty) SQLite database file at *path*."""
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
sqlite3.connect(str(path)).close()
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_check_database_health_without_engine(monkeypatch):
|
||||
"""backend=memory (no engine) must report not_configured, never unreachable."""
|
||||
monkeypatch.setattr("app.gateway.health.get_engine", lambda: None)
|
||||
|
||||
assert await check_database_health() == DATABASE_NOT_CONFIGURED
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_check_database_health_reachable(monkeypatch):
|
||||
monkeypatch.setattr("app.gateway.health.get_engine", lambda: _FakeEngine())
|
||||
|
||||
assert await check_database_health() == DATABASE_OK
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_check_database_health_unreachable(monkeypatch):
|
||||
monkeypatch.setattr("app.gateway.health.get_engine", lambda: _FakeEngine(unreachable=True))
|
||||
|
||||
assert await check_database_health() == DATABASE_UNREACHABLE
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_readiness_payload_ready_when_database_ok_and_memory_backend(monkeypatch):
|
||||
monkeypatch.setattr("app.gateway.health.get_engine", lambda: _FakeEngine())
|
||||
|
||||
status_code, payload = await readiness_payload(CheckpointerConfig(type="memory"))
|
||||
|
||||
assert status_code == 200
|
||||
assert payload["status"] == "ready"
|
||||
assert payload["database"] == DATABASE_OK
|
||||
assert payload["checkpointer"] == DATABASE_NOT_CONFIGURED
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_readiness_payload_degraded_when_database_unreachable(monkeypatch):
|
||||
monkeypatch.setattr("app.gateway.health.get_engine", lambda: _FakeEngine(unreachable=True))
|
||||
|
||||
status_code, payload = await readiness_payload(CheckpointerConfig(type="memory"))
|
||||
|
||||
assert status_code == 503
|
||||
assert payload["status"] == "degraded"
|
||||
assert payload["database"] == DATABASE_UNREACHABLE
|
||||
assert payload["checkpointer"] == DATABASE_NOT_CONFIGURED
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_readiness_payload_ready_when_nothing_configured(monkeypatch):
|
||||
"""backend=memory end to end must stay ready with not_configured results."""
|
||||
monkeypatch.setattr("app.gateway.health.get_engine", lambda: None)
|
||||
|
||||
status_code, payload = await readiness_payload(CheckpointerConfig(type="memory"))
|
||||
|
||||
assert status_code == 200
|
||||
assert payload["status"] == "ready"
|
||||
assert payload["database"] == DATABASE_NOT_CONFIGURED
|
||||
assert payload["checkpointer"] == DATABASE_NOT_CONFIGURED
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_readiness_payload_degraded_when_checkpointer_unreachable_but_database_ok(tmp_path, monkeypatch):
|
||||
"""A healthy ORM engine must not mask an unreachable legacy checkpointer backend."""
|
||||
monkeypatch.setattr("app.gateway.health.get_engine", lambda: _FakeEngine())
|
||||
config = CheckpointerConfig(type="sqlite", connection_string=str(tmp_path / "missing" / "checkpoints.db"))
|
||||
|
||||
status_code, payload = await readiness_payload(config)
|
||||
|
||||
assert status_code == 503
|
||||
assert payload["status"] == "degraded"
|
||||
assert payload["database"] == DATABASE_OK
|
||||
assert payload["checkpointer"] == DATABASE_UNREACHABLE
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_readiness_payload_fails_closed_without_startup_snapshot(monkeypatch):
|
||||
"""No startup config snapshot must degrade readiness, never report ready."""
|
||||
monkeypatch.setattr("app.gateway.health.get_engine", lambda: _FakeEngine())
|
||||
|
||||
status_code, payload = await readiness_payload(None)
|
||||
|
||||
assert status_code == 503
|
||||
assert payload["status"] == "degraded"
|
||||
assert payload["database"] == DATABASE_OK
|
||||
assert payload["checkpointer"] == DATABASE_UNREACHABLE
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_readiness_probes_run_concurrently(monkeypatch):
|
||||
"""Slow-but-healthy probes must not add their budgets together."""
|
||||
|
||||
async def _slow_ok(*args) -> str:
|
||||
await asyncio.sleep(0.35)
|
||||
return DATABASE_OK
|
||||
|
||||
monkeypatch.setattr(health_module, "check_database_health", _slow_ok)
|
||||
monkeypatch.setattr(health_module, "_probe_checkpointer_backend", _slow_ok)
|
||||
|
||||
started = time.perf_counter()
|
||||
status_code, payload = await readiness_payload(CheckpointerConfig(type="memory"))
|
||||
elapsed = time.perf_counter() - started
|
||||
|
||||
assert status_code == 200
|
||||
assert payload["database"] == DATABASE_OK
|
||||
assert payload["checkpointer"] == DATABASE_OK
|
||||
# Two sequential 0.35s probes would take ~0.7s; concurrent ones finish
|
||||
# within a single probe window.
|
||||
assert elapsed < 0.6
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_readiness_payload_enforces_endpoint_deadline(monkeypatch):
|
||||
"""A probe ignoring its own budget must trip the endpoint-wide deadline."""
|
||||
|
||||
async def _hanging(*args) -> str:
|
||||
await asyncio.sleep(30)
|
||||
return DATABASE_OK
|
||||
|
||||
monkeypatch.setattr(health_module, "check_database_health", _hanging)
|
||||
monkeypatch.setattr(health_module, "_probe_checkpointer_backend", _hanging)
|
||||
monkeypatch.setattr(health_module, "_READINESS_DEADLINE_SECONDS", 0.05)
|
||||
|
||||
status_code, payload = await readiness_payload(CheckpointerConfig(type="memory"))
|
||||
|
||||
assert status_code == 503
|
||||
assert payload["status"] == "degraded"
|
||||
assert payload["database"] == DATABASE_UNREACHABLE
|
||||
assert payload["checkpointer"] == DATABASE_UNREACHABLE
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_concurrent_readiness_requests_do_not_open_concurrent_probe_connections(monkeypatch):
|
||||
"""Public /health/ready must serialize connection-opening probes.
|
||||
|
||||
An unauthenticated thundering herd must never translate into an unbounded
|
||||
number of new database connections (e.g. past PostgreSQL
|
||||
max_connections): at most one probe connection may be in flight at a time
|
||||
per process.
|
||||
"""
|
||||
active = 0
|
||||
max_active = 0
|
||||
|
||||
async def _tracked_probe(conn_string: str | None) -> str:
|
||||
nonlocal active, max_active
|
||||
active += 1
|
||||
max_active = max(max_active, active)
|
||||
try:
|
||||
await asyncio.sleep(0.05)
|
||||
return DATABASE_OK
|
||||
finally:
|
||||
active -= 1
|
||||
|
||||
monkeypatch.setattr(health_module, "_probe_sqlite_backend", _tracked_probe)
|
||||
monkeypatch.setattr("app.gateway.health.get_engine", lambda: None)
|
||||
config = CheckpointerConfig(type="sqlite", connection_string=":memory:")
|
||||
|
||||
results = await asyncio.gather(*(readiness_payload(config) for _ in range(8)))
|
||||
|
||||
assert [status_code for status_code, _ in results] == [200] * 8
|
||||
assert max_active == 1
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_probe_checkpointer_memory_reports_not_configured():
|
||||
assert await _probe_checkpointer_backend(CheckpointerConfig(type="memory")) == DATABASE_NOT_CONFIGURED
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_probe_checkpointer_sqlite_reachable(tmp_path):
|
||||
db_path = tmp_path / "checkpoints.db"
|
||||
_create_sqlite_file(db_path)
|
||||
|
||||
result = await _probe_checkpointer_backend(CheckpointerConfig(type="sqlite", connection_string=str(db_path)))
|
||||
|
||||
assert result == DATABASE_OK
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_probe_checkpointer_sqlite_file_uri_reachable(tmp_path):
|
||||
db_path = tmp_path / "checkpoints.db"
|
||||
_create_sqlite_file(db_path)
|
||||
|
||||
result = await _probe_checkpointer_backend(CheckpointerConfig(type="sqlite", connection_string=pathlib.Path(db_path).as_uri()))
|
||||
|
||||
assert result == DATABASE_OK
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_probe_checkpointer_sqlite_missing_file_stays_missing_and_unreachable(tmp_path):
|
||||
"""The probe must never create a missing SQLite file (regression)."""
|
||||
missing = tmp_path / "checkpoints.db"
|
||||
|
||||
result = await _probe_checkpointer_backend(CheckpointerConfig(type="sqlite", connection_string=str(missing)))
|
||||
|
||||
assert result == DATABASE_UNREACHABLE
|
||||
assert not missing.exists()
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_probe_checkpointer_sqlite_unreachable_when_parent_missing(tmp_path):
|
||||
missing_parent = tmp_path / "does-not-exist" / "checkpoints.db"
|
||||
|
||||
result = await _probe_checkpointer_backend(CheckpointerConfig(type="sqlite", connection_string=str(missing_parent)))
|
||||
|
||||
assert result == DATABASE_UNREACHABLE
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
@pytest.mark.parametrize(
|
||||
"conn_string",
|
||||
[":memory:", "file:memdb1?mode=memory&cache=shared", "file::memory:?cache=shared"],
|
||||
)
|
||||
async def test_probe_checkpointer_sqlite_in_memory_is_not_configured(conn_string):
|
||||
"""In-memory SQLite has no external state, mirroring the memory backend."""
|
||||
result = await _probe_checkpointer_backend(CheckpointerConfig(type="sqlite", connection_string=conn_string))
|
||||
|
||||
assert result == DATABASE_NOT_CONFIGURED
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_probe_checkpointer_postgres_without_psycopg_is_unreachable(monkeypatch):
|
||||
monkeypatch.setitem(sys.modules, "psycopg", None)
|
||||
|
||||
result = await _probe_checkpointer_backend(
|
||||
CheckpointerConfig(
|
||||
type="postgres",
|
||||
connection_string="postgresql://user:pass@localhost:5432/deerflow",
|
||||
)
|
||||
)
|
||||
|
||||
assert result == DATABASE_UNREACHABLE
|
||||
|
||||
|
||||
def test_resolve_checkpointer_config_passes_through_resolution(monkeypatch):
|
||||
resolved = CheckpointerConfig(type="memory")
|
||||
monkeypatch.setattr(
|
||||
"deerflow.runtime.checkpointer.provider._resolve_checkpointer_config",
|
||||
lambda app_config: resolved,
|
||||
)
|
||||
|
||||
assert resolve_checkpointer_config(object()) is resolved
|
||||
|
||||
|
||||
def test_resolve_checkpointer_config_failure_fails_closed(monkeypatch):
|
||||
"""A resolution failure must surface as None, never as a memory default."""
|
||||
|
||||
def _raise(app_config):
|
||||
raise RuntimeError("broken checkpointer config")
|
||||
|
||||
monkeypatch.setattr(
|
||||
"deerflow.runtime.checkpointer.provider._resolve_checkpointer_config",
|
||||
_raise,
|
||||
)
|
||||
|
||||
assert resolve_checkpointer_config(object()) is None
|
||||
@ -155,10 +155,14 @@ spec:
|
||||
name: http
|
||||
readinessProbe:
|
||||
httpGet:
|
||||
path: /health
|
||||
path: /health/ready
|
||||
port: http
|
||||
initialDelaySeconds: 5
|
||||
periodSeconds: 10
|
||||
# /health/ready bounds both probes to a 3s endpoint deadline;
|
||||
# timeoutSeconds must stay above that bound or Kubernetes aborts
|
||||
# slow-but-healthy responses at its 1s default.
|
||||
timeoutSeconds: 5
|
||||
livenessProbe:
|
||||
httpGet:
|
||||
path: /health
|
||||
|
||||
@ -150,7 +150,9 @@ services:
|
||||
redis:
|
||||
condition: service_healthy
|
||||
healthcheck:
|
||||
test: ["CMD", "python", "-c", "import urllib.request; response = urllib.request.urlopen('http://127.0.0.1:8001/health', timeout=3); raise SystemExit(0 if response.status == 200 else 1)"]
|
||||
# /health/ready runs both probes concurrently within a 3s endpoint
|
||||
# deadline; the client timeout must stay above that bound.
|
||||
test: ["CMD", "python", "-c", "import urllib.request; response = urllib.request.urlopen('http://127.0.0.1:8001/health/ready', timeout=5); raise SystemExit(0 if response.status == 200 else 1)"]
|
||||
interval: 5s
|
||||
timeout: 5s
|
||||
retries: 30
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user