diff --git a/agentex/database/migrations/alembic/versions/2026_08_03_1500_add_task_states_771e95623724.py b/agentex/database/migrations/alembic/versions/2026_08_03_1500_add_task_states_771e95623724.py new file mode 100644 index 00000000..48ea95e1 --- /dev/null +++ b/agentex/database/migrations/alembic/versions/2026_08_03_1500_add_task_states_771e95623724.py @@ -0,0 +1,73 @@ +"""add task_states + +Revision ID: 771e95623724 +Revises: a1b2c3d4e5f6 +Create Date: 2026-08-03 15:00:00.000000 + +Creates the task_states table: the optional PostgreSQL backend for task state +(selected per deployment via TASK_STATE_STORAGE_PHASE; MongoDB remains the +default). Schema-only on a brand-new table, so creation is instant and holds +no lock against live traffic; nothing reads or writes the table until the +Postgres task-state repository lands. + +The (task_id, agent_id) index is deliberately absent: its shape is the open +write-semantics decision (a unique constraint backing an atomic upsert, or a +plain compound index mirroring MongoDB) and it ships with that decision. +""" + +from collections.abc import Sequence + +import sqlalchemy as sa +from alembic import op +from sqlalchemy.dialects.postgresql import JSONB + +# revision identifiers, used by Alembic. +revision: str = "771e95623724" +down_revision: str | None = "a1b2c3d4e5f6" +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + + +def upgrade() -> None: + op.create_table( + "task_states", + sa.Column("id", sa.String(), nullable=False), + sa.Column("task_id", sa.String(), nullable=False), + sa.Column("agent_id", sa.String(), nullable=False), + sa.Column("state", JSONB(), nullable=False), + sa.Column( + "created_at", + sa.DateTime(timezone=True), + server_default=sa.text("now()"), + nullable=False, + ), + sa.Column( + "updated_at", + sa.DateTime(timezone=True), + server_default=sa.text("now()"), + nullable=False, + ), + # Both FKs are bare (no ON DELETE action), deliberately: MongoDB has no + # cascades, so state deletion is application-driven on both backends, + # and the FK alone prevents orphaned rows. This matches the sibling + # child tables of tasks (task_agents, events), whose bare FKs already + # reject the existing DELETE /tasks path for any task with children; + # task_states behaves identically rather than cascading on one backend + # only. + sa.ForeignKeyConstraint(["task_id"], ["tasks.id"]), + sa.ForeignKeyConstraint(["agent_id"], ["agents.id"]), + sa.PrimaryKeyConstraint("id"), + ) + # The index targets the table created in this same migration, so it holds + # no write-blocking lock against live traffic (the table has no rows yet). + op.create_index( + "ix_task_states_agent_id", + "task_states", + ["agent_id"], + unique=False, + ) + + +def downgrade() -> None: + op.drop_index("ix_task_states_agent_id", table_name="task_states") + op.drop_table("task_states") diff --git a/agentex/src/adapters/orm.py b/agentex/src/adapters/orm.py index e5f7b139..95a95963 100644 --- a/agentex/src/adapters/orm.py +++ b/agentex/src/adapters/orm.py @@ -147,6 +147,49 @@ class AgentTaskTrackerORM(BaseORM): ) +class TaskStateORM(BaseORM): + """Task state on PostgreSQL: one row per (task, agent) pair. + + Optional storage backend for task state (selected per deployment via + TASK_STATE_STORAGE_PHASE); MongoDB remains the default. Nothing reads or + writes this table until the Postgres task-state repository lands. + """ + + __tablename__ = "task_states" + + id = Column(String, primary_key=True, default=orm_id) + # Both FKs are bare (no ON DELETE action), deliberately: MongoDB has no + # cascades, so state deletion is application-driven on both backends, and + # the FK alone prevents orphaned rows. This matches the sibling child + # tables of tasks (task_agents, events), whose bare FKs already reject the + # existing DELETE /tasks path for any task with children; task_states + # behaves identically rather than cascading on one backend only. + task_id = Column(String, ForeignKey("tasks.id"), nullable=False) + agent_id = Column(String, ForeignKey("agents.id"), nullable=False) + state = Column(JSONB, nullable=False) + # NOT NULL with a server default: the repository must OMIT unset (None) + # timestamps when constructing rows. SQLAlchemy renders an explicitly + # assigned None as a literal NULL, which violates the constraint instead + # of falling back to the default (StateEntity defaults these to None). + created_at = Column( + DateTime(timezone=True), server_default=func.now(), nullable=False + ) + updated_at = Column( + DateTime(timezone=True), + server_default=func.now(), + onupdate=func.now(), + nullable=False, + ) + + __table_args__ = ( + # The (task_id, agent_id) index is deliberately absent: its shape is + # the open write-semantics decision (a unique constraint backing an + # atomic upsert, or a plain compound index mirroring MongoDB) and it + # ships with the repository that reads this table. + Index("ix_task_states_agent_id", "agent_id"), + ) + + class SpanORM(BaseORM): __tablename__ = "spans" id = Column(String, primary_key=True, default=orm_id) # Using UUIDs for IDs