# Copyright (c) Microsoft Corporation. # Licensed under the MIT license. from __future__ import annotations import ast import asyncio import logging import uuid from abc import ABC, abstractmethod from contextlib import asynccontextmanager from copy import deepcopy from dataclasses import dataclass from enum import Enum from typing import TYPE_CHECKING, Any, Generic, TypeVar from pyrit.common import default_values from pyrit.common.logger import logger from pyrit.exceptions import clear_execution_context, get_execution_context from pyrit.exceptions.retry_collector import ( RetryCollector, clear_retry_collector, set_retry_collector, ) from pyrit.models import StrategyResultT if TYPE_CHECKING: from collections.abc import AsyncIterator, MutableMapping StrategyContextT = TypeVar("StrategyContextT", bound="StrategyContext") class _StrategyRuntimeError(RuntimeError): """RuntimeError subclass for strategy execution failures.""" @dataclass class StrategyContext(ABC): # noqa: B024 """Base class for all strategy contexts.""" def duplicate(self: StrategyContextT) -> StrategyContextT: """ Create a deep copy of the context. Returns: StrategyContext: A deep copy of the context. """ return deepcopy(self) class StrategyEvent(Enum): """Enumeration of all strategy lifecycle events.""" # Strategy lifecycle events # Validation Events ON_PRE_VALIDATE = "on_pre_validate" ON_POST_VALIDATE = "on_post_validate" # Setup Events ON_PRE_SETUP = "on_pre_setup" ON_POST_SETUP = "on_post_setup" # Execution Events ON_PRE_EXECUTE = "on_pre_execute" ON_POST_EXECUTE = "on_post_execute" # Teardown Events ON_PRE_TEARDOWN = "on_pre_teardown" ON_POST_TEARDOWN = "on_post_teardown" # Error Handling Events ON_ERROR = "on_error" @dataclass class StrategyEventData(Generic[StrategyContextT, StrategyResultT]): """Data passed to event observers.""" # The event that triggered this data event: StrategyEvent # Strategy information strategy_name: str strategy_id: str # Context and result of the strategy context: StrategyContextT result: StrategyResultT | None = None # Optional error if the event is related to an error error: Exception | None = None class StrategyEventHandler(ABC, Generic[StrategyContextT, StrategyResultT]): """ Abstract base class for strategy event handlers. Generic parameters: StrategyContextT: The specific context type (must be a subclass of StrategyContext) StrategyResultT: The specific result type (must be a subclass of StrategyResult) """ @abstractmethod async def on_event_async(self, event_data: StrategyEventData[StrategyContextT, StrategyResultT]) -> None: """ Handle a strategy event. Args: event_data: Data about the event that occurred. """ # ``logging.LoggerAdapter`` only became subscriptable at runtime on Python 3.11, but PyRIT # supports 3.10. Parameterize it for type checkers only and use the bare class at runtime. if TYPE_CHECKING: _LoggerAdapter = logging.LoggerAdapter[logging.Logger] else: _LoggerAdapter = logging.LoggerAdapter class StrategyLogAdapter(_LoggerAdapter): """ Custom logger adapter that adds strategy information to log messages. """ _STRATEGY_NAME_KEY = "strategy_name" _STRATEGY_ID_KEY = "strategy_id" def process(self, msg: Any, kwargs: MutableMapping[str, Any]) -> tuple[str, MutableMapping[str, Any]]: """ Add strategy context to each log message. Returns: tuple: The modified log message and keyword arguments. """ if not self.extra: return msg, kwargs # Extract strategy information from extra strategy_name = self.extra.get(self._STRATEGY_NAME_KEY) strategy_id = self.extra.get(self._STRATEGY_ID_KEY) if strategy_name and strategy_id: strategy_info = f"[{strategy_name} (ID: {str(strategy_id)})]" return f"{strategy_info} {msg}", kwargs return msg, kwargs class Strategy(ABC, Generic[StrategyContextT, StrategyResultT]): """ Abstract base class for strategies with enforced lifecycle management. Ensures a consistent execution flow: validate -> setup -> execute -> teardown. The teardown phase is guaranteed to run even if exceptions occur. Subclasses must implement: _validate_context(): Validate context _setup_async(): Initialize resources _perform_async(): Execute the logic _teardown_async(): Clean up resources """ def __init__( self, *, context_type: type[StrategyContextT], event_handler: StrategyEventHandler[StrategyContextT, StrategyResultT] | None = None, logger: logging.Logger = logger, ) -> None: """ Initialize the strategy with a context type and logger. Args: context_type (type[StrategyContextT]): The type of context this strategy will use. event_handler (StrategyEventHandler[StrategyContextT, StrategyResultT] | None): An optional event handler for strategy events. logger (logging.Logger): The logger to use for this strategy. """ self._id = uuid.uuid4() self._context_type = context_type self._event_handlers: dict[str, StrategyEventHandler[StrategyContextT, StrategyResultT]] = {} if event_handler is not None: self._register_event_handler(event_handler) self._logger = StrategyLogAdapter( logger, { StrategyLogAdapter._STRATEGY_NAME_KEY: self.__class__.__name__, StrategyLogAdapter._STRATEGY_ID_KEY: str(self._id)[:8], }, ) self._memory_labels: dict[str, str] = ast.literal_eval( default_values.get_non_required_value(env_var_name="GLOBAL_MEMORY_LABELS") or "{}" ) def _register_event_handler(self, event_handler: StrategyEventHandler[StrategyContextT, StrategyResultT]) -> None: """ Register an event handler for strategy events. Args: event_handler: The event handler to register. """ self._event_handlers[event_handler.__class__.__name__] = event_handler @abstractmethod def _validate_context(self, *, context: StrategyContextT) -> None: """ Validate the strategy context before execution. This method should be implemented by subclasses to validate that the context is suitable for the strategy. Args: context (StrategyContextT): The context to validate. Raises: Exception: If the context is invalid for this strategy. """ @abstractmethod async def _setup_async(self, *, context: StrategyContextT) -> None: """ Set up the phase before executing the strategy. This method should be implemented by subclasses to prepare any necessary state or resources. This method is guaranteed to be called before the strategy execution and after context validation. Args: context (StrategyContextT): The context for the strategy. """ @abstractmethod async def _perform_async(self, *, context: StrategyContextT) -> StrategyResultT: """ Core implementation to be defined by subclasses. This contains the actual strategy logic that subclasses must implement. Args: context (StrategyContextT): The context for the strategy. Returns: StrategyResultT: The result of the strategy execution. """ @abstractmethod async def _teardown_async(self, *, context: StrategyContextT) -> None: """ Teardown phase after executing the strategy. This method should be implemented by subclasses to clean up any resources or state. This method is guaranteed to be called even if exceptions occur during execution. Args: context (StrategyContextT): The context for the strategy. """ async def _handle_event_async( self, *, event: StrategyEvent, context: StrategyContextT, result: StrategyResultT | None = None, error: Exception | None = None, ) -> None: """ Handle a strategy event by notifying all registered event handlers. Args: event (StrategyEvent): The event that occurred. context (StrategyContextT): The context for the strategy. result (StrategyResultT | None): The result of the strategy execution, if applicable. error (Exception | None): An error that occurred during execution, if applicable. """ event_data = StrategyEventData( event=event, strategy_name=self.__class__.__name__, strategy_id=str(self._id), context=context, result=result, error=error, ) # Dispatch events to all handlers in parallel if self._event_handlers: tasks = [ asyncio.create_task(handler.on_event_async(event_data)) for handler in self._event_handlers.values() ] await asyncio.gather(*tasks, return_exceptions=True) @asynccontextmanager async def _execution_context_async(self, context: StrategyContextT) -> AsyncIterator[None]: """ Manage the complete lifecycle of a strategy execution as an async context manager. This method provides a context manager that ensures proper setup and teardown of strategy resources, regardless of whether the strategy completes successfully or encounters an error. Args: context (StrategyContextT): The execution context containing configuration and state for the strategy. Yields: None: Control is yielded back to the caller after setup is complete. """ try: # Notify pre-setup event await self._handle_event_async(event=StrategyEvent.ON_PRE_SETUP, context=context) await self._setup_async(context=context) # Notify post-setup event await self._handle_event_async(event=StrategyEvent.ON_POST_SETUP, context=context) yield finally: # Notify pre-teardown event await self._handle_event_async(event=StrategyEvent.ON_PRE_TEARDOWN, context=context) await self._teardown_async(context=context) # Notify post-teardown event await self._handle_event_async(event=StrategyEvent.ON_POST_TEARDOWN, context=context) async def execute_with_context_async(self, *, context: StrategyContextT) -> StrategyResultT: """ Execute strategy with complete lifecycle management. Enforces: validate -> setup -> execute -> teardown. Args: context (StrategyContextT): The context for the strategy, containing configuration and state. Returns: StrategyResultT: The result of the strategy execution, including outcome and reason. Raises: ValueError: If the context validation fails. RuntimeError: If the strategy execution fails. """ # Validation phase # This is a critical step to ensure the context is suitable for the strategy try: # Notify pre-validation event await self._handle_event_async(event=StrategyEvent.ON_PRE_VALIDATE, context=context) self._validate_context(context=context) # Notify post-validation event await self._handle_event_async(event=StrategyEvent.ON_POST_VALIDATE, context=context) except Exception as e: raise ValueError(f"Strategy context validation failed for {self.__class__.__name__}: {str(e)}") from e # Execution with lifecycle management # This uses an async context manager to ensure setup and teardown are handled correctly retry_collector_started = False try: async with self._execution_context_async(context): await self._handle_event_async(event=StrategyEvent.ON_PRE_EXECUTE, context=context) # Set up RetryCollector in the parent task so it is visible to # Tenacity callbacks that fire during _perform_async. Event # handlers run in child tasks (asyncio.create_task) which # inherit a *copy* of the parent's ContextVar — setting the # collector here ensures both the execution path and the event # handlers can see it. collector = RetryCollector() set_retry_collector(collector) retry_collector_started = True result = await self._perform_async(context=context) await self._handle_event_async(event=StrategyEvent.ON_POST_EXECUTE, context=context, result=result) return result except Exception as e: # Notify error event await self._handle_event_async(event=StrategyEvent.ON_ERROR, context=context, error=e) # Build enhanced error message with execution context if available # Note: The context is preserved on exception by ExecutionContextManager exec_context = get_execution_context() if exec_context: error_details = exec_context.get_exception_details() # Extract the root cause exception for better diagnostics root_cause: BaseException = e while root_cause.__cause__ is not None: root_cause = root_cause.__cause__ # Include root cause type and message if different from the immediate exception if root_cause is not e: root_cause_info = f"\n\nRoot cause: {type(root_cause).__name__}: {str(root_cause)}" else: root_cause_info = "" error_message = ( f"Strategy execution failed for {exec_context.component_role.value} " f"in {self.__class__.__name__}: {str(e)}{root_cause_info}\n\nDetails:\n{error_details}" ) # Clear the context now that we've read it clear_execution_context() else: error_message = f"Strategy execution failed for {self.__class__.__name__}: {str(e)}" runtime_error = _StrategyRuntimeError(error_message) raise runtime_error from e finally: if retry_collector_started: clear_retry_collector() async def execute_async(self, **kwargs: Any) -> StrategyResultT: """ Execute the strategy asynchronously with the given keyword arguments. Returns: StrategyResultT: The result of the strategy execution. """ context = self._context_type(**kwargs) return await self.execute_with_context_async(context=context)