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().

Author: vineetvora-scale

agentex/src/config/environment_variables.py

@@ -80,8 +80,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"
@@ -115,6 +117,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
@@ -318,6 +344,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
 

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."

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

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()