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

agentex/docker-compose.yml

@@ -164,6 +164,10 @@ services:
       - REDIS_URL=redis://agentex-redis:6379
       - MONGODB_URI=mongodb://agentex-mongodb:27017
       - MONGODB_DATABASE_NAME=agentex
+      # Storage backend per document store; only mongodb is selectable until
+      # the Postgres repositories land (startup rejects anything else).
+      - 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 +242,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

@@ -68,6 +68,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):
@@ -76,6 +78,19 @@ class Environment(str, Enum):
     PROD = "production"
 
 
+class StoragePhase(str, Enum):
+    """Which backend serves a document store (task state, task messages).
+
+    POSTGRES becomes selectable per store once that store's Postgres
+    repository lands; until then refresh() rejects it at startup. 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
 
 
@@ -103,6 +118,30 @@ def _parse_bool_env(key: EnvVarKeys, default: bool) -> bool:
     )
 
 
+def _validate_storage_phases(environment_variables: EnvironmentVariables) -> None:
+    """
+    The Postgres storage path lands store-by-store; until a store's repository
+    exists, selecting its phase must fail here — at process startup, in the
+    API and Temporal workers alike — rather than on first use, where it would
+    surface as request-time 500s or a crash inside worker wiring.
+    """
+    for key, phase in (
+        (
+            EnvVarKeys.TASK_STATE_STORAGE_PHASE,
+            environment_variables.TASK_STATE_STORAGE_PHASE,
+        ),
+        (
+            EnvVarKeys.TASK_MESSAGE_STORAGE_PHASE,
+            environment_variables.TASK_MESSAGE_STORAGE_PHASE,
+        ),
+    ):
+        if phase is not StoragePhase.MONGODB:
+            raise ValueError(
+                f"{key.value}={phase.value!r} is not supported yet; "
+                "only 'mongodb' is currently available"
+            )
+
+
 class EnvironmentVariables(BaseModel):
     ENVIRONMENT: str | None = Environment.DEV
     OPENAI_API_KEY: str | None
@@ -171,6 +210,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:
@@ -293,7 +345,14 @@ 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
+            ),
         )
+        _validate_storage_phases(environment_variables)
         refreshed_environment_variables = environment_variables
         return refreshed_environment_variables
 

agentex/src/domain/repositories/task_state_repository.py

@@ -1,16 +1,76 @@
-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 TaskStateRepository(MongoDBCRUDRepository[StateEntity]):
+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], TaskStateRepositoryProtocol
+):
     """Repository for managing task states in MongoDB."""
 
     COLLECTION_NAME = "task_states"
@@ -47,4 +107,29 @@ 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)
+    # Unreachable through env config (refresh() rejects unimplemented phases
+    # at startup); defense in depth for configs constructed another way.
+    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
@@ -113,7 +113,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,90 @@
+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", "mongodb")
+    monkeypatch.setenv("TASK_MESSAGE_STORAGE_PHASE", "mongodb")
+
+    env = EnvironmentVariables.refresh(force_refresh=True)
+
+    assert env.TASK_STATE_STORAGE_PHASE == StoragePhase.MONGODB
+    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(
+    "key", ["TASK_STATE_STORAGE_PHASE", "TASK_MESSAGE_STORAGE_PHASE"]
+)
+def test_unimplemented_phase_is_rejected_at_refresh(monkeypatch, key):
+    """postgres is a valid phase value but has no repository yet: the config
+    refresh (process startup, API and Temporal workers alike) must fail with
+    the variable named — not the first state request or worker wiring."""
+    monkeypatch.setenv(key, "postgres")
+
+    with pytest.raises(ValueError, match=key):
+        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.delenv("TASK_STATE_STORAGE_PHASE", raising=False)
+    monkeypatch.delenv("TASK_MESSAGE_STORAGE_PHASE", raising=False)
+    env = EnvironmentVariables.refresh(force_refresh=True)
+
+    # Unimplemented phases cannot come from the environment (refresh rejects
+    # them at startup), so pin the derived property on updated copies.
+    patched = env.model_copy(
+        update={
+            "TASK_STATE_STORAGE_PHASE": StoragePhase(state_phase),
+            "TASK_MESSAGE_STORAGE_PHASE": StoragePhase(message_phase),
+        }
+    )
+
+    assert patched.mongodb_required is expected