Merge repaired #2376 parent into #2377

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
Copilot-Session: 5d02c2d5-b499-4f78-a04d-03bffa750817

Author: romanlutz

pyproject.toml

@@ -222,6 +222,12 @@ include = ["pyrit/prompt_target/hugging_face/**"]
 [tool.ty.overrides.rules]
 invalid-argument-type = "ignore"
 
+# Historical Alembic revisions are immutable; retain their defensive rowcount checks.
+[[tool.ty.overrides]]
+include = ["pyrit/memory/alembic/versions/1b3d5f7a9c2e_persist_scored_expectation.py"]
+[tool.ty.overrides.rules]
+redundant-condition-strict = "ignore"
+
 # Unit tests intentionally exercise runtime validation failures, mutable test doubles,
 # and optional fields after fixtures have populated them. Keep library code strict while
 # avoiding hundreds of per-assertion ignores in tests.

pyrit/auth/azure_auth.py

@@ -206,9 +206,9 @@ async def async_token_provider() -> str:  # pyrit-async-suffix-exempt
             str: The token string from the synchronous provider.
         """
         result = api_key()
-        if inspect.isawaitable(result):
-            return await result  # type: ignore[ty:invalid-return-type]
-        return result
+        if isinstance(result, str):
+            return result
+        return await result
 
     return async_token_provider
 

pyrit/backend/mappers/attack_mappers.py

@@ -262,15 +262,23 @@ def _resolve_summary_timestamps(ar: AttackResult) -> tuple[datetime, datetime]:
     Returns:
         A ``(created_at, updated_at)`` tuple.
     """
-    created_str = ar.metadata.get("created_at")
+    return _resolve_timestamps(created_str=ar.metadata.get("created_at"), timestamp=ar.timestamp)
+
+
+def _resolve_timestamps(*, created_str: str | None, timestamp: datetime | None) -> tuple[datetime, datetime]:
+    """
+    Resolve display times, retaining fallbacks for unpersisted mutable results.
+
+    Returns:
+        tuple[datetime, datetime]: Creation and last-update timestamps.
+    """
     if created_str:
         created_at = datetime.fromisoformat(created_str)
-    elif ar.timestamp is not None:
-        created_at = ar.timestamp
+    elif timestamp is not None:
+        created_at = timestamp
     else:
         created_at = datetime.now(UTC)
-    updated_at = ar.timestamp if ar.timestamp is not None else created_at
-    return created_at, updated_at
+    return created_at, timestamp if timestamp is not None else created_at
 
 
 async def _summary_last_response_async(piece: MessagePiece | None) -> MessagePieceView | None:

pyrit/backend/services/configuration_file_service.py

@@ -8,7 +8,7 @@
 import os
 import tempfile
 from collections.abc import AsyncGenerator
-from contextlib import asynccontextmanager
+from contextlib import ExitStack, asynccontextmanager
 from pathlib import Path
 from urllib.parse import urlparse
 
@@ -103,9 +103,8 @@ def _validate_configuration_content(content: str) -> None:
 
 def _replace_local_config_file(*, path: Path, content: str) -> None:
     """Atomically replace a local configuration file."""
-    temporary_path: Path | None = None
-    try:
-        path.parent.mkdir(parents=True, exist_ok=True)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    with ExitStack() as cleanup:
         with tempfile.NamedTemporaryFile(
             mode="w",
             encoding="utf-8",
@@ -114,12 +113,10 @@ def _replace_local_config_file(*, path: Path, content: str) -> None:
             suffix=".tmp",
             delete=False,
         ) as temporary_file:
-            temporary_file.write(content)
             temporary_path = Path(temporary_file.name)
+            cleanup.callback(temporary_path.unlink, missing_ok=True)
+            temporary_file.write(content)
         os.replace(temporary_path, path)
-    finally:
-        if temporary_path is not None:
-            temporary_path.unlink(missing_ok=True)
 
 
 class ConfigurationFileService:

pyrit/backend/services/scenario_progress_read_model.py

@@ -364,10 +364,19 @@ def calculate_progress_counts(
     @staticmethod
     def _result_order_key(attack_result: AttackResult) -> tuple[datetime, str]:
         """Return a deterministic chronological key for one hydrated result attempt."""
-        timestamp = attack_result.timestamp
-        if not isinstance(timestamp, datetime):
-            timestamp = datetime.min.replace(tzinfo=UTC)
-        return timestamp, str(attack_result.attack_result_id)
+        return ScenarioProgressReadModel._timestamp_order_key(attack_result.timestamp), str(
+            attack_result.attack_result_id
+        )
+
+    @staticmethod
+    def _timestamp_order_key(timestamp: object) -> datetime:
+        """
+        Normalize potentially malformed timestamps from mutable result objects.
+
+        Returns:
+            datetime: The timestamp or a stable earliest-time fallback.
+        """
+        return timestamp if isinstance(timestamp, datetime) else datetime.min.replace(tzinfo=UTC)
 
     @staticmethod
     def total_retry_pressure(*, attempts_per_unit: Iterable[int], persisted_retries: Iterable[int]) -> int:

pyrit/backend/services/scenario_run_service.py

@@ -893,10 +893,14 @@ def _schedule_terminalization_retry(self, *, active: _ActiveTask) -> None:
         self._handoff_retry_tasks.add(retry_task)
         retry_task.add_done_callback(self._handoff_retry_tasks.discard)
 
+    def _can_retry_active_run(self, *, scenario_result_id: str) -> bool:
+        """Return whether retry work may continue for the active run."""
+        return not self._stopping and self._active_scenario_result_id == scenario_result_id
+
     async def _retry_handoff_async(self, *, scenario_result_id: str) -> None:
         """Retry scheduler handoff with bounded exponential delay until it succeeds or shutdown begins."""
         delay = _SCHEDULER_RETRY_INITIAL_SECONDS
-        while not self._stopping and self._active_scenario_result_id == scenario_result_id:
+        while self._can_retry_active_run(scenario_result_id=scenario_result_id):
             await asyncio.sleep(delay)
             try:
                 await self._handoff_scheduler_async(scenario_result_id=scenario_result_id)
@@ -909,11 +913,11 @@ async def _retry_handoff_async(self, *, scenario_result_id: str) -> None:
     async def _retry_terminalization_async(self, *, active: _ActiveTask) -> None:
         """Retry a failed cancellation transition, then perform the terminal handoff."""
         delay = _SCHEDULER_RETRY_INITIAL_SECONDS
-        while not self._stopping and self._active_scenario_result_id == active.scenario_result_id:
+        while self._can_retry_active_run(scenario_result_id=active.scenario_result_id):
             await asyncio.sleep(delay)
             try:
                 async with self._scheduler_lock:
-                    if self._stopping or self._active_scenario_result_id != active.scenario_result_id:
+                    if not self._can_retry_active_run(scenario_result_id=active.scenario_result_id):
                         return
                     await asyncio.to_thread(
                         self._memory.try_update_scenario_run_state,
@@ -1447,6 +1451,16 @@ def _load_started_at(*, scenario_result: ScenarioResult) -> datetime | None:
             return None
         return started_at if started_at.tzinfo is not None else None
 
+    @staticmethod
+    def _identifier_techniques(scenario_identifier: ScenarioIdentifier | None) -> list[str]:
+        """
+        Read techniques when legacy persisted metadata has an identifier.
+
+        Returns:
+            list[str]: Stored techniques or an empty list.
+        """
+        return list(scenario_identifier.techniques or []) if scenario_identifier is not None else []
+
     @staticmethod
     def _safe_run_metadata(
         *,
@@ -1793,10 +1807,8 @@ def get_run_progress_from_storage(
         target, datasets_used, scenario_parameters = self._safe_run_metadata(scenario_identifier=scenario_identifier)
         if plan is not None:
             techniques_used = list(dict.fromkeys(group.display_group for group in plan.atomic_groups))
-        elif scenario_identifier is not None:
-            techniques_used = list(scenario_identifier.techniques or [])
         else:
-            techniques_used = []
+            techniques_used = self._identifier_techniques(scenario_identifier)
         return ScenarioRunProgress(
             run=ScenarioProgressHeader(
                 scenario_result_id=scenario_result_id,

pyrit/cli/_server_launcher.py

@@ -825,18 +825,7 @@ async def start_async(
             _logger.info("Backend launcher PID: %d (logs: %s)", process.pid, self._log_path)
             await self._wait_for_readiness_async(plan=plan, host=host, startup_state=startup_state)
         finally:
-            if startup_state.phase is not _StartupPhase.READY and startup_state.process is not None:
-                if self._process is None:
-                    self._process = startup_state.process
-                    self._listener_pid = startup_state.process.pid
-                    self._port = port if startup_state.pid_record_written else None
-                try:
-                    await startup_state.cleanup_async()
-                finally:
-                    if startup_state.cleanup_succeeded:
-                        self._clear_process_state()
-                    elif startup_state.cleanup_succeeded is False:
-                        _logger.warning("Failed to stop backend launcher process %d", startup_state.process.pid)
+            await self._cleanup_failed_startup_async(startup_state=startup_state, port=port)
 
         if startup_state.timed_out:
             cleanup_message = (
@@ -856,6 +845,22 @@ async def start_async(
             )
         return plan.base_url
 
+    async def _cleanup_failed_startup_async(self, *, startup_state: _ServerStartupState, port: int) -> None:
+        """Clean up even when spawning raised before the launcher stored its process."""
+        if startup_state.phase is _StartupPhase.READY or startup_state.process is None:
+            return
+        if self._process is None:
+            self._process = startup_state.process
+            self._listener_pid = startup_state.process.pid
+            self._port = port if startup_state.pid_record_written else None
+        try:
+            await startup_state.cleanup_async()
+        finally:
+            if startup_state.cleanup_succeeded:
+                self._clear_process_state()
+            elif startup_state.cleanup_succeeded is False:
+                _logger.warning("Failed to stop backend launcher process %d", startup_state.process.pid)
+
     def _clear_process_state(self) -> None:
         """Clear process state after the owned backend exits."""
         self._process = None

pyrit/common/apply_defaults.py

@@ -293,7 +293,7 @@ def wrapper(self: object, *args: object, **kwargs: object) -> T:
                         f"Either pass the parameter explicitly or register a default using set_default_value()."
                     )
                 # If None was explicitly passed and parameter has REQUIRED_VALUE as default, also raise
-                elif param_value is None:
+                else:
                     # Check if the parameter's default in the signature is REQUIRED_VALUE
                     param_obj = sig.parameters.get(param_name)
                     if param_obj and isinstance(param_obj.default, _RequiredValueSentinel):

pyrit/common/net_utility.py

@@ -53,7 +53,7 @@ def extract_url_parameters(url: str) -> dict[str, str]:
     parsed_url = urlparse(url)
     url_params = parse_qs(parsed_url.query, keep_blank_values=True)
     # Flatten params (parse_qs returns lists)
-    return {k: v[0] if isinstance(v, list) and len(v) > 0 else "" for k, v in url_params.items()}
+    return {k: v[0] if v else "" for k, v in url_params.items()}
 
 
 def remove_url_parameters(url: str) -> str:

pyrit/common/text_helper.py

@@ -1,7 +1,17 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT license.
 
-from typing import IO, Any
+from typing import IO, Any, TypeGuard
+
+
+def is_non_empty_string(value: object) -> TypeGuard[str]:
+    """
+    Check whether an untrusted value is a string containing non-whitespace text.
+
+    Returns:
+        bool: Whether the value is a non-empty string.
+    """
+    return isinstance(value, str) and bool(value.strip())
 
 
 def read_txt(file: IO[Any]) -> list[dict[str, str]]: