diff --git a/backend/packages/harness/deerflow/agents/middlewares/pii_redaction_middleware.py b/backend/packages/harness/deerflow/agents/middlewares/pii_redaction_middleware.py index 81ead6d4ec..b1176a9359 100644 --- a/backend/packages/harness/deerflow/agents/middlewares/pii_redaction_middleware.py +++ b/backend/packages/harness/deerflow/agents/middlewares/pii_redaction_middleware.py @@ -55,6 +55,7 @@ import hmac import logging import re from collections.abc import Awaitable, Callable, Sequence +from copy import deepcopy from dataclasses import dataclass from dataclasses import replace as dc_replace from typing import override @@ -291,6 +292,31 @@ class _Redactor: return replace +def _independent_copy(value: object, *, what: str) -> object: + """Copy *value* so the rebuilt message shares no mutable state with the original. + + ``deepcopy`` keeps nested mutable blocks and metadata independent, which is + what the rewrite needs: ``_redact_content`` retains references to the blocks + it leaves untouched, and ``model_copy`` alone is shallow. It can raise for an + exotic value a caller placed in a block, and a copy failure must never skip + the rewrite — that would hand raw PII to the model. Fall back to a shallow + per-container copy, which still detaches the containers the rewrite touches. + """ + try: + return deepcopy(value) + except Exception: + logger.warning( + "PII redaction could not deep-copy %s; using a shallow copy", + what, + exc_info=True, + ) + if isinstance(value, list): + return [dict(item) if isinstance(item, dict) else item for item in value] + if isinstance(value, dict): + return dict(value) + return value + + def _redact_content(content: object, redactor: _Redactor) -> tuple[object, bool]: """Redact *content*, preserving its shape. Returns ``(content, changed)``. @@ -370,11 +396,14 @@ class PiiRedactionMiddleware(AgentMiddleware[AgentState]): continue if not changed_msg: continue - messages[index] = HumanMessage( - content=content, - id=msg.id, - name=msg.name, - additional_kwargs=dict(msg.additional_kwargs or {}), + messages[index] = msg.model_copy( + update={ + "content": _independent_copy(content, what="message content"), + "additional_kwargs": dict(msg.additional_kwargs or {}), + # model_copy is shallow, so the preserved metadata would + # otherwise be the same dict as the original message's. + "response_metadata": _independent_copy(dict(msg.response_metadata or {}), what="response_metadata"), + }, ) changed = True updates = {} diff --git a/backend/tests/test_pii_redaction_middleware.py b/backend/tests/test_pii_redaction_middleware.py index 53c37581b9..b2ac89bd44 100644 --- a/backend/tests/test_pii_redaction_middleware.py +++ b/backend/tests/test_pii_redaction_middleware.py @@ -248,6 +248,55 @@ class TestModelCallBoundary: messages, _ = _run_model_call(_make_middleware(), [original]) assert messages[0].additional_kwargs["custom"] == "v" + def test_redacted_user_message_preserves_other_fields(self): + # The rewrite must not drop fields it does not touch: a hand-built + # constructor call silently loses response_metadata and every other + # field outside content/id/name/additional_kwargs. The tool-result + # seam in this same middleware already rebuilds with model_copy for + # exactly that reason. + original = HumanMessage( + content="alice@example.com", + id="msg-1", + name="user", + additional_kwargs={"custom": "v"}, + response_metadata={"source": "gateway"}, + ) + messages, _ = _run_model_call(_make_middleware(), [original]) + assert messages[0].content == EMAIL_ALICE + assert messages[0].id == "msg-1" + assert messages[0].name == "user" + assert messages[0].additional_kwargs["custom"] == "v" + assert messages[0].response_metadata == {"source": "gateway"} + + def test_preserved_metadata_is_not_shared_with_the_original(self): + # model_copy is shallow: without an explicit copy the rewritten message + # and the original would hold the same response_metadata dict, so + # writing metadata on the model-facing message would mutate the + # request message retained in thread state. + original = HumanMessage( + content="alice@example.com", + response_metadata={"source": "gateway", "nested": {"k": "v"}}, + ) + messages, request = _run_model_call(_make_middleware(), [original]) + assert messages[0].response_metadata == original.response_metadata + + messages[0].response_metadata["source"] = "mutated" + assert request.messages[0].response_metadata["source"] == "gateway" + + def test_redaction_survives_a_content_block_that_cannot_be_copied(self): + # A caller can put a value whose copy raises into a block. The rewrite + # must still happen: letting the copy failure escape would fail the + # whole request open and hand raw PII to the model. + class NoCopy: + def __deepcopy__(self, memo): + raise RuntimeError("cannot copy this block value") + + original = HumanMessage(content=[{"type": "text", "text": "reach me at alice@example.com", "meta": NoCopy()}]) + messages, _ = _run_model_call(_make_middleware(), [original]) + + assert EMAIL_ALICE in messages[0].content[0]["text"] + assert "alice@example.com" not in messages[0].content[0]["text"] + def test_ai_message_untouched(self): ai = AIMessage("contact alice@example.com") messages, _ = _run_model_call(_make_middleware(), [ai]) @@ -300,6 +349,17 @@ class TestModelCallBoundary: # The original message object is untouched. assert original.content[0] == "reach me at alice@example.com" + def test_rebuilt_user_message_does_not_alias_the_original_content(self): + # _redact_content appends untouched blocks by reference, so the + # rebuilt message must copy the container: without the copy, mutating + # a non-text block on the redacted message would write through to the + # message kept in thread state. + image_block = {"type": "image_url", "image_url": {"url": "https://example.com/a.png"}} + original = HumanMessage(["alice@example.com", image_block]) + messages, _ = _run_model_call(_make_middleware(), [original]) + messages[0].content[1]["image_url"]["url"] = "https://example.com/mutated.png" + assert original.content[1]["image_url"]["url"] == "https://example.com/a.png" + # --------------------------------------------------------------------------- # Tool boundary