Skip to content

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

Draft
vineetvora-scale wants to merge 1 commit into
mainfrom
vineetvora/task-state-storage-phase-seam
Draft

feat(states): add storage-phase config and task-state repository selector#380
vineetvora-scale wants to merge 1 commit into
mainfrom
vineetvora/task-state-storage-phase-seam

Conversation

@vineetvora-scale

@vineetvora-scale vineetvora-scale commented Jul 28, 2026

Copy link
Copy Markdown
Contributor

Summary

First PR in a short series adding an optional PostgreSQL backend for task state storage, selected per deployment by a config flag. This PR lands the configuration and the selection seam only — no behavior change: the default phase is mongodb, and existing deployments continue to run exactly as before. The dual-write/dual-read phases are reserved for later in the series (enum values exist, selector rejects them).

What's in this PR

1. Storage-phase settings (src/config/environment_variables.py)

  • TASK_STATE_STORAGE_PHASE and TASK_MESSAGE_STORAGE_PHASE: mongodb | dual_write | dual_read | postgres, default mongodb.
  • Values are validated against a StoragePhase enum — a typo fails at startup instead of silently routing to a backend the operator didn't choose.
  • Derived mongodb_required property (true unless every store's phase is postgres) for later use in conditional Mongo startup/readiness.

2. TaskStateRepositoryProtocol (src/domain/repositories/task_state_repository.py)

  • Structural interface capturing 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.
  • The docstring records the behavioral contract a backend must honor (string ids, caller-supplied timestamp passthrough, ItemDoesNotExist/DuplicateItemError, the Mongo-shaped $in list filter). The MongoDB repository satisfies it unchanged.

3. Phase selector wired into the seam

  • get_task_state_repository() resolves the repository for the configured phase; DTaskStateRepository now points at it, so every DI consumer follows automatically.
  • Construction is lazy per branch: only the mongodb branch touches a Mongo handle, which is what will allow running without a MongoDB connection once no store needs one.
  • Non-mongodb phases raise NotImplementedError until the Postgres repository lands in a follow-up.

4. Temporal factories routed through the selector (src/temporal/task_retention_factory.py, src/temporal/scheduled_agent_run_factory.py)

  • Both previously constructed TaskStateRepository(mongodb_database) by hand, bypassing DI. Without this change, a postgres-phase deployment's workers would silently keep using MongoDB while the API used Postgres.

Plus docker-compose pass-throughs for both variables (defaulting to mongodb) on the API and worker services.

Review focus

  • The Protocol's method set: it's derived from actual call sites, notably the retention service's find_by_field/delete_by_field/batch_create — these are part of the contract the future Postgres repository must implement.
  • The two factory changes are the correctness-critical part of this PR (worker/API consistency).
  • Protocol location (repositories module) and shape (Protocol vs the existing CRUDRepository ABC): the ABC's list() takes no filter/pagination arguments and lacks the task-state-specific methods, and a Protocol requires no change to the MongoDB repository's class hierarchy.

Testing

  • New unit tests: env parsing/validation and the mongodb_required matrix; selector resolution for all phases, including a test pinning that non-mongodb branches never touch Mongo wiring; both Temporal factories resolving through the selector.
  • Full config/selector/factory unit suites pass locally (42 tests). openapi.yaml regenerated with no diff; ruff clean.

Follow-ups (separate PRs)

  1. TaskStateORM + a schema change creating the new (empty) task_states table — no data is moved.
  2. TaskStatePostgresRepository + repository tests parametrized over both backends; selector's postgres branch goes live.
  3. Task messages, conditional Mongo startup, and dual-write/dual-read phases later in the series.

…ctor

Introduce TASK_STATE_STORAGE_PHASE and TASK_MESSAGE_STORAGE_PHASE settings
(mongodb|dual_write|dual_read|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. Phases other than mongodb raise
NotImplementedError until the Postgres repository lands, and the mongodb
default keeps existing deployments unchanged.
@vineetvora-scale
vineetvora-scale force-pushed the vineetvora/task-state-storage-phase-seam branch from 6f216c5 to 0be2190 Compare July 28, 2026 20:39
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant