RFC: Deepen Agent Lifecycle & Persistence Hub

[parthashirolkar]

Problem

The agent lifecycle and persistence initialization in apps/backend/src/rerai_agent/graph.py is shallow, tightly coupled, and creates cascading import-time side effects that make the codebase hard to test and unsafe under async frameworks.

Shallow modules & tight coupling

• graph.py executes graph = build_persisted_graph() at import time, eagerly opening Postgres/SQLite connections via ExitStack.
• tools/config.py calls load_project_env() at import time and immediately reads os.environ["OPENROUTER_API_KEY"].
• Six pieces of global mutable state (_checkpointer, _store, _persistence_stack, _async_persistence_stack, _async_checkpointer, _async_store, _udcpr_store_initialized) leak across imports and make parallel tests impossible.
• The atexit handler _close_persistence() calls asyncio.run(_async_persistence_stack.aclose()), which crashes under pytest-asyncio and Uvicorn because an event loop is already running.
• runtime.py calls build_persisted_graph_async() during FastAPI lifespan, but a sync version was already built at import time — so two persistence stacks are opened, but only one is managed.

Integration risk

Any change to persistence (e.g., adding Redis, switching to async-only SQLite) requires touching graph.py, runtime.py, and every test that imports the module. The seam between graph construction and persistence connection is non-existent — they are the same function.

Why this matters

• pytest collection fails if env vars are missing.
• Tests cannot run in parallel because globals collide.
• The atexit crash is a production reliability bug waiting to happen.
• New developers cannot reason about when DB connections are opened or closed.

Proposed Interface

Introduce a single deep module — AgentHub — that hides all graph construction, persistence lifecycle, env loading, and UDCPR vector-store initialization behind a small, explicit async interface.

Interface signature

from __future__ import annotations

from collections.abc import Sequence
from contextlib import AbstractAsyncContextManager
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Callable

from langchain_core.embeddings import Embeddings
from langchain_core.language_models import BaseChatModel
from langgraph.checkpoint.base import BaseCheckpointSaver
from langgraph.store.base import BaseStore
from pydantic import SecretStr


class CompiledGraph(Protocol):
    async def astream(self, input: Any, config: dict[str, Any], **kwargs: Any) -> Any: ...
    async def ainvoke(self, input: Any, config: Any | None = None, **kwargs: Any) -> Any: ...
    def get_input_jsonschema(self) -> dict[str, Any] | None: ...
    def get_output_jsonschema(self) -> dict[str, Any] | None: ...
    def config_schema(self) -> Any: ...
    def get_context_jsonschema(self) -> dict[str, Any] | None: ...


@dataclass(frozen=True, slots=True)
class AgentHubConfig:
    database_uri: str | None = None
    openrouter_api_key: SecretStr | None = None
    chat_model: str = "nvidia/nemotron-3-super-120b-a12b:free"
    subagent_model: str = "nvidia/nemotron-3-nano-30b-a3b:free"
    embedding_model: str = "nvidia/llama-nemotron-embed-vl-1b-v2:free"
    openrouter_base_url: str = "https://openrouter.ai/api/v1"
    memory_files: Sequence[str | Path] = field(default_factory=_default_memory_files)
    skills_dir: str | Path | None = None
    system_prompt: str = SYSTEM_PROMPT
    tools: Sequence[Callable] = field(default_factory=_default_tools)
    subagents: Sequence[Any] = field(default_factory=_default_subagents)
    backend_factory: Callable[[Any], Any] | None = None
    interrupt_on: Sequence[str] | None = None
    setup_db: bool = True
    init_udcpr_store: bool = True

    @classmethod
    def from_env(cls, *, env_loader: Callable[[], None] | None = None) -> AgentHubConfig:
        """Reads .env and os.environ. THE ONLY PLACE ENV IS TOUCHED."""
        ...

    @classmethod
    def for_testing(
        cls,
        *,
        database_uri: str = "sqlite://:memory:",
        skip_udcpr: bool = True,
    ) -> AgentHubConfig:
        """Sensible defaults for tests. No env read."""
        ...


class AgentHub(AbstractAsyncContextManager):
    graph: CompiledGraph
    config: AgentHubConfig
    checkpointer: BaseCheckpointSaver | None
    store: BaseStore | None

    @classmethod
    async def production(
        cls,
        *,
        database_uri: str | None = None,
        config: AgentHubConfig | None = None,
        backend_factory: Callable[[Any], Any] | None = None,
        interrupt_on: Sequence[str] | None = None,
    ) -> AgentHub:
        """One-liner for production:
        1. Loads env (if config not provided).
        2. Builds real LLM/embeddings.
        3. Connects to Postgres/SQLite.
        4. Initializes UDCPR store.
        5. Returns a ready-to-run hub.
        """
        ...

    @classmethod
    async def testing(
        cls,
        *,
        database_uri: str = "sqlite://:memory:",
        fake_llm: BaseChatModel | None = None,
        fake_embeddings: Embeddings | None = None,
        skip_udcpr: bool = True,
        config: AgentHubConfig | None = None,
    ) -> AgentHub:
        """One-liner for tests:
        1. Uses in-memory SQLite (or no persistence if URI is empty).
        2. Injects fake LLM/embeddings when provided.
        3. Skips UDCPR by default.
        4. Fast setup / teardown.
        """
        ...

    @classmethod
    def build(
        cls,
        *,
        config: AgentHubConfig,
        llm: BaseChatModel | None = None,
        embeddings: Embeddings | None = None,
        checkpointer: BaseCheckpointSaver | None = None,
        store: BaseStore | None = None,
    ) -> AgentHub:
        """Inject any dependency. Returns an UN-INITIALIZED hub.
        Caller MUST await hub.setup(). No env is read.
        """
        ...

    async def setup(self) -> None:
        """Idempotent.
        - Resolves missing deps (LLM, checkpointer, store) from config.
        - Builds the deepagents graph.
        - Runs DB setup() if configured.
        - Initializes UDCPR store once per instance.
        """
        ...

    async def close(self) -> None:
        """Idempotent. Closes checkpointer, store, and any connection pools.
        NEVER calls asyncio.run(). Safe under pytest-asyncio and Uvicorn.
        """
        ...

    async def __aenter__(self) -> AgentHub:
        await self.setup()
        return self

    async def __aexit__(self, exc_type: Any, exc: Any, tb: Any) -> None:
        await self.close()

Usage example — production

# rerai_api/runtime.py
class BackendRuntime:
    async def setup(self) -> None:
        self.hub = await AgentHub.production(database_uri=self.database_uri)
        self.graph = self.hub.graph
        self.run_manager = RunManager(self.metadata_store, self.hub.graph)

Usage example — test

# tests/conftest.py
import pytest
from rerai_agent.hub import AgentHub

@pytest.fixture
async def agent_hub():
    async with AgentHub.testing(fake_llm=FakeChatModel()) as hub:
        yield hub

async def test_permit_query(agent_hub: AgentHub):
    result = await agent_hub.graph.ainvoke({"query": "Gat 123, Haveli"})
    assert "permit" in result

Usage example — advanced caller

from dataclasses import replace

config = AgentHubConfig.from_env()
config = replace(config, tools=[my_mock_gis_tool, *config.tools])

hub = AgentHub.build(config=config, llm=my_custom_llm)
await hub.setup()

What complexity it hides

Hidden concern	Previously exposed	Now hidden
Env loading	config.py called load_project_env() at import	Only AgentHubConfig.from_env() reads env; happens inside production(), never at import
Global mutable state	_checkpointer, _store, _persistence_stack, etc.	Private instance attributes on AgentHub
Persistence detection	String parsing for postgresql:// vs sqlite://	URI-to-adapter mapping inside setup()
Sync vs async DB	Two sets of globals: _checkpointer / _async_checkpointer	Single internal adapter; setup() is always async
DB schema setup	store.setup(); checkpointer.setup() scattered in callers	Executed once inside setup() if config.setup_db
UDCPR vector store	Global _udcpr_store_initialized boolean	Instance-level flag; executed once inside setup()
atexit crash	atexit.register(_close_persistence) calling asyncio.run()	Gone. Explicit await hub.close() or context manager
LLM / embeddings construction	get_chat_model(), get_embeddings() imported from tool modules	Built from AgentHubConfig unless injected
Tool registry assembly	ALL_TOOLS and ALL_SUBAGENTS imported and passed manually	Default lists baked into config factories
Graph compilation	create_deep_agent(...) with 8+ args	Single internal call driven by config

Dependency Strategy

Concern	Category	How it is handled
Env / config reading	Local-substitutable	AgentHubConfig is the boundary. from_env() is the only adapter that touches os.environ / .env files. Tests bypass env entirely via for_testing() or direct dataclass construction.
DB savers / stores	Ports & Adapters (LangGraph built-ins)	BaseCheckpointSaver and BaseStore are the ports. Default adapter resolves Postgres vs SQLite async savers/stores, runs .setup(), and manages an internal AsyncExitStack. Tests inject None for in-memory SQLite or pass a mock BaseCheckpointSaver via AgentHub.build().
LLM / embeddings	Local-substitutable	Default factories wrap OpenRouter (get_chat_model, get_embeddings). Tests pass fake_llm=FakeChatModel() via AgentHub.testing().
Tool registry	In-process	Defaults are the current ALL_TOOLS / ALL_SUBAGENTS lists, resolved lazily inside setup(). Tests override with empty or mocked lists via AgentHubConfig replacement.
UDCPR vector store	True external (Mock boundary)	Chroma Cloud is external. The hub treats init_udcpr_store() as an injected initializer. Tests pass skip_udcpr=True or inject a no-op.

Testing Strategy

New boundary tests to write

1. Lifecycle boundary tests (test_hub_lifecycle.py)

   • AgentHub.production() initializes successfully with a real SQLite file.
   • AgentHub.testing() initializes successfully with :memory: SQLite.
   • AgentHub.build() + fake LLM + no persistence initializes without env vars.
   • setup() is idempotent: calling twice does not open duplicate connections.
   • close() is idempotent: calling twice does not raise.
   • close() does NOT call asyncio.run() — assert no crash under pytest-asyncio.
   • Context-manager (async with) correctly sets up and tears down.

2. Graph construction boundary tests (test_hub_graph.py)

   • hub.graph is not None after setup().
   • hub.graph.astream() accepts input and yields chunks (using a fake LLM).
   • Custom tool injection via AgentHubConfig replacement works end-to-end.
   • Custom subagent injection works end-to-end.

3. Config boundary tests (test_hub_config.py)

   • AgentHubConfig.from_env() reads the expected keys from .env.
   • AgentHubConfig.for_testing() produces a config with no env access.
   • Missing required env keys raise a clear error at from_env() time, not at import.

Old tests to delete

• Any tests that monkeypatch rerai_agent.graph._checkpointer or _store globals.
• Any tests that rely on from rerai_agent.graph import graph being pre-built.
• The atexit-related workaround tests (if any exist).

Test environment needs

• pytest-asyncio (already in use).
• A FakeChatModel and FakeEmbeddings fixture (already partially present in test_fastapi_backend.py — extract and reuse).
• Temporary SQLite file paths for lifecycle tests (use tmp_path fixture).

Implementation Recommendations

What the module should own

• Graph construction: create_deep_agent() call with tools, subagents, memory, skills, system prompt, backend, interrupt_on.
• Persistence lifecycle: URI parsing, async saver/store selection, AsyncExitStack management, .setup() calls.
• Env loading: .env discovery and os.environ reading — once, inside AgentHubConfig.from_env().
• UDCPR vector store initialization: Lazy, once-per-instance, no globals.
• Clean teardown: Explicit close() that never calls asyncio.run().

What it should hide

• ExitStack / AsyncExitStack details.
• Postgres vs SQLite driver selection.
• deepagents internal parameter shapes.
• OpenRouterEmbeddings encoding_format workaround.
• CHAT_MODEL, SUBAGENT_MODEL, EMBEDDING_MODEL hardcoded strings (live inside AgentHubConfig defaults).
• ALL_TOOLS, ALL_SUBAGENTS, SYSTEM_PROMPT, MEMORY_FILE, SKILLS_DIR module-level constants (become private defaults).

What it should expose

• AgentHub.production() — the default production entry point.
• AgentHub.testing() — the default test entry point.
• AgentHub.build() — the full-injection escape hatch.
• AgentHubConfig.from_env() — the only env-touching code path.
• AgentHubConfig.for_testing() — the only test-default config path.
• AgentHub.graph — the compiled graph, ready for RunManager.
• AgentHub.close() — safe async teardown.

How callers should migrate

1. Replace import-time graph usage:

   # OLD
   from rerai_agent.graph import graph
   # NEW
   hub = await AgentHub.production()
   graph = hub.graph

2. Replace BackendRuntime.setup():

   # OLD
   from rerai_agent.graph import build_persisted_graph_async
   self.graph = await build_persisted_graph_async(database_uri=self.database_uri)
   # NEW
   self.hub = await AgentHub.production(database_uri=self.database_uri)
   self.graph = self.hub.graph

3. Replace root app.py / agent.py shims:
   These files currently exist solely to load .env before importing modules that touch env at import time. After the refactor, .env loading moves into AgentHubConfig.from_env(), which is called inside the FastAPI lifespan. The root shims can be deleted or reduced to a single re-export.

4. Replace atexit cleanup:
   Remove atexit.register(_close_persistence) entirely. FastAPI lifespan and test fixtures become the only teardown paths.

5. Update tests:
   Replace FakeGraph patterns with AgentHub.testing(fake_llm=FakeChatModel()). Tests no longer need to avoid importing rerai_agent.graph.

[parthashirolkar]

Fixed in the backend-refactor branch (merged).\n\n- Removed all global mutable state (checkpointer, store, persistence stacks, UDCPR init flag).\n- Introduced AgentHub with explicit async lifecycle (setup/close, async context manager).\n- Production, testing, and full-DI build paths.\n- Removed atexit handler and the asyncio.run() crash.\n- Lazy env loading so pytest collection works without env vars.