FIX: Reject concurrent SDK scenario runs and make attempt accounting atomic (#3748)

pyrit/memory/memory_interface.py

@@ -7158,6 +7158,32 @@ def _execute_update_scenario_metadata_fields(
             flag_modified(entry, "scenario_metadata")
             session.commit()
 
+    def _execute_increment_scenario_number_tries(
+        self,
+        *,
+        scenario_result_id: str,
+    ) -> int:
+        """
+        Atomically increment and return the number_tries counter for a scenario result.
+
+        Args:
+            scenario_result_id (str): The ID of the scenario result.
+
+        Returns:
+            int: The new number_tries value after incrementing.
+
+        Raises:
+            ValueError: If the scenario result is not found in memory.
+        """
+        with closing(self._get_session()) as session:
+            entry = session.query(ScenarioResultEntry).filter_by(id=scenario_result_id).first()
+            if not entry:
+                raise ValueError(f"Scenario result with ID {scenario_result_id} not found in memory")
+            entry.number_tries += 1
+            new_tries = entry.number_tries
+            session.commit()
+            return new_tries
+
     def _execute_get_scenario_result_header(self, *, scenario_result_id: str) -> ScenarioResult | None:
         """Return one ScenarioResult header without hydrating linked attack results."""
         with closing(self._get_session()) as session:
@@ -11747,6 +11773,51 @@ async def update_scenario_metadata_fields_async(
             self._execute_update_scenario_metadata_fields, scenario_result_id=scenario_result_id, fields=fields
         )
 
+    def increment_scenario_number_tries(
+        self,
+        *,
+        scenario_result_id: str,
+    ) -> int:
+        """
+        Use ``increment_scenario_number_tries_async``.
+
+        This synchronous API is deprecated and can block the caller.
+
+        Returns:
+            int: The new number_tries value after incrementing.
+        """
+        print_deprecation_message(
+            old_item="MemoryInterface.increment_scenario_number_tries",
+            new_item="MemoryInterface.increment_scenario_number_tries_async",
+            removed_in="1.4.0",
+        )
+        return self._execute_increment_scenario_number_tries(
+            scenario_result_id=scenario_result_id,
+        )
+
+    @legacy_sync_override(lambda: MemoryInterface.increment_scenario_number_tries)
+    async def increment_scenario_number_tries_async(
+        self,
+        *,
+        scenario_result_id: str,
+    ) -> int:
+        """
+        Atomically increment and return the number_tries counter for a scenario result.
+
+        Args:
+            scenario_result_id (str): The ID of the scenario result to update.
+
+        Returns:
+            int: The new number_tries value after incrementing.
+
+        Raises:
+            ValueError: If the scenario result is not found.
+        """
+        return await self._run_database_operation_async(
+            self._execute_increment_scenario_number_tries,
+            scenario_result_id=scenario_result_id,
+        )
+
     def get_scenario_result_header(self, *, scenario_result_id: str) -> ScenarioResult | None:
         """
         Use ``get_scenario_result_header_async``.

pyrit/scenario/core/scenario.py

@@ -11,6 +11,7 @@
 import ast
 import asyncio
 import logging
+import threading
 import uuid
 from abc import ABC, abstractmethod
 from collections.abc import Mapping, Sequence
@@ -163,6 +164,16 @@ class Scenario(ABC):
     #: Maximum number of resolved concrete techniques; None leaves the selection unrestricted.
     MAX_CONCRETE_TECHNIQUES: ClassVar[int | None] = None
 
+    #: Active scenario result IDs currently executing in this process.
+    _active_runs: ClassVar[set[str]] = set()
+    _active_runs_lock: ClassVar[threading.Lock] = threading.Lock()
+
+    @classmethod
+    def _clear_active_runs(cls) -> None:
+        """Clear active scenario runs. Intended for test isolation."""
+        with cls._active_runs_lock:
+            cls._active_runs.clear()
+
     def __init_subclass__(cls, **kwargs: Any) -> None:
         """
         Enforce the keyword-only constructor contract on subclasses.
@@ -1705,57 +1716,68 @@ async def run_async(self) -> ScenarioResult:
         # Type narrowing: create local variable that type checker knows is non-None
         scenario_result_id: str = self._scenario_result_id
 
-        # Implement retry logic
-        last_exception = None
-        for retry_attempt in range(self._max_retries + 1):  # +1 for initial attempt
-            try:
-                return await self._execute_scenario_async()
-            except asyncio.CancelledError:
+        with Scenario._active_runs_lock:
+            if scenario_result_id in Scenario._active_runs:
+                raise RuntimeError(
+                    f"Scenario '{self._name}' execution is already in progress for run '{scenario_result_id}'."
+                )
+            Scenario._active_runs.add(scenario_result_id)
+
+        try:
+            # Implement retry logic
+            last_exception = None
+            for retry_attempt in range(self._max_retries + 1):  # +1 for initial attempt
                 try:
-                    (
-                        await self._memory.update_scenario_run_state_async(
-                            scenario_result_id=scenario_result_id,
-                            scenario_run_state=ScenarioRunState.CANCELLED,
-                            error_message="Scenario run was cancelled",
-                            error_type="CancelledError",
+                    return await self._execute_scenario_async()
+                except asyncio.CancelledError:
+                    try:
+                        (
+                            await self._memory.update_scenario_run_state_async(
+                                scenario_result_id=scenario_result_id,
+                                scenario_run_state=ScenarioRunState.CANCELLED,
+                                error_message="Scenario run was cancelled",
+                                error_type="CancelledError",
+                            )
                         )
+                    except Exception:
+                        logger.exception(f"Failed to persist cancellation state for scenario '{self._name}'")
+                    raise
+                except Exception as e:
+                    last_exception = e
+
+                    # Get current scenario to check number of tries
+                    scenario_results = await self._memory.get_scenario_results_async(
+                        scenario_result_ids=[scenario_result_id]
                     )
-                except Exception:
-                    logger.exception(f"Failed to persist cancellation state for scenario '{self._name}'")
-                raise
-            except Exception as e:
-                last_exception = e
-
-                # Get current scenario to check number of tries
-                scenario_results = await self._memory.get_scenario_results_async(
-                    scenario_result_ids=[scenario_result_id]
-                )
-                current_tries = scenario_results[0].number_tries if scenario_results else retry_attempt + 1
+                    current_tries = scenario_results[0].number_tries if scenario_results else retry_attempt + 1
 
-                # Check if we have more retries available
-                remaining_retries = self._max_retries - retry_attempt
+                    # Check if we have more retries available
+                    remaining_retries = self._max_retries - retry_attempt
 
-                if remaining_retries > 0:
+                    if remaining_retries > 0:
+                        logger.error(
+                            f"Scenario '{self._name}' failed on attempt {current_tries} with error: {str(e)}. "
+                            f"Retrying... ({remaining_retries} retries remaining)",
+                            exc_info=True,
+                        )
+                        # Continue to next iteration for retry
+                        continue
+                    # No more retries, log final failure
                     logger.error(
-                        f"Scenario '{self._name}' failed on attempt {current_tries} with error: {str(e)}. "
-                        f"Retrying... ({remaining_retries} retries remaining)",
+                        f"Scenario '{self._name}' failed after {current_tries} attempts "
+                        f"(initial + {self._max_retries} retries) with error: {str(e)}. Giving up.",
                         exc_info=True,
                     )
-                    # Continue to next iteration for retry
-                    continue
-                # No more retries, log final failure
-                logger.error(
-                    f"Scenario '{self._name}' failed after {current_tries} attempts "
-                    f"(initial + {self._max_retries} retries) with error: {str(e)}. Giving up.",
-                    exc_info=True,
-                )
-                (await self._mark_scenario_failed_async(scenario_result_id=scenario_result_id, error=e))
-                raise
+                    (await self._mark_scenario_failed_async(scenario_result_id=scenario_result_id, error=e))
+                    raise
 
-        # This should never be reached, but just in case
-        if last_exception:
-            raise last_exception
-        raise RuntimeError(f"Scenario '{self._name}' completed unexpectedly without result")
+            # This should never be reached, but just in case
+            if last_exception:
+                raise last_exception
+            raise RuntimeError(f"Scenario '{self._name}' completed unexpectedly without result")
+        finally:
+            with Scenario._active_runs_lock:
+                Scenario._active_runs.discard(scenario_result_id)
 
     async def _execute_scenario_async(self) -> ScenarioResult:
         """
@@ -1781,15 +1803,9 @@ async def _execute_scenario_async(self) -> ScenarioResult:
             raise ValueError("self._scenario_result_id is not initialized")
         scenario_result_id: str = self._scenario_result_id
 
-        # Increment number_tries at the start of each run
-        scenario_results = await self._memory.get_scenario_results_async(scenario_result_ids=[scenario_result_id])
-        if scenario_results:
-            current_scenario = scenario_results[0]
-            current_scenario.number_tries += 1
-            await self._memory.update_scenario_result_async(scenario_result=current_scenario)
-            logger.info(f"Scenario '{self._name}' attempt #{current_scenario.number_tries}")
-        else:
-            raise ValueError(f"Scenario result with ID {scenario_result_id} not found")
+        # Increment number_tries atomically at the start of each run
+        current_tries = await self._memory.increment_scenario_number_tries_async(scenario_result_id=scenario_result_id)
+        logger.info(f"Scenario '{self._name}' attempt #{current_tries}")
 
         # Mark scenario as in progress
         (

tests/unit/memory/memory_interface/test_interface_scenario_results.py

@@ -1,11 +1,17 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT license.
 
+import asyncio
 from datetime import UTC, datetime, timedelta
 from uuid import uuid4
 
 import pytest
-from unit.mocks import get_mock_scorer_identifier, make_scenario_result, run_memory_session_async
+from unit.mocks import (
+    get_mock_scorer_identifier,
+    get_mock_target_identifier,
+    make_scenario_result,
+    run_memory_session_async,
+)
 
 from pyrit.memory import MemoryInterface
 from pyrit.memory.memory_models import (
@@ -1172,3 +1178,70 @@ async def test_add_scenario_result_without_objective_target_raises(sqlite_instan
         (await sqlite_instance.add_scenario_results_to_memory_async(scenario_results=[scenario]))
 
     assert (await sqlite_instance.get_scenario_results_async(scenario_name="Targetless Scenario")) == []
+
+
+async def test_increment_scenario_number_tries_async(sqlite_instance: MemoryInterface):
+    """increment_scenario_number_tries_async atomically increments and returns the new try count."""
+    scenario = make_scenario_result(
+        scenario_name="Tries Scenario",
+        scenario_version=1,
+        objective_target_identifier=get_mock_target_identifier(),
+        objective_scorer_identifier=get_mock_scorer_identifier(),
+        attack_results={},
+    )
+    await sqlite_instance.add_scenario_results_to_memory_async(scenario_results=[scenario])
+    scenario_id = str(scenario.id)
+
+    first = await sqlite_instance.increment_scenario_number_tries_async(scenario_result_id=scenario_id)
+    assert first == 1
+
+    second = await sqlite_instance.increment_scenario_number_tries_async(scenario_result_id=scenario_id)
+    assert second == 2
+
+    stored = (await sqlite_instance.get_scenario_results_async(scenario_result_ids=[scenario_id]))[0]
+    assert stored.number_tries == 2
+
+
+async def test_increment_scenario_number_tries_sync(sqlite_instance: MemoryInterface):
+    """Synchronous increment_scenario_number_tries works and updates try count."""
+    scenario = make_scenario_result(
+        scenario_name="Sync Tries Scenario",
+        scenario_version=1,
+        objective_target_identifier=get_mock_target_identifier(),
+        objective_scorer_identifier=get_mock_scorer_identifier(),
+        attack_results={},
+    )
+    await sqlite_instance.add_scenario_results_to_memory_async(scenario_results=[scenario])
+    scenario_id = str(scenario.id)
+
+    with pytest.deprecated_call():
+        first = sqlite_instance.increment_scenario_number_tries(scenario_result_id=scenario_id)
+    assert first == 1
+
+
+async def test_increment_scenario_number_tries_missing_id_raises(sqlite_instance: MemoryInterface):
+    """increment_scenario_number_tries_async raises ValueError when scenario_result_id does not exist."""
+    missing_id = str(uuid4())
+    with pytest.raises(ValueError, match=f"Scenario result with ID {missing_id} not found in memory"):
+        await sqlite_instance.increment_scenario_number_tries_async(scenario_result_id=missing_id)
+
+
+async def test_increment_scenario_number_tries_concurrent(sqlite_instance: MemoryInterface):
+    """Concurrent increments are serialized by the transaction and never lose counts."""
+    scenario = make_scenario_result(
+        scenario_name="Concurrent Tries Scenario",
+        scenario_version=1,
+        objective_target_identifier=get_mock_target_identifier(),
+        objective_scorer_identifier=get_mock_scorer_identifier(),
+        attack_results={},
+    )
+    await sqlite_instance.add_scenario_results_to_memory_async(scenario_results=[scenario])
+    scenario_id = str(scenario.id)
+
+    results = await asyncio.gather(
+        *[sqlite_instance.increment_scenario_number_tries_async(scenario_result_id=scenario_id) for _ in range(10)]
+    )
+    assert sorted(results) == list(range(1, 11))
+
+    stored = (await sqlite_instance.get_scenario_results_async(scenario_result_ids=[scenario_id]))[0]
+    assert stored.number_tries == 10