diff --git a/agentex/docker-compose.yml b/agentex/docker-compose.yml index 517f6579..8a2523df 100644 --- a/agentex/docker-compose.yml +++ b/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:-} diff --git a/agentex/src/config/environment_variables.py b/agentex/src/config/environment_variables.py index 97c06b64..d49d0a25 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,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 diff --git a/agentex/src/domain/repositories/task_state_repository.py b/agentex/src/domain/repositories/task_state_repository.py index 9cc0381e..2c515e25 100644 --- a/agentex/src/domain/repositories/task_state_repository.py +++ b/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) +] 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..2448a8fa --- /dev/null +++ b/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 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..c81a91ec --- /dev/null +++ b/agentex/tests/unit/repositories/test_task_state_repository_selector.py @@ -0,0 +1,105 @@ +from types import SimpleNamespace +from typing import get_args +from unittest.mock import MagicMock + +import pytest +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, + TaskStateRepository, + TaskStateRepositoryProtocol, + 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) + + +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): + _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.""" + _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="postgres"): + get_task_state_repository() + + 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 + 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