fix(rbac+estimation): drop legacy OrgToUser.role and harden every estimate_bud_dates caller (#181)

* fix(rbac+estimation): repair count_active_by_role GROUP BY and harden every estimate_bud_dates caller

The dynamic-role refactor introduced count_active_by_role with a
parameterised CASE in both SELECT and GROUP BY. SQLAlchemy compiled
each call to _effective_role_case() as a fresh expression, so the two
CASEs emitted distinct bind parameters and Postgres rejected the query
with "column roles.scope_type must appear in the GROUP BY clause". That
GroupingError cascaded through get_role_capacity into estimate_bud_dates
and crashed every caller. The original failure was then masked as
InFailedSQLTransactionError on the next innocent statement because every
"except Exception" around estimation swallowed the error without
rolling back, leaving the connection in an aborted state.

Two-part fix:

1. Repository: hoist the CASE into a single labelled expression so SELECT
   and GROUP BY reference the same object. Postgres accepts GROUP BY by
   output column alias, sidestepping the expression-identity trap.

2. Every estimate_bud_dates caller that ran on a shared session is now
   wrapped in db.begin_nested() (Postgres SAVEPOINT). A query failure
   inside the estimator only rolls back its own writes; the outer txn
   stays alive and prior flushed writes survive. Same session, same
   connection — no row-lock conflict with the outer's pending writes.
   Each except now logs with exc_info=True so the real traceback is not
   lost again.

   - handle_prd_result: explicit flush of linked-feature inserts before
     the savepoint so an autoflush inside it cannot scope them to the
     savepoint and lose them on rollback.
   - handle_tech_arch_result: existing flush already anchored
     tech_spec_md + impacted_repos; only the swallow shape changed.
   - handle_testing_result: moved the trailing flush above the savepoint
     for the same anchoring reason.
   - check_all_prs_merged: create_agent_task_for_stage already commits
     internally, so no extra flush; SAVEPOINT still isolates the webhook's
     trailing commit from estimator failure.
   - code_review_override endpoint: outer was already committed earlier
     in the handler, so estimation safely runs in a fresh AsyncSessionLocal
     without a row-lock conflict.

Inline imports in the touched handlers were hoisted to top of file per
the project's import policy.

Verified end-to-end against live Postgres on the failing bud:
- count_active_by_role / get_role_capacity / estimate_bud_dates all
  succeed.
- Forced SQL error inside db.begin_nested() raises cleanly; outer txn
  stays valid; subsequent outer query sees the pre-savepoint flush.
- Full handler flow (outer flush -> SAVEPOINT estimate incl. LLM call
  -> outer rollback) completes in 5.7s with no deadlock.
- All four triggers (prd_completed, tech_arch_completed,
  testing_completed, prs_merged) pass.

ruff + mypy clean on the four modified files.

Signed-off-by: Arun Rajkumar <mickyarunr@gmail.com>

* chore(format): apply ruff format on the estimation-resilience edits

The PR-181 commit passed `ruff check` locally but not `ruff format
--check` — CI's separate format gate flagged three files. Pure
whitespace: ruff collapses multi-line calls / logger.warning args
that fit on one line. No behavioural change. lint, mypy, and the
behaviour-verifying live tests from the original commit still pass.

Signed-off-by: Arun Rajkumar <mickyarunr@gmail.com>

---------

Signed-off-by: Arun Rajkumar <mickyarunr@gmail.com>

Author: mickyarun

backend/app/api/v1/bud.py

@@ -34,6 +34,7 @@
 from app.api.v1.bud_versions import router as versions_router
 from app.api.v1.bud_workflows import router as workflows_router
 from app.core.deps import get_current_user, get_db, get_user_permissions, require_permissions
+from app.database import AsyncSessionLocal
 from app.models.bud import (
     BUDDesignStatus,
     BUDDocument,
@@ -92,6 +93,7 @@
 )
 from app.services.bud_assignment_actions import assign_bud, unassign_bud
 from app.services.bud_edit_policy import assert_section_editable
+from app.services.bud_estimation import estimate_bud_dates
 from app.services.bud_timeline import record_event
 from app.services.job_queue import JOB_BUD_AGENT, create_job
 
@@ -1126,13 +1128,8 @@ async def _bg_estimate(
     actor_name: str,
 ) -> None:
     """Run estimation in the background with its own DB session."""
-    from app.database import AsyncSessionLocal
-    from app.services.bud_estimation import estimate_bud_dates
-
     try:
         async with AsyncSessionLocal() as db:
-            from app.repositories.bud import BUDRepository
-
             bud_repo = BUDRepository(db, org_id=org_id)
             bud = await bud_repo.get_by_id(bud_id)
             if bud is None:
@@ -1274,23 +1271,29 @@ async def override_code_review(
             bud_id=str(bud_id),
         )
 
-    # Refresh estimates so dashboards reflect the new phase.
+    # Refresh estimates so dashboards reflect the new phase. Run in a
+    # fresh session so a DB failure inside the estimator can't poison
+    # the request session — the trailing _bud_response below queries
+    # the same session and would otherwise 500 with the confusing
+    # InFailedSQLTransactionError instead of the real cause.  The
+    # status transition was already committed above, so a fresh
+    # refetch sees status=testing without any mirroring.
     try:
-        from app.services.bud_estimation import estimate_bud_dates
-
-        refreshed_for_est = await bud_repo.get_by_id(bud_id)
-        if refreshed_for_est is not None:
-            await estimate_bud_dates(
-                db,
-                current_user.org_id,
-                refreshed_for_est,
-                trigger="code_review_override",
-            )
-            await db.commit()
+        async with AsyncSessionLocal() as est_db:
+            est_bud = await BUDRepository(est_db, org_id=current_user.org_id).get_by_id(bud_id)
+            if est_bud is not None:
+                await estimate_bud_dates(
+                    est_db,
+                    current_user.org_id,
+                    est_bud,
+                    trigger="code_review_override",
+                )
+                await est_db.commit()
     except Exception:
         logger.warning(
             "code_review_override_estimation_failed",
             bud_id=str(bud_id),
+            exc_info=True,
         )
 
     refreshed = await bud_repo.get_by_id(bud_id)

backend/app/repositories/user.py

@@ -284,13 +284,18 @@ async def count_active_by_role(self, org_id: uuid.UUID) -> dict[UserRole, int]:
         ``role_id``-less rows fall out because they have no pool semantics.
         """
         base_role = aliased(Role)
+        # The CASE must be the SAME expression object in SELECT and GROUP BY;
+        # calling ``_effective_role_case`` twice would emit two distinct
+        # parameterised CASEs and Postgres would reject the GROUP BY as not
+        # covering the SELECT's underlying ``roles.scope_type`` reference.
+        role_expr = _effective_role_case(base_role).label("effective_role")
         stmt = (
-            select(_effective_role_case(base_role), func.count())
+            select(role_expr, func.count())
             .select_from(OrgToUser)
             .outerjoin(Role, Role.id == OrgToUser.role_id)
             .outerjoin(base_role, base_role.id == Role.base_role_id)
             .where(OrgToUser.org_id == org_id)
-            .group_by(_effective_role_case(base_role))
+            .group_by(role_expr)
         )
         result = await self._db.execute(stmt)
         counts: dict[UserRole, int] = {}

backend/app/services/agent_result_handlers.py

@@ -30,9 +30,13 @@
 from app.models.bud_feature_link import BUDFeatureLinkSource
 from app.models.bud_version import BUDEditSource
 from app.repositories import bud_version as bud_version_repo
+from app.repositories.bud import BUDRepository
 from app.repositories.bud_feature_link import BUDFeatureLinkRepository
+from app.repositories.tracked_repository import TrackedRepoRepository
+from app.services.bud_estimation import estimate_bud_dates
 from app.services.bud_timeline import record_event
 from app.services.json_parser import parse_json_response, strip_insight_blocks
+from app.services.notification_service import send_lifecycle_notification
 
 logger = structlog.get_logger(__name__)
 
@@ -96,18 +100,25 @@ async def handle_prd_result(
     TechPlanner, Code Reviewer, Tester) inherit the grounding.
     """
     linked_count = await _persist_pm_linked_features(bud_id, org_id, output, db)
-
-    # Generate initial delivery estimates now that PRD content exists
-    try:
-        from app.repositories.bud import BUDRepository
-        from app.services.bud_estimation import estimate_bud_dates
-
-        bud_repo = BUDRepository(db, org_id=org_id)
-        bud = await bud_repo.get_by_id(bud_id)
-        if bud:
-            await estimate_bud_dates(db, org_id, bud, trigger="prd_completed")
-    except Exception:
-        logger.warning("estimation_failed_after_prd", bud_id=str(bud_id))
+    # Flush the linked-feature inserts to the outer transaction proper
+    # before opening the savepoint below — otherwise an autoflush triggered
+    # inside the savepoint would scope those inserts to it, and an
+    # estimation failure would roll the links back too.
+    await db.flush()
+
+    # Generate initial delivery estimates now that PRD content exists.
+    # Wrap in a SAVEPOINT so a query failure inside the estimator only
+    # rolls back the estimator's own writes — the outer transaction stays
+    # alive, its prior writes survive, and the next log_agent_activity
+    # call doesn't trip InFailedSQLTransactionError.
+    bud_repo = BUDRepository(db, org_id=org_id)
+    bud = await bud_repo.get_by_id(bud_id)
+    if bud is not None:
+        try:
+            async with db.begin_nested():
+                await estimate_bud_dates(db, org_id, bud, trigger="prd_completed")
+        except Exception:
+            logger.warning("estimation_failed_after_prd", bud_id=str(bud_id), exc_info=True)
 
     return {
         "section": "requirements_md",
@@ -254,9 +265,6 @@ async def handle_tech_arch_result(
     We parse that JSON to determine which repos actually need changes,
     then strip the JSON block before storing the markdown.
     """
-    from app.repositories.bud import BUDRepository
-    from app.repositories.tracked_repository import TrackedRepoRepository
-
     bud_repo = BUDRepository(db, org_id=org_id)
     bud = await bud_repo.get_by_id(bud_id)
     if bud:
@@ -292,17 +300,20 @@ async def handle_tech_arch_result(
         # todos crystallize from the *approved* plan at the DEVELOPMENT-phase
         # transition — see services/bud_development.on_bud_development_started.
 
-        # Re-estimate with richer context (now has tech spec + impacted repos)
+        # Re-estimate with richer context (now has tech spec + impacted repos).
+        # Wrap in a SAVEPOINT so a query failure inside the estimator only
+        # rolls back the estimator's own writes — the outer transaction
+        # stays alive (preserving the tech_spec_md / impacted_repos writes
+        # flushed just above), and the next log_agent_activity call doesn't
+        # trip InFailedSQLTransactionError. Same-session also means the
+        # estimator sees the in-memory bud writes without a refetch + mirror.
         try:
-            from app.services.bud_estimation import estimate_bud_dates
-
-            await estimate_bud_dates(db, org_id, bud, trigger="tech_arch_completed")
+            async with db.begin_nested():
+                await estimate_bud_dates(db, org_id, bud, trigger="tech_arch_completed")
         except Exception:
-            logger.warning("estimation_failed_after_tech_arch", bud_id=str(bud_id))
+            logger.warning("estimation_failed_after_tech_arch", bud_id=str(bud_id), exc_info=True)
 
     if bud and bud.assignee_id:
-        from app.services.notification_service import send_lifecycle_notification
-
         bud_ref = f"BUD-{bud.bud_number:03d}"
         send_lifecycle_notification(
             org_id=str(org_id),
@@ -476,14 +487,20 @@ async def handle_testing_result(
             f"- **{manual_count}** manual test cases\n\n"
             f"{parsed_data['test_execution_plan']}"
         )
-        # Re-estimate with QA test case context
+        # Flush the QA writes to the outer transaction proper BEFORE the
+        # savepoint — otherwise an autoflush triggered inside the savepoint
+        # would scope them to it, and an estimation failure would roll the
+        # qa_* / test_plan_md back too.
+        await db.flush()
+        # Re-estimate with QA test case context. Wrap in a SAVEPOINT so a
+        # query failure inside the estimator only rolls back the estimator's
+        # own writes — the QA writes above and the subsequent notification
+        # all stay on the live outer transaction.
         try:
-            from app.services.bud_estimation import estimate_bud_dates
-
-            await estimate_bud_dates(db, org_id, bud, trigger="testing_completed")
+            async with db.begin_nested():
+                await estimate_bud_dates(db, org_id, bud, trigger="testing_completed")
         except Exception:
-            logger.warning("estimation_failed_after_testing", bud_id=str(bud_id))
-        await db.flush()
+            logger.warning("estimation_failed_after_testing", bud_id=str(bud_id), exc_info=True)
 
     if bud and bud.assignee_id:
         from app.services.notification_service import send_lifecycle_notification

backend/app/services/pr_auto_transition.py

@@ -31,6 +31,7 @@
 from app.repositories.bud import BUDRepository
 from app.repositories.bud_todo import BUDTodoRepository
 from app.repositories.pull_request import PullRequestRepository
+from app.services.bud_estimation import estimate_bud_dates
 from app.services.bud_timeline import record_event
 
 logger = structlog.get_logger(__name__)
@@ -178,10 +179,13 @@ async def check_all_prs_merged(
 
     await create_agent_task_for_stage(bud, "testing", org_id, db, force=True)
 
+    # create_agent_task_for_stage above already committed the status
+    # transition, timeline events, and new agent-task row. The SAVEPOINT
+    # below exists purely so a query failure inside the estimator doesn't
+    # poison the webhook's trailing flushes with InFailedSQLTransaction.
     try:
-        from app.services.bud_estimation import estimate_bud_dates
-
-        await estimate_bud_dates(db, org_id, bud, trigger="prs_merged")
+        async with db.begin_nested():
+            await estimate_bud_dates(db, org_id, bud, trigger="prs_merged")
     except Exception:
-        logger.warning("estimation_failed_after_prs_merged", bud_id=str(bud_id))
+        logger.warning("estimation_failed_after_prs_merged", bud_id=str(bud_id), exc_info=True)
     logger.info("auto_transition_to_testing", bud_id=str(bud_id))