From 857150f7823c5b1c3955a6b96f39410ce4b6e4d9 Mon Sep 17 00:00:00 2001 From: Vineet Vora Date: Tue, 28 Jul 2026 12:01:40 -0400 Subject: [PATCH 1/4] 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. --- agentex/docker-compose.yml | 5 ++ agentex/src/config/environment_variables.py | 32 +++++++ .../repositories/task_state_repository.py | 87 ++++++++++++++++++- .../temporal/scheduled_agent_run_factory.py | 6 +- .../src/temporal/task_retention_factory.py | 6 +- .../unit/config/test_storage_phase_env.py | 68 +++++++++++++++ .../test_task_state_repository_selector.py | 69 +++++++++++++++ .../temporal/test_factory_storage_selector.py | 38 ++++++++ 8 files changed, 304 insertions(+), 7 deletions(-) create mode 100644 agentex/tests/unit/config/test_storage_phase_env.py create mode 100644 agentex/tests/unit/repositories/test_task_state_repository_selector.py create mode 100644 agentex/tests/unit/temporal/test_factory_storage_selector.py diff --git a/agentex/docker-compose.yml b/agentex/docker-compose.yml index 517f6579..741806a9 100644 --- a/agentex/docker-compose.yml +++ b/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:-} diff --git a/agentex/src/config/environment_variables.py b/agentex/src/config/environment_variables.py index 97c06b64..75c3f6a5 100644 --- a/agentex/src/config/environment_variables.py +++ b/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,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 @@ -171,6 +184,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,6 +319,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 diff --git a/agentex/src/domain/repositories/task_state_repository.py b/agentex/src/domain/repositories/task_state_repository.py index 9cc0381e..18fd7e22 100644 --- a/agentex/src/domain/repositories/task_state_repository.py +++ b/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) +] diff --git a/agentex/src/temporal/scheduled_agent_run_factory.py b/agentex/src/temporal/scheduled_agent_run_factory.py index 121f6e4b..aebe244a 100644 --- a/agentex/src/temporal/scheduled_agent_run_factory.py +++ b/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 ) diff --git a/agentex/src/temporal/task_retention_factory.py b/agentex/src/temporal/task_retention_factory.py index 0f9290f6..d7cb8efa 100644 --- a/agentex/src/temporal/task_retention_factory.py +++ b/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 ) diff --git a/agentex/tests/unit/config/test_storage_phase_env.py b/agentex/tests/unit/config/test_storage_phase_env.py new file mode 100644 index 00000000..3359b68e --- /dev/null +++ b/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 diff --git a/agentex/tests/unit/repositories/test_task_state_repository_selector.py b/agentex/tests/unit/repositories/test_task_state_repository_selector.py new file mode 100644 index 00000000..506844aa --- /dev/null +++ b/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 diff --git a/agentex/tests/unit/temporal/test_factory_storage_selector.py b/agentex/tests/unit/temporal/test_factory_storage_selector.py new file mode 100644 index 00000000..5933173b --- /dev/null +++ b/agentex/tests/unit/temporal/test_factory_storage_selector.py @@ -0,0 +1,38 @@ +"""The Temporal factories build repositories by hand, outside FastAPI's Depends +DI. These tests pin that both factories resolve the task-state repository +through the shared phase selector instead of constructing the Mongo repository +directly — otherwise a Postgres deployment's workers would silently keep +writing task state to MongoDB.""" + +from unittest.mock import MagicMock + +import pytest +from src.temporal import scheduled_agent_run_factory, task_retention_factory + + +@pytest.mark.unit +def test_retention_factory_resolves_state_repo_through_selector(monkeypatch): + sentinel_repository = MagicMock() + selector = MagicMock(return_value=sentinel_repository) + monkeypatch.setattr(task_retention_factory, "get_task_state_repository", selector) + + use_case = task_retention_factory.build_task_retention_use_case(MagicMock()) + + selector.assert_called_once_with() + assert use_case.retention_service.task_state_repository is sentinel_repository + + +@pytest.mark.unit +def test_scheduled_run_factory_resolves_state_repo_through_selector(monkeypatch): + sentinel_repository = MagicMock() + selector = MagicMock(return_value=sentinel_repository) + monkeypatch.setattr( + scheduled_agent_run_factory, "get_task_state_repository", selector + ) + + use_case = scheduled_agent_run_factory.build_acp_use_case_for_principal( + MagicMock(), {"user_id": "u1", "account_id": "a1"} + ) + + selector.assert_called_once_with() + assert use_case.task_service.task_state_repository is sentinel_repository From c08fca8336b33b1353b37472378b9fdcf2422c8c Mon Sep 17 00:00:00 2001 From: Vineet Vora Date: Fri, 31 Jul 2026 15:37:50 -0400 Subject: [PATCH 2/4] fix(states): reject unimplemented storage phases at config refresh An accepted-but-unimplemented phase previously surfaced only on first use: request-time 500s from the API and a crash inside Temporal worker wiring. Fail at process startup instead, with the misconfigured variable named in the error. The selector keeps its NotImplementedError as defense in depth for configs constructed outside refresh(). --- agentex/src/config/environment_variables.py | 31 +++++++++++++++-- .../repositories/task_state_repository.py | 2 ++ .../unit/config/test_storage_phase_env.py | 34 +++++++++++++++---- .../test_task_state_repository_selector.py | 21 +++++++++--- 4 files changed, 76 insertions(+), 12 deletions(-) diff --git a/agentex/src/config/environment_variables.py b/agentex/src/config/environment_variables.py index 75c3f6a5..d49d0a25 100644 --- a/agentex/src/config/environment_variables.py +++ b/agentex/src/config/environment_variables.py @@ -81,8 +81,10 @@ class Environment(str, Enum): 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. + 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" @@ -116,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 @@ -326,6 +352,7 @@ def refresh(cls, force_refresh: bool = False) -> EnvironmentVariables | None: EnvVarKeys.TASK_MESSAGE_STORAGE_PHASE, StoragePhase.MONGODB ), ) + _validate_storage_phases(environment_variables) refreshed_environment_variables = environment_variables return refreshed_environment_variables diff --git a/agentex/src/domain/repositories/task_state_repository.py b/agentex/src/domain/repositories/task_state_repository.py index 18fd7e22..be8f8953 100644 --- a/agentex/src/domain/repositories/task_state_repository.py +++ b/agentex/src/domain/repositories/task_state_repository.py @@ -120,6 +120,8 @@ def get_task_state_repository() -> TaskStateRepositoryProtocol: 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." diff --git a/agentex/tests/unit/config/test_storage_phase_env.py b/agentex/tests/unit/config/test_storage_phase_env.py index 3359b68e..2448a8fa 100644 --- a/agentex/tests/unit/config/test_storage_phase_env.py +++ b/agentex/tests/unit/config/test_storage_phase_env.py @@ -24,12 +24,12 @@ def test_storage_phases_default_to_mongodb(monkeypatch): @pytest.mark.unit def test_storage_phases_parse_from_environment(monkeypatch): - monkeypatch.setenv("TASK_STATE_STORAGE_PHASE", "postgres") + 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.POSTGRES + assert env.TASK_STATE_STORAGE_PHASE == StoragePhase.MONGODB assert env.TASK_MESSAGE_STORAGE_PHASE == StoragePhase.MONGODB @@ -47,6 +47,20 @@ def test_storage_phase_rejects_unknown_values(monkeypatch, raw): 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"), @@ -60,9 +74,17 @@ def test_storage_phase_rejects_unknown_values(monkeypatch, raw): 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) - + monkeypatch.delenv("TASK_STATE_STORAGE_PHASE", raising=False) + monkeypatch.delenv("TASK_MESSAGE_STORAGE_PHASE", raising=False) env = EnvironmentVariables.refresh(force_refresh=True) - assert env.mongodb_required is expected + # 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 diff --git a/agentex/tests/unit/repositories/test_task_state_repository_selector.py b/agentex/tests/unit/repositories/test_task_state_repository_selector.py index 506844aa..deff02f3 100644 --- a/agentex/tests/unit/repositories/test_task_state_repository_selector.py +++ b/agentex/tests/unit/repositories/test_task_state_repository_selector.py @@ -3,7 +3,8 @@ from unittest.mock import MagicMock import pytest -from src.config.environment_variables import EnvironmentVariables +from src.config import environment_variables as environment_variables_module +from src.config.environment_variables import EnvironmentVariables, StoragePhase from src.domain.repositories import task_state_repository as selector_module from src.domain.repositories.task_state_repository import ( DTaskStateRepository, @@ -27,6 +28,19 @@ def _set_phase(monkeypatch, phase: str | None): EnvironmentVariables.refresh(force_refresh=True) +def _force_cached_phase(monkeypatch, phase: StoragePhase): + """Unimplemented phases are rejected by refresh() itself (the startup + guard), so exercise the selector's own defense-in-depth branch by patching + the cached config directly.""" + monkeypatch.delenv("TASK_STATE_STORAGE_PHASE", raising=False) + env = EnvironmentVariables.refresh(force_refresh=True) + monkeypatch.setattr( + environment_variables_module, + "refreshed_environment_variables", + env.model_copy(update={"TASK_STATE_STORAGE_PHASE": phase}), + ) + + @pytest.mark.unit @pytest.mark.parametrize("phase", [None, "mongodb"]) def test_selector_returns_mongo_repository_for_mongodb_phase(monkeypatch, phase): @@ -48,14 +62,13 @@ def test_selector_returns_mongo_repository_for_mongodb_phase(monkeypatch, phase) 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) + _force_cached_phase(monkeypatch, StoragePhase.POSTGRES) 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): + with pytest.raises(NotImplementedError, match="postgres"): get_task_state_repository() global_dependencies.assert_not_called() From 74356710d620e99bc8e453be60162637c2269c7f Mon Sep 17 00:00:00 2001 From: Vineet Vora Date: Mon, 3 Aug 2026 12:03:35 -0400 Subject: [PATCH 3/4] docs(compose): correct storage-phase comment, postgres is not selectable yet --- agentex/docker-compose.yml | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/agentex/docker-compose.yml b/agentex/docker-compose.yml index 741806a9..8a2523df 100644 --- a/agentex/docker-compose.yml +++ b/agentex/docker-compose.yml @@ -164,7 +164,8 @@ 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`. + # 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 From 07433d06f24be957ba42852c6ed60d82fceb9853 Mon Sep 17 00:00:00 2001 From: Vineet Vora Date: Mon, 3 Aug 2026 16:34:17 -0400 Subject: [PATCH 4/4] refactor(states): make the Mongo repository explicitly implement the Protocol Declares TaskStateRepositoryProtocol as a base of TaskStateRepository so conformance is visible at the class and checkable by IDEs. Because an explicit subclass silently inherits the protocol's placeholder methods, a test pins that every required method is a real implementation. --- .../repositories/task_state_repository.py | 4 +++- .../test_task_state_repository_selector.py | 23 +++++++++++++++++++ 2 files changed, 26 insertions(+), 1 deletion(-) diff --git a/agentex/src/domain/repositories/task_state_repository.py b/agentex/src/domain/repositories/task_state_repository.py index be8f8953..2c515e25 100644 --- a/agentex/src/domain/repositories/task_state_repository.py +++ b/agentex/src/domain/repositories/task_state_repository.py @@ -68,7 +68,9 @@ async def get_by_task_and_agent( ) -> StateEntity | None: ... -class TaskStateRepository(MongoDBCRUDRepository[StateEntity]): +class TaskStateRepository( + MongoDBCRUDRepository[StateEntity], TaskStateRepositoryProtocol +): """Repository for managing task states in MongoDB.""" COLLECTION_NAME = "task_states" diff --git a/agentex/tests/unit/repositories/test_task_state_repository_selector.py b/agentex/tests/unit/repositories/test_task_state_repository_selector.py index deff02f3..c81a91ec 100644 --- a/agentex/tests/unit/repositories/test_task_state_repository_selector.py +++ b/agentex/tests/unit/repositories/test_task_state_repository_selector.py @@ -9,6 +9,7 @@ from src.domain.repositories.task_state_repository import ( DTaskStateRepository, TaskStateRepository, + TaskStateRepositoryProtocol, get_task_state_repository, ) @@ -74,6 +75,28 @@ def test_selector_rejects_unimplemented_phases_lazily(monkeypatch): global_dependencies.assert_not_called() +@pytest.mark.unit +def test_mongo_repository_implements_every_protocol_method(): + """The Mongo repo explicitly subclasses the Protocol, so a missing method + would silently inherit the protocol's `...` placeholder (returning None) + instead of raising AttributeError. Pin that every required method is a + real implementation.""" + for name in ( + "create", + "batch_create", + "get", + "update", + "delete", + "list", + "find_by_field", + "delete_by_field", + "get_by_task_and_agent", + ): + assert getattr(TaskStateRepository, name) is not getattr( + TaskStateRepositoryProtocol, name + ), f"{name} is still the protocol placeholder — not implemented" + + @pytest.mark.unit def test_di_seam_resolves_through_selector(): """DTaskStateRepository is the seam every consumer (use case, authorization