FIX: Keep the backend responsive while starting scenario runs

pyrit/backend/routes/scenarios.py

@@ -146,7 +146,9 @@ async def start_scenario_run(request: RunScenarioRequest) -> ScenarioRunSummary:
     """
     Start a new scenario run as a background task.
 
-    Returns immediately with a scenario_result_id that can be polled for status.
+    Initialization runs eagerly so configuration errors surface here, then the run
+    itself continues in the background. Returns a scenario_result_id that can be
+    polled for status.
 
     Args:
         request: Scenario run configuration.

pyrit/memory/memory_interface.py

@@ -10,7 +10,7 @@
 import re
 import uuid
 import weakref
-from collections.abc import Iterator, Mapping, MutableSequence, Sequence
+from collections.abc import Collection, Iterator, Mapping, MutableSequence, Sequence
 from contextlib import closing
 from dataclasses import dataclass
 from datetime import datetime, timezone
@@ -3886,6 +3886,66 @@ def update_scenario_run_state(
 
         logger.info(f"Updated scenario {scenario_result_id} state to '{scenario_run_state.value}'")
 
+    def try_update_scenario_run_state(
+        self,
+        *,
+        scenario_result_id: str,
+        expected_states: Collection[ScenarioRunState],
+        scenario_run_state: ScenarioRunState,
+        error_message: str | None = None,
+        error_type: str | None = None,
+    ) -> bool:
+        """
+        Update the run state only when the stored state is one of ``expected_states``.
+
+        The compare and the write are a single UPDATE so a run that reached a terminal state
+        on another thread is not overwritten. A read followed by
+        :meth:`update_scenario_run_state` cannot give that guarantee because scenario
+        preparation and cancellation run on different threads.
+
+        Args:
+            scenario_result_id (str): The ID of the scenario result to update.
+            expected_states (Collection[ScenarioRunState]): States the row may currently be in.
+            scenario_run_state (ScenarioRunState): The new state for the scenario.
+            error_message (str | None): Optional scenario-level error message.
+            error_type (str | None): Optional exception class name.
+
+        Returns:
+            bool: True if the row was updated, False if it was missing or in another state.
+
+        Raises:
+            ValueError: If ``expected_states`` is empty.
+        """
+        if not expected_states:
+            raise ValueError("expected_states must not be empty")
+
+        values: dict[str, Any] = {
+            "scenario_run_state": scenario_run_state.value,
+            "error_message": error_message,
+            "error_type": error_type,
+        }
+        if scenario_run_state in (
+            ScenarioRunState.COMPLETED,
+            ScenarioRunState.FAILED,
+            ScenarioRunState.CANCELLED,
+        ):
+            values["completion_time"] = datetime.now(tz=timezone.utc)
+
+        with closing(self.get_session()) as session:
+            updated_rows = (
+                session.query(ScenarioResultEntry)
+                .filter(
+                    ScenarioResultEntry.id == scenario_result_id,
+                    ScenarioResultEntry.scenario_run_state.in_([state.value for state in expected_states]),
+                )
+                .update(values, synchronize_session=False)
+            )
+            session.commit()
+
+        if updated_rows:
+            logger.info(f"Updated scenario {scenario_result_id} state to '{scenario_run_state.value}'")
+        return bool(updated_rows)
+
     def update_scenario_metadata(
         self,
         *,

pyrit/memory/sqlite_memory.py

@@ -2,6 +2,8 @@
 # Licensed under the MIT license.
 
 import logging
+import threading
+import weakref
 from collections.abc import Sequence
 from contextlib import closing
 from datetime import datetime
@@ -72,6 +74,11 @@ def __init__(
             self.db_path = Path(db_path or Path(DB_DATA_PATH, self.DEFAULT_DB_FILE_NAME)).resolve()
         self.results_path = str(DB_DATA_PATH)
 
+        # An in-memory database shares a single DBAPI connection across every thread (see
+        # ``_create_engine``), so concurrent sessions would interleave on it. Serialize session
+        # lifetimes for that backend only; file-backed databases get a connection per checkout.
+        self._connection_lock: threading.RLock | None = threading.RLock() if self.db_path == ":memory:" else None
+
         self.engine = self._create_engine(has_echo=verbose)
         self.SessionFactory = sessionmaker(bind=self.engine)
         if not skip_schema_migration:
@@ -285,10 +292,43 @@ def get_session(self) -> Session:
         """
         Provide a SQLAlchemy session for transactional operations.
 
+        For an in-memory database every session borrows the same DBAPI connection, so the
+        session is handed out under a lock that is only released when it is closed. That keeps
+        a whole transaction, not just a single statement, isolated from the other threads.
+
         Returns:
             Session: A SQLAlchemy session bound to the engine.
         """
-        return self.SessionFactory()
+        session = self.SessionFactory()
+        connection_lock = self._connection_lock
+        if connection_lock is None:
+            return session
+
+        connection_lock.acquire()
+        close_session = session.close
+        released = False
+
+        def release_once() -> None:
+            # Also runs if the session is discarded without being closed, so one caller that
+            # forgets cannot leave the lock held and stall every other thread forever.
+            nonlocal released
+            if released:
+                return
+            released = True
+            try:
+                connection_lock.release()
+            except RuntimeError:
+                logger.warning("An in-memory session was discarded by a thread that did not open it.")
+
+        def close_and_release() -> None:
+            try:
+                close_session()
+            finally:
+                release_once()
+
+        session.close = close_and_release  # type: ignore[ty:invalid-assignment]
+        weakref.finalize(session, release_once)
+        return session
 
     def print_schema(self) -> None:
         """