diff --git a/backend/packages/harness/deerflow/config/database_config.py b/backend/packages/harness/deerflow/config/database_config.py index bcd869069..000ebe482 100644 --- a/backend/packages/harness/deerflow/config/database_config.py +++ b/backend/packages/harness/deerflow/config/database_config.py @@ -63,6 +63,16 @@ class DatabaseConfig(BaseModel): default=5, description="Connection pool size for the app ORM engine (postgres only).", ) + pool_recycle: int = Field( + default=300, + gt=0, + description="Seconds before app ORM PostgreSQL connections are recycled.", + ) + command_timeout: float | None = Field( + default=30, + gt=0, + description="Timeout in seconds for app ORM PostgreSQL commands. Set to null to disable the command timeout.", + ) # -- Derived helpers (not user-configured) -- diff --git a/backend/packages/harness/deerflow/persistence/engine.py b/backend/packages/harness/deerflow/persistence/engine.py index 0e12f57ef..e6ac0e17a 100644 --- a/backend/packages/harness/deerflow/persistence/engine.py +++ b/backend/packages/harness/deerflow/persistence/engine.py @@ -16,12 +16,38 @@ import logging from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, async_sessionmaker, create_async_engine +# Recycle pooled Postgres connections before stale idle sockets can hang +# pool_pre_ping. The command timeout bounds stalled ORM queries independently. +POSTGRES_POOL_RECYCLE_SECONDS = 300 +POSTGRES_COMMAND_TIMEOUT_SECONDS = 30 + def _json_serializer(obj: object) -> str: """JSON serializer with ensure_ascii=False for Chinese character support.""" return json.dumps(obj, ensure_ascii=False) +def _postgres_engine_kwargs( + *, + echo: bool, + pool_size: int, + pool_recycle: int = POSTGRES_POOL_RECYCLE_SECONDS, + command_timeout: float | None = POSTGRES_COMMAND_TIMEOUT_SECONDS, +) -> dict[str, object]: + """Build the shared SQLAlchemy engine options for PostgreSQL.""" + connect_args = {} + if command_timeout is not None: + connect_args["command_timeout"] = command_timeout + return { + "echo": echo, + "pool_size": pool_size, + "pool_pre_ping": True, + "pool_recycle": pool_recycle, + "connect_args": connect_args, + "json_serializer": _json_serializer, + } + + logger = logging.getLogger(__name__) _engine: AsyncEngine | None = None @@ -61,6 +87,8 @@ async def init_engine( url: str = "", echo: bool = False, pool_size: int = 5, + pool_recycle: int = POSTGRES_POOL_RECYCLE_SECONDS, + command_timeout: float | None = POSTGRES_COMMAND_TIMEOUT_SECONDS, sqlite_dir: str = "", ) -> None: """Create the async engine and session factory, then auto-create tables. @@ -70,6 +98,8 @@ async def init_engine( url: SQLAlchemy async URL (for sqlite/postgres). echo: Echo SQL to log. pool_size: Postgres connection pool size. + pool_recycle: Seconds before Postgres connections are recycled. + command_timeout: Timeout in seconds for app ORM Postgres commands, or None to disable. sqlite_dir: Directory to create for SQLite (ensured to exist). """ global _engine, _session_factory @@ -133,10 +163,12 @@ async def init_engine( elif backend == "postgres": _engine = create_async_engine( url, - echo=echo, - pool_size=pool_size, - pool_pre_ping=True, - json_serializer=_json_serializer, + **_postgres_engine_kwargs( + echo=echo, + pool_size=pool_size, + pool_recycle=pool_recycle, + command_timeout=command_timeout, + ), ) else: raise ValueError(f"Unknown persistence backend: {backend!r}") @@ -162,7 +194,15 @@ async def init_engine( await _auto_create_postgres_db(url) # Rebuild engine against the now-existing database await _engine.dispose() - _engine = create_async_engine(url, echo=echo, pool_size=pool_size, pool_pre_ping=True, json_serializer=_json_serializer) + _engine = create_async_engine( + url, + **_postgres_engine_kwargs( + echo=echo, + pool_size=pool_size, + pool_recycle=pool_recycle, + command_timeout=command_timeout, + ), + ) _session_factory = async_sessionmaker(_engine, expire_on_commit=False) await bootstrap_schema(_engine, backend=backend) else: @@ -181,6 +221,8 @@ async def init_engine_from_config(config) -> None: url=config.app_sqlalchemy_url, echo=config.echo_sql, pool_size=config.pool_size, + pool_recycle=config.pool_recycle, + command_timeout=config.command_timeout, sqlite_dir=config.sqlite_dir if config.backend == "sqlite" else "", ) diff --git a/backend/tests/test_persistence_engine_postgres_config.py b/backend/tests/test_persistence_engine_postgres_config.py new file mode 100644 index 000000000..843919db0 --- /dev/null +++ b/backend/tests/test_persistence_engine_postgres_config.py @@ -0,0 +1,225 @@ +"""Tests for hardened PostgreSQL async engine configuration.""" + +from __future__ import annotations + +import asyncio +import sys +from time import monotonic +from types import ModuleType +from unittest.mock import AsyncMock, MagicMock, call, patch + +import pytest + +from deerflow.config.database_config import DatabaseConfig +from deerflow.persistence import engine as engine_mod + + +def test_postgres_engine_kwargs_include_connection_hardening() -> None: + kwargs = engine_mod._postgres_engine_kwargs(echo=False, pool_size=5) + + assert kwargs["echo"] is False + assert kwargs["pool_size"] == 5 + assert kwargs["pool_pre_ping"] is True + assert kwargs["pool_recycle"] == engine_mod.POSTGRES_POOL_RECYCLE_SECONDS + assert kwargs["connect_args"]["command_timeout"] == engine_mod.POSTGRES_COMMAND_TIMEOUT_SECONDS + assert kwargs["json_serializer"] is engine_mod._json_serializer + + +def test_database_command_timeout_defaults_to_30_seconds() -> None: + config = DatabaseConfig() + + assert config.command_timeout == 30 + + +def test_database_pool_recycle_defaults_to_300_seconds() -> None: + config = DatabaseConfig() + + assert config.pool_recycle == 300 + + +def test_postgres_engine_kwargs_preserve_caller_values() -> None: + kwargs = engine_mod._postgres_engine_kwargs(echo=True, pool_size=20, pool_recycle=120, command_timeout=90) + + assert kwargs["echo"] is True + assert kwargs["pool_size"] == 20 + assert kwargs["pool_recycle"] == 120 + assert kwargs["connect_args"] == {"command_timeout": 90} + + +def test_postgres_engine_kwargs_allow_command_timeout_opt_out() -> None: + config = DatabaseConfig(command_timeout=None) + kwargs = engine_mod._postgres_engine_kwargs(echo=False, pool_size=5, command_timeout=config.command_timeout) + + assert config.command_timeout is None + assert kwargs["connect_args"] == {} + + +@pytest.mark.asyncio +async def test_configured_command_timeout_ends_stalled_command() -> None: + config = DatabaseConfig( + backend="postgres", + postgres_url="postgresql://user:password@localhost/deerflow", + command_timeout=0.01, + ) + + class _StalledAsyncpgEngine: + def __init__(self, command_timeout: float) -> None: + self.command_timeout = command_timeout + + async def checkout(self) -> None: + async with asyncio.timeout(self.command_timeout): + await asyncio.Event().wait() + + async def dispose(self) -> None: + return None + + def _create_engine(_url: str, **kwargs) -> _StalledAsyncpgEngine: + assert kwargs["pool_pre_ping"] is True + return _StalledAsyncpgEngine(kwargs["connect_args"]["command_timeout"]) + + bootstrap_schema = AsyncMock() + + with ( + patch.dict(sys.modules, {"asyncpg": ModuleType("asyncpg")}), + patch.object(engine_mod, "create_async_engine", side_effect=_create_engine), + patch.object(engine_mod, "async_sessionmaker", return_value=MagicMock()), + patch("deerflow.persistence.bootstrap.bootstrap_schema", new=bootstrap_schema), + ): + try: + await engine_mod.init_engine_from_config(config) + engine = engine_mod.get_engine() + assert isinstance(engine, _StalledAsyncpgEngine) + started_at = monotonic() + with pytest.raises(TimeoutError): + await engine.checkout() + elapsed = monotonic() - started_at + + assert engine.command_timeout == config.command_timeout + assert elapsed < 1 + finally: + await engine_mod.close_engine() + + +@pytest.mark.asyncio +async def test_init_engine_from_config_preserves_longer_command_timeout_override() -> None: + config = DatabaseConfig( + backend="postgres", + postgres_url="postgresql://user:password@localhost/deerflow", + pool_recycle=120, + command_timeout=90, + ) + mock_engine = MagicMock() + mock_engine.dispose = AsyncMock() + bootstrap_schema = AsyncMock() + + with ( + patch.dict(sys.modules, {"asyncpg": ModuleType("asyncpg")}), + patch.object(engine_mod, "create_async_engine", return_value=mock_engine) as create_engine, + patch.object(engine_mod, "async_sessionmaker", return_value=MagicMock()), + patch("deerflow.persistence.bootstrap.bootstrap_schema", new=bootstrap_schema), + ): + try: + await engine_mod.init_engine_from_config(config) + + kwargs = create_engine.call_args.kwargs + assert kwargs["connect_args"]["command_timeout"] == 90 + assert kwargs["pool_recycle"] == 120 + finally: + await engine_mod.close_engine() + + +@pytest.mark.asyncio +async def test_init_engine_postgres_uses_hardened_kwargs() -> None: + url = "postgresql+asyncpg://user:password@localhost/deerflow" + mock_engine = MagicMock() + mock_engine.dispose = AsyncMock() + bootstrap_schema = AsyncMock() + + with ( + patch.dict(sys.modules, {"asyncpg": ModuleType("asyncpg")}), + patch.object(engine_mod, "create_async_engine", return_value=mock_engine) as create_engine, + patch.object(engine_mod, "async_sessionmaker", return_value=MagicMock()), + patch("deerflow.persistence.bootstrap.bootstrap_schema", new=bootstrap_schema), + ): + try: + await engine_mod.init_engine(backend="postgres", url=url, echo=True, pool_size=12) + + create_engine.assert_called_once_with(url, **engine_mod._postgres_engine_kwargs(echo=True, pool_size=12)) + bootstrap_schema.assert_awaited_once_with(mock_engine, backend="postgres") + finally: + await engine_mod.close_engine() + + +@pytest.mark.asyncio +async def test_init_engine_postgres_retry_uses_hardened_kwargs() -> None: + url = "postgresql+asyncpg://user:password@localhost/deerflow" + initial_engine = MagicMock() + initial_engine.dispose = AsyncMock() + retry_engine = MagicMock() + retry_engine.dispose = AsyncMock() + bootstrap_schema = AsyncMock(side_effect=[Exception("database does not exist"), None]) + auto_create = AsyncMock() + + with ( + patch.dict(sys.modules, {"asyncpg": ModuleType("asyncpg")}), + patch.object(engine_mod, "create_async_engine", side_effect=[initial_engine, retry_engine]) as create_engine, + patch.object(engine_mod, "async_sessionmaker", return_value=MagicMock()), + patch.object(engine_mod, "_auto_create_postgres_db", new=auto_create), + patch("deerflow.persistence.bootstrap.bootstrap_schema", new=bootstrap_schema), + ): + try: + await engine_mod.init_engine(backend="postgres", url=url, echo=False, pool_size=8) + + kwargs = engine_mod._postgres_engine_kwargs(echo=False, pool_size=8) + assert create_engine.call_args_list == [call(url, **kwargs), call(url, **kwargs)] + auto_create.assert_awaited_once_with(url) + initial_engine.dispose.assert_awaited_once() + assert bootstrap_schema.await_args_list == [call(initial_engine, backend="postgres"), call(retry_engine, backend="postgres")] + finally: + await engine_mod.close_engine() + + +@pytest.mark.asyncio +async def test_init_engine_sqlite_omits_postgres_kwargs_and_keeps_wal_listener(tmp_path) -> None: + url = f"sqlite+aiosqlite:///{tmp_path / 'deerflow.db'}" + mock_engine = MagicMock() + mock_engine.sync_engine = object() + mock_engine.dispose = AsyncMock() + bootstrap_schema = AsyncMock() + registered: dict[str, object] = {} + + def _capture_listener(target, event_name): + assert target is mock_engine.sync_engine + assert event_name == "connect" + + def _decorator(fn): + registered["listener"] = fn + return fn + + return _decorator + + with ( + patch.object(engine_mod, "create_async_engine", return_value=mock_engine) as create_engine, + patch.object(engine_mod, "async_sessionmaker", return_value=MagicMock()), + patch("sqlalchemy.event.listens_for", new=_capture_listener), + patch("deerflow.persistence.bootstrap.bootstrap_schema", new=bootstrap_schema), + ): + try: + await engine_mod.init_engine(backend="sqlite", url=url, echo=True, sqlite_dir=str(tmp_path)) + + create_engine.assert_called_once_with(url, echo=True, json_serializer=engine_mod._json_serializer) + + cursor = MagicMock() + dbapi_connection = MagicMock() + dbapi_connection.cursor.return_value = cursor + listener = registered["listener"] + listener(dbapi_connection, None) + assert [entry.args[0] for entry in cursor.execute.call_args_list] == [ + "PRAGMA journal_mode=WAL;", + "PRAGMA synchronous=NORMAL;", + "PRAGMA foreign_keys=ON;", + "PRAGMA busy_timeout=30000;", + ] + cursor.close.assert_called_once_with() + finally: + await engine_mod.close_engine() diff --git a/config.example.yaml b/config.example.yaml index aa706921c..23b8a7a93 100644 --- a/config.example.yaml +++ b/config.example.yaml @@ -1662,6 +1662,11 @@ skill_evolution: database: backend: sqlite sqlite_dir: .deer-flow/data + # Recycle app ORM PostgreSQL connections before the environment's idle cutoff. + pool_recycle: 300 + # App ORM PostgreSQL command timeout. Set to null to disable it or raise it + # for intentionally long commands. + command_timeout: 30 # ============================================================================ # Run Events Configuration diff --git a/deploy/helm/deer-flow/README.md b/deploy/helm/deer-flow/README.md index 4405af16a..2fe040b76 100644 --- a/deploy/helm/deer-flow/README.md +++ b/deploy/helm/deer-flow/README.md @@ -137,6 +137,8 @@ config: | database: backend: postgres postgres_url: $DATABASE_URL + pool_recycle: 300 + command_timeout: 30 checkpointer: type: postgres connection_string: $DATABASE_URL diff --git a/deploy/helm/deer-flow/values.yaml b/deploy/helm/deer-flow/values.yaml index 47b988cd2..aab205b34 100644 --- a/deploy/helm/deer-flow/values.yaml +++ b/deploy/helm/deer-flow/values.yaml @@ -265,6 +265,8 @@ config: | # resolved by the harness before the config is instantiated. backend: postgres postgres_url: $DATABASE_URL + pool_recycle: 300 + command_timeout: 30 # The LangGraph Store (cross-thread memory + thread list) reads this legacy # `checkpointer:` section — it does NOT fall back to `database:` the way the