diff --git a/docs/user_doc/shields_guide.md b/docs/user_doc/shields_guide.md index 71bd84194..15f6ea4fb 100644 --- a/docs/user_doc/shields_guide.md +++ b/docs/user_doc/shields_guide.md @@ -95,6 +95,22 @@ for a complete example. Invalid regex patterns are rejected at configuration load time. +In a [compacted conversation](conversation_compaction.md), Lightspeed Core +Stack builds the model input and stores the turn itself. It applies the +redaction shields selected for the request to the new query there, so that the +model receives and the conversation stores the same redacted query as outside +compacted mode. The user messages and the summaries it replays from the +conversation, and the user messages it hands to the summarizer, go through the +same rules each time they are sent, so a replacement must not itself match a +rule. + +Not covered: answers of the model are stored and replayed as the model +produced them, in both modes, and the summarizer gets them as stored. Text +that is already stored is not rewritten, so +`GET /v1/conversations/{conversation_id}` returns it as stored. The recursive +re-summarization sends the stored summaries as they are, and `/v1/responses` +replays a compacted conversation as stored. + ## granite_guardian IBM Granite Guardian screening with configurable risk definitions, thresholds, diff --git a/src/app/endpoints/a2a.py b/src/app/endpoints/a2a.py index c64718dc2..0a8d85f26 100644 --- a/src/app/endpoints/a2a.py +++ b/src/app/endpoints/a2a.py @@ -76,6 +76,7 @@ from utils.pydantic_ai_helpers import build_agent, captured_output_items from utils.query import extract_provider_and_model_from_model_id from utils.responses import prepare_responses_params +from utils.shields import request_redactor from utils.suid import normalize_conversation_id from version import __version__ @@ -241,6 +242,7 @@ async def _compact_a2a_request( responses_params, configuration.inference, configuration.compaction, + redact=request_redactor(configuration.shields), ) return compaction.params, PendingTurn.for_request( client, compaction.params, compaction.original_input diff --git a/src/app/endpoints/query.py b/src/app/endpoints/query.py index 538bd5b1e..7b8357e8a 100644 --- a/src/app/endpoints/query.py +++ b/src/app/endpoints/query.py @@ -60,7 +60,7 @@ maybe_get_topic_summary, prepare_responses_params, ) -from utils.shields import validate_shield_ids_override +from utils.shields import request_redactor, validate_shield_ids_override from utils.suid import normalize_conversation_id from utils.types import Responses from utils.vector_search import build_rag_context @@ -239,6 +239,7 @@ async def _handle_query_with_tracing( responses_params, configuration.inference, configuration.compaction, + redact=request_redactor(configuration.shields, query_request.shield_ids), cache=configured_conversation_cache(), user_id=user_id, skip_user_id_check=_skip_userid_check, diff --git a/src/app/endpoints/streaming_query.py b/src/app/endpoints/streaming_query.py index 3c8a11de7..d73d05ea5 100644 --- a/src/app/endpoints/streaming_query.py +++ b/src/app/endpoints/streaming_query.py @@ -81,7 +81,7 @@ extract_vector_store_ids_from_tools, prepare_responses_params, ) -from utils.shields import validate_shield_ids_override +from utils.shields import request_redactor, validate_shield_ids_override from utils.streaming_sse import ( http_exception_stream_event, stream_compaction_event, @@ -431,6 +431,9 @@ async def generate_response_with_compaction( configuration.inference, configuration.compaction, emit_events=True, + redact=request_redactor( + configuration.shields, context.query_request.shield_ids + ), cache=configured_conversation_cache(), user_id=context.user_id, skip_user_id_check=context.skip_userid_check, diff --git a/src/utils/conversation_compaction.py b/src/utils/conversation_compaction.py index af4ba903e..208d348a2 100644 --- a/src/utils/conversation_compaction.py +++ b/src/utils/conversation_compaction.py @@ -80,6 +80,7 @@ extract_message_text, is_message_item, ) +from utils.types import TextRedactor logger = get_logger(__name__) @@ -184,7 +185,7 @@ class CompactionResult: or the cache), so the ``conversation`` parameter is omitted. This is True whether the summary was created this request or reused from a prior one. Drives ``context_status``. - original_input: The new user query exactly as it arrived (before the + original_input: The new user query after redaction, if any (before the explicit-input rewrite). Populated only in compacted mode (where ``compacted`` is True); ``None`` otherwise. In compacted mode the caller hands it to ``utils.pending_turn.PendingTurn``, which @@ -438,6 +439,33 @@ def _query_input_message(original_input: ResponseInput) -> list[Any]: return list(original_input) +def _redact_user_messages( + items: list[Any], redact: Optional[TextRedactor] +) -> list[Any]: + """Return conversation items with the text of their user messages redacted. + + User messages are the ones :func:`_verbatim_input_message` replays with + the user role; other messages are left as stored, as the redaction + capability leaves them. A message the redaction changes is replaced by a + text-only copy and every other item is the object that was passed in, so + the stored conversation is not rewritten. Without a redaction the items + are returned as they are. + """ + if redact is None: + return items + redacted_items: list[Any] = [] + for item in items: + replacement = item + role = getattr(item, "role", "user") + if is_message_item(item) and role not in ("system", "developer", "assistant"): + text = extract_message_text(item) + redacted_text = redact(text) + if redacted_text != text: + replacement = OpenAIResponseMessage(role="user", content=redacted_text) + redacted_items.append(replacement) + return redacted_items + + def _build_explicit_input( summaries: list[str], recent_items: list[Any], @@ -701,6 +729,7 @@ async def apply_compaction( # pylint: disable=too-many-arguments,too-many-posit compaction_config: CompactionConfiguration, emit_events: bool = False, encoding_name: str = DEFAULT_ENCODING_NAME, + redact: Optional[TextRedactor] = None, cache: Optional[Cache] = None, user_id: str = "", skip_user_id_check: bool = False, @@ -725,6 +754,13 @@ async def apply_compaction( # pylint: disable=too-many-arguments,too-many-posit compaction_config: Compaction tuning (enabled, threshold, buffer, ...). emit_events: Whether to yield CompactionStartedEvent before summarizing. encoding_name: tiktoken encoding name for estimation/summarization. + redact: The redaction of the shields the request selects + (``utils.shields.request_redactor``), or ``None``. It is applied + to the new query, so that a compacted turn sends and stores the + text the redaction capability sends outside compacted mode. It + is also applied to the summaries and the user turns replayed from + the conversation and to the user turns handed to the summarizer, + which may have been stored unredacted. cache: Conversation cache, the preferred summary store and the home of the persisted recursive fold. ``None`` (or a non-persisting backend) falls back to marker-only summaries with no folding. @@ -755,6 +791,10 @@ async def apply_compaction( # pylint: disable=too-many-arguments,too-many-posit summaries, cached_summaries, recent_items = _load_compaction_state( items, cache, user_id, conversation_id, skip_user_id_check ) + # An item list (the /v1/responses form) is left alone: that endpoint + # passes no redaction. + if redact is not None and isinstance(original_input, str): + original_input = redact(original_input) context_window = resolve_context_window(model, inference_config) estimated = _estimate_total_tokens( @@ -778,7 +818,7 @@ async def apply_compaction( # pylint: disable=too-many-arguments,too-many-posit summary = await summarize_chunk( client, model, - old_items, + _redact_user_messages(old_items, redact), summarized_through_turn=_covered_item_count(items, keep_items), encoding_name=encoding_name, ) @@ -813,6 +853,13 @@ async def apply_compaction( # pylint: disable=too-many-arguments,too-many-posit yield CompactionResult(params, compacted=False) return + if redact is not None: + # Summaries and replayed user turns go out as user messages and may + # have been stored unredacted. Not earlier: the kept turns are found + # among the stored items by identity (_covered_item_count). + summaries = [redact(text) for text in summaries] + recent_items = _redact_user_messages(recent_items, redact) + # Compacted mode: lightspeed owns the context. Build explicit input and # stop passing the conversation parameter to inference. yield _compacted_result(params, summaries, recent_items, original_input) @@ -824,6 +871,7 @@ async def apply_compaction_blocking( # pylint: disable=too-many-arguments,too-m inference_config: InferenceConfiguration, compaction_config: CompactionConfiguration, encoding_name: str = DEFAULT_ENCODING_NAME, + redact: Optional[TextRedactor] = None, cache: Optional[Cache] = None, user_id: str = "", skip_user_id_check: bool = False, @@ -842,6 +890,7 @@ async def apply_compaction_blocking( # pylint: disable=too-many-arguments,too-m compaction_config, emit_events=False, encoding_name=encoding_name, + redact=redact, cache=cache, user_id=user_id, skip_user_id_check=skip_user_id_check, diff --git a/src/utils/shields.py b/src/utils/shields.py index 0834c8dda..7fe883330 100644 --- a/src/utils/shields.py +++ b/src/utils/shields.py @@ -36,6 +36,7 @@ from pydantic_ai_lightspeed.capabilities.redaction._capability import ( PiiRedactionCapability, ) +from pydantic_ai_lightspeed.capabilities.redaction.core import redact_text from utils.agents.error_handler import map_agent_inference_error from utils.input_sanitization import sanitize_input from utils.otel_tracing import ( @@ -44,6 +45,7 @@ set_span_attributes, shield_span_attributes, ) +from utils.types import TextRedactor logger = get_logger(__name__) tracer = trace.get_tracer(__name__) @@ -215,3 +217,43 @@ def get_shields_for_request( raise HTTPException(**response.model_dump()) return [shield for shield in shields if shield.name in requested] + + +def request_redactor( + shields: list[ShieldConfiguration], + shield_ids: Optional[list[str]] = None, +) -> Optional[TextRedactor]: + """Build the text redaction of the shields a request selects. + + The rules of the selected redaction shields are applied one after another + in configuration order, the order in which their capabilities run on an + agent. It is for the places where lightspeed-stack builds or stores text + that the redaction capability does not see. + + Parameters: + shields: Configured LCS shields. + shield_ids: The shield selection of the request, with the meaning it + has for ``get_shields_for_request``. + + Returns: + A function that redacts a text, or ``None`` when the selection holds + no redaction rule. + + Raises: + HTTPException: 404 if shield_ids names a shield that is not configured. + + """ + patterns = [ + compiled + for shield in get_shields_for_request(shields, shield_ids) + if isinstance(shield.config, RedactionConfig) + for compiled in shield.config.compiled_patterns + ] + if not patterns: + return None + + def redact(text: str) -> str: + """Apply the selected redaction rules to a text.""" + return redact_text(text, patterns).content + + return redact diff --git a/src/utils/types.py b/src/utils/types.py index 44b49700b..81a76560b 100644 --- a/src/utils/types.py +++ b/src/utils/types.py @@ -1,5 +1,6 @@ """Common types for the project.""" +from collections.abc import Callable from re import Pattern from typing import Any, TypeVar, cast @@ -9,6 +10,9 @@ CompiledPatterns = list[tuple[Pattern[str], str]] +TextRedactor = Callable[[str], str] +"""Rewrites a text with the redaction rules selected for a request.""" + T = TypeVar("T") diff --git a/tests/integration/endpoints/test_compaction_redaction.py b/tests/integration/endpoints/test_compaction_redaction.py new file mode 100644 index 000000000..9d0855864 --- /dev/null +++ b/tests/integration/endpoints/test_compaction_redaction.py @@ -0,0 +1,177 @@ +"""Integration tests for PII redaction on compacted conversations (LCORE-4593).""" + +from collections.abc import Awaitable, Callable +from typing import Any, Optional + +import pytest +from a2a.server.agent_execution import RequestContext +from fastapi import Request +from ogx_api.openai_responses import OpenAIResponseMessage +from pytest_mock import MockerFixture +from sqlalchemy.orm import Session + +from app.endpoints.a2a import A2AAgentExecutor +from app.endpoints.query import query_endpoint_handler +from app.endpoints.streaming_query import streaming_query_endpoint_handler +from authentication.interface import AuthTuple +from configuration import AppConfig +from models.api.requests import QueryRequest +from models.common.responses.responses_api_params import ResponsesApiParams +from models.config import RedactionConfig, RedactionRule, RedactionShieldConfiguration +from tests.integration.conftest import InMemoryConversationStore +from tests.integration.endpoints._compaction_helpers import ( + CONV_ID_LLAMA, + DEFAULT_MODEL_RESPONSE, + EXISTING_CONV_ID, + TEST_MODEL, + create_existing_conversation, + enable_compaction, + marker, + msg, +) + +RAW_QUERY = "Mail the report to jane@example.com" +REDACTED_QUERY = "Mail the report to [EMAIL]" + + +@pytest.fixture(name="store") +def store_fixture( + test_config: AppConfig, + mock_conversation_store: InMemoryConversationStore, + test_auth: AuthTuple, + patch_db_session: Session, +) -> InMemoryConversationStore: + """Configure a redaction shield and store an already compacted conversation. + + The context window is far larger than the conversation, so no summarization + runs: the stored marker alone puts the conversation into compacted mode. + """ + enable_compaction(test_config, context_window=100_000) + test_config.configuration.shields = [ + RedactionShieldConfiguration( + name="pii-redaction", + provider_id="redaction", + config=RedactionConfig( + rules=[RedactionRule(pattern=r"\S+@\S+", replacement="[EMAIL]")] + ), + ) + ] + create_existing_conversation(patch_db_session, test_auth[0]) + mock_conversation_store.store[CONV_ID_LLAMA] = [ + msg("user", "old question"), + msg("assistant", "old answer"), + marker("earlier summary", covers=2), + ] + return mock_conversation_store + + +@pytest.fixture(name="ask") +def ask_fixture( + mock_query_agent: Any, + mock_streaming_query_agent: Any, + test_request: Request, + test_auth: AuthTuple, +) -> Callable[..., Awaitable[Any]]: + """Return a function that sends RAW_QUERY and returns the params the agent got.""" + + async def ask(streaming: bool, shield_ids: Optional[list[str]] = None) -> Any: + """Send RAW_QUERY to /v1/streaming_query or to /v1/query.""" + agent = mock_streaming_query_agent if streaming else mock_query_agent + agent.model.last_output_items = [ + OpenAIResponseMessage(role="assistant", content=DEFAULT_MODEL_RESPONSE) + ] + query_request = QueryRequest( + query=RAW_QUERY, conversation_id=EXISTING_CONV_ID, shield_ids=shield_ids + ) + if streaming: + response = await streaming_query_endpoint_handler( + test_request, query_request, auth=test_auth, mcp_headers={} + ) + async for _ in response.body_iterator: + pass + else: + await query_endpoint_handler( + test_request, query_request, auth=test_auth, mcp_headers={} + ) + return agent.build_agent_mock.call_args[0][1] + + return ask + + +def _stored_texts(store: InMemoryConversationStore) -> list[str]: + """Return the text of every item stored for the conversation.""" + return [str(item.content) for item in store.store[CONV_ID_LLAMA]] + + +@pytest.mark.parametrize("streaming", [False, True], ids=["query", "streaming_query"]) +async def test_compacted_turn_is_redacted( + store: InMemoryConversationStore, + ask: Callable[..., Awaitable[Any]], + streaming: bool, +) -> None: + """The model input and the stored turn carry the redacted query.""" + params = await ask(streaming) + + assert params.omit_conversation is True + assert params.input[-1].content == REDACTED_QUERY + assert _stored_texts(store)[3:] == [REDACTED_QUERY, DEFAULT_MODEL_RESPONSE] + + +@pytest.mark.parametrize("streaming", [False, True], ids=["query", "streaming_query"]) +async def test_compacted_turn_honours_shield_ids( + store: InMemoryConversationStore, + ask: Callable[..., Awaitable[Any]], + streaming: bool, +) -> None: + """A request that deselects the shield sends and stores the query as it arrived.""" + params = await ask(streaming, shield_ids=[]) + + assert params.input[-1].content == RAW_QUERY + assert _stored_texts(store)[3:] == [RAW_QUERY, DEFAULT_MODEL_RESPONSE] + + +@pytest.mark.parametrize("streaming", [False, True], ids=["query", "streaming_query"]) +async def test_not_compacted_turn_is_unchanged( + store: InMemoryConversationStore, + ask: Callable[..., Awaitable[Any]], + streaming: bool, +) -> None: + """Without a summary the agent gets the raw query and its capability redacts it.""" + store.store[CONV_ID_LLAMA] = [msg("user", "hi"), msg("assistant", "hello")] + + params = await ask(streaming) + + assert params.omit_conversation is False + assert params.input == RAW_QUERY + assert _stored_texts(store) == ["hi", "hello"] # OGX stores this turn + + +async def test_a2a_compacted_turn_is_redacted( + store: InMemoryConversationStore, mocker: MockerFixture +) -> None: + """A2A, which runs every configured shield, sends and stores the redacted query.""" + mocker.patch( + "app.endpoints.a2a.prepare_responses_params", + new=mocker.AsyncMock( + return_value=ResponsesApiParams( + input=RAW_QUERY, + model=TEST_MODEL, + conversation=CONV_ID_LLAMA, + store=True, + stream=True, + ) + ), + ) + build_agent = mocker.patch("app.endpoints.a2a.build_agent") + context = mocker.MagicMock(spec=RequestContext) + context.get_user_input.return_value = RAW_QUERY + context.message = None + executor = A2AAgentExecutor(auth_token="test-token") + + await executor._process_task_streaming( # pylint: disable=protected-access + context, mocker.AsyncMock(), "task-4593", "ctx-4593" + ) + + assert build_agent.call_args[0][1].input[-1].content == REDACTED_QUERY + # The mocked agent captures no output items, so only the query is stored. + assert _stored_texts(store)[3:] == [REDACTED_QUERY] diff --git a/tests/unit/app/endpoints/test_streaming_query.py b/tests/unit/app/endpoints/test_streaming_query.py index bb7a6b86c..6fb57a84f 100644 --- a/tests/unit/app/endpoints/test_streaming_query.py +++ b/tests/unit/app/endpoints/test_streaming_query.py @@ -941,6 +941,7 @@ async def fake_apply_compaction( return_value=None, ) mock_config = mocker.Mock() + mock_config.shields = [] mocker.patch("app.endpoints.streaming_query.configuration", mock_config) async def inner_generator() -> AsyncIterator[str]: diff --git a/tests/unit/utils/test_conversation_compaction_redaction.py b/tests/unit/utils/test_conversation_compaction_redaction.py new file mode 100644 index 000000000..5572b3343 --- /dev/null +++ b/tests/unit/utils/test_conversation_compaction_redaction.py @@ -0,0 +1,171 @@ +"""Unit tests for redaction in compacted mode (LCORE-4593).""" + +from typing import Any, Optional, cast + +import pytest +from ogx_api.openai_responses import OpenAIResponseMessage +from pytest_mock import MockerFixture + +from models.common.responses.responses_api_params import ResponsesApiParams +from models.compaction import ConversationSummary +from models.config import ( + CompactionConfiguration, + InferenceConfiguration, + RedactionConfig, + RedactionRule, + RedactionShieldConfiguration, +) +from pydantic_ai_lightspeed.ogx._model import _model_settings_from_responses_params +from utils import conversation_compaction as cc +from utils.shields import request_redactor +from utils.types import TextRedactor + +MODEL = "openai/gpt-4o-mini" +RAW_QUERY = "mail me at jane@example.com" +REDACTED_QUERY = "mail me at [EMAIL]" +SHIELD = RedactionShieldConfiguration( + name="pii-redaction", + provider_id="redaction", + config=RedactionConfig( + rules=[RedactionRule(pattern=r"\S+@\S+", replacement="[EMAIL]")] + ), +) + + +def _msg(role: str, text: str) -> OpenAIResponseMessage: + """Build a typed OGX message item for tests.""" + return OpenAIResponseMessage(role=cast("Any", role), content=text) + + +# A conversation compacted once, with a buffer of two turns: the marker covers +# the first turn and sits behind the two turns that compaction kept. +HISTORY = [ + _msg("user", "q1"), + _msg("assistant", "a1"), + _msg("user", "ask bob@example.org"), + _msg("assistant", "bob@example.org is out"), + _msg("user", "then ask ann@example.net"), + _msg("assistant", "ann@example.net replied"), + _msg( + "user", f"{cc.MARKER_SENTINEL} {cc.MARKER_COVERS_PREFIX}2] met eve@example.com" + ), + _msg("user", "q4"), + _msg("assistant", "a4"), +] + + +async def _compact( + mocker: MockerFixture, + items: list[Any], + redact: Optional[TextRedactor], + window: int = 1_000_000, + **compaction: Any, +) -> cc.CompactionResult: + """Apply compaction to RAW_QUERY over *items*, by default far below the trigger.""" + mocker.patch.object( + cc, "get_all_conversation_items", mocker.AsyncMock(return_value=items) + ) + return await cc.apply_compaction_blocking( + client=mocker.AsyncMock(), + params=ResponsesApiParams( + input=RAW_QUERY, + model=MODEL, + conversation="conv_abc123", + store=True, + stream=False, + ), + inference_config=InferenceConfiguration(context_windows={MODEL: window}), + compaction_config=CompactionConfiguration(enabled=True, **compaction), + redact=redact, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("shield_ids", "expected_query"), + [(None, REDACTED_QUERY), ([], RAW_QUERY)], + ids=["shield selected", "shield deselected"], +) +async def test_compacted_turn_carries_the_redacted_query( + mocker: MockerFixture, shield_ids: Optional[list[str]], expected_query: str +) -> None: + """A compacted turn sends and stores the query as the selected shields redact it. + + The explicit input, the ``extra_body`` override built from it (which + replaces the input on the wire) and the text the endpoints store all carry + the redacted query. Without a redaction shield they carry it as it arrived. + """ + items = [ + _msg("user", "old q"), + _msg("assistant", "old a"), + _msg("user", f"{cc.MARKER_SENTINEL} {cc.MARKER_COVERS_PREFIX}2] earlier"), + ] + + result = await _compact(mocker, items, request_redactor([SHIELD], shield_ids)) + + assert result.compacted is True + assert isinstance(result.params.input, list) + assert result.params.input[-1].content == expected_query + settings = _model_settings_from_responses_params(result.params) + override = cast("dict[str, Any]", settings["extra_body"])["input"] + assert override[-1]["content"] == expected_query + assert result.original_input == expected_query + + +@pytest.mark.asyncio +async def test_replayed_user_turns_and_summaries_are_redacted( + mocker: MockerFixture, +) -> None: + """What is replayed as user messages is redacted; other turns are not.""" + items = [*HISTORY, _msg("system", "ops@example.com owns the report")] + result = await _compact(mocker, items, request_redactor([SHIELD])) + + assert [message.content for message in result.params.input] == [ + "Summary of earlier conversation:\nmet [EMAIL]", + "ask [EMAIL]", + "bob@example.org is out", + "then ask [EMAIL]", + "ann@example.net replied", + "q4", + "a4", + "ops@example.com owns the report", + REDACTED_QUERY, + ] + + +@pytest.mark.asyncio +async def test_summarizer_gets_redacted_user_turns(mocker: MockerFixture) -> None: + """The summarizer is handed redacted user turns, and the boundary stays put. + + The first kept turn is looked up in the stored items by identity + (LCORE-4219), so the kept turns must not be replaced by redacted copies + before that: with the older marker among them the fallback count is 5. + """ + summary = ConversationSummary( + summary_text="second summary", + summarized_through_turn=4, + token_count=2, + created_at="2026-10-07T00:00:00Z", + model_used=MODEL, + ) + summarize = mocker.patch.object( + cc, "summarize_chunk", mocker.AsyncMock(return_value=summary) + ) + mocker.patch.object(cc, "_write_summary_marker", mocker.AsyncMock()) + + await _compact( + mocker, + HISTORY, + request_redactor([SHIELD]), + window=1000, + threshold_ratio=0.01, + token_floor=0, + buffer_turns=2, + ) + + summarized = summarize.await_args.args[2] + assert [item.content for item in summarized] == [ + "ask [EMAIL]", + "bob@example.org is out", + ] + assert summarize.await_args.kwargs["summarized_through_turn"] == 4 diff --git a/tests/unit/utils/test_shields.py b/tests/unit/utils/test_shields.py index 451bdd776..a6503b798 100644 --- a/tests/unit/utils/test_shields.py +++ b/tests/unit/utils/test_shields.py @@ -14,11 +14,15 @@ from models.config import ( QuestionValidityConfig, QuestionValidityShieldConfiguration, + RedactionConfig, + RedactionRule, + RedactionShieldConfiguration, ShieldConfiguration, ) from utils.otel_tracing import SpanAttributes, SpanEvents from utils.shields import ( get_shields_for_request, + request_redactor, run_shield_moderation_v2, validate_shield_ids_override, ) @@ -33,6 +37,19 @@ def _shield_config(name: str) -> QuestionValidityShieldConfiguration: ) +def _redaction_shield( + name: str, pattern: str, replacement: str +) -> RedactionShieldConfiguration: + """Build a redaction shield with one rule for tests.""" + return RedactionShieldConfiguration( + name=name, + provider_id="redaction", + config=RedactionConfig( + rules=[RedactionRule(pattern=pattern, replacement=replacement)] + ), + ) + + class TestValidateShieldIdsOverride: """Tests for validate_shield_ids_override function.""" @@ -403,3 +420,51 @@ def test_raises_404_when_multiple_requested_shields_not_configured(self) -> None assert "Shields" in detail["response"] assert "missing-1" in detail["cause"] assert "missing-2" in detail["cause"] + + +class TestRequestRedactor: + """Tests for request_redactor function.""" + + def test_returns_none_when_no_redaction_rule_is_selected(self) -> None: + """There is no redactor when the selection holds no redaction rule.""" + email = _redaction_shield("email", r"\S+@\S+", "[EMAIL]") + no_rules = RedactionShieldConfiguration( + name="no-rules", provider_id="redaction", config=RedactionConfig() + ) + validity = _shield_config("validity") + + assert request_redactor([validity, no_rules]) is None + assert request_redactor([email], shield_ids=[]) is None + assert request_redactor([email, validity], shield_ids=["validity"]) is None + + def test_redacts_with_all_configured_shields_by_default(self) -> None: + """Without shield_ids every configured redaction shield applies.""" + redact = request_redactor( + [ + _shield_config("validity"), + _redaction_shield("email", r"\S+@\S+", "[EMAIL]"), + ] + ) + + assert redact is not None + assert redact("mail jane@example.com now") == "mail [EMAIL] now" + + def test_applies_shields_in_configuration_order(self) -> None: + """Rules run in configuration order, whatever the order of shield_ids.""" + first = _redaction_shield("first", "secret", "token") + second = _redaction_shield("second", "token", "[GONE]") + + redact = request_redactor([first, second], shield_ids=["second", "first"]) + + assert redact is not None + assert redact("my secret") == "my [GONE]" + + def test_raises_404_when_requested_shield_not_configured(self) -> None: + """An unknown shield name is the same 404 as on the agent path.""" + with pytest.raises(HTTPException) as exc_info: + request_redactor( + [_redaction_shield("email", r"\S+@\S+", "[EMAIL]")], + shield_ids=["missing-shield"], + ) + + assert exc_info.value.status_code == status.HTTP_404_NOT_FOUND