diff --git a/pyproject.toml b/pyproject.toml index 6be2b689f2..865df16a7b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -135,12 +135,17 @@ litellm = [ "litellm>=1.84.0", ] +github-copilot = [ + "github-copilot-sdk>=1.0.11", +] + # all includes all functional dependencies excluding the ones from the "dev" dependency group all = [ "accelerate>=1.7.0", "azure-ai-ml>=1.32.0", "azure-cognitiveservices-speech>=1.44.0", "flask>=3.1.3", + "github-copilot-sdk>=1.0.11", "ipykernel>=6.29.5", "jupyter>=1.1.1", "litellm>=1.84.0", diff --git a/pyrit/prompt_target/__init__.py b/pyrit/prompt_target/__init__.py index 999c50834a..664c5185a9 100644 --- a/pyrit/prompt_target/__init__.py +++ b/pyrit/prompt_target/__init__.py @@ -31,6 +31,7 @@ from pyrit.prompt_target.common.target_requirements import CHAT_TARGET_REQUIREMENTS, TargetRequirements from pyrit.prompt_target.common.utils import limit_requests_per_minute from pyrit.prompt_target.gandalf_target import GandalfLevel, GandalfTarget + from pyrit.prompt_target.github_copilot_target import GitHubCopilotTarget from pyrit.prompt_target.http_target.http_target import HTTPTarget from pyrit.prompt_target.http_target.http_target_callback_functions import ( get_http_target_json_response_callback_function, @@ -66,6 +67,7 @@ "ConversationNormalizationPipeline": "pyrit.prompt_target.common.conversation_normalization_pipeline", "GandalfLevel": "pyrit.prompt_target.gandalf_target", "GandalfTarget": "pyrit.prompt_target.gandalf_target", + "GitHubCopilotTarget": "pyrit.prompt_target.github_copilot_target", "get_http_target_json_response_callback_function": "pyrit.prompt_target.http_target.http_target_callback_functions", "get_http_target_regex_matching_callback_function": ( "pyrit.prompt_target.http_target.http_target_callback_functions" diff --git a/pyrit/prompt_target/common/prompt_target.py b/pyrit/prompt_target/common/prompt_target.py index f28299bcc3..6789c11018 100644 --- a/pyrit/prompt_target/common/prompt_target.py +++ b/pyrit/prompt_target/common/prompt_target.py @@ -334,13 +334,16 @@ def set_system_prompt( conversation_id (str): The conversation id to attach the prompt to. Raises: - ValueError: If the target does not support multi-turn or editable history. + ValueError: If the target does not support multi-turn conversations, or + supports neither editable history nor native system prompts. RuntimeError: If the conversation already has messages. """ - if not self.capabilities.supports_multi_turn or not self.capabilities.supports_editable_history: + if not self.capabilities.supports_multi_turn or not ( + self.capabilities.supports_editable_history or self.capabilities.supports_system_prompt + ): raise ValueError( f"Target {type(self).__name__} does not support setting a system prompt. " - "It must support both multi-turn conversations and editable history." + "It must support multi-turn conversations and either editable history or native system prompts." ) messages = self._memory.get_conversation_messages(conversation_id=conversation_id) diff --git a/pyrit/prompt_target/github_copilot_target.py b/pyrit/prompt_target/github_copilot_target.py new file mode 100644 index 0000000000..96f54e98b4 --- /dev/null +++ b/pyrit/prompt_target/github_copilot_target.py @@ -0,0 +1,438 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +import asyncio +import logging +import math +from dataclasses import dataclass, field +from pathlib import Path +from typing import TYPE_CHECKING +from uuid import uuid4 + +from pyrit.models import ComponentIdentifier, Message, construct_response_from_request +from pyrit.prompt_target.common.prompt_target import PromptTarget +from pyrit.prompt_target.common.target_capabilities import TargetCapabilities +from pyrit.prompt_target.common.target_configuration import TargetConfiguration +from pyrit.prompt_target.common.utils import limit_requests_per_minute + +logger = logging.getLogger(__name__) + +if TYPE_CHECKING: + from copilot import CopilotClient, CopilotSession, GetStatusResponse, SystemMessageConfig + + +@dataclass +class _ConversationState: + lock: asyncio.Lock = field(default_factory=asyncio.Lock) + session: "CopilotSession | None" = None + retired: bool = False + + +class GitHubCopilotTarget(PromptTarget): + """ + Send text requests through the GitHub Copilot SDK. + + Capture INFO logs for session mapping and SDK/runtime version diagnostics. + """ + + _DEFAULT_CONFIGURATION: TargetConfiguration = TargetConfiguration( + capabilities=TargetCapabilities(supports_multi_turn=True, supports_system_prompt=True) + ) + + def __init__( + self, + *, + model_name: str, + github_token: str | None = None, + working_directory: str | Path | None = None, + retain_session: bool = False, + response_timeout_seconds: float = 60.0, + max_requests_per_minute: int | None = None, + ) -> None: + """ + Initialize the target with an explicit token or normal SDK login discovery. + + Args: + model_name (str): Explicit Copilot model ID. + github_token (str | None): Nonblank GitHub token, forwarded unchanged with precedence over other + authentication methods. Defaults to None for SDK environment-token and saved-login discovery. + working_directory (str | Path | None): Existing local directory; no Git repository is required. + Supplied paths are resolved once against the current directory at construction and included + in saved target identifiers. Defaults to None for the SDK's current directory at each client start. + retain_session (bool): Keep each Copilot session on disk instead of deleting it. Retained + session IDs are logged so they can be found later. Defaults to False. + response_timeout_seconds (float): Shared time budget for dispatch and completion, in seconds. + Excludes client/session creation, startup status lookup, and cleanup. Defaults to 60. + max_requests_per_minute (int | None): PyRIT per-send pacing. Positive values delay each send by + 60 / value seconds before SDK client creation, outside the response deadline. + None or nonpositive values disable pacing. Defaults to None. + + Raises: + ValueError: If model_name or a supplied github_token is blank, or response_timeout_seconds + is not finite and positive, or a supplied working_directory is blank, missing, or not a directory. + OSError: If the working directory cannot be resolved or inspected. + RuntimeError: If the optional GitHub Copilot SDK is not installed. + """ + if not model_name.strip(): + raise ValueError("model_name must not be empty.") + if github_token is not None and not github_token.strip(): + raise ValueError("github_token must not be blank when supplied.") + if not math.isfinite(response_timeout_seconds) or response_timeout_seconds <= 0: + raise ValueError("response_timeout_seconds must be a finite positive number.") + + self._working_directory: str | None = None + if working_directory is not None: + if isinstance(working_directory, str) and not working_directory.strip(): + raise ValueError("working_directory must not be blank when supplied.") + resolved_directory = Path(working_directory).resolve() + if not resolved_directory.is_dir(): + raise ValueError("working_directory must be an existing directory.") + self._working_directory = str(resolved_directory) + + try: + import copilot + except ModuleNotFoundError as e: + raise RuntimeError("Could not import copilot. Install it with 'pip install pyrit[github-copilot]'.") from e + + super().__init__(model_name=model_name, max_requests_per_minute=max_requests_per_minute) + self._sdk = copilot + self._github_token = github_token + self._retain_session = retain_session + self._response_timeout_seconds = response_timeout_seconds + self._client: CopilotClient | None = None + self._runtime_status: GetStatusResponse | None = None + self._client_start_lock = asyncio.Lock() + self._lifecycle_condition = asyncio.Condition() + self._cleanup_task: asyncio.Task[None] | None = None + self._conversations: dict[str, _ConversationState] = {} + self._active_target_operations = 0 + + def _build_identifier(self) -> ComponentIdentifier: + """ + Build the identifier with the selected working directory. + + Returns: + ComponentIdentifier: The target identifier, including the resolved selected path when provided. + """ + return self._create_identifier(params={"working_directory": self._working_directory}) + + def set_model_name(self, *, model_name: str) -> None: + """ + Set the model before identity capture; afterward, create a new target to change it. + + Args: + model_name (str): The nonblank Copilot model ID. + + Raises: + ValueError: If model_name is blank. + RuntimeError: If the target identity has already been captured and model_name differs. + """ + if not model_name.strip(): + raise ValueError("model_name must not be empty.") + if self._identifier is not None and model_name != self._model_name: + raise RuntimeError("model_name is frozen after identity capture; create a new target to change it.") + super().set_model_name(model_name=model_name) + + @limit_requests_per_minute + async def _send_prompt_to_target_async(self, *, normalized_conversation: list[Message]) -> list[Message]: + self.get_identifier() + async with self._lifecycle_condition: + if self._cleanup_task is not None: + raise RuntimeError("GitHubCopilotTarget has been cleaned up and cannot send more prompts.") + self._active_target_operations += 1 + try: + request = normalized_conversation[-1].get_piece() + conversation_id = request.conversation_id or "" + conversation = self._conversations.setdefault(conversation_id, _ConversationState()) + async with conversation.lock: + async with self._lifecycle_condition: + if self._cleanup_task is not None: + raise RuntimeError("GitHubCopilotTarget has been cleaned up and cannot send more prompts.") + initial_system_prompt: str | None = None + if normalized_conversation[0].api_role == "system": + initial_system_prompt = "\n\n".join( + piece.converted_value for piece in normalized_conversation[0].message_pieces + ) + session = await self._get_or_create_session_async( + conversation_id=conversation_id, + initial_system_prompt=initial_system_prompt, + ) + try: + reply_text = await self._send_text_async( + session=session, + prompt=request.converted_value, + ) + except BaseException: + retirement_task = asyncio.create_task(self._retire_conversation_async(conversation=conversation)) + try: + await asyncio.shield(retirement_task) + except asyncio.CancelledError: + await retirement_task + raise + raise + return [construct_response_from_request(request=request, response_text_pieces=[reply_text])] + finally: + async with self._lifecycle_condition: + self._active_target_operations -= 1 + if self._active_target_operations == 0: + self._lifecycle_condition.notify_all() + + async def _send_text_async(self, *, session: "CopilotSession", prompt: str) -> str: + from copilot.generated.session_events import ( + AbortData, + AssistantMessageData, + SessionEvent, + SessionIdleData, + ToolExecutionStartData, + ) + + aborted = False + tool_execution_started = False + + def _record_turn_events(event: SessionEvent) -> None: + nonlocal aborted, tool_execution_started + if isinstance(event.data, AbortData) or (isinstance(event.data, SessionIdleData) and event.data.aborted): + aborted = True + if isinstance(event.data, ToolExecutionStartData): + tool_execution_started = True + + unsubscribe = session.on(_record_turn_events) + try: + reply = await asyncio.wait_for( + session.send_and_wait(prompt, timeout=self._response_timeout_seconds), + timeout=self._response_timeout_seconds, + ) + finally: + unsubscribe() + + if aborted: + raise RuntimeError("Copilot turn was aborted.") + if tool_execution_started: + raise RuntimeError("Copilot turn reported tool execution.") + if ( + reply is None + or reply.agent_id is not None + or not isinstance(reply.data, AssistantMessageData) + or not isinstance(reply.data.content, str) + or not reply.data.content + ): + raise ValueError("Copilot did not return a non-empty root assistant text reply.") + return reply.data.content + + async def cleanup_target_async(self) -> None: + """ + Stop accepting target work, drain active operations, and release owned SDK resources. + + Cleanup is terminal and idempotent. Retained sessions are preserved while the shared + client is always stopped. Cleanup attempts every owned resource before surfacing failures. + """ + async with self._lifecycle_condition: + if self._cleanup_task is None: + self._cleanup_task = asyncio.create_task(self._cleanup_owned_resources_async()) + cleanup_task = self._cleanup_task + await asyncio.shield(cleanup_task) + + async def reset_conversation_async(self, *, conversation_id: str) -> None: + """ + Release one established Copilot conversation without stopping the shared client. + + Unknown or already released conversations are no-ops. Established conversations + become retired before release, so a later send cannot silently create a new native + history for the same PyRIT conversation ID. + + Args: + conversation_id (str): The PyRIT conversation ID to release. + + Raises: + asyncio.CancelledError: If the caller is cancelled while waiting for cleanup. + """ + async with self._lifecycle_condition: + conversation = self._conversations.get(conversation_id) + if conversation is None or (conversation.retired and conversation.session is None): + return + cleanup_task = self._cleanup_task + if cleanup_task is None: + self._active_target_operations += 1 + + if cleanup_task is not None: + try: + await asyncio.shield(cleanup_task) + except asyncio.CancelledError: + raise + except Exception: + async with self._lifecycle_condition: + selected_conversation_released = conversation.retired and conversation.session is None + if not selected_conversation_released: + raise + return + + try: + async with conversation.lock: + await self._retire_conversation_async(conversation=conversation) + finally: + async with self._lifecycle_condition: + self._active_target_operations -= 1 + if self._active_target_operations == 0: + self._lifecycle_condition.notify_all() + + async def _get_or_create_session_async( + self, + *, + conversation_id: str, + initial_system_prompt: str | None, + ) -> "CopilotSession": + async with self._lifecycle_condition: + conversation = self._conversations[conversation_id] + if conversation.retired: + raise RuntimeError( + f"Copilot conversation {conversation_id} was retired and cannot accept further sends; " + "use a new conversation ID." + ) + existing_session = conversation.session + if existing_session is not None: + return existing_session + + client = await self._get_or_start_client_async() + session_id = str(uuid4()) + status = self._runtime_status + if status is None: + raise RuntimeError("Copilot runtime status is unavailable after client startup.") + logger.info( + "Attempting Copilot session creation: pyrit_conversation_id=%s requested_sdk_session_id=%s " + "sdk_version=%s runtime_version=%s runtime_protocol_version=%s retain_session=%s remote_mode=OFF", + conversation_id, + session_id, + self._sdk.__version__, + status.version, + status.protocol_version, + self._retain_session, + ) + + try: + system_message: SystemMessageConfig + if initial_system_prompt is None: + system_message = { + "mode": "customize", + "sections": { + "environment_context": {"action": "remove"}, + "custom_instructions": {"action": "remove"}, + }, + } + else: + system_message = {"mode": "replace", "content": initial_system_prompt} + session = await client.create_session( + session_id=session_id, + model=self._model_name, + system_message=system_message, + remote_session=self._sdk.RemoteSessionMode.OFF, + available_tools=[], + skip_custom_instructions=True, + instruction_directories=[], + enable_host_git_operations=False, + enable_config_discovery=False, + organization_custom_instructions="", + enable_on_demand_instruction_discovery=False, + infinite_sessions={"enabled": False}, + memory={"enabled": False}, + enable_session_store=False, + enable_file_hooks=False, + ) + async with self._lifecycle_condition: + conversation.session = session + return session + except BaseException: + await self._cleanup_allocated_session_async(client=client, session_id=session_id) + raise + + async def _get_or_start_client_async(self) -> "CopilotClient": + async with self._client_start_lock: + if self._client is not None: + return self._client + + client = await asyncio.to_thread( + self._sdk.CopilotClient, + github_token=self._github_token, + working_directory=self._working_directory, + ) + try: + await client.start() + self._runtime_status = await client.get_status() + except BaseException as error: + try: + await client.stop() + except BaseException as cleanup_error: + if isinstance(error, asyncio.CancelledError): + raise error from cleanup_error + raise BaseExceptionGroup( + "Copilot client startup and cleanup failed", + [error, cleanup_error], + ) from error + raise + + async with self._lifecycle_condition: + self._client = client + return client + + async def _retire_conversation_async(self, *, conversation: _ConversationState) -> None: + async with self._lifecycle_condition: + session = conversation.session + client = self._client + if session is None or client is None: + return + conversation.retired = True + await self._release_session_async(client=client, conversation=conversation) + + async def _cleanup_allocated_session_async(self, *, client: "CopilotClient", session_id: str) -> None: + if await client.get_session_metadata(session_id) is None: + return + if self._retain_session: + logger.info("Retaining Copilot session %s as requested; delete it manually.", session_id) + else: + await client.delete_session(session_id) + + async def _release_session_async( + self, + *, + client: "CopilotClient", + conversation: _ConversationState, + ) -> None: + session = conversation.session + if session is None: + return + session_id = session.session_id + if self._retain_session: + await session.disconnect() + logger.info("Retaining Copilot session %s as requested; delete it manually.", session_id) + else: + await client.delete_session(session_id) + async with self._lifecycle_condition: + if conversation.session is session: + conversation.session = None + conversation.retired = True + + async def _cleanup_owned_resources_async(self) -> None: + async with self._lifecycle_condition: + await self._lifecycle_condition.wait_for(lambda: self._active_target_operations == 0) + conversations = [ + conversation for conversation in self._conversations.values() if conversation.session is not None + ] + client = self._client + self._client = None + self._runtime_status = None + + errors: list[BaseException] = [] + if client is not None: + for conversation in conversations: + try: + await self._release_session_async(client=client, conversation=conversation) + except BaseException as error: + errors.append(error) + try: + await client.stop() + except BaseException as error: + errors.append(error) + + if len(errors) == 1: + raise errors[0] + if errors: + raise BaseExceptionGroup("Copilot target cleanup failed", errors) diff --git a/pyrit/score/llm_scoring.py b/pyrit/score/llm_scoring.py index fecbaed396..d554f42a3d 100644 --- a/pyrit/score/llm_scoring.py +++ b/pyrit/score/llm_scoring.py @@ -13,6 +13,7 @@ EmptyResponseException, InvalidJsonException, ScorerLLMResponseBlockedException, + pyrit_json_retry, ) from pyrit.models import ( Acquisition, @@ -83,6 +84,7 @@ async def _run_llm_scoring_async( category: Sequence[str] | str | None = None, objective: str | None = None, normalizer: PromptNormalizer | None = None, + fresh_conversation_per_attempt: bool = False, observation_metadata: Mapping[str, str] | None = None, requires_message_piece_evidence: bool = False, judgment_replay_identifier: Mapping[str, object] | None = None, @@ -93,13 +95,13 @@ async def _run_llm_scoring_async( This is the shared LLM evaluation mechanism: it optionally sets a system prompt on the target, sends the value to be scored (forwarding ``response_handler.json_response_config`` so targets that support structured output can enforce it), and delegates parsing and validation to - ``response_handler``. The round-trip is routed through a ``PromptNormalizer`` via - ``send_json_with_retry_async`` so the scorer's question and the target's answer are persisted - to memory (a full audit trail, and a real conversation an attack can link as a SCORE-type - related conversation) and so JSON retries roll memory back to a clean baseline between attempts - instead of replaying the target's own malformed reply. It is intentionally stateless and - independent of any particular ``Scorer`` so that scorers can compose it without inheriting LLM - machinery. + ``response_handler``. The round-trip uses a ``PromptNormalizer`` so the scorer's question and + the target's answer are persisted to memory (a full audit trail, and a real conversation an + attack can link as a SCORE-type related conversation). The default editable-history path rolls + memory back between JSON attempts; ``fresh_conversation_per_attempt`` keeps malformed judge + exchanges in separate conversations instead of replaying native history. It is intentionally + stateless and independent of any particular ``Scorer`` so scorers can compose it without + inheriting LLM machinery. The round-trip owns only the transport; the ``ResponseHandler`` owns the response contract — the optional response schema and turning raw text into a validated ``UnvalidatedScore``. @@ -127,8 +129,10 @@ async def _run_llm_scoring_async( objective (str | None): Transitional objective context for direct helper callers. Defaults to None. normalizer (PromptNormalizer | None): Normalizer used to send the scoring round-trip - and whose memory is rolled back between JSON retries. Injectable for testing; - defaults to a fresh ``PromptNormalizer()`` when not supplied. + and resolve scorer evidence. Injectable for testing; defaults to a fresh + ``PromptNormalizer()`` when not supplied. + fresh_conversation_per_attempt (bool): Use a new conversation for each JSON retry when + target history cannot be rolled back. Defaults to False. observation_metadata (Mapping[str, str] | None): Scorer-specific state required to reconstruct the response parser during replay. Defaults to None. requires_message_piece_evidence (bool): Whether the rendered request reads fields that a @@ -193,7 +197,7 @@ async def _run_llm_scoring_async( observation_scorable is not None and scored_evidence_digest is not None and has_required_evidence ) - if system_prompt is not None: + if system_prompt is not None and not fresh_conversation_per_attempt: chat_target.set_system_prompt( system_prompt=system_prompt, conversation_id=conversation_id, @@ -273,17 +277,51 @@ def _parse(response: Message) -> UnvalidatedScore: objective=expectation.objective if expectation else None, ) - # Route the round-trip through the normalizer so the scorer Q&A is persisted and JSON retries - # replay on a clean history. + # Editable targets retry on a rolled-back conversation; non-editable judges opt into fresh sessions. try: - unvalidated_score = await send_json_with_retry_async( - normalizer=resolved_normalizer, - target=chat_target, - message=scorer_llm_request, - conversation_id=conversation_id, - parse=_parse, - on_response=_capture_response, - ) + if fresh_conversation_per_attempt: + first_attempt = True + + @pyrit_json_retry + async def _fresh_attempt_async() -> UnvalidatedScore: + nonlocal first_attempt + attempt_conversation_id = conversation_id if first_attempt else str(uuid.uuid4()) + attempt_message = scorer_llm_request if first_attempt else scorer_llm_request.duplicate() + if system_prompt is not None: + chat_target.set_system_prompt( + system_prompt=system_prompt, + conversation_id=attempt_conversation_id, + ) + first_attempt = False + response = await resolved_normalizer.send_prompt_async( + message=attempt_message, + conversation_id=attempt_conversation_id, + target=chat_target, + ) + if not response: + raise ValueError(f"No response received for conversation ID: {attempt_conversation_id}") + _capture_response(response) + try: + return _parse(response) + except InvalidJsonException: + try: + await chat_target.reset_conversation_async(conversation_id=attempt_conversation_id) + except Exception as reset_error: + raise RuntimeError( + "Could not release the malformed judge session; refusing another retry." + ) from reset_error + raise + + unvalidated_score: UnvalidatedScore = await _fresh_attempt_async() + else: + unvalidated_score = await send_json_with_retry_async( + normalizer=resolved_normalizer, + target=chat_target, + message=scorer_llm_request, + conversation_id=conversation_id, + parse=_parse, + on_response=_capture_response, + ) except ScorerLLMResponseBlockedException as error: if terminal_response is not None and can_collect_observation and _has_observation_collection(): observation = _build_judgment_observation( diff --git a/pyrit/score/scorer.py b/pyrit/score/scorer.py index b58096f69b..fc8f10b7a0 100644 --- a/pyrit/score/scorer.py +++ b/pyrit/score/scorer.py @@ -7,6 +7,7 @@ import asyncio import logging from abc import abstractmethod +from dataclasses import replace from typing import TYPE_CHECKING, Any, ClassVar, TypeVar, cast, final, overload from pyrit.common.deprecation import print_deprecation_message @@ -32,7 +33,8 @@ ScoringExpectation, ) from pyrit.prompt_target.batch_helper import batch_task_async -from pyrit.prompt_target.common.target_requirements import TargetRequirements +from pyrit.prompt_target.common.target_capabilities import CapabilityName +from pyrit.prompt_target.common.target_requirements import CHAT_TARGET_REQUIREMENTS, TargetRequirements from pyrit.score.observation.execution import ( NonReplayableObservationError, _observation_collection, @@ -62,6 +64,18 @@ ConditionT = TypeVar("ConditionT", bound=Condition) +class _SelfContainedJudgeTargetRequirements(TargetRequirements): + def validate(self, *, target: PromptTarget) -> None: + requirements = CHAT_TARGET_REQUIREMENTS + if not target.capabilities.supports_editable_history: + requirements = replace( + requirements, + required=requirements.required - {CapabilityName.EDITABLE_HISTORY}, + native_required=requirements.native_required | {CapabilityName.SYSTEM_PROMPT}, + ) + requirements.validate(target=target) + + async def _legacy_score_scorable_async( self: Scorer, *, diff --git a/pyrit/score/true_false/self_ask_question_answer_scorer.py b/pyrit/score/true_false/self_ask_question_answer_scorer.py index 803b32344a..95d4353a2b 100644 --- a/pyrit/score/true_false/self_ask_question_answer_scorer.py +++ b/pyrit/score/true_false/self_ask_question_answer_scorer.py @@ -7,6 +7,7 @@ from pyrit.common.path import SCORER_SEED_PROMPT_PATH from pyrit.models import AnswerMatches, ScoringExpectation +from pyrit.prompt_target import CHAT_TARGET_REQUIREMENTS from pyrit.score.llm_scoring import _run_llm_scoring_async from pyrit.score.scorer_prompt_validator import ScorerPromptValidator from pyrit.score.true_false.self_ask_true_false_scorer import ( @@ -52,6 +53,7 @@ class SelfAskQuestionAnswerScorer(SelfAskTrueFalseScorer): _JUDGMENT_PROMPT = ( "{objective}\n\nThe target response is as follows. Evaluate if the response is correct:\n{response}" ) + TARGET_REQUIREMENTS = CHAT_TARGET_REQUIREMENTS def __init__( self, diff --git a/pyrit/score/true_false/self_ask_refusal_scorer.py b/pyrit/score/true_false/self_ask_refusal_scorer.py index f7b4621d11..8b4846242a 100644 --- a/pyrit/score/true_false/self_ask_refusal_scorer.py +++ b/pyrit/score/true_false/self_ask_refusal_scorer.py @@ -17,10 +17,11 @@ ScoringExpectation, SeedPrompt, ) -from pyrit.prompt_target import CHAT_TARGET_REQUIREMENTS, PromptTarget +from pyrit.prompt_target import PromptTarget from pyrit.score.llm_scoring import _parse_judgment_observation, _run_llm_scoring_async from pyrit.score.observation.execution import _ObservationEvidence from pyrit.score.response_handler import JsonSchemaResponseHandler, ResponseHandler, TrueFalseResponseHandler +from pyrit.score.scorer import _SelfContainedJudgeTargetRequirements from pyrit.score.scorer_prompt_validator import ScorerPromptValidator from pyrit.score.true_false.true_false_score_aggregator import ( TrueFalseAggregatorFunc, @@ -76,7 +77,7 @@ class SelfAskRefusalScorer(MessageTrueFalseScorer): ) _DEFAULT_VALIDATOR: ScorerPromptValidator = ScorerPromptValidator() - TARGET_REQUIREMENTS = CHAT_TARGET_REQUIREMENTS + TARGET_REQUIREMENTS = _SelfContainedJudgeTargetRequirements() def __init__( self, @@ -93,8 +94,9 @@ def __init__( Initialize the SelfAskRefusalScorer. Args: - chat_target (PromptTarget | None): The chat target used for scoring. Must satisfy - CHAT_TARGET_REQUIREMENTS. + chat_target (PromptTarget | None): The chat target used for scoring. Must support + multi-turn conversations and either editable history or native system prompts. + Non-editable targets use fresh conversations when malformed JSON is retried. system_prompt (SeedPrompt | str | None): The refusal-detection system prompt. A ``SeedPrompt`` (e.g. loaded from a ``RefusalScorerPaths`` YAML) is used verbatim and may carry a ``response_json_schema``; a ``str`` is used as-is; ``None`` falls back to @@ -248,6 +250,7 @@ async def _score_piece_async(self, message_piece: MessagePiece, *, objective: st scorer_identifier=self.get_identifier(), judgment_replay_identifier=self._get_judgment_replay_identifier(), category=self._score_category, + fresh_conversation_per_attempt=not self._prompt_target.capabilities.supports_editable_history, ) score = unvalidated_score.to_score(score_value=unvalidated_score.raw_score_value, score_type="true_false") diff --git a/pyrit/score/true_false/self_ask_true_false_scorer.py b/pyrit/score/true_false/self_ask_true_false_scorer.py index 95292cf350..9b9c33410a 100644 --- a/pyrit/score/true_false/self_ask_true_false_scorer.py +++ b/pyrit/score/true_false/self_ask_true_false_scorer.py @@ -20,10 +20,11 @@ SeedPrompt, UnvalidatedScore, ) -from pyrit.prompt_target import CHAT_TARGET_REQUIREMENTS, PromptTarget +from pyrit.prompt_target import PromptTarget from pyrit.score.llm_scoring import _parse_judgment_observation, _run_llm_scoring_async from pyrit.score.observation.execution import _ObservationEvidence from pyrit.score.response_handler import JsonSchemaResponseHandler, ResponseHandler, TrueFalseResponseHandler +from pyrit.score.scorer import _SelfContainedJudgeTargetRequirements from pyrit.score.scorer_prompt_validator import ScorerPromptValidator from pyrit.score.system_prompt import _render_system_prompt_template from pyrit.score.true_false.true_false_score_aggregator import ( @@ -153,7 +154,7 @@ class SelfAskTrueFalseScorer(MessageTrueFalseScorer): _DEFAULT_VALIDATOR: ScorerPromptValidator = ScorerPromptValidator( supported_data_types=["text", "image_path"], ) - TARGET_REQUIREMENTS = CHAT_TARGET_REQUIREMENTS + TARGET_REQUIREMENTS = _SelfContainedJudgeTargetRequirements() def __init__( self, @@ -169,8 +170,9 @@ def __init__( Initialize the SelfAskTrueFalseScorer. Args: - chat_target (PromptTarget | None): The chat target used for scoring. Must satisfy - CHAT_TARGET_REQUIREMENTS. + chat_target (PromptTarget | None): The chat target used for scoring. Must support multi-turn + conversations and either editable history or native system prompts. Noneditable targets + are supported for text scoring only. system_prompt (SeedPrompt | str | None): The scoring system prompt. A ``SeedPrompt`` (e.g. rendered via ``render_true_false_system_prompt``) is used verbatim and may carry a ``response_json_schema``; a ``str`` is used as-is; ``None`` falls back to the @@ -305,9 +307,17 @@ async def _score_piece_async(self, message_piece: MessagePiece, *, objective: st The category is configured from the TrueFalseQuestionPath. The score_value is True or False based on which description fits best. Metadata can be configured to provide additional information. + + Raises: + ValueError: If non-text scoring uses a target without editable history. """ # Build scoring prompt - for non-text content, extra context about objective is sent as a prepended text piece is_non_text = message_piece.converted_value_data_type != "text" + if is_non_text and not self._prompt_target.capabilities.supports_editable_history: + raise ValueError( + "non-text scoring requires editable history; fresh-conversation retries support text only." + ) + if is_non_text: prepended_text = f"objective: {objective}\nresponse:" scoring_value = message_piece.converted_value @@ -328,6 +338,9 @@ async def _score_piece_async(self, message_piece: MessagePiece, *, objective: st judgment_replay_identifier=self._get_judgment_replay_identifier(), prepended_text=prepended_text, category=self._score_category, + fresh_conversation_per_attempt=( + scoring_data_type == "text" and not self._prompt_target.capabilities.supports_editable_history + ), ) return [self._convert_score(unvalidated_score)] diff --git a/tests/unit/mocks.py b/tests/unit/mocks.py index 5a6051979d..d84fe44687 100644 --- a/tests/unit/mocks.py +++ b/tests/unit/mocks.py @@ -205,8 +205,14 @@ class MockPromptTarget(PromptTarget): prompt_sent: list[str] - def __init__(self, *, id=None, rpm=None) -> None: # noqa: A002 - super().__init__(max_requests_per_minute=rpm) + def __init__( + self, + *, + id=None, # noqa: A002 + rpm=None, + custom_configuration: TargetConfiguration | None = None, + ) -> None: + super().__init__(max_requests_per_minute=rpm, custom_configuration=custom_configuration) self.id = id self.prompt_sent = [] diff --git a/tests/unit/prompt_target/target/test_github_copilot_target.py b/tests/unit/prompt_target/target/test_github_copilot_target.py new file mode 100644 index 0000000000..8aac0c28a3 --- /dev/null +++ b/tests/unit/prompt_target/target/test_github_copilot_target.py @@ -0,0 +1,1584 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +from __future__ import annotations + +import asyncio +import logging +import threading +from contextlib import suppress +from dataclasses import replace +from datetime import UTC, datetime +from functools import partial +from typing import TYPE_CHECKING, Any +from unittest.mock import AsyncMock, NonCallableMagicMock, call, create_autospec, patch +from uuid import UUID, uuid4 + +import pytest +from unit.mocks import store_message + +from pyrit.models import Message, MessagePiece, MessageScorable, ScoringExpectation +from pyrit.prompt_normalizer import PromptNormalizer +from pyrit.prompt_target import GitHubCopilotTarget +from pyrit.score import SelfAskTrueFalseScorer, TrueFalseQuestion + +if TYPE_CHECKING: + from collections.abc import Callable, Iterator + from pathlib import Path + + from copilot.generated.session_events import SessionEvent + + from pyrit.memory import MemoryInterface + +TARGET_LOGGER = "pyrit.prompt_target.github_copilot_target" + + +@pytest.fixture +def sdk() -> Any: + return pytest.importorskip("copilot") + + +@pytest.fixture +def client(sdk: Any) -> Iterator[NonCallableMagicMock]: + from copilot.client import GetStatusResponse + + session = _make_sdk_session(sdk=sdk, session_id="sdk-session-id") + session.send_and_wait.return_value = _assistant_reply("HELLO") + client = create_autospec(sdk.CopilotClient, instance=True) + assert isinstance(client, NonCallableMagicMock) + client.create_session.return_value = session + client.get_status.return_value = GetStatusResponse(version="6.5.4", protocol_version=3) + with patch.object(sdk, "CopilotClient", return_value=client): + yield client + + +def _assistant_reply(text: str) -> SessionEvent: + from copilot.generated.session_events import AssistantMessageData, SessionEvent, SessionEventType + + return SessionEvent( + id=uuid4(), + timestamp=datetime.now(UTC), + type=SessionEventType.ASSISTANT_MESSAGE, + data=AssistantMessageData(content=text, message_id="sdk-reply"), + ) + + +def _mock_session_storage(*, client: NonCallableMagicMock, sessions: set[str]) -> None: + from copilot import SessionMetadata + + async def get_session_metadata_async(session_id: str) -> SessionMetadata | None: + if session_id not in sessions: + return None + return SessionMetadata( + session_id=session_id, + start_time=datetime(2026, 1, 1, tzinfo=UTC), + modified_time=datetime(2026, 1, 1, tzinfo=UTC), + is_remote=False, + ) + + client.get_session_metadata.side_effect = get_session_metadata_async + client.delete_session.side_effect = sessions.remove + + +def _make_sdk_session( + *, + sdk: Any, + session_id: str, +) -> NonCallableMagicMock: + session = create_autospec(sdk.CopilotSession, instance=True) + assert isinstance(session, NonCallableMagicMock) + session.session_id = session_id + return session + + +def _user_message( + *, + original_value: str, + converted_value: str | None = None, + conversation_id: str | None = None, +) -> Message: + return MessagePiece( + role="user", + conversation_id=conversation_id, + original_value=original_value, + converted_value=original_value if converted_value is None else converted_value, + ).to_message() + + +async def _send_normalized_async( + *, + target: GitHubCopilotTarget, + original_value: str, + converted_value: str | None = None, + conversation_id: str | None = None, +) -> Message: + return await PromptNormalizer().send_prompt_async( + message=_user_message( + original_value=original_value, + converted_value=converted_value, + ), + conversation_id=conversation_id, + target=target, + ) + + +def _message_state(*, memory: MemoryInterface, conversation_id: str) -> list[tuple[str, str, str, str | None, str]]: + return [ + ( + message.get_piece().role, + message.get_piece().original_value, + message.get_piece().converted_value, + message.get_piece().conversation_id, + message.get_piece().response_error, + ) + for message in memory.get_conversation_messages(conversation_id=conversation_id) + ] + + +def _message_roles_and_errors(*, memory: MemoryInterface, conversation_id: str) -> list[tuple[str, str]]: + return [ + (message.get_piece().role, message.get_piece().response_error) + for message in memory.get_conversation_messages(conversation_id=conversation_id) + ] + + +def _message_values_and_errors(*, memory: MemoryInterface, conversation_id: str) -> list[tuple[str, str, str]]: + return [ + (message.get_piece().role, message.get_piece().converted_value, message.get_piece().response_error) + for message in memory.get_conversation_messages(conversation_id=conversation_id) + ] + + +def _expected_session_configuration(*, system_message: dict[str, Any]) -> dict[str, Any]: + from copilot.generated.rpc import RemoteSessionMode + + return { + "model": "gpt-5-mini", + "system_message": system_message, + "remote_session": RemoteSessionMode.OFF, + "available_tools": [], + "skip_custom_instructions": True, + "instruction_directories": [], + "enable_host_git_operations": False, + "enable_config_discovery": False, + "organization_custom_instructions": "", + "enable_on_demand_instruction_discovery": False, + "infinite_sessions": {"enabled": False}, + "memory": {"enabled": False}, + "enable_session_store": False, + "enable_file_hooks": False, + } + + +async def _cancel_tasks_async(*tasks: asyncio.Task[Any] | None) -> None: + for task in tasks: + if task is not None and not task.done(): + task.cancel() + await asyncio.gather(*[task for task in tasks if task is not None], return_exceptions=True) + + +def _assert_no_resource_release(client: NonCallableMagicMock) -> None: + client.delete_session.assert_not_awaited() + client.stop.assert_not_awaited() + + +@pytest.mark.usefixtures("patch_central_database") +@pytest.mark.parametrize("retain_session", [False, True], ids=["delete", "retain"]) +async def test_normalizer_round_trip_and_retention_async( + *, + sdk: Any, + client: NonCallableMagicMock, + sqlite_instance: MemoryInterface, + caplog: pytest.LogCaptureFixture, + retain_session: bool, +) -> None: + conversation_id = str(uuid4()) + target = GitHubCopilotTarget(model_name="gpt-5-mini", retain_session=retain_session) + with patch.object(sdk, "__version__", "9.8.7"), caplog.at_level(logging.INFO, logger=TARGET_LOGGER): + response = await _send_normalized_async( + target=target, + original_value="Original text before conversion.", + converted_value="Reply exactly HELLO.", + conversation_id=conversation_id, + ) + _assert_no_resource_release(client=client) + await target.cleanup_target_async() + + assert response.get_piece().converted_value == "HELLO" + assert _message_state(memory=sqlite_instance, conversation_id=conversation_id) == [ + ("user", "Original text before conversion.", "Reply exactly HELLO.", conversation_id, "none"), + ("assistant", "HELLO", "HELLO", conversation_id, "none"), + ] + session = client.create_session.return_value + session.send_and_wait.assert_awaited_once_with("Reply exactly HELLO.", timeout=60.0) + session.on.return_value.assert_called_once_with() + client.create_session.assert_awaited_once() + client.get_session_metadata.assert_not_awaited() + sdk.CopilotClient.assert_called_once_with(github_token=None, working_directory=None) + requested_id = client.create_session.await_args.kwargs["session_id"] + assert str(UUID(requested_id)) == requested_id + assert requested_id != "sdk-session-id" + records = [r.getMessage() for r in caplog.records if r.name == TARGET_LOGGER and r.levelno == logging.INFO] + assert len(records) == (2 if retain_session else 1) + assert records[0].startswith("Attempting Copilot session creation:") + for field in ( + f"pyrit_conversation_id={conversation_id}", + f"requested_sdk_session_id={requested_id}", + "sdk_version=9.8.7", + "runtime_version=6.5.4", + "runtime_protocol_version=3", + f"retain_session={retain_session}", + "remote_mode=OFF", + ): + assert field in records[0] + if retain_session: + client.delete_session.assert_not_awaited() + assert records[1] == "Retaining Copilot session sdk-session-id as requested; delete it manually." + else: + client.delete_session.assert_awaited_once_with("sdk-session-id") + client.stop.assert_awaited_once() + + +@pytest.mark.usefixtures("patch_central_database") +@pytest.mark.parametrize( + "initial_system_prompt", + [pytest.param(None, id="default-customize"), pytest.param("initial system instructions", id="initial-replacement")], +) +async def test_normalizer_continues_native_session_across_turns_async( + *, + sdk: Any, + client: NonCallableMagicMock, + sqlite_instance: MemoryInterface, + initial_system_prompt: str | None, +) -> None: + conversation_id = "native-two-turn-conversation" + session = client.create_session.return_value + session.session_id = "sdk-session-id" + session.send_and_wait.side_effect = [_assistant_reply("FIRST"), _assistant_reply("SECOND")] + target = GitHubCopilotTarget(model_name="gpt-5-mini") + + if initial_system_prompt is not None: + target.set_system_prompt(system_prompt=initial_system_prompt, conversation_id=conversation_id) + + first_response = await _send_normalized_async( + target=target, + original_value="first original", + converted_value="first prepared", + conversation_id=conversation_id, + ) + assert first_response.get_piece().converted_value == "FIRST" + + configuration = dict(client.create_session.await_args.kwargs) + requested_session_id = configuration.pop("session_id") + assert str(UUID(requested_session_id)) == requested_session_id + assert configuration == _expected_session_configuration( + system_message=( + {"mode": "replace", "content": initial_system_prompt} + if initial_system_prompt is not None + else { + "mode": "customize", + "sections": { + "environment_context": {"action": "remove"}, + "custom_instructions": {"action": "remove"}, + }, + } + ) + ) + if initial_system_prompt is not None: + with pytest.raises(RuntimeError, match="Conversation already exists"): + target.set_system_prompt(system_prompt="different system instructions", conversation_id=conversation_id) + + second_response = await _send_normalized_async( + target=target, + original_value="second original", + converted_value="second prepared", + conversation_id=conversation_id, + ) + assert second_response.get_piece().converted_value == "SECOND" + assert session.send_and_wait.await_args_list == [ + call("first prepared", timeout=60.0), + call("second prepared", timeout=60.0), + ] + client.create_session.assert_awaited_once() + expected_messages = [ + ("user", "first original", "first prepared", conversation_id, "none"), + ("assistant", "FIRST", "FIRST", conversation_id, "none"), + ("user", "second original", "second prepared", conversation_id, "none"), + ("assistant", "SECOND", "SECOND", conversation_id, "none"), + ] + if initial_system_prompt is not None: + expected_messages.insert(0, ("system", initial_system_prompt, initial_system_prompt, conversation_id, "none")) + assert _message_state(memory=sqlite_instance, conversation_id=conversation_id) == expected_messages + _assert_no_resource_release(client=client) + await target.cleanup_target_async() + client.delete_session.assert_awaited_once_with(session.session_id) + client.stop.assert_awaited_once() + + +@pytest.mark.usefixtures("patch_central_database") +async def test_distinct_conversations_progress_on_shared_client_async( + *, + sdk: Any, + client: NonCallableMagicMock, +) -> None: + session_a = _make_sdk_session(sdk=sdk, session_id="sdk-session-a") + session_b = _make_sdk_session(sdk=sdk, session_id="sdk-session-b") + first_send_started = asyncio.Event() + release_first_send = asyncio.Event() + + async def send_a_async(*_args: Any, **_kwargs: Any) -> Any: + first_send_started.set() + await release_first_send.wait() + return _assistant_reply("A") + + session_a.send_and_wait.side_effect = send_a_async + session_b.send_and_wait.return_value = _assistant_reply("B") + client.create_session.side_effect = [session_a, session_b] + target = GitHubCopilotTarget(model_name="gpt-5-mini") + first_task = asyncio.create_task( + target.send_prompt_async( + message=_user_message( + conversation_id="conversation-a", + original_value="a original", + converted_value="a prepared", + ) + ) + ) + second_task: asyncio.Task[list[Message]] | None = None + try: + await asyncio.wait_for(first_send_started.wait(), timeout=2.0) + second_task = asyncio.create_task( + target.send_prompt_async( + message=_user_message( + conversation_id="conversation-b", + original_value="b original", + converted_value="b prepared", + ) + ) + ) + second_response = await asyncio.wait_for(second_task, timeout=2.0) + assert not release_first_send.is_set() + assert (second_response[0].get_piece().conversation_id, second_response[0].get_piece().converted_value) == ( + "conversation-b", + "B", + ) + + release_first_send.set() + first_response = await asyncio.wait_for(first_task, timeout=2.0) + assert (first_response[0].get_piece().conversation_id, first_response[0].get_piece().converted_value) == ( + "conversation-a", + "A", + ) + _assert_no_resource_release(client=client) + await target.cleanup_target_async() + finally: + release_first_send.set() + await _cancel_tasks_async(first_task, second_task) + with suppress(Exception): + await asyncio.wait_for(target.cleanup_target_async(), timeout=2.0) + + sdk.CopilotClient.assert_called_once_with(github_token=None, working_directory=None) + client.start.assert_awaited_once() + client.get_status.assert_awaited_once() + assert client.create_session.await_count == 2 + requested_session_ids = [entry.kwargs["session_id"] for entry in client.create_session.await_args_list] + assert all(str(UUID(session_id)) == session_id for session_id in requested_session_ids) + assert len(set(requested_session_ids)) == 2 + session_a.send_and_wait.assert_awaited_once_with("a prepared", timeout=60.0) + session_b.send_and_wait.assert_awaited_once_with("b prepared", timeout=60.0) + assert client.delete_session.await_args_list == [call(session_a.session_id), call(session_b.session_id)] + client.stop.assert_awaited_once() + + +@pytest.mark.usefixtures("patch_central_database") +async def test_reset_conversation_releases_only_requested_session_async( + *, + sdk: Any, + client: NonCallableMagicMock, + sqlite_instance: MemoryInterface, +) -> None: + session_a = _make_sdk_session(sdk=sdk, session_id="sdk-session-a") + session_a.send_and_wait.side_effect = [_assistant_reply("A1"), _assistant_reply("A_REOPENED")] + session_b = _make_sdk_session(sdk=sdk, session_id="sdk-session-b") + session_b.send_and_wait.side_effect = [_assistant_reply("B1"), _assistant_reply("B2")] + client.create_session.side_effect = [session_a, session_b] + target = GitHubCopilotTarget(model_name="gpt-5-mini", retain_session=True) + + await _send_normalized_async( + target=target, + original_value="a original", + converted_value="a prepared", + conversation_id="conversation-a", + ) + await _send_normalized_async( + target=target, + original_value="b original", + converted_value="b prepared", + conversation_id="conversation-b", + ) + memory_before_reset = { + conversation_id: _message_values_and_errors(memory=sqlite_instance, conversation_id=conversation_id) + for conversation_id in ("conversation-a", "conversation-b") + } + + await target.reset_conversation_async(conversation_id="conversation-a") + assert session_a.disconnect.await_count == 1 + _assert_no_resource_release(client=client) + session_b.disconnect.assert_not_awaited() + memory_after_reset = { + conversation_id: _message_values_and_errors(memory=sqlite_instance, conversation_id=conversation_id) + for conversation_id in ("conversation-a", "conversation-b") + } + assert memory_after_reset == memory_before_reset + + await target.reset_conversation_async(conversation_id="conversation-a") + assert session_a.disconnect.await_count == 1 + + with pytest.raises(Exception, match="Error sending prompt with conversation ID:") as reset_error: + await _send_normalized_async( + target=target, + original_value="a retry original", + converted_value="a retry prepared", + conversation_id="conversation-a", + ) + assert isinstance(reset_error.value.__cause__, RuntimeError) + assert "retired" in str(reset_error.value.__cause__).lower() + + response_b = await _send_normalized_async( + target=target, + original_value="b second original", + converted_value="b second prepared", + conversation_id="conversation-b", + ) + assert response_b.get_piece().converted_value == "B2" + assert client.create_session.await_count == 2 + session_a.send_and_wait.assert_awaited_once_with("a prepared", timeout=60.0) + session_b.send_and_wait.assert_has_awaits( + [ + call("b prepared", timeout=60.0), + call("b second prepared", timeout=60.0), + ] + ) + client.delete_session.assert_not_awaited() + await target.cleanup_target_async() + client.delete_session.assert_not_awaited() + client.stop.assert_awaited_once() + + +@pytest.mark.usefixtures("patch_central_database") +async def test_reset_waits_for_in_progress_cleanup_release_async( + *, + client: NonCallableMagicMock, +) -> None: + session = client.create_session.return_value + target = GitHubCopilotTarget(model_name="gpt-5-mini", retain_session=True) + await _send_normalized_async( + target=target, + original_value="a original", + converted_value="a prepared", + conversation_id="conversation-a", + ) + disconnect_started = asyncio.Event() + release_disconnect = asyncio.Event() + + async def disconnect_async() -> None: + disconnect_started.set() + await release_disconnect.wait() + + session.disconnect.side_effect = disconnect_async + cleanup_task = asyncio.create_task(target.cleanup_target_async()) + reset_task: asyncio.Task[None] | None = None + try: + await asyncio.wait_for(disconnect_started.wait(), timeout=2.0) + reset_task = asyncio.create_task(target.reset_conversation_async(conversation_id="conversation-a")) + await asyncio.sleep(0) + assert not reset_task.done() + + release_disconnect.set() + await asyncio.wait_for(cleanup_task, timeout=2.0) + await asyncio.wait_for(reset_task, timeout=2.0) + session.disconnect.assert_awaited_once() + client.delete_session.assert_not_awaited() + client.stop.assert_awaited_once() + finally: + release_disconnect.set() + await _cancel_tasks_async(cleanup_task, reset_task) + + +@pytest.mark.usefixtures("patch_central_database") +async def test_reset_waits_for_conversation_creation_async( + *, + client: NonCallableMagicMock, +) -> None: + session = client.create_session.return_value + session.session_id = "sdk-session-a" + session.send_and_wait.return_value = _assistant_reply("FIRST") + create_started = asyncio.Event() + release_creation = asyncio.Event() + + async def create_session_async(*_args: Any, **_kwargs: Any) -> Any: + create_started.set() + await release_creation.wait() + return session + + client.create_session.side_effect = create_session_async + target = GitHubCopilotTarget(model_name="gpt-5-mini") + conversation_id = "provisioning-conversation" + first_task = asyncio.create_task( + target.send_prompt_async( + message=_user_message( + conversation_id=conversation_id, + original_value="first original", + converted_value="first prepared", + ) + ) + ) + reset_task: asyncio.Task[None] | None = None + try: + await asyncio.wait_for(create_started.wait(), timeout=2.0) + reset_task = asyncio.create_task(target.reset_conversation_async(conversation_id=conversation_id)) + await asyncio.sleep(0) + assert not reset_task.done() + + release_creation.set() + first_response = await asyncio.wait_for(first_task, timeout=2.0) + await asyncio.wait_for(reset_task, timeout=2.0) + assert first_response[0].get_piece().converted_value == "FIRST" + session.send_and_wait.assert_awaited_once_with("first prepared", timeout=60.0) + client.create_session.assert_awaited_once() + client.delete_session.assert_awaited_once_with(session.session_id) + client.stop.assert_not_awaited() + + await target.cleanup_target_async() + client.stop.assert_awaited_once() + finally: + release_creation.set() + await _cancel_tasks_async(first_task, reset_task) + with suppress(Exception): + await asyncio.wait_for(target.cleanup_target_async(), timeout=2.0) + + +@pytest.mark.usefixtures("patch_central_database") +@pytest.mark.parametrize( + "reset_before_cleanup", + [pytest.param(True, id="explicit-reset"), pytest.param(False, id="whole-cleanup-release")], +) +async def test_repeated_reset_does_not_join_unrelated_cleanup_async( + *, + sdk: Any, + client: NonCallableMagicMock, + reset_before_cleanup: bool, +) -> None: + session_a = _make_sdk_session(sdk=sdk, session_id="sdk-session-a") + session_a.send_and_wait.return_value = _assistant_reply("A") + session_b = _make_sdk_session(sdk=sdk, session_id="sdk-session-b") + session_b.send_and_wait.return_value = _assistant_reply("B") + client.create_session.side_effect = [session_a, session_b] + target = GitHubCopilotTarget(model_name="gpt-5-mini", retain_session=True) + + for conversation_id, prompt in (("conversation-a", "a prepared"), ("conversation-b", "b prepared")): + await target.send_prompt_async( + message=_user_message( + conversation_id=conversation_id, + original_value=prompt, + ) + ) + if reset_before_cleanup: + await target.reset_conversation_async(conversation_id="conversation-a") + session_a.disconnect.assert_awaited_once() + + disconnect_started = asyncio.Event() + release_disconnect = asyncio.Event() + + async def disconnect_b_async() -> None: + disconnect_started.set() + await release_disconnect.wait() + raise RuntimeError("B disconnect failed") + + session_b.disconnect.side_effect = disconnect_b_async + cleanup_task = asyncio.create_task(target.cleanup_target_async()) + reset_task: asyncio.Task[None] | None = None + try: + await asyncio.wait_for(disconnect_started.wait(), timeout=2.0) + if not reset_before_cleanup: + session_a.disconnect.assert_awaited_once() + reset_task = asyncio.create_task(target.reset_conversation_async(conversation_id="conversation-a")) + await asyncio.sleep(0) + assert reset_task.done() + await reset_task + assert session_a.disconnect.await_count == 1 + client.stop.assert_not_awaited() + + release_disconnect.set() + with pytest.raises(RuntimeError, match="B disconnect failed"): + await asyncio.wait_for(cleanup_task, timeout=2.0) + session_b.disconnect.assert_awaited_once() + client.stop.assert_awaited_once() + finally: + release_disconnect.set() + await _cancel_tasks_async(cleanup_task, reset_task) + + +@pytest.mark.usefixtures("patch_central_database") +@pytest.mark.parametrize( + "failed_conversation", + [ + pytest.param("conversation-a", id="selected-session-release-fails"), + pytest.param("conversation-b", id="unrelated-session-release-fails"), + ], +) +async def test_reset_during_cleanup_propagates_only_selected_session_failure_async( + *, + sdk: Any, + client: NonCallableMagicMock, + failed_conversation: str, +) -> None: + session_a = _make_sdk_session(sdk=sdk, session_id="sdk-session-a") + session_a.send_and_wait.return_value = _assistant_reply("A") + session_b = _make_sdk_session(sdk=sdk, session_id="sdk-session-b") + session_b.send_and_wait.return_value = _assistant_reply("B") + client.create_session.side_effect = [session_a, session_b] + target = GitHubCopilotTarget(model_name="gpt-5-mini", retain_session=True) + + for conversation_id, prompt in (("conversation-a", "A"), ("conversation-b", "B")): + await target.send_prompt_async( + message=_user_message(conversation_id=conversation_id, original_value=prompt), + ) + + a_release_started = asyncio.Event() + release_a = asyncio.Event() + a_release_finished = asyncio.Event() + b_release_started = asyncio.Event() + release_b = asyncio.Event() + a_failure = RuntimeError("selected session A release failed") + b_failure = RuntimeError("unrelated session B release failed") + + async def disconnect_a_async() -> None: + a_release_started.set() + try: + await release_a.wait() + if failed_conversation == "conversation-a": + raise a_failure + finally: + a_release_finished.set() + + async def disconnect_b_async() -> None: + b_release_started.set() + await release_b.wait() + if failed_conversation == "conversation-b": + raise b_failure + + session_a.disconnect.side_effect = disconnect_a_async + session_b.disconnect.side_effect = disconnect_b_async + cleanup_task = asyncio.create_task(target.cleanup_target_async()) + reset_task: asyncio.Task[None] | None = None + try: + await asyncio.wait_for(a_release_started.wait(), timeout=2.0) + reset_task = asyncio.create_task(target.reset_conversation_async(conversation_id="conversation-a")) + await asyncio.sleep(0) + assert not reset_task.done() + + release_a.set() + await asyncio.wait_for(a_release_finished.wait(), timeout=2.0) + await asyncio.wait_for(b_release_started.wait(), timeout=2.0) + assert a_release_finished.is_set() + release_b.set() + + expected_cleanup_failure = a_failure if failed_conversation == "conversation-a" else b_failure + with pytest.raises(RuntimeError) as cleanup_error: + await asyncio.wait_for(cleanup_task, timeout=2.0) + assert cleanup_error.value is expected_cleanup_failure + + if failed_conversation == "conversation-a": + with pytest.raises(RuntimeError) as reset_error: + await asyncio.wait_for(reset_task, timeout=2.0) + assert reset_error.value is a_failure + else: + await asyncio.wait_for(reset_task, timeout=2.0) + + assert client.create_session.await_count == 2 + session_a.send_and_wait.assert_awaited_once() + session_b.send_and_wait.assert_awaited_once() + session_a.disconnect.assert_awaited_once() + session_b.disconnect.assert_awaited_once() + client.stop.assert_awaited_once() + finally: + release_a.set() + release_b.set() + tasks = {cleanup_task} + if reset_task is not None: + tasks.add(reset_task) + await asyncio.gather(*tasks, return_exceptions=True) + + +@pytest.mark.usefixtures("patch_central_database") +async def test_normalizer_rejects_retired_conversation_but_allows_fresh_async( + *, + sdk: Any, + client: NonCallableMagicMock, + sqlite_instance: MemoryInterface, +) -> None: + session_a = _make_sdk_session(sdk=sdk, session_id="sdk-session-a") + session_a.send_and_wait.side_effect = TimeoutError("ambiguous mock send") + + session_b = _make_sdk_session(sdk=sdk, session_id="sdk-session-b") + session_b.send_and_wait.return_value = _assistant_reply("FRESH") + client.create_session.side_effect = [session_a, session_b] + target = GitHubCopilotTarget(model_name="gpt-5-mini") + + conversation_a = "retired-conversation" + with pytest.raises(Exception, match="Error sending prompt with conversation ID:") as first_error: + await _send_normalized_async( + target=target, + original_value="ambiguous original", + converted_value="ambiguous prepared", + conversation_id=conversation_a, + ) + assert isinstance(first_error.value.__cause__, TimeoutError) + assert str(first_error.value.__cause__) == "ambiguous mock send" + + with pytest.raises(Exception, match="Error sending prompt with conversation ID:") as retired_error: + await _send_normalized_async( + target=target, + original_value="retry original", + converted_value="retry prepared", + conversation_id=conversation_a, + ) + assert isinstance(retired_error.value.__cause__, RuntimeError) + assert "retired" in str(retired_error.value.__cause__).lower() + + conversation_b = "fresh-conversation" + response_b = await _send_normalized_async( + target=target, + original_value="fresh original", + converted_value="fresh prepared", + conversation_id=conversation_b, + ) + assert response_b.get_piece().converted_value == "FRESH" + assert _message_roles_and_errors(memory=sqlite_instance, conversation_id=conversation_a) == [ + ("user", "none"), + ("assistant", "processing"), + ("user", "none"), + ("assistant", "processing"), + ] + assert _message_values_and_errors(memory=sqlite_instance, conversation_id=conversation_b) == [ + ("user", "fresh prepared", "none"), + ("assistant", "FRESH", "none"), + ] + + sdk.CopilotClient.assert_called_once_with(github_token=None, working_directory=None) + client.start.assert_awaited_once() + client.get_status.assert_awaited_once() + assert client.create_session.await_count == 2 + requested_session_ids = [entry.kwargs["session_id"] for entry in client.create_session.await_args_list] + assert all(str(UUID(session_id)) == session_id for session_id in requested_session_ids) + assert len(set(requested_session_ids)) == 2 + session_a.send_and_wait.assert_awaited_once_with("ambiguous prepared", timeout=60.0) + session_b.send_and_wait.assert_awaited_once_with("fresh prepared", timeout=60.0) + assert client.delete_session.await_args_list == [call("sdk-session-a")] + client.stop.assert_not_awaited() + await target.cleanup_target_async() + assert client.delete_session.await_args_list == [call("sdk-session-a"), call("sdk-session-b")] + client.stop.assert_awaited_once() + + +@pytest.mark.usefixtures("patch_central_database") +async def test_self_ask_true_false_uses_fresh_copilot_session_after_invalid_json_async( + *, + sdk: Any, + client: NonCallableMagicMock, +) -> None: + failed_session = _make_sdk_session(sdk=sdk, session_id="malformed-judge-session") + failed_session.send_and_wait.return_value = _assistant_reply("malformed judge response") + successful_session = _make_sdk_session(sdk=sdk, session_id="valid-judge-session") + released_session_ids: list[str] = [] + judge_json = '{"score_value":true,"description":"Correct","rationale":"Paris is the capital of France."}' + + async def delete_session_async(session_id: str) -> None: + released_session_ids.append(session_id) + + async def send_valid_judge_reply_async(prompt: str, *, timeout: float) -> SessionEvent: + assert released_session_ids == [failed_session.session_id] + return _assistant_reply(judge_json) + + successful_session.send_and_wait.side_effect = send_valid_judge_reply_async + client.create_session.side_effect = [failed_session, successful_session] + client.delete_session.side_effect = delete_session_async + target = GitHubCopilotTarget(model_name="gpt-5-mini") + answer = "Paris is the capital of France." + saved_answer = store_message( + MessagePiece( + role="assistant", + conversation_id=str(uuid4()), + original_value=answer, + ).to_message() + ) + input_piece = saved_answer.get_piece() + input_scorable = MessageScorable.from_message(saved_answer) + + try: + question = TrueFalseQuestion( + category="capital correctness", + true_description="The response correctly identifies Paris as the capital of France.", + false_description="The response does not correctly identify Paris as the capital of France.", + ) + scorer = SelfAskTrueFalseScorer.from_question(chat_target=target, question=question) + scores = await scorer.score_async( + scorable=input_scorable, + expectation=ScoringExpectation(objective="Name France's capital"), + ) + + assert len(scores) == 1 + assert scores[0].get_value() is True + + failed_session.send_and_wait.assert_awaited_once() + successful_session.send_and_wait.assert_awaited_once() + client.create_session.assert_awaited() + assert client.create_session.await_count == 2 + requested_session_ids = [entry.kwargs["session_id"] for entry in client.create_session.await_args_list] + assert all(str(UUID(session_id)) == session_id for session_id in requested_session_ids) + assert len(set(requested_session_ids)) == 2 + session_configurations = [entry.kwargs for entry in client.create_session.await_args_list] + assert session_configurations[0]["system_message"] == session_configurations[1]["system_message"] + assert session_configurations[0]["system_message"]["mode"] == "replace" + finally: + await target.cleanup_target_async() + + assert client.delete_session.await_args_list == [ + call(failed_session.session_id), + call(successful_session.session_id), + ] + client.stop.assert_awaited_once() + + +@pytest.mark.usefixtures("patch_central_database") +async def test_cancelled_send_keeps_retirement_owned_until_cleanup_async( + *, + client: NonCallableMagicMock, +) -> None: + session = client.create_session.return_value + session.send_and_wait.side_effect = TimeoutError("ambiguous mock send") + delete_started = asyncio.Event() + release_delete = asyncio.Event() + delete_tasks: list[asyncio.Task[Any]] = [] + delete_count = 0 + + async def delete_session_async(session_id: str) -> None: + nonlocal delete_count + delete_count += 1 + current_task = asyncio.current_task() + if current_task is not None: + delete_tasks.append(current_task) + if delete_count == 1: + delete_started.set() + await release_delete.wait() + + client.delete_session.side_effect = delete_session_async + target = GitHubCopilotTarget(model_name="gpt-5-mini") + send_task = asyncio.create_task( + target.send_prompt_async( + message=_user_message( + conversation_id="cancelled-retirement", + original_value="ambiguous request", + ) + ) + ) + cleanup_started = asyncio.Event() + + async def cleanup_target_for_test_async() -> None: + cleanup_started.set() + await target.cleanup_target_async() + + cleanup_task: asyncio.Task[None] | None = None + try: + await asyncio.wait_for(delete_started.wait(), timeout=2.0) + send_task.cancel() + cleanup_task = asyncio.create_task(cleanup_target_for_test_async()) + await asyncio.wait_for(cleanup_started.wait(), timeout=2.0) + + assert not send_task.done() + assert not cleanup_task.done() + client.delete_session.assert_awaited_once_with(session.session_id) + client.stop.assert_not_awaited() + + release_delete.set() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(asyncio.shield(send_task), timeout=2.0) + await asyncio.wait_for(cleanup_task, timeout=2.0) + client.delete_session.assert_awaited_once_with(session.session_id) + client.stop.assert_awaited_once() + assert all(task.done() for task in delete_tasks) + finally: + release_delete.set() + tasks = {send_task, *delete_tasks} + if cleanup_task is not None: + tasks.add(cleanup_task) + await asyncio.wait_for(asyncio.gather(*tasks, return_exceptions=True), timeout=2.0) + + +@pytest.mark.usefixtures("patch_central_database") +async def test_cleanup_rejects_queued_turn_and_drains_active_send_async( + *, + client: NonCallableMagicMock, +) -> None: + session = client.create_session.return_value + first_send_started = asyncio.Event() + release_first_send = asyncio.Event() + + async def send_and_wait_async(*_args: Any, **_kwargs: Any) -> Any: + first_send_started.set() + await release_first_send.wait() + return _assistant_reply("FIRST") + + session.send_and_wait.side_effect = send_and_wait_async + target = GitHubCopilotTarget(model_name="gpt-5-mini") + conversation_id = "queued-cleanup-conversation" + first_task = asyncio.create_task( + target.send_prompt_async( + message=_user_message(conversation_id=conversation_id, original_value="first"), + ) + ) + queued_task: asyncio.Task[list[Message]] | None = None + cleanup_task: asyncio.Task[None] | None = None + try: + await asyncio.wait_for(first_send_started.wait(), timeout=2.0) + queued_task = asyncio.create_task( + target.send_prompt_async( + message=_user_message(conversation_id=conversation_id, original_value="queued"), + ) + ) + await asyncio.sleep(0) + cleanup_task = asyncio.create_task(target.cleanup_target_async()) + await asyncio.sleep(0) + + with pytest.raises(RuntimeError, match="cleaned up"): + await asyncio.wait_for( + target.send_prompt_async( + message=_user_message( + conversation_id="fresh-after-cleanup", + original_value="fresh", + ) + ), + timeout=2.0, + ) + _assert_no_resource_release(client=client) + + release_first_send.set() + first_response = await asyncio.wait_for(first_task, timeout=2.0) + assert first_response[0].get_piece().converted_value == "FIRST" + assert queued_task is not None + with pytest.raises(RuntimeError, match="cleaned up"): + await asyncio.wait_for(queued_task, timeout=2.0) + assert cleanup_task is not None + await asyncio.wait_for(cleanup_task, timeout=2.0) + finally: + release_first_send.set() + await _cancel_tasks_async(first_task, queued_task, cleanup_task) + + session.send_and_wait.assert_awaited_once_with("first", timeout=60.0) + client.create_session.assert_awaited_once() + client.delete_session.assert_awaited_once_with(session.session_id) + client.stop.assert_awaited_once() + + +@pytest.mark.usefixtures("patch_central_database", "sdk") +def test_target_advertises_native_text_only_capabilities() -> None: + capabilities = GitHubCopilotTarget(model_name="gpt-4o").capabilities + assert capabilities.supports_multi_turn is True + assert capabilities.supports_system_prompt is True + assert capabilities.supports_multi_message_pieces is False + assert capabilities.input_modalities == frozenset({frozenset({"text"})}) + assert capabilities.output_modalities == frozenset({frozenset({"text"})}) + + +@pytest.mark.usefixtures("patch_central_database") +@pytest.mark.parametrize( + ("capture_before", "requested_model", "expected_model", "expected_error"), + [ + pytest.param(True, "gpt-5.4", "gpt-5-mini", RuntimeError, id="different-after-capture"), + pytest.param(False, "gpt-5.4", "gpt-5.4", None, id="different-before-capture"), + pytest.param(True, "gpt-5-mini", "gpt-5-mini", None, id="same-after-capture"), + pytest.param(False, " ", "gpt-5-mini", ValueError, id="blank-before-capture"), + pytest.param(True, " ", "gpt-5-mini", ValueError, id="blank-after-capture"), + ], +) +def test_set_model_name_boundaries( + *, + sdk: Any, + client: NonCallableMagicMock, + capture_before: bool, + requested_model: str, + expected_model: str, + expected_error: type[Exception] | None, +) -> None: + target = GitHubCopilotTarget(model_name="gpt-5-mini") + captured_identifier = target.get_identifier() if capture_before else None + + if expected_error is None: + target.set_model_name(model_name=requested_model) + else: + with pytest.raises(expected_error): + target.set_model_name(model_name=requested_model) + + identity = target.get_identifier() + assert identity.params["model_name"] == expected_model + if captured_identifier is not None: + assert identity == captured_identifier + sdk.CopilotClient.assert_not_called() + client.start.assert_not_awaited() + client.create_session.assert_not_awaited() + + +@pytest.mark.usefixtures("patch_central_database") +async def test_direct_send_captures_model_identity_before_startup_async( + *, + client: NonCallableMagicMock, +) -> None: + startup_entered = asyncio.Event() + release_startup = asyncio.Event() + + async def start_async() -> None: + startup_entered.set() + await release_startup.wait() + + client.start.side_effect = start_async + target = GitHubCopilotTarget(model_name="gpt-5-mini") + first_task = asyncio.create_task( + target.send_prompt_async( + message=_user_message( + conversation_id="direct-model-capture", + original_value="first original", + converted_value="first prepared", + ) + ) + ) + try: + await asyncio.wait_for(startup_entered.wait(), timeout=2.0) + with pytest.raises(RuntimeError, match="new target"): + target.set_model_name(model_name="gpt-5.4") + + release_startup.set() + response = await asyncio.wait_for(first_task, timeout=2.0) + assert response[0].get_piece().converted_value == "HELLO" + assert client.create_session.await_args.kwargs["model"] == "gpt-5-mini" + with pytest.raises(RuntimeError, match="new target"): + target.set_model_name(model_name="gpt-5.4") + assert target.get_identifier().params["model_name"] == "gpt-5-mini" + await target.cleanup_target_async() + finally: + release_startup.set() + await _cancel_tasks_async(first_task) + with suppress(Exception): + await asyncio.wait_for(target.cleanup_target_async(), timeout=2.0) + + +@pytest.mark.usefixtures("patch_central_database") +@pytest.mark.parametrize("failure_stage", ["start", "status", "send", "stop"]) +async def test_normalizer_surfaces_lifecycle_failures_async( + *, + client: NonCallableMagicMock, + sqlite_instance: MemoryInterface, + failure_stage: str, +) -> None: + from copilot.client import StopError + + session = client.create_session.return_value + error = ( + ExceptionGroup("SDK shutdown failed", [StopError(message="Synthetic shutdown failure")]) + if failure_stage == "stop" + else RuntimeError(f"Synthetic {failure_stage} failure") + ) + target = GitHubCopilotTarget(model_name="gpt-5-mini") + conversation_id = str(uuid4()) + + if failure_stage == "stop": + response = await _send_normalized_async( + target=target, + original_value="Reply exactly HELLO.", + conversation_id=conversation_id, + ) + assert response.get_piece().response_error == "none" + client.stop.side_effect = error + with pytest.raises(Exception) as exc_info: + await target.cleanup_target_async() + assert exc_info.value is error + client.start.assert_awaited_once() + client.get_status.assert_awaited_once() + client.create_session.assert_awaited_once() + session.send_and_wait.assert_awaited_once() + client.delete_session.assert_awaited_once_with("sdk-session-id") + client.stop.assert_awaited_once() + assert _message_roles_and_errors(memory=sqlite_instance, conversation_id=conversation_id) == [ + ("user", "none"), + ("assistant", "none"), + ] + return + + operation = { + "start": client.start, + "status": client.get_status, + "send": session.send_and_wait, + }[failure_stage] + operation.side_effect = error + with pytest.raises(Exception, match="Error sending prompt with conversation ID:") as exc_info: + await _send_normalized_async( + target=target, + original_value="Reply exactly HELLO.", + conversation_id=conversation_id, + ) + + assert exc_info.value.__cause__ is error + assert _message_roles_and_errors(memory=sqlite_instance, conversation_id=conversation_id) == [ + ("user", "none"), + ("assistant", "processing"), + ] + client.start.assert_awaited_once() + client.get_session_metadata.assert_not_awaited() + if failure_stage == "start": + client.get_status.assert_not_awaited() + client.stop.assert_awaited_once() + else: + client.get_status.assert_awaited_once() + if failure_stage in ("start", "status"): + client.create_session.assert_not_awaited() + session.send_and_wait.assert_not_awaited() + client.delete_session.assert_not_awaited() + else: + client.create_session.assert_awaited_once() + session.send_and_wait.assert_awaited_once() + client.delete_session.assert_awaited_once_with("sdk-session-id") + session.on.return_value.assert_called_once_with() + client.stop.assert_not_awaited() + session.send.assert_not_awaited() + await target.cleanup_target_async() + if failure_stage == "send": + client.stop.assert_awaited_once() + + +@pytest.mark.usefixtures("patch_central_database") +async def test_normalizer_surfaces_dispatch_timeout_and_cleans_up_without_replay_async( + *, + sdk: Any, + client: NonCallableMagicMock, +) -> None: + session = client.create_session.return_value + target = GitHubCopilotTarget(model_name="gpt-5-mini", response_timeout_seconds=0.01) + + async def stall_send_async(*_args: Any, **_kwargs: Any) -> None: + await asyncio.Event().wait() + + session.send.side_effect = stall_send_async + session.send_and_wait.side_effect = partial(sdk.CopilotSession.send_and_wait, session) + with pytest.raises(Exception, match="Error sending prompt with conversation ID:") as exc_info: + # A bare watchdog TimeoutError must not satisfy the normalizer-wrapped failure. + await asyncio.wait_for( + _send_normalized_async(target=target, original_value="Reply exactly HELLO."), + timeout=2.0, + ) + assert isinstance(exc_info.value.__cause__, TimeoutError) + session.send.assert_awaited_once() + client.delete_session.assert_awaited_once_with("sdk-session-id") + client.stop.assert_not_awaited() + await target.cleanup_target_async() + client.stop.assert_awaited_once() + + +@pytest.mark.usefixtures("patch_central_database") +@pytest.mark.parametrize("invalid_reply", ["empty", "subagent", "absent", "non-assistant"]) +async def test_normalizer_rejects_invalid_reply_async( + *, + client: NonCallableMagicMock, + sqlite_instance: MemoryInterface, + invalid_reply: str, +) -> None: + from copilot.generated.session_events import SessionEventType, SessionIdleData + + reply = _assistant_reply("HELLO") + session = client.create_session.return_value + session.send_and_wait.return_value = { + "empty": _assistant_reply(""), + "subagent": replace(reply, agent_id="sdk-subagent-id"), + "absent": None, + "non-assistant": replace(reply, type=SessionEventType.SESSION_IDLE, data=SessionIdleData(aborted=False)), + }[invalid_reply] + conversation_id = str(uuid4()) + target = GitHubCopilotTarget(model_name="gpt-5-mini") + with pytest.raises(Exception, match="Error sending prompt with conversation ID:") as exc_info: + await _send_normalized_async( + target=target, + original_value="Reply exactly HELLO.", + conversation_id=conversation_id, + ) + assert isinstance(exc_info.value.__cause__, ValueError) + assert _message_roles_and_errors(memory=sqlite_instance, conversation_id=conversation_id) == [ + ("user", "none"), + ("assistant", "processing"), + ] + session.send_and_wait.assert_awaited_once() + session.on.return_value.assert_called_once_with() + client.delete_session.assert_awaited_once_with("sdk-session-id") + client.stop.assert_not_awaited() + await target.cleanup_target_async() + client.stop.assert_awaited_once() + + +@pytest.mark.usefixtures("patch_central_database") +@pytest.mark.parametrize( + ("event_stream", "expected_error"), + [("aborted-idle", "abort"), ("abort-then-idle", "abort"), ("tool-then-idle", "tool")], + ids=["aborted-idle", "abort-then-idle", "tool-then-idle"], +) +async def test_normalizer_rejects_unsafe_events_async( + *, + sdk: Any, + client: NonCallableMagicMock, + sqlite_instance: MemoryInterface, + event_stream: str, + expected_error: str, +) -> None: + from copilot.generated.session_events import AbortData, SessionEventType, SessionIdleData, ToolExecutionStartData + + session = client.create_session.return_value + reply = _assistant_reply("HELLO") + events = { + "aborted-idle": [ + reply, + replace(reply, id=uuid4(), type=SessionEventType.SESSION_IDLE, data=SessionIdleData(aborted=True)), + ], + "abort-then-idle": [ + reply, + replace( + reply, + id=uuid4(), + type=SessionEventType.ABORT, + data=AbortData.from_dict({"reason": "user_initiated"}), + ), + replace(reply, id=uuid4(), type=SessionEventType.SESSION_IDLE, data=SessionIdleData(aborted=None)), + ], + "tool-then-idle": [ + replace( + reply, + id=uuid4(), + type=SessionEventType.TOOL_EXECUTION_START, + data=ToolExecutionStartData( + tool_call_id="synthetic-tool-call", tool_name="benign_test_tool", arguments={"text": "HELLO"} + ), + ), + reply, + replace(reply, id=uuid4(), type=SessionEventType.SESSION_IDLE, data=SessionIdleData(aborted=False)), + ], + }[event_stream] + handlers: list[Callable[[SessionEvent], None]] = [] + + def subscribe(handler: Callable[[SessionEvent], None]) -> Callable[[], None]: + handlers.append(handler) + return partial(handlers.remove, handler) + + async def send_events_async(*_args: Any, **_kwargs: Any) -> str: + for event in events: + for handler in tuple(handlers): + handler(event) + return "sdk-request-id" + + session.on.side_effect = subscribe + session.send.side_effect = send_events_async + session.send_and_wait.side_effect = partial(sdk.CopilotSession.send_and_wait, session) + conversation_id = str(uuid4()) + target = GitHubCopilotTarget(model_name="gpt-5-mini", response_timeout_seconds=1.0) + with pytest.raises(Exception, match="Error sending prompt with conversation ID:") as exc_info: + await _send_normalized_async( + target=target, + original_value="Reply exactly HELLO.", + conversation_id=conversation_id, + ) + assert isinstance(exc_info.value.__cause__, RuntimeError) + assert expected_error in str(exc_info.value.__cause__).lower() + assert _message_roles_and_errors(memory=sqlite_instance, conversation_id=conversation_id) == [ + ("user", "none"), + ("assistant", "processing"), + ] + session.send.assert_awaited_once() + client.delete_session.assert_awaited_once_with("sdk-session-id") + client.stop.assert_not_awaited() + await target.cleanup_target_async() + client.stop.assert_awaited_once() + assert not handlers + + +@pytest.mark.usefixtures("patch_central_database") +@pytest.mark.parametrize("retain_session", [False, True], ids=["delete", "retain"]) +async def test_normalizer_cleans_up_partial_creation_async( + *, + client: NonCallableMagicMock, + sqlite_instance: MemoryInterface, + caplog: pytest.LogCaptureFixture, + retain_session: bool, +) -> None: + session = client.create_session.return_value + sessions = {"unrelated-session-id"} + allocated_session_id = "" + creation_error = RuntimeError("Copilot post-create options update failed") + + async def create_session_async(*, session_id: str = "sdk-generated-session-id", **_kwargs: Any) -> None: + nonlocal allocated_session_id + allocated_session_id = session_id + sessions.add(session_id) + raise creation_error + + client.create_session.side_effect = create_session_async + _mock_session_storage(client=client, sessions=sessions) + conversation_id = str(uuid4()) + target = GitHubCopilotTarget(model_name="gpt-5-mini", retain_session=retain_session) + with caplog.at_level(logging.INFO, logger=TARGET_LOGGER): + with pytest.raises(Exception, match="Error sending prompt with conversation ID:") as exc_info: + await _send_normalized_async( + target=target, + original_value="Reply exactly HELLO.", + conversation_id=conversation_id, + ) + + assert exc_info.value.__cause__ is creation_error + assert allocated_session_id == client.create_session.await_args.kwargs["session_id"] + assert str(UUID(allocated_session_id)) == allocated_session_id + assert allocated_session_id != "sdk-session-id" + client.create_session.assert_awaited_once() + client.get_session_metadata.assert_awaited_once_with(allocated_session_id) + session.send_and_wait.assert_not_awaited() + session.send.assert_not_awaited() + client.stop.assert_not_awaited() + await target.cleanup_target_async() + client.stop.assert_awaited_once() + assert _message_roles_and_errors(memory=sqlite_instance, conversation_id=conversation_id) == [ + ("user", "none"), + ("assistant", "processing"), + ] + retained_logs = [ + r.getMessage() + for r in caplog.records + if r.name == TARGET_LOGGER + and r.levelno == logging.INFO + and r.getMessage().startswith("Retaining Copilot session ") + ] + if retain_session: + client.delete_session.assert_not_awaited() + assert sessions == {"unrelated-session-id", allocated_session_id} + assert retained_logs == [f"Retaining Copilot session {allocated_session_id} as requested; delete it manually."] + else: + client.delete_session.assert_awaited_once_with(allocated_session_id) + assert sessions == {"unrelated-session-id"} + assert not retained_logs + + +@pytest.mark.usefixtures("patch_central_database") +async def test_normalizer_deletes_owned_session_when_creation_is_cancelled_after_allocation_async( + client: NonCallableMagicMock, +) -> None: + session = client.create_session.return_value + sessions = {"unrelated-session-id"} + allocated = asyncio.Event() + + async def create_session_async(*, session_id: str = "sdk-generated-session-id", **_kwargs: Any) -> None: + sessions.add(session_id) + allocated.set() + await asyncio.Event().wait() + + client.create_session.side_effect = create_session_async + _mock_session_storage(client=client, sessions=sessions) + target = GitHubCopilotTarget(model_name="gpt-5-mini") + request_task = asyncio.create_task(_send_normalized_async(target=target, original_value="Reply exactly HELLO.")) + try: + await asyncio.wait_for(allocated.wait(), timeout=2.0) + request_task.cancel() + with pytest.raises(asyncio.CancelledError): + await request_task + finally: + await _cancel_tasks_async(request_task) + + client.create_session.assert_awaited_once() + allocated_id = client.create_session.await_args.kwargs["session_id"] + client.get_session_metadata.assert_awaited_once_with(allocated_id) + client.delete_session.assert_awaited_once_with(allocated_id) + session.send_and_wait.assert_not_awaited() + session.send.assert_not_awaited() + client.stop.assert_not_awaited() + await target.cleanup_target_async() + client.stop.assert_awaited_once() + assert sessions == {"unrelated-session-id"} + + +@pytest.mark.usefixtures("patch_central_database") +async def test_normalizer_stops_owned_client_when_startup_is_cancelled_async( + *, + sdk: Any, + client: NonCallableMagicMock, +) -> None: + session = client.create_session.return_value + owned_resources: set[str] = set() + allocated = asyncio.Event() + original_cancellation: asyncio.CancelledError | None = None + + async def start_async() -> None: + nonlocal original_cancellation + owned_resources.add("owned-runtime") + allocated.set() + try: + await asyncio.Event().wait() + except asyncio.CancelledError as error: + original_cancellation = error + raise + + async def stop_async() -> None: + owned_resources.remove("owned-runtime") + + client.start.side_effect = start_async + client.stop.side_effect = stop_async + target = GitHubCopilotTarget(model_name="gpt-5-mini") + request_task = asyncio.create_task(_send_normalized_async(target=target, original_value="Reply exactly HELLO.")) + try: + await asyncio.wait_for(allocated.wait(), timeout=2.0) + request_task.cancel() + with pytest.raises(asyncio.CancelledError) as exc_info: + await request_task + finally: + await _cancel_tasks_async(request_task) + + assert exc_info.value is original_cancellation + sdk.CopilotClient.assert_called_once() + client.start.assert_awaited_once() + client.get_status.assert_not_awaited() + client.create_session.assert_not_awaited() + client.get_session_metadata.assert_not_awaited() + client.delete_session.assert_not_awaited() + session.send.assert_not_awaited() + session.send_and_wait.assert_not_awaited() + client.stop.assert_awaited_once() + assert owned_resources == set() + await target.cleanup_target_async() + + +@pytest.mark.usefixtures("patch_central_database") +@pytest.mark.parametrize( + ("github_token", "use_working_directory", "max_requests_per_minute"), + [ + pytest.param(" dummy-github-token ", False, None, id="token-only"), + pytest.param(None, True, None, id="directory-only"), + pytest.param(None, False, 30, id="throttle-only"), + pytest.param(" dummy-github-token ", True, 30, id="all-options"), + ], +) +async def test_normalizer_forwards_options_without_exposing_token_async( + *, + sdk: Any, + client: NonCallableMagicMock, + caplog: pytest.LogCaptureFixture, + tmp_path: Path, + github_token: str | None, + use_working_directory: bool, + max_requests_per_minute: int | None, +) -> None: + with ( + caplog.at_level(logging.DEBUG, logger=TARGET_LOGGER), + patch.object(asyncio, "sleep", new_callable=AsyncMock) as mock_sleep, + ): + target = GitHubCopilotTarget( + model_name="gpt-5-mini", + github_token=github_token, + working_directory=tmp_path if use_working_directory else None, + max_requests_per_minute=max_requests_per_minute, + ) + response = await _send_normalized_async(target=target, original_value="Reply exactly HELLO.") + await target.cleanup_target_async() + sdk.CopilotClient.assert_called_once_with( + github_token=github_token, working_directory=str(tmp_path) if use_working_directory else None + ) + assert response.get_piece().converted_value == "HELLO" + client.create_session.return_value.send_and_wait.assert_awaited_once() + if max_requests_per_minute is not None: + mock_sleep.assert_awaited_once_with(2.0) + assert target.get_identifier().params["max_requests_per_minute"] == 30 + else: + mock_sleep.assert_not_awaited() + if use_working_directory: + assert target.get_identifier().params["working_directory"] == str(tmp_path) + assert "dummy-github-token" not in target.get_identifier().model_dump_json() + assert "dummy-github-token" not in caplog.text + + +def test_init_without_copilot_sdk_reports_installation_guidance() -> None: + with patch.dict("sys.modules", {"copilot": None}): + with pytest.raises(RuntimeError, match=r"pip install pyrit\[github-copilot\]"): + GitHubCopilotTarget(model_name="gpt-5-mini") + + +@pytest.mark.parametrize( + ("overrides", "field"), + [ + pytest.param({"model_name": " "}, "model_name", id="blank-model"), + pytest.param({"github_token": " "}, "github_token", id="blank-token"), + pytest.param({"response_timeout_seconds": 0}, "response_timeout_seconds", id="zero-timeout"), + pytest.param({"response_timeout_seconds": -1}, "response_timeout_seconds", id="negative-timeout"), + pytest.param({"response_timeout_seconds": float("inf")}, "response_timeout_seconds", id="infinite-timeout"), + pytest.param({"response_timeout_seconds": float("nan")}, "response_timeout_seconds", id="nan-timeout"), + pytest.param({"working_directory": " "}, "working_directory", id="blank-directory"), + ], +) +def test_init_rejects_invalid_options_before_sdk_import(*, overrides: dict[str, Any], field: str) -> None: + with patch.dict("sys.modules", {"copilot": None}), pytest.raises(ValueError, match=field): + GitHubCopilotTarget(**{"model_name": "gpt-5-mini", **overrides}) + + +@pytest.mark.parametrize("path_kind", ["missing", "file"]) +def test_init_rejects_non_directory_before_sdk_import(*, tmp_path: Path, path_kind: str) -> None: + path = tmp_path / "not-a-directory" + if path_kind == "file": + path.write_text("local test fixture", encoding="utf-8") + with patch.dict("sys.modules", {"copilot": None}), pytest.raises(ValueError, match="working_directory"): + GitHubCopilotTarget(model_name="gpt-5-mini", working_directory=path) + + +@pytest.mark.usefixtures("patch_central_database") +async def test_normalizer_keeps_event_loop_responsive_during_client_construction_async( + *, sdk: Any, client: NonCallableMagicMock +) -> None: + loop = asyncio.get_running_loop() + constructor_entered = asyncio.Event() + constructor_finished = asyncio.Event() + release = threading.Event() + released_while_constructing = False + + def construct_client(*, github_token: str | None, working_directory: str | None) -> NonCallableMagicMock: + nonlocal released_while_constructing + try: + loop.call_soon_threadsafe(constructor_entered.set) + # The bound lets a blocked event loop escape without satisfying the responsiveness assertion. + released_while_constructing = release.wait(timeout=5.0) + return client + finally: + loop.call_soon_threadsafe(constructor_finished.set) + + async def release_constructor_async() -> None: + await constructor_entered.wait() + release.set() + + sdk.CopilotClient.side_effect = construct_client + conversation_id = str(uuid4()) + target = GitHubCopilotTarget(model_name="gpt-5-mini") + request_task = asyncio.create_task( + _send_normalized_async( + target=target, + original_value="Reply exactly HELLO.", + conversation_id=conversation_id, + ) + ) + release_task = asyncio.create_task(release_constructor_async()) + try: + await asyncio.wait_for(asyncio.gather(request_task, release_task), timeout=10.0) + finally: + release.set() + await _cancel_tasks_async(request_task, release_task) + await asyncio.wait_for(constructor_finished.wait(), timeout=5.0) + + await target.cleanup_target_async() + client.stop.assert_awaited_once() + client.delete_session.assert_awaited_once_with("sdk-session-id") + assert released_while_constructing is True diff --git a/tests/unit/prompt_target/target/test_prompt_target.py b/tests/unit/prompt_target/target/test_prompt_target.py index 453870e382..6abaf85b4e 100644 --- a/tests/unit/prompt_target/target/test_prompt_target.py +++ b/tests/unit/prompt_target/target/test_prompt_target.py @@ -12,7 +12,14 @@ from pyrit.executor.attack.core.attack_strategy import AttackStrategy from pyrit.memory.memory_interface import MemoryInterface -from pyrit.models import ChatMessageRole, ComponentIdentifier, Message, MessagePiece, flatten_to_message_pieces +from pyrit.models import ( + ChatMessageRole, + ComponentIdentifier, + Conversation, + Message, + MessagePiece, + flatten_to_message_pieces, +) from pyrit.prompt_target import OpenAIChatTarget from pyrit.prompt_target.common.target_capabilities import ( CapabilityHandlingPolicy, @@ -81,6 +88,109 @@ async def test_set_system_prompt_adds_memory( assert chats[0].api_role == "system" +@pytest.mark.usefixtures("patch_central_database") +@pytest.mark.parametrize( + ("supports_multi_turn", "supports_editable_history", "supports_system_prompt", "policy", "expected_error"), + [ + pytest.param(True, False, True, None, None, id="native-system-without-editable-history"), + pytest.param( + True, + True, + False, + CapabilityHandlingPolicy(behaviors={CapabilityName.SYSTEM_PROMPT: UnsupportedCapabilityBehavior.ADAPT}), + None, + id="editable-history-without-native-system", + ), + pytest.param(False, True, True, None, ValueError, id="without-multi-turn-editable"), + pytest.param(False, False, True, None, ValueError, id="without-multi-turn-native-system"), + pytest.param(True, False, False, None, ValueError, id="without-editable-or-native-system"), + ], +) +def test_set_system_prompt_capability_admission_and_nonmutation( + *, + sqlite_instance: MemoryInterface, + supports_multi_turn: bool, + supports_editable_history: bool, + supports_system_prompt: bool, + policy: CapabilityHandlingPolicy | None, + expected_error: type[ValueError] | None, +) -> None: + conversation_id = "system-prompt-capability-conversation" + target = _make_identifier_target( + capabilities=TargetCapabilities( + supports_multi_turn=supports_multi_turn, + supports_editable_history=supports_editable_history, + supports_system_prompt=supports_system_prompt, + ), + policy=policy, + ) + + if expected_error is None: + target.set_system_prompt(system_prompt="be concise", conversation_id=conversation_id) + stored = sqlite_instance.get_conversation_messages(conversation_id=conversation_id) + assert len(stored) == 1 + piece = stored[0].get_piece() + assert piece.api_role == "system" + assert piece.converted_value == "be concise" + assert piece.conversation_id == conversation_id + else: + with pytest.raises( + expected_error, + match="It must support multi-turn conversations and either editable history or native system prompts.", + ): + target.set_system_prompt(system_prompt="be concise", conversation_id=conversation_id) + assert sqlite_instance.get_conversation_messages(conversation_id=conversation_id) == [] + assert sqlite_instance.get_target_identifiers(identifier_hashes=[target.get_identifier().hash]) == [] + + +@pytest.mark.usefixtures("patch_central_database") +@pytest.mark.parametrize( + ("existing_role", "existing_content"), + [ + pytest.param("system", "be concise", id="repeated-system-prompt"), + pytest.param("user", "existing user message", id="existing-user-message"), + ], +) +def test_set_system_prompt_rejects_nonempty_conversation_without_mutation( + *, + sqlite_instance: MemoryInterface, + existing_role: str, + existing_content: str, +) -> None: + conversation_id = "nonempty-system-prompt-conversation" + target = _make_identifier_target( + capabilities=TargetCapabilities( + supports_multi_turn=True, + supports_editable_history=False, + supports_system_prompt=True, + ) + ) + if existing_role == "system": + target.set_system_prompt(system_prompt=existing_content, conversation_id=conversation_id) + else: + sqlite_instance.add_conversation_to_memory( + conversation=Conversation(conversation_id=conversation_id, target_identifier=target.get_identifier()) + ) + sqlite_instance.add_message_to_memory( + request=MessagePiece( + role="user", + conversation_id=conversation_id, + original_value=existing_content, + converted_value=existing_content, + ).to_message() + ) + + with pytest.raises(RuntimeError, match="Conversation already exists"): + target.set_system_prompt(system_prompt="be expansive", conversation_id=conversation_id) + + stored = sqlite_instance.get_conversation_messages(conversation_id=conversation_id) + assert len(stored) == 1 + piece = stored[0].get_piece() + assert piece.api_role == existing_role + assert piece.converted_value == existing_content + assert piece.conversation_id == conversation_id + + async def test_send_prompt_with_system_calls_chat_complete( azure_openai_target: OpenAIChatTarget, openai_response_json: dict, diff --git a/tests/unit/score/test_scorer.py b/tests/unit/score/test_scorer.py index c16ad97b8c..58c4321ce9 100644 --- a/tests/unit/score/test_scorer.py +++ b/tests/unit/score/test_scorer.py @@ -3,27 +3,35 @@ import asyncio import uuid +from dataclasses import replace from textwrap import dedent from unittest.mock import AsyncMock, MagicMock, patch import pytest -from unit.mocks import get_mock_target_identifier, store_message +from unit.mocks import MockPromptTarget, get_mock_target_identifier, store_message +import pyrit.score.scorer as scorer_module from pyrit.exceptions import InvalidJsonException, remove_markdown_json from pyrit.memory import CentralMemory, MemoryInterface from pyrit.models import ( + Acquisition, AnswerMatches, ChatMessageRole, ComponentIdentifier, ContentScorable, Message, MessagePiece, + PromptDataType, Scorable, Score, + ScorerTargetResponsePayload, ScoreStatus, ScoringExpectation, ) -from pyrit.prompt_target import PromptTarget +from pyrit.prompt_target import CapabilityName, PromptTarget +from pyrit.prompt_target.common import target_requirements as target_requirements_module +from pyrit.prompt_target.common.target_capabilities import TargetCapabilities +from pyrit.prompt_target.common.target_configuration import TargetConfiguration from pyrit.score import ( FloatScaleScorer, FloatScaleThresholdScorer, @@ -34,7 +42,10 @@ MessageTrueFalseScorer, Scorer, ScorerPromptValidator, + SelfAskRefusalScorer, + SelfAskTrueFalseScorer, TrueFalseInverterScorer, + TrueFalseQuestion, TrueFalseScorer, ) from pyrit.score.llm_scoring import _run_llm_scoring_async @@ -122,6 +133,269 @@ def validate_return_scores(self, scores: list[Score]): assert all(s.score_value in ["true", "false"] for s in scores if s.status != ScoreStatus.UNDETERMINED) +def _make_mock_judge_target(*, editable_history: bool = False) -> MockPromptTarget: + return MockPromptTarget( + custom_configuration=TargetConfiguration( + capabilities=TargetCapabilities( + supports_multi_turn=True, + supports_multi_message_pieces=True, + supports_system_prompt=True, + supports_editable_history=editable_history, + ) + ) + ) + + +@pytest.mark.usefixtures("patch_central_database") +@pytest.mark.parametrize( + "scorer_type", + [ + pytest.param(SelfAskRefusalScorer, id="refusal"), + pytest.param(SelfAskTrueFalseScorer, id="true-false"), + ], +) +@pytest.mark.parametrize( + ( + "extra_required", + "extra_native_required", + "extra_input_modalities", + "extra_output_modalities", + "rejected_routes", + "error_fragments", + ), + [ + pytest.param( + frozenset({CapabilityName.STREAMING_AUDIO}), + frozenset(), + frozenset(), + frozenset(), + frozenset({"editable", "native"}), + ("supports_streaming_audio",), + id="additional-required-capability", + ), + pytest.param( + frozenset(), + frozenset({CapabilityName.EDITABLE_HISTORY}), + frozenset(), + frozenset(), + frozenset({"native"}), + ("natively support 'supports_editable_history'",), + id="native-required-editable-history", + ), + pytest.param( + frozenset(), + frozenset(), + frozenset({frozenset({"audio_path"})}), + frozenset({frozenset({"audio_path"})}), + frozenset({"editable", "native"}), + ("input modality {audio_path}", "output modality {audio_path}"), + id="additional-input-output-modalities", + ), + ], +) +def test_self_ask_scorers_preserve_shared_target_requirements( + *, + scorer_type: type[SelfAskRefusalScorer] | type[SelfAskTrueFalseScorer], + extra_required: frozenset[CapabilityName], + extra_native_required: frozenset[CapabilityName], + extra_input_modalities: frozenset[frozenset[PromptDataType]], + extra_output_modalities: frozenset[frozenset[PromptDataType]], + rejected_routes: frozenset[str], + error_fragments: tuple[str, ...], +) -> None: + shared_requirements = target_requirements_module.CHAT_TARGET_REQUIREMENTS + shared_fields = ( + shared_requirements.required, + shared_requirements.native_required, + shared_requirements.required_input_modalities, + shared_requirements.required_output_modalities, + ) + targets = ( + ("editable", _make_mock_judge_target(editable_history=True)), + ("native", _make_mock_judge_target()), + ) + + for _, target in targets: + scorer_type(chat_target=target) + + future_requirements = replace( + shared_requirements, + required=shared_requirements.required | extra_required, + native_required=shared_requirements.native_required | extra_native_required, + required_input_modalities=shared_requirements.required_input_modalities | extra_input_modalities, + required_output_modalities=shared_requirements.required_output_modalities | extra_output_modalities, + ) + with patch.object(scorer_module, "CHAT_TARGET_REQUIREMENTS", future_requirements): + for route, target in targets: + if route in rejected_routes: + with pytest.raises(ValueError) as exc_info: + scorer_type(chat_target=target) + assert all(fragment in str(exc_info.value) for fragment in error_fragments) + else: + scorer_type(chat_target=target) + + assert target_requirements_module.CHAT_TARGET_REQUIREMENTS is shared_requirements + assert ( + shared_requirements.required, + shared_requirements.native_required, + shared_requirements.required_input_modalities, + shared_requirements.required_output_modalities, + ) == shared_fields + + +@pytest.mark.usefixtures("patch_central_database") +@pytest.mark.parametrize( + ( + "scorer_kind", + "answer", + "judge_json", + "expected_score", + "expected_rationale", + "expected_judge_prompt", + "expected_system_prompt_fragment", + ), + [ + pytest.param( + "refusal", + "I cannot answer that question.", + '{"score_value":"True","description":"Refusal detected","rationale":"The response explicitly refuses."}', + True, + "The response explicitly refuses.", + "conversation_objective: Name France's capital\nresponse_to_evaluate_input: I cannot answer that question.", + None, + id="refusal", + ), + pytest.param( + "refusal", + "France's capital is Paris.", + '{"score_value":"False","description":"No refusal detected",' + '"rationale":"The response answers the objective."}', + False, + "The response answers the objective.", + "conversation_objective: Name France's capital\nresponse_to_evaluate_input: France's capital is Paris.", + None, + id="answer", + ), + pytest.param( + "true_false", + "Paris is the capital of France.", + '{"score_value":true,"description":"Correct","rationale":"Paris is the capital of France."}', + True, + "Paris is the capital of France.", + "objective: Name France's capital\nresponse: Paris is the capital of France.", + "The response correctly identifies Paris as the capital of France.", + id="true-false", + ), + ], +) +async def test_self_ask_scorers_use_fresh_conversations_for_noneditable_text_judgments_async( + *, + sqlite_instance: MemoryInterface, + scorer_kind: str, + answer: str, + judge_json: str, + expected_score: bool, + expected_rationale: str, + expected_judge_prompt: str, + expected_system_prompt_fragment: str | None, +) -> None: + target = _make_mock_judge_target() + sent_conversations: list[list[Message]] = [] + sent_conversation_ids: list[str] = [] + reset_conversation_ids: list[str] = [] + sent_response_piece_ids: list[uuid.UUID] = [] + + async def send_judge_reply_async(*, normalized_conversation: list[Message]) -> list[Message]: + sent_conversations.append(normalized_conversation) + conversation_id = normalized_conversation[-1].get_piece().conversation_id + sent_conversation_ids.append(conversation_id) + if len(sent_conversation_ids) == 2: + assert reset_conversation_ids == [sent_conversation_ids[0]] + response_text = "malformed judge response" if len(sent_conversation_ids) == 1 else judge_json + response = MessagePiece( + role="assistant", + original_value=response_text, + conversation_id=conversation_id, + ).to_message() + sent_response_piece_ids.append(response.get_piece().id) + return [response] + + async def reset_conversation_async(*, conversation_id: str) -> None: + reset_conversation_ids.append(conversation_id) + + target_send = AsyncMock(side_effect=send_judge_reply_async) + reset = AsyncMock(side_effect=reset_conversation_async) + with ( + patch.object(target, "_send_prompt_to_target_async", new=target_send), + patch.object(target, "reset_conversation_async", new=reset), + ): + saved_answer = store_message( + MessagePiece( + role="assistant", + conversation_id=str(uuid.uuid4()), + original_value=answer, + ).to_message() + ) + input_piece = saved_answer.get_piece() + input_scorable = MessageScorable.from_message(saved_answer) + + if scorer_kind == "refusal": + scorer = SelfAskRefusalScorer(chat_target=target) + else: + question = TrueFalseQuestion( + category="capital correctness", + true_description="The response correctly identifies Paris as the capital of France.", + false_description="The response does not correctly identify Paris as the capital of France.", + ) + scorer = SelfAskTrueFalseScorer.from_question(chat_target=target, question=question) + + expectation = ScoringExpectation(objective="Name France's capital") + scores = await scorer.score_async(scorable=input_scorable, expectation=expectation) + + assert len(scores) == 1 + assert scores[0].get_value() is expected_score + assert scores[0].score_rationale == expected_rationale + assert scores[0].message_piece_id == input_piece.id + assert scores[0].scorable == input_scorable + assert len(scores[0].observation_ids) == 1 + observation = sqlite_instance.get_observations(observation_ids=scores[0].observation_ids)[0] + assert observation.acquisition is Acquisition.COMPLETE + assert isinstance(observation.payload, ScorerTargetResponsePayload) + assert observation.payload.message_piece_ids == (sent_response_piece_ids[1],) + assert target_send.await_count == 2 + assert len(set(sent_conversation_ids)) == 2 + assert reset_conversation_ids == [sent_conversation_ids[0]] + + user_pieces = [conversation[-1].get_piece() for conversation in sent_conversations] + assert all(piece.original_value == expected_judge_prompt for piece in user_pieces) + assert user_pieces[0].converted_value == user_pieces[1].converted_value + assert "The response should conform to the following JSON schema:" in user_pieces[0].converted_value + assert '"score_value"' in user_pieces[0].converted_value + assert '"rationale"' in user_pieces[0].converted_value + + memory_pieces = sqlite_instance.get_message_pieces() + system_pieces = [ + piece for piece in memory_pieces if piece.role == "system" and piece.conversation_id in sent_conversation_ids + ] + assert len(system_pieces) == 2 + assert system_pieces[0].original_value == system_pieces[1].original_value + if expected_system_prompt_fragment is not None: + assert expected_system_prompt_fragment in system_pieces[0].original_value + assert any( + piece.role == "assistant" + and piece.original_value == "malformed judge response" + and piece.conversation_id == sent_conversation_ids[0] + for piece in memory_pieces + ) + stored_answer = sqlite_instance.get_message_pieces(prompt_ids=[input_piece.id]) + assert len(stored_answer) == 1 + assert stored_answer[0].original_value == answer + replayed_scores = await scorer.score_observation_async(observation=observation, expectation=expectation) + assert len(replayed_scores) == 1 + assert replayed_scores[0].get_value() is expected_score + assert target_send.await_count == 2 + + class SelectiveValidator(ScorerPromptValidator): """Validator that only supports text pieces, not images.""" @@ -241,6 +515,101 @@ async def test_scorer_score_value_with_llm_exception_display_prompt_id(patch_cen ) +@pytest.mark.usefixtures("patch_central_database") +async def test_fresh_llm_scoring_bounds_invalid_json_retries_and_resets_each_attempt( + sqlite_instance: MemoryInterface, +) -> None: + target = _make_mock_judge_target() + attempted_conversation_ids: list[str] = [] + reset_conversation_ids: list[str] = [] + + async def send_invalid_json_async(*, normalized_conversation: list[Message]) -> list[Message]: + conversation_id = normalized_conversation[-1].get_piece().conversation_id + attempted_conversation_ids.append(conversation_id) + return [ + MessagePiece( + role="assistant", + original_value=BAD_JSON, + conversation_id=conversation_id, + ).to_message() + ] + + async def reset_async(*, conversation_id: str) -> None: + reset_conversation_ids.append(conversation_id) + + target_send = AsyncMock(side_effect=send_invalid_json_async) + reset = AsyncMock(side_effect=reset_async) + scorer = MockScorer() + with ( + patch.object(target, "_send_prompt_to_target_async", new=target_send), + patch.object(target, "reset_conversation_async", new=reset), + pytest.raises(InvalidJsonException), + ): + await _run_llm_scoring_async( + chat_target=target, + response_handler=JsonSchemaResponseHandler(), + scorer_identifier=scorer.get_identifier(), + system_prompt="Judge this answer.", + value="The answer to judge.", + data_type="text", + scored_prompt_id="saved-answer-id", + objective="Name France's capital", + fresh_conversation_per_attempt=True, + ) + + assert target_send.await_count == 2 + assert len(set(attempted_conversation_ids)) == 2 + assert reset_conversation_ids == attempted_conversation_ids + + +@pytest.mark.usefixtures("patch_central_database") +async def test_fresh_llm_scoring_does_not_retry_when_conversation_reset_fails( + sqlite_instance: MemoryInterface, +) -> None: + target = _make_mock_judge_target() + reset_failure = RuntimeError("native session release failed") + + async def send_invalid_json_async(*, normalized_conversation: list[Message]) -> list[Message]: + conversation_id = normalized_conversation[-1].get_piece().conversation_id + return [ + MessagePiece( + role="assistant", + original_value=BAD_JSON, + conversation_id=conversation_id, + ).to_message() + ] + + async def fail_reset_async(*, conversation_id: str) -> None: + raise reset_failure + + target_send = AsyncMock(side_effect=send_invalid_json_async) + reset = AsyncMock(side_effect=fail_reset_async) + scorer = MockScorer() + with ( + patch.object(target, "_send_prompt_to_target_async", new=target_send), + patch.object(target, "reset_conversation_async", new=reset), + pytest.raises(Exception, match="Error scoring prompt with original prompt ID: saved-answer-id") as exc_info, + ): + await _run_llm_scoring_async( + chat_target=target, + response_handler=JsonSchemaResponseHandler(), + scorer_identifier=scorer.get_identifier(), + system_prompt="Judge this answer.", + value="The answer to judge.", + data_type="text", + scored_prompt_id="saved-answer-id", + objective="Name France's capital", + fresh_conversation_per_attempt=True, + ) + + reset_error = exc_info.value.__cause__ + assert isinstance(reset_error, RuntimeError) + assert "Could not release the malformed judge session" in str(reset_error) + assert reset_error.__cause__ is reset_failure + target_send.assert_awaited_once() + reset.assert_awaited_once() + + async def test_scorer_send_chat_target_async_good_response(good_json, patch_central_database): chat_target = MagicMock(PromptTarget) chat_target.get_identifier.return_value = get_mock_target_identifier("MockChatTarget") diff --git a/tests/unit/score/test_self_ask_question_answer_scorer.py b/tests/unit/score/test_self_ask_question_answer_scorer.py index 83d124607f..b7810dfa1b 100644 --- a/tests/unit/score/test_self_ask_question_answer_scorer.py +++ b/tests/unit/score/test_self_ask_question_answer_scorer.py @@ -4,7 +4,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest -from unit.mocks import store_message +from unit.mocks import MockPromptTarget, store_message from pyrit.memory import MemoryInterface from pyrit.models import ( @@ -19,6 +19,8 @@ UnvalidatedScore, ) from pyrit.prompt_target import PromptTarget +from pyrit.prompt_target.common.target_capabilities import TargetCapabilities +from pyrit.prompt_target.common.target_configuration import TargetConfiguration from pyrit.score import ( MessageScorable, NonReplayableObservationError, @@ -74,6 +76,21 @@ async def test_score_async_returns_score_from_unvalidated(mock_chat_target): assert scores[0].get_value() is True +@pytest.mark.usefixtures("patch_central_database") +def test_question_answer_scorer_keeps_editable_history_requirement() -> None: + target = MockPromptTarget( + custom_configuration=TargetConfiguration( + capabilities=TargetCapabilities( + supports_multi_turn=True, + supports_system_prompt=True, + ) + ) + ) + + with pytest.raises(ValueError, match="supports_editable_history"): + SelfAskQuestionAnswerScorer(chat_target=target) + + @pytest.mark.parametrize("objective", [None, "What is the capital of France?"]) async def test_typed_answer_supplies_judge_ground_truth_async( mock_chat_target: MagicMock, objective: str | None diff --git a/tests/unit/score/test_self_ask_refusal.py b/tests/unit/score/test_self_ask_refusal.py index 29c49833d9..838e9665f8 100644 --- a/tests/unit/score/test_self_ask_refusal.py +++ b/tests/unit/score/test_self_ask_refusal.py @@ -20,7 +20,12 @@ MessagePiece, SeedPrompt, ) -from pyrit.score import JsonSchemaResponseHandler, MessageScorable, RefusalScorerPaths, SelfAskRefusalScorer +from pyrit.score import ( + JsonSchemaResponseHandler, + MessageScorable, + RefusalScorerPaths, + SelfAskRefusalScorer, +) @pytest.fixture diff --git a/tests/unit/score/test_self_ask_true_false.py b/tests/unit/score/test_self_ask_true_false.py index 001bfd7814..f3e88097a1 100644 --- a/tests/unit/score/test_self_ask_true_false.py +++ b/tests/unit/score/test_self_ask_true_false.py @@ -1,16 +1,19 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. +from pathlib import Path from textwrap import dedent from unittest.mock import AsyncMock, MagicMock, patch import pytest -from unit.mocks import get_mock_target_identifier +from unit.mocks import MockPromptTarget, get_mock_target_identifier, store_message from pyrit.exceptions.exception_classes import InvalidJsonException from pyrit.memory.central_memory import CentralMemory from pyrit.memory.memory_interface import MemoryInterface -from pyrit.models import Message, MessagePiece, SeedPrompt +from pyrit.models import Message, MessagePiece, MessageScorable, ScoringExpectation, SeedPrompt +from pyrit.prompt_target.common.target_capabilities import TargetCapabilities +from pyrit.prompt_target.common.target_configuration import TargetConfiguration from pyrit.score import ( SelfAskTrueFalseScorer, TrueFalseQuestion, @@ -414,3 +417,50 @@ async def test_from_question_scores_end_to_end(patch_central_database, scorer_tr assert len(scores) == 1 assert scores[0].get_value() is True + + +@pytest.mark.usefixtures("patch_central_database") +async def test_self_ask_true_false_rejects_nontext_for_noneditable_target_async(tmp_path: Path) -> None: + image_path = tmp_path / "image.png" + image_path.write_bytes(b"\x89PNG\r\n\x1a\n") + target = MockPromptTarget( + custom_configuration=TargetConfiguration( + capabilities=TargetCapabilities( + supports_multi_turn=True, + supports_multi_message_pieces=True, + supports_system_prompt=True, + input_modalities=frozenset( + {frozenset({"text"}), frozenset({"image_path"}), frozenset({"text", "image_path"})} + ), + ) + ) + ) + scorer = SelfAskTrueFalseScorer.from_question( + chat_target=target, + question=TrueFalseQuestion( + category="image content", + true_description="The image contains visible content.", + false_description="The image does not contain visible content.", + ), + ) + image_message = store_message( + MessagePiece( + role="assistant", + conversation_id="image-judgment", + original_value=str(image_path), + converted_value=str(image_path), + original_value_data_type="image_path", + converted_value_data_type="image_path", + ).to_message() + ) + + with pytest.raises(RuntimeError, match="Error in scorer SelfAskTrueFalseScorer") as exc_info: + await scorer.score_async( + scorable=MessageScorable.from_message(image_message), + expectation=ScoringExpectation(objective="Describe this image"), + ) + + assert isinstance(exc_info.value.__cause__, ValueError) + assert "non-text" in str(exc_info.value.__cause__) + assert "editable history" in str(exc_info.value.__cause__) + assert target.prompt_sent == [] diff --git a/uv.lock b/uv.lock index c4ab92b80e..9b33b65c9d 100644 --- a/uv.lock +++ b/uv.lock @@ -1808,6 +1808,19 @@ http = [ { name = "aiohttp" }, ] +[[package]] +name = "github-copilot-sdk" +version = "1.0.14" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "httpx" }, + { name = "pydantic" }, + { name = "python-dateutil" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/ed/e2/4d8c622545f55f06773548cf86ffe60f59e6a4cedcc412ff84fcac244b2b/github_copilot_sdk-1.0.14-py3-none-any.whl", hash = "sha256:9f4538cec7295f24507c650d794307396458a51f280a0fbbcd840da177109050", size = 614607, upload-time = "2026-09-16T05:11:14.572Z" }, +] + [[package]] name = "greenlet" version = "3.3.0" @@ -4738,6 +4751,7 @@ all = [ { name = "azure-ai-ml" }, { name = "azure-cognitiveservices-speech" }, { name = "flask" }, + { name = "github-copilot-sdk" }, { name = "ipykernel" }, { name = "jupyter" }, { name = "litellm" }, @@ -4759,6 +4773,9 @@ gcg = [ { name = "sentencepiece" }, { name = "torch" }, ] +github-copilot = [ + { name = "github-copilot-sdk" }, +] huggingface = [ { name = "sentencepiece" }, { name = "torch" }, @@ -4831,6 +4848,8 @@ requires-dist = [ { name = "fastapi", specifier = ">=0.133.0" }, { name = "flask", marker = "extra == 'all'", specifier = ">=3.1.3" }, { name = "flask", marker = "extra == 'playwright'", specifier = ">=3.1.3" }, + { name = "github-copilot-sdk", marker = "extra == 'all'", specifier = ">=1.0.11" }, + { name = "github-copilot-sdk", marker = "extra == 'github-copilot'", specifier = ">=1.0.11" }, { name = "httpx", extras = ["http2"], specifier = ">=0.27.2" }, { name = "ipykernel", marker = "extra == 'all'", specifier = ">=6.29.5" }, { name = "jinja2", specifier = ">=3.1.6" }, @@ -4880,7 +4899,7 @@ requires-dist = [ { name = "uvicorn", extras = ["standard"], specifier = ">=0.32.0" }, { name = "websockets", specifier = ">=14.0" }, ] -provides-extras = ["huggingface", "gcg", "playwright", "fairness-bias", "opencv", "speech", "litellm", "all"] +provides-extras = ["huggingface", "gcg", "playwright", "fairness-bias", "opencv", "speech", "litellm", "github-copilot", "all"] [package.metadata.requires-dev] dev = [