Subject keying rides DBOS workflow attributes; durable_runs shrinks

backend/druks/build/extension.py

@@ -1,7 +1,6 @@
 from typing import TYPE_CHECKING
 
 from pydantic import BaseModel, Field
-from sqlalchemy import select
 
 from druks.agents import Agent
 from druks.build.contracts import (
@@ -17,10 +16,9 @@
 from druks.build.models import WorkItem
 from druks.build.scoping.contracts import ScopeBriefOutput
 from druks.db import db_session
-from druks.durable.enums import RunState
 from druks.events import Event, FeedItem
 from druks.extensions import Extension
-from druks.workflows import Run, SubjectActivity, get_run_phase
+from druks.workflows import Run, RunState, SubjectActivity, get_run_phase
 
 if TYPE_CHECKING:
     from druks.build.schemas import WorkItemSummary
@@ -219,13 +217,8 @@ async def subject_activity(cls, subject_id: str) -> SubjectActivity | None:
         assert item is not None  # the read-side resolves the summary before the activity
         run = item.get_build_run()
         if not run:
-            newest_scope = (
-                select(Run)
-                .where(Run.kind == Scope.kind, Run.subject["id"].as_integer() == item.id)
-                .order_by(Run.created_at.desc())
-                .limit(1)
-            )
-            run = db_session().scalars(newest_scope).first()
+            scope_runs = Run.list_for_subject("work_item", str(item.id), kind=Scope.kind)
+            run = scope_runs[0] if scope_runs else None
         # Only while a run is actually RUNNING — a run parked on a gate isn't working.
         if not run or run.state != RunState.RUNNING.value:
             return None

backend/druks/build/scoping/workflows.py

@@ -1,14 +1,10 @@
 import logging
 
-from sqlalchemy import select
-
 from druks.build.extension import Build
 from druks.build.models import Project, ProjectRepo, WorkItem
 from druks.build.scoping.contracts import ScopeBriefOutput
-from druks.db import db_session
-from druks.durable.enums import RunState
 from druks.ticketing.datastructures import Ticket
-from druks.workflows import Gate, Run, Workflow
+from druks.workflows import Gate, Run, RunState, Workflow
 
 logger = logging.getLogger(__name__)
 
@@ -57,12 +53,8 @@ async def dispatch(cls, *, ticket: Ticket) -> str | None:
 
     @classmethod
     def parked_for(cls, work_item_id: int) -> Run | None:
-        stmt = select(Run).where(
-            Run.kind == cls.kind,
-            Run.state == RunState.PENDING_INPUT.value,
-            Run.subject["id"].as_integer() == work_item_id,
-        )
-        return db_session().scalars(stmt).first()
+        runs = Run.list_for_subject("work_item", str(work_item_id), kind=cls.kind)
+        return next((run for run in runs if run.state == RunState.PENDING_INPUT.value), None)
 
     async def get_prompt_context(self, **context: object) -> dict[str, object]:
         # Everything the agent needs beyond the ticket it fetches itself: where

backend/druks/durable/dbos_state.py

@@ -1,6 +1,7 @@
 from datetime import timedelta
 
 import sqlalchemy as sa
+from sqlalchemy.dialects.postgresql import JSONB
 
 from druks.durable.enums import RunState
 
@@ -14,19 +15,48 @@
 # its executor destroyed) and reads orphaned rather than scheduled forever.
 _MISSING_STATUS_GRACE = timedelta(minutes=5)
 
-# Read-only handle on the columns the derivation needs; DBOS owns and migrates
-# the real table. NOTE: DBOS workflow_status GC/retention must stay off while
-# durable_runs keeps history — a purged workflow row derives as orphaned once
-# past the grace window.
+# Read-only handle on the columns the derivation and subject keying need; DBOS
+# owns and migrates the real table. NOTE: DBOS workflow_status GC/retention must
+# stay off — a purged workflow row derives as orphaned once past the grace
+# window and falls out of its subject's timeline.
 workflow_status = sa.Table(
     "workflow_status",
     sa.MetaData(schema=DBOS_SYSTEM_SCHEMA),
     sa.Column("workflow_uuid", sa.String, primary_key=True),
     sa.Column("status", sa.String),
     sa.Column("updated_at", sa.BigInteger),
+    sa.Column("attributes", JSONB),
 )
 
 
+def extension_expression(run_id: sa.ColumnElement) -> sa.ColumnElement:
+    return (
+        sa.select(workflow_status.c.attributes["extension"].as_string())
+        .where(workflow_status.c.workflow_uuid == run_id)
+        .correlate_except(workflow_status)
+        .scalar_subquery()
+    )
+
+
+def subject_filter(
+    run_id: sa.ColumnElement, subject_type: str, subject_id: str
+) -> sa.ColumnElement:
+    # "This run is about that subject" — the predicate every runs-for-a-subject
+    # query composes, reading the attributes stamped at start(). A fresh alias
+    # per call keeps it independent of the state/updated_at subqueries, which
+    # claim the bare workflow_status table via correlate_except.
+    ws = workflow_status.alias()
+    return (
+        sa.select(ws.c.workflow_uuid)
+        .where(
+            ws.c.workflow_uuid == run_id,
+            ws.c.attributes["subject_type"].as_string() == subject_type,
+            ws.c.attributes["subject_id"].as_string() == subject_id,
+        )
+        .exists()
+    )
+
+
 def state_expression(
     run_id: sa.ColumnElement, input_gate: sa.ColumnElement, created_at: sa.ColumnElement
 ) -> sa.ColumnElement:

backend/druks/durable/models.py

@@ -11,7 +11,12 @@
 
 from druks.core.models import Uuid7Pk
 from druks.database import db_session, get_session
-from druks.durable.dbos_state import state_expression, updated_at_expression
+from druks.durable.dbos_state import (
+    extension_expression,
+    state_expression,
+    subject_filter,
+    updated_at_expression,
+)
 from druks.durable.enums import ACTIVE_STATES
 from druks.harnesses.artifacts import normalize_token_usage
 from druks.models import Base
@@ -44,19 +49,15 @@ class Run(Base):
     # error, so read-sides tell e.g. a gate timeout from a crash without parsing
     # `failure`. Empty/None for a crash.
     failure_code: Mapped[str | None] = mapped_column(default=None)
-    # run()'s validated input, dumped to JSON at start(); {} for an input-less run.
-    input: Mapped[dict[str, Any]] = mapped_column(JSONB, default=dict)
-    # What the run is about ({type, id}); None for a framework cron. Routing
-    # metadata, not body input — its own column so events and reads key off it.
-    subject: Mapped[dict[str, Any] | None] = mapped_column(JSONB, default=None)
-    # The extension that launched the run, stamped onto its events.
-    extension: Mapped[str | None] = mapped_column(default=None)
     created_at: Mapped[datetime] = mapped_column(default=Base.utc_now)
     # Derived, never stored: DBOS owns whether the workflow is live — every
     # SELECT (session.get() on a not-yet-loaded instance included) computes it
     # fresh; an already-loaded instance keeps what it read until expired.
     # Read-only; an operator ends a run through cancel().
     state: Mapped[str] = column_property(state_expression(id, input_gate, created_at))
+    # The extension that launched the run, read back from the attributes
+    # stamped at start(); None for a framework cron.
+    extension: Mapped[str | None] = column_property(extension_expression(id))
     # When the run last changed — the newest of creation, the parked ask, and
     # DBOS's status write.
     updated_at: Mapped[datetime] = column_property(
@@ -68,27 +69,14 @@ def is_active(self) -> bool:
         return self.state in {s.value for s in ACTIVE_STATES}
 
     @classmethod
-    def create_row(
-        cls,
-        engine,
-        *,
-        workflow_id: str,
-        kind: str,
-        input: dict[str, Any],
-        subject: dict[str, Any] | None = None,
-        extension: str | None = None,
-    ) -> None:
+    def create_row(cls, engine, *, workflow_id: str, kind: str) -> None:
         # Own committed transaction (not the caller's request txn) so the row
         # exists before the running workflow's first lifecycle event. Idempotent:
         # a scheduled run creates its row inside the (replayable) body, and a
         # start that races its own retry must not double-insert.
         with get_session(engine) as session:
             session.execute(
-                pg_insert(cls)
-                .values(
-                    id=workflow_id, kind=kind, input=input, subject=subject, extension=extension
-                )
-                .on_conflict_do_nothing()
+                pg_insert(cls).values(id=workflow_id, kind=kind).on_conflict_do_nothing()
             )
             session.commit()
 
@@ -97,17 +85,21 @@ def get(cls, workflow_id: str) -> "Run | None":
         return db_session().get(cls, workflow_id)
 
     @classmethod
-    def list_for_subject(cls, subject_type: str, subject_id: str) -> list["Run"]:
+    def list_for_subject(
+        cls, subject_type: str, subject_id: str, kind: str | None = None
+    ) -> list["Run"]:
         # Every run about this subject (stamped at start), newest first — a
-        # subject's lifecycle spans many runs, so its status is theirs aggregated.
+        # subject's lifecycle spans many runs, so its status is theirs
+        # aggregated. ``kind`` narrows to one workflow's runs; per (kind,
+        # subject) the queue dedup makes runs strictly sequential, so newest
+        # first holds within a kind too.
         stmt = (
             select(cls)
-            .where(
-                cls.subject["type"].as_string() == subject_type,
-                cls.subject["id"].as_string() == subject_id,
-            )
+            .where(subject_filter(cls.id, subject_type, subject_id))
             .order_by(cls.updated_at.desc())
         )
+        if kind:
+            stmt = stmt.where(cls.kind == kind)
         return list(db_session().scalars(stmt))
 
     def get_ask(self) -> dict[str, Any]:
@@ -150,16 +142,17 @@ def get_rendered_ask(self) -> dict[str, Any]:
         # actions; url is an optional gate-author-declared view-link.
         return {"body": ask["label"], "actions": None, "deep_link": ask.get("url")}
 
-    def create_park_notification(self, destination_id: str) -> str:
+    def create_park_notification(self, destination_id: str, subject: dict[str, Any]) -> str:
         # Create the notification for the round this run just parked on — the
-        # caller enqueues delivery. run_id + run_parked_at snapshot the round
-        # so a click on an old button can be refused once the run re-parks.
+        # caller supplies the run's subject and enqueues delivery. run_id +
+        # run_parked_at snapshot the round so a click on an old button can be
+        # refused once the run re-parks.
         rendered = self.get_rendered_ask()
         notification = Notification.create(
             destination_id=destination_id,
             reason="gate.parked",
             body=rendered["body"],
-            subject=self.subject,
+            subject=subject,
             actions=rendered["actions"],
             run_id=self.id,
             run_parked_at=self.input_requested_at,
@@ -362,11 +355,7 @@ def by_run(cls, run_ids: list[str]) -> dict[str, list["AgentCall"]]:
     def list_for_subject(cls, subject_type: str, subject_id: str) -> list["AgentCall"]:
         stmt = (
             select(cls)
-            .join(Run, cls.run_id == Run.id)
-            .where(
-                Run.subject["type"].as_string() == subject_type,
-                Run.subject["id"].as_string() == subject_id,
-            )
+            .where(subject_filter(cls.run_id, subject_type, subject_id))
             .order_by(cls.created_at, cls.id)
         )
         return list(db_session().scalars(stmt))