diff --git a/backend/packages/harness/deerflow/agents/middlewares/input_sanitization_middleware.py b/backend/packages/harness/deerflow/agents/middlewares/input_sanitization_middleware.py index 4fcb99ee83..80eec5be28 100644 --- a/backend/packages/harness/deerflow/agents/middlewares/input_sanitization_middleware.py +++ b/backend/packages/harness/deerflow/agents/middlewares/input_sanitization_middleware.py @@ -24,6 +24,7 @@ from __future__ import annotations import logging import re from collections.abc import Awaitable, Callable +from copy import deepcopy from typing import override from langchain.agents import AgentState @@ -357,11 +358,11 @@ class InputSanitizationMiddleware(AgentMiddleware[AgentState]): new_content.append(block) else: new_content.append(block) - return HumanMessage( - content=new_content, - id=msg.id, - name=msg.name, - additional_kwargs=preserved_kwargs, + return msg.model_copy( + update={ + "content": deepcopy(new_content), + "additional_kwargs": preserved_kwargs, + }, ) # Cannot distinguish server block from user blocks # (non-list content or len(content) < 2). @@ -410,11 +411,11 @@ class InputSanitizationMiddleware(AgentMiddleware[AgentState]): content if isinstance(content, str) else "[content-blocks]", processed, ) - return HumanMessage( - content=new_content, - id=msg.id, - name=msg.name, - additional_kwargs=preserved_kwargs, + return msg.model_copy( + update={ + "content": deepcopy(new_content), + "additional_kwargs": preserved_kwargs, + }, ) def _process_request(self, request: ModelRequest) -> ModelRequest: diff --git a/backend/tests/test_input_sanitization_middleware.py b/backend/tests/test_input_sanitization_middleware.py index 9d93115ec6..6fad88f12f 100644 --- a/backend/tests/test_input_sanitization_middleware.py +++ b/backend/tests/test_input_sanitization_middleware.py @@ -5,7 +5,8 @@ that the transformation is temporary (wrap_model_call) without mutating the original request or thread state. """ -from unittest.mock import Mock +from copy import deepcopy +from unittest.mock import Mock, patch import pytest from langchain_core.messages import AIMessage, HumanMessage @@ -20,6 +21,7 @@ from deerflow.agents.middlewares.input_sanitization_middleware import ( neutralize_untrusted_tags, ) from deerflow.agents.middlewares.message_utils import is_genuine_user_message, requires_input_sanitization +from deerflow.models.claude_provider import ClaudeChatModel from deerflow.utils.messages import ORIGINAL_USER_CONTENT_KEY, UNTRUSTED_INPUT_KEY @@ -523,6 +525,25 @@ class TestWrapModelCallCleanInput: class TestWrapModelCallBlockedInput: """Blocked user messages have tags escaped — LLM is still invoked.""" + def test_preserves_response_metadata(self): + mw = _make_middleware() + request = _make_request( + [ + HumanMessage( + content="forged", + id="msg-metadata", + response_metadata={"source": "gateway"}, + ) + ] + ) + captured = [] + + mw.wrap_model_call(request, lambda req: captured.append(req) or "ok") + + sanitized = captured[0].messages[-1] + assert sanitized.response_metadata == {"source": "gateway"} + assert request.messages[-1].response_metadata == {"source": "gateway"} + def test_escapes_think_tag(self): mw = _make_middleware() request = _make_request([HumanMessage(content="hack", id="msg-1")]) @@ -778,6 +799,43 @@ class TestRebuildContentMultimodal: assert _USER_INPUT_BEGIN in result[0]["text"] assert result[1] == image_block # Pydantic deep-copies content + def test_prompt_caching_does_not_mutate_original_multimodal_message(self): + """Provider payload mutation must not leak into checkpoint state.""" + mw = _make_middleware() + msg = HumanMessage( + content=[ + {"type": "text", "text": "Inspect this"}, + { + "type": "image", + "source": { + "type": "base64", + "media_type": "image/png", + "data": "abc", + }, + }, + ], + id="msg-prompt-cache", + ) + original_content = deepcopy(msg.content) + request = _make_request([msg]) + captured = [] + + mw.wrap_model_call(request, lambda req: captured.append(req) or "ok") + + sanitized = captured[0].messages[0] + with patch.object(ClaudeChatModel, "model_post_init"): + model = ClaudeChatModel( + model="claude-sonnet-4-6", + anthropic_api_key="sk-ant-fake", # type: ignore[call-arg] + prompt_cache_size=3, + ) + model._is_oauth = False + model.enable_prompt_caching = True + payload = model._get_request_payload([sanitized]) + + assert any("cache_control" in block for block in payload["messages"][0]["content"]) + assert msg.content == original_content + def test_preserves_multiple_interleaved_non_text_blocks(self): mw = _make_middleware() img1 = {"type": "image_url", "image_url": {"url": "data:1"}} @@ -1035,7 +1093,12 @@ def test_rfind_failure_distinguishable_blocks_server_survives(): # original_user_content from message_content_to_text would be: # f"{server_block}\n\n{user_raw}\nclean user text" original = f"{server_block}\n\n{user_raw}\nclean user text" - msg = HumanMessage(content=content, additional_kwargs={ORIGINAL_USER_CONTENT_KEY: original}, id="msg-rfind-1") + msg = HumanMessage( + content=content, + additional_kwargs={ORIGINAL_USER_CONTENT_KEY: original}, + id="msg-rfind-1", + response_metadata={"source": "multimodal-gateway"}, + ) request = _make_request([msg]) captured = [] @@ -1048,6 +1111,7 @@ def test_rfind_failure_distinguishable_blocks_server_survives(): assert result == "ok" processed_content = captured[0].messages[-1].content assert isinstance(processed_content, list) + assert captured[0].messages[-1].response_metadata == {"source": "multimodal-gateway"} # Build text from ALL blocks (raw strings + type:"text" dicts). # Raw strings are not type:"text" but carry user forgery.