feat(states): add storage-phase config and task-state repository selector

Introduce TASK_STATE_STORAGE_PHASE and TASK_MESSAGE_STORAGE_PHASE settings
(mongodb|postgres, default mongodb) with a derived mongodb_required
property, a TaskStateRepositoryProtocol capturing the contract callers
rely on through the DTaskStateRepository seam, and a
get_task_state_repository() selector that resolves the backend by phase.

The DI seam and both Temporal factories (task retention, scheduled agent
runs) now construct the repository through the selector, so a phase switch
applies to API handlers and workers alike. The postgres phase raises
NotImplementedError until the Postgres repository lands, and the mongodb
default keeps existing deployments unchanged.

Author: vineetvora-scale

agentex/docker-compose.yml

@@ -164,6 +164,9 @@ services:
       - REDIS_URL=redis://agentex-redis:6379
       - MONGODB_URI=mongodb://agentex-mongodb:27017
       - MONGODB_DATABASE_NAME=agentex
+      # Storage backend per document store (mongodb|postgres), e.g. `TASK_STATE_STORAGE_PHASE=postgres ./dev.sh`.
+      - TASK_STATE_STORAGE_PHASE=${TASK_STATE_STORAGE_PHASE:-mongodb}
+      - TASK_MESSAGE_STORAGE_PHASE=${TASK_MESSAGE_STORAGE_PHASE:-mongodb}
       - WATCHFILES_FORCE_POLLING=true
       - ENABLE_HEALTH_CHECK_WORKFLOW=true
       # Disabled by default; enable when testing, e.g. `ENABLE_AGENT_RUN_SCHEDULES=true ./dev.sh`.
@@ -238,6 +241,8 @@ services:
       - REDIS_URL=redis://agentex-redis:6379
       - MONGODB_URI=mongodb://agentex-mongodb:27017
       - MONGODB_DATABASE_NAME=agentex
+      - TASK_STATE_STORAGE_PHASE=${TASK_STATE_STORAGE_PHASE:-mongodb}
+      - TASK_MESSAGE_STORAGE_PHASE=${TASK_MESSAGE_STORAGE_PHASE:-mongodb}
       - AGENTEX_SERVER_TASK_QUEUE=agentex-server
       - RETENTION_CLEANUP_ENABLED=${RETENTION_CLEANUP_ENABLED:-false}
       - RETENTION_CLEANUP_AGENT_ALLOWLIST=${RETENTION_CLEANUP_AGENT_ALLOWLIST:-}

agentex/src/config/environment_variables.py

@@ -67,6 +67,8 @@ class EnvVarKeys(str, Enum):
     RETENTION_CLEANUP_MAX_IN_FLIGHT = "RETENTION_CLEANUP_MAX_IN_FLIGHT"
     RETENTION_CLEANUP_DRY_RUN = "RETENTION_CLEANUP_DRY_RUN"
     RETENTION_CLEANUP_STALE_RUNNING_DAYS = "RETENTION_CLEANUP_STALE_RUNNING_DAYS"
+    TASK_STATE_STORAGE_PHASE = "TASK_STATE_STORAGE_PHASE"
+    TASK_MESSAGE_STORAGE_PHASE = "TASK_MESSAGE_STORAGE_PHASE"
 
 
 class Environment(str, Enum):
@@ -75,6 +77,17 @@ class Environment(str, Enum):
     PROD = "production"
 
 
+class StoragePhase(str, Enum):
+    """Which backend serves a document store (task state, task messages).
+
+    A future data-migration effort would define its own additional phases;
+    new values are additive, not breaking.
+    """
+
+    MONGODB = "mongodb"
+    POSTGRES = "postgres"
+
+
 refreshed_environment_variables = None
 
 
@@ -166,6 +179,19 @@ class EnvironmentVariables(BaseModel):
     # are treated as abandoned and become eligible for cleanup. 0 disables the
     # override (RUNNING tasks are never cleaned), preserving prior behavior.
     RETENTION_CLEANUP_STALE_RUNNING_DAYS: int = 0
+    # Storage backend per document store. The mongodb default keeps existing
+    # deployments unchanged; postgres serves that store from the relational
+    # database instead.
+    TASK_STATE_STORAGE_PHASE: StoragePhase = StoragePhase.MONGODB
+    TASK_MESSAGE_STORAGE_PHASE: StoragePhase = StoragePhase.MONGODB
+
+    @property
+    def mongodb_required(self) -> bool:
+        """True while any document store still needs a MongoDB connection."""
+        return not (
+            self.TASK_STATE_STORAGE_PHASE == StoragePhase.POSTGRES
+            and self.TASK_MESSAGE_STORAGE_PHASE == StoragePhase.POSTGRES
+        )
 
     @classmethod
     def refresh(cls, force_refresh: bool = False) -> EnvironmentVariables | None:
@@ -285,6 +311,12 @@ def refresh(cls, force_refresh: bool = False) -> EnvironmentVariables | None:
             RETENTION_CLEANUP_STALE_RUNNING_DAYS=int(
                 os.environ.get(EnvVarKeys.RETENTION_CLEANUP_STALE_RUNNING_DAYS, "0")
             ),
+            TASK_STATE_STORAGE_PHASE=os.environ.get(
+                EnvVarKeys.TASK_STATE_STORAGE_PHASE, StoragePhase.MONGODB
+            ),
+            TASK_MESSAGE_STORAGE_PHASE=os.environ.get(
+                EnvVarKeys.TASK_MESSAGE_STORAGE_PHASE, StoragePhase.MONGODB
+            ),
         )
         refreshed_environment_variables = environment_variables
         return refreshed_environment_variables

agentex/src/domain/repositories/task_state_repository.py

@@ -1,15 +1,73 @@
-from typing import Annotated
+import builtins
+from typing import Annotated, Any, Protocol
 
 import pymongo
 from fastapi import Depends
 from src.adapters.crud_store.adapter_mongodb import MongoDBCRUDRepository
-from src.config.dependencies import DMongoDBDatabase
+from src.config.dependencies import DMongoDBDatabase, GlobalDependencies
+from src.config.environment_variables import EnvironmentVariables, StoragePhase
 from src.domain.entities.states import StateEntity
 from src.utils.logging import make_logger
 
 logger = make_logger(__name__)
 
 
+class TaskStateRepositoryProtocol(Protocol):
+    """Contract every task-state storage backend must satisfy.
+
+    Covers everything callers actually invoke through the DTaskStateRepository
+    seam: the states use case and authorization shortcuts (create / get /
+    update / delete / list), the retention service (find_by_field /
+    delete_by_field / batch_create), and get_by_task_and_agent.
+
+    Behavioral requirements beyond the signatures: `.id` is presented as a
+    string; `create` honors caller-supplied created_at/updated_at and only
+    falls back to server time when absent; missing rows raise ItemDoesNotExist
+    and duplicates raise DuplicateItemError; and `list` accepts the
+    Mongo-shaped filter dict the states use case builds (plain equality plus
+    `{"$in": [...]}` for the authorized-task allow-list).
+    """
+
+    async def create(self, item: StateEntity) -> StateEntity: ...
+
+    async def batch_create(
+        self, items: builtins.list[StateEntity]
+    ) -> builtins.list[StateEntity]: ...
+
+    async def get(
+        self, id: str | None = None, name: str | None = None
+    ) -> StateEntity | None: ...
+
+    async def update(self, item: StateEntity) -> StateEntity: ...
+
+    async def delete(self, id: str | None = None, name: str | None = None) -> None: ...
+
+    async def list(
+        self,
+        filters: dict[str, Any] | None = None,
+        limit: int | None = None,
+        page_number: int | None = None,
+        order_by: str | None = None,
+        order_direction: str | None = None,
+    ) -> builtins.list[StateEntity]: ...
+
+    async def find_by_field(
+        self,
+        field_name: str,
+        field_value: Any,
+        limit: int | None = None,
+        page_number: int | None = None,
+        sort_by: dict[str, int] | None = None,
+        filters: dict[str, Any] | None = None,
+    ) -> builtins.list[StateEntity]: ...
+
+    async def delete_by_field(self, field_name: str, field_value: Any) -> int: ...
+
+    async def get_by_task_and_agent(
+        self, task_id: str, agent_id: str
+    ) -> StateEntity | None: ...
+
+
 class TaskStateRepository(MongoDBCRUDRepository[StateEntity]):
     """Repository for managing task states in MongoDB."""
 
@@ -47,4 +105,27 @@ async def get_by_task_and_agent(
         return self._deserialize(doc) if doc else None
 
 
-DTaskStateRepository = Annotated[TaskStateRepository, Depends(TaskStateRepository)]
+def get_task_state_repository() -> TaskStateRepositoryProtocol:
+    """Select the task-state repository for the configured storage phase.
+
+    This is the single construction point for task-state repositories: the
+    FastAPI seam below resolves through it, and so must every site that builds
+    the repository by hand outside Depends (the Temporal factories), so that a
+    phase switch applies to request handlers and workers alike.
+
+    Each branch constructs its repository lazily — the Postgres phase must
+    never touch a Mongo handle, which is what allows the MongoDB connection to
+    be absent entirely once no store needs it.
+    """
+    phase = EnvironmentVariables.refresh().TASK_STATE_STORAGE_PHASE
+    if phase == StoragePhase.MONGODB:
+        return TaskStateRepository(GlobalDependencies().mongodb_database)
+    raise NotImplementedError(
+        f"TASK_STATE_STORAGE_PHASE={phase.value!r} is not implemented yet; "
+        "only 'mongodb' is currently supported."
+    )
+
+
+DTaskStateRepository = Annotated[
+    TaskStateRepositoryProtocol, Depends(get_task_state_repository)
+]

agentex/src/temporal/scheduled_agent_run_factory.py

@@ -37,7 +37,7 @@
 from src.domain.repositories.event_repository import EventRepository
 from src.domain.repositories.task_message_repository import TaskMessageRepository
 from src.domain.repositories.task_repository import TaskRepository
-from src.domain.repositories.task_state_repository import TaskStateRepository
+from src.domain.repositories.task_state_repository import get_task_state_repository
 from src.domain.services.agent_acp_service import AgentACPService
 from src.domain.services.authorization_service import AuthorizationService
 from src.domain.services.task_message_service import TaskMessageService
@@ -98,7 +98,9 @@ def build_acp_use_case_for_principal(
     task_repository = TaskRepository(rw_session_maker, ro_session_maker)
     event_repository = EventRepository(rw_session_maker, ro_session_maker)
 
-    task_state_repository = TaskStateRepository(global_dependencies.mongodb_database)
+    # Resolved through the shared selector (not constructed by hand) so the
+    # worker honors TASK_STATE_STORAGE_PHASE the same way the API does.
+    task_state_repository = get_task_state_repository()
     task_message_repository = TaskMessageRepository(
         global_dependencies.mongodb_database
     )

agentex/src/temporal/task_retention_factory.py

@@ -18,7 +18,7 @@
 from src.domain.repositories.event_repository import EventRepository
 from src.domain.repositories.task_message_repository import TaskMessageRepository
 from src.domain.repositories.task_repository import TaskRepository
-from src.domain.repositories.task_state_repository import TaskStateRepository
+from src.domain.repositories.task_state_repository import get_task_state_repository
 from src.domain.services.task_message_service import TaskMessageService
 from src.domain.services.task_retention_service import TaskRetentionService
 from src.domain.use_cases.task_retention_use_case import TaskRetentionUseCase
@@ -41,7 +41,9 @@ def build_task_retention_use_case(
     task_message_repository = TaskMessageRepository(
         global_dependencies.mongodb_database
     )
-    task_state_repository = TaskStateRepository(global_dependencies.mongodb_database)
+    # Resolved through the shared selector (not constructed by hand) so the
+    # worker honors TASK_STATE_STORAGE_PHASE the same way the API does.
+    task_state_repository = get_task_state_repository()
     task_message_service = TaskMessageService(
         message_repository=task_message_repository
     )

agentex/tests/unit/config/test_storage_phase_env.py

@@ -0,0 +1,68 @@
+import pytest
+from pydantic import ValidationError
+from src.config.environment_variables import EnvironmentVariables, StoragePhase
+
+
+@pytest.fixture(autouse=True)
+def _reset_env_cache():
+    """Drop the forced-refresh cache so later tests re-read a clean environment."""
+    yield
+    EnvironmentVariables.clear_cache()
+
+
+@pytest.mark.unit
+def test_storage_phases_default_to_mongodb(monkeypatch):
+    monkeypatch.delenv("TASK_STATE_STORAGE_PHASE", raising=False)
+    monkeypatch.delenv("TASK_MESSAGE_STORAGE_PHASE", raising=False)
+
+    env = EnvironmentVariables.refresh(force_refresh=True)
+
+    assert env.TASK_STATE_STORAGE_PHASE == StoragePhase.MONGODB
+    assert env.TASK_MESSAGE_STORAGE_PHASE == StoragePhase.MONGODB
+    assert env.mongodb_required is True
+
+
+@pytest.mark.unit
+def test_storage_phases_parse_from_environment(monkeypatch):
+    monkeypatch.setenv("TASK_STATE_STORAGE_PHASE", "postgres")
+    monkeypatch.setenv("TASK_MESSAGE_STORAGE_PHASE", "mongodb")
+
+    env = EnvironmentVariables.refresh(force_refresh=True)
+
+    assert env.TASK_STATE_STORAGE_PHASE == StoragePhase.POSTGRES
+    assert env.TASK_MESSAGE_STORAGE_PHASE == StoragePhase.MONGODB
+
+
+@pytest.mark.unit
+@pytest.mark.parametrize(
+    "raw", ["mongo", "POSTGRES", "true", "", "dual_write", "dual_read"]
+)
+def test_storage_phase_rejects_unknown_values(monkeypatch, raw):
+    # Fail loud on typos rather than silently falling back to a backend the
+    # operator did not choose. dual_write/dual_read are deliberately not
+    # defined: a data-migration effort would add its own phases.
+    monkeypatch.setenv("TASK_STATE_STORAGE_PHASE", raw)
+
+    with pytest.raises(ValidationError):
+        EnvironmentVariables.refresh(force_refresh=True)
+
+
+@pytest.mark.unit
+@pytest.mark.parametrize(
+    ("state_phase", "message_phase", "expected"),
+    [
+        ("mongodb", "mongodb", True),
+        ("postgres", "mongodb", True),
+        ("mongodb", "postgres", True),
+        ("postgres", "postgres", False),
+    ],
+)
+def test_mongodb_required_unless_every_phase_is_postgres(
+    monkeypatch, state_phase, message_phase, expected
+):
+    monkeypatch.setenv("TASK_STATE_STORAGE_PHASE", state_phase)
+    monkeypatch.setenv("TASK_MESSAGE_STORAGE_PHASE", message_phase)
+
+    env = EnvironmentVariables.refresh(force_refresh=True)
+
+    assert env.mongodb_required is expected

agentex/tests/unit/repositories/test_task_state_repository_selector.py

@@ -0,0 +1,69 @@
+from types import SimpleNamespace
+from typing import get_args
+from unittest.mock import MagicMock
+
+import pytest
+from src.config.environment_variables import EnvironmentVariables
+from src.domain.repositories import task_state_repository as selector_module
+from src.domain.repositories.task_state_repository import (
+    DTaskStateRepository,
+    TaskStateRepository,
+    get_task_state_repository,
+)
+
+
+@pytest.fixture(autouse=True)
+def _reset_env_cache():
+    """Drop the forced-refresh cache so later tests re-read a clean environment."""
+    yield
+    EnvironmentVariables.clear_cache()
+
+
+def _set_phase(monkeypatch, phase: str | None):
+    if phase is None:
+        monkeypatch.delenv("TASK_STATE_STORAGE_PHASE", raising=False)
+    else:
+        monkeypatch.setenv("TASK_STATE_STORAGE_PHASE", phase)
+    EnvironmentVariables.refresh(force_refresh=True)
+
+
+@pytest.mark.unit
+@pytest.mark.parametrize("phase", [None, "mongodb"])
+def test_selector_returns_mongo_repository_for_mongodb_phase(monkeypatch, phase):
+    _set_phase(monkeypatch, phase)
+    mongo_db = MagicMock()
+    monkeypatch.setattr(
+        selector_module,
+        "GlobalDependencies",
+        lambda: SimpleNamespace(mongodb_database=mongo_db),
+    )
+
+    repository = get_task_state_repository()
+
+    assert isinstance(repository, TaskStateRepository)
+    assert repository.db is mongo_db
+
+
+@pytest.mark.unit
+def test_selector_rejects_unimplemented_phases_lazily(monkeypatch):
+    """Unimplemented phases fail loud, and without touching any Mongo wiring:
+    the mongodb branch must be the only one that needs a Mongo handle."""
+    phase = "postgres"
+    _set_phase(monkeypatch, phase)
+    global_dependencies = MagicMock(
+        side_effect=AssertionError("selector must not touch Mongo wiring")
+    )
+    monkeypatch.setattr(selector_module, "GlobalDependencies", global_dependencies)
+
+    with pytest.raises(NotImplementedError, match=phase):
+        get_task_state_repository()
+
+    global_dependencies.assert_not_called()
+
+
+@pytest.mark.unit
+def test_di_seam_resolves_through_selector():
+    """DTaskStateRepository is the seam every consumer (use case, authorization
+    shortcuts, services) resolves; it must point at the phase selector."""
+    depends_marker = get_args(DTaskStateRepository)[1]
+    assert depends_marker.dependency is get_task_state_repository