Fix run cancellation and live activity stability

src/run_state_store.py

@@ -131,6 +131,50 @@ class RunStateStore:
             finished_at=time.time(),
         )
 
+    def finish_active_for_session(
+        self,
+        account_key: str,
+        session_id: str,
+        *,
+        status: str,
+        stage: str | None = None,
+    ) -> list[str]:
+        now = time.time()
+        with self._connect() as conn:
+            rows = conn.execute(
+                """
+                select run_id, started_at
+                from run_states
+                where account_key = ?
+                  and session_id = ?
+                  and status in ('queued', 'running')
+                """,
+                (account_key, session_id),
+            ).fetchall()
+            run_ids: list[str] = []
+            for row in rows:
+                run_ids.append(row['run_id'])
+                started_at = row['started_at']
+                elapsed_ms = (
+                    max(0, int((now - float(started_at)) * 1000))
+                    if isinstance(started_at, (int, float))
+                    else None
+                )
+                conn.execute(
+                    """
+                    update run_states
+                    set status = ?,
+                        current_stage = coalesce(?, current_stage),
+                        cancellable = 0,
+                        elapsed_ms = coalesce(?, elapsed_ms),
+                        finished_at = ?,
+                        updated_at = ?
+                    where run_id = ?
+                    """,
+                    (status, stage, elapsed_ms, now, now, row['run_id']),
+                )
+            return run_ids
+
     def record_event(self, run_id: str, event: dict[str, Any]) -> None:
         recorded_at = event.get('recorded_at')
         if not isinstance(recorded_at, (int, float)):