From 3c50add420ba4012762cad97e45ec56beb50ad1e Mon Sep 17 00:00:00 2001 From: Maxim Svistunov Date: Wed, 7 Oct 2026 23:56:38 +0200 Subject: [PATCH 1/3] LCORE-4593: add the text redaction of the shields a request selects The redaction shield rewrites the user prompt as an agent capability, so it only reaches text that goes through the agent. Compacted mode builds the model input and stores the turn outside the agent and needs the same rules as a plain function. request_redactor() returns that function for a request: the rules of the redaction shields the request selects, applied in configuration order, which is the order the capabilities run in. It returns None when the selection holds no redaction rule and raises the 404 of get_shields_for_request for an unknown shield name. TextRedactor is the type of the returned function. Nothing calls it yet, so no behaviour changes. Unit tests cover the None cases, the default selection, the rule order and the 404. --- src/utils/shields.py | 42 +++++++++++++++++++++ src/utils/types.py | 4 ++ tests/unit/utils/test_shields.py | 65 ++++++++++++++++++++++++++++++++ 3 files changed, 111 insertions(+) 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/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 From 5def42b333c1ed454527351d76fc5904fd8b6a10 Mon Sep 17 00:00:00 2001 From: Maxim Svistunov Date: Thu, 8 Oct 2026 00:12:14 +0200 Subject: [PATCH 2/3] LCORE-4593: redact the query of a compacted turn With a redaction shield and compaction enabled, the query of a compacted turn reached the model and the conversation store unredacted. The redaction capability rewrites pydantic-ai's messages, but in compacted mode the request input is the explicit list built by apply_compaction, sent as an extra_body override that replaces the message-derived input, and the endpoints store CompactionResult.original_input themselves. Both were built from the raw query. apply_compaction and apply_compaction_blocking take the redaction of the request (request_redactor) and apply it to the new query, so the explicit input and original_input carry the redacted text. /v1/query, /v1/streaming_query and A2A pass it; /v1/responses does not, its redaction shield blocks instead of rewriting. A request that is not compacted is returned untouched, and without a selected redaction shield nothing changes. Two side effects. On /v1/query and /v1/streaming_query the agent prompt of a compacted turn is taken from the explicit input, so it is now the redacted query and the other shields judge the text the model gets. An unknown shield name fails before compaction instead of after it. Tested at unit level (explicit input, the override built from it and original_input, with the shield selected and deselected) and with integration tests on the three entry points: model input and stored turn of a compacted conversation, per-request shield_ids, and the unchanged request of a conversation that is not compacted. The integration tests of the compacted turn fail without the fix. --- docs/user_doc/shields_guide.md | 7 + src/app/endpoints/a2a.py | 2 + src/app/endpoints/query.py | 3 +- src/app/endpoints/streaming_query.py | 5 +- src/utils/conversation_compaction.py | 14 +- .../endpoints/test_compaction_redaction.py | 177 ++++++++++++++++++ .../app/endpoints/test_streaming_query.py | 1 + .../test_conversation_compaction_redaction.py | 90 +++++++++ 8 files changed, 296 insertions(+), 3 deletions(-) create mode 100644 tests/integration/endpoints/test_compaction_redaction.py create mode 100644 tests/unit/utils/test_conversation_compaction_redaction.py diff --git a/docs/user_doc/shields_guide.md b/docs/user_doc/shields_guide.md index 71bd84194..d1b44f7c2 100644 --- a/docs/user_doc/shields_guide.md +++ b/docs/user_doc/shields_guide.md @@ -95,6 +95,13 @@ 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. Answers of the model are stored and replayed as the model +produced them, in both modes. Text that is already stored is not rewritten. + ## 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..258d1d633 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 @@ -701,6 +702,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 +727,10 @@ 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. 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 +761,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( @@ -824,6 +834,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 +853,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/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..abe32aaa1 --- /dev/null +++ b/tests/unit/utils/test_conversation_compaction_redaction.py @@ -0,0 +1,90 @@ +"""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.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) + + +async def _compact( + mocker: MockerFixture, items: list[Any], redact: Optional[TextRedactor] +) -> cc.CompactionResult: + """Apply compaction to RAW_QUERY over *items*, 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: 1_000_000}), + compaction_config=CompactionConfiguration(enabled=True), + 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 From ad38391e3887d92dbc8ee81164dc5d0da04038dd Mon Sep 17 00:00:00 2001 From: Maxim Svistunov Date: Thu, 8 Oct 2026 00:16:27 +0200 Subject: [PATCH 3/3] LCORE-4593: redact the user turns and summaries compaction replays Conversations hold user turns that were stored unredacted: compacted turns written before the previous commit, and turns of requests that deselected the shield. Compacted mode replayed them as stored, handed them to the summarizer as stored, and replayed the summaries made from them. apply_compaction now applies the redaction of the request to the user turns and the summaries it puts into the explicit input, and to the user turns it hands to the summarizer. Summaries are covered because they are sent as user messages. Assistant turns are left as stored, as the redaction capability leaves them. Nothing stored is rewritten, and the recursive fold still sends the stored summaries. The replayed turns are redacted only after the summary boundary is taken: the first kept turn is found among the stored items by identity, and a redacted copy would move the boundary past a turn that is still verbatim. Unit tests cover the explicit input (summary and user turns redacted, assistant and system turns untouched), the items handed to the summarizer, and the boundary on a conversation with an older marker among the kept turns. The shields guide says what is redacted in compacted mode and what is not. --- docs/user_doc/shields_guide.md | 13 ++- src/utils/conversation_compaction.py | 41 ++++++++- .../test_conversation_compaction_redaction.py | 89 ++++++++++++++++++- 3 files changed, 135 insertions(+), 8 deletions(-) diff --git a/docs/user_doc/shields_guide.md b/docs/user_doc/shields_guide.md index d1b44f7c2..15f6ea4fb 100644 --- a/docs/user_doc/shields_guide.md +++ b/docs/user_doc/shields_guide.md @@ -99,8 +99,17 @@ 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. Answers of the model are stored and replayed as the model -produced them, in both modes. Text that is already stored is not rewritten. +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 diff --git a/src/utils/conversation_compaction.py b/src/utils/conversation_compaction.py index 258d1d633..208d348a2 100644 --- a/src/utils/conversation_compaction.py +++ b/src/utils/conversation_compaction.py @@ -439,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], @@ -730,7 +757,10 @@ async def apply_compaction( # pylint: disable=too-many-arguments,too-many-posit 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. + 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. @@ -788,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, ) @@ -823,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) diff --git a/tests/unit/utils/test_conversation_compaction_redaction.py b/tests/unit/utils/test_conversation_compaction_redaction.py index abe32aaa1..5572b3343 100644 --- a/tests/unit/utils/test_conversation_compaction_redaction.py +++ b/tests/unit/utils/test_conversation_compaction_redaction.py @@ -7,6 +7,7 @@ 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, @@ -36,10 +37,31 @@ def _msg(role: str, text: str) -> OpenAIResponseMessage: 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] + mocker: MockerFixture, + items: list[Any], + redact: Optional[TextRedactor], + window: int = 1_000_000, + **compaction: Any, ) -> cc.CompactionResult: - """Apply compaction to RAW_QUERY over *items*, far below the trigger.""" + """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) ) @@ -52,8 +74,8 @@ async def _compact( store=True, stream=False, ), - inference_config=InferenceConfiguration(context_windows={MODEL: 1_000_000}), - compaction_config=CompactionConfiguration(enabled=True), + inference_config=InferenceConfiguration(context_windows={MODEL: window}), + compaction_config=CompactionConfiguration(enabled=True, **compaction), redact=redact, ) @@ -88,3 +110,62 @@ async def test_compacted_turn_carries_the_redacted_query( 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