From 724ab9098dfbf6bb4ebbf5de5733381dc64257b9 Mon Sep 17 00:00:00 2001
From: pedrommaiaa <100535780+pedrommaiaa@users.noreply.github.com>
Date: Thu, 4 Jun 2026 20:25:44 -0400
Subject: [PATCH] perf: avoid broad message prep deepcopies

(cherry picked from commit 030746c56010ce9f343140a16dfe66a7966b0705)
---
 agent/prompt_caching.py                       | 50 +++++++++--
 agent/transports/chat_completions.py          | 69 +++++++++++----
 .../model-providers/qwen-oauth/__init__.py    | 54 ++++++++----
 tests/agent/test_prompt_caching.py            | 48 +++++++++++
 .../agent/transports/test_chat_completions.py | 84 ++++++++++++++++++-
 tests/providers/test_provider_profiles.py     | 63 ++++++++++++++
 tests/run_agent/test_provider_parity.py       | 10 +++
 7 files changed, 341 insertions(+), 37 deletions(-)

diff --git a/agent/prompt_caching.py b/agent/prompt_caching.py
index 9a2fdf4ccce..34d34143127 100644
--- a/agent/prompt_caching.py
+++ b/agent/prompt_caching.py
@@ -8,7 +8,6 @@ single session.
 Pure functions -- no class state, no AIAgent dependency.
 """
 
-import copy
 from typing import Any, Dict, List
 
 
@@ -81,6 +80,41 @@ def _build_marker(ttl: str) -> Dict[str, str]:
     return marker
 
 
+def _copy_message_for_api_mutation(msg: Dict[str, Any]) -> Dict[str, Any]:
+    """Copy message containers that later API-retry recovery may mutate."""
+    copied = dict(msg)
+
+    content = msg.get("content")
+    if isinstance(content, list) and content:
+        copied_content = []
+        for part in content:
+            if isinstance(part, dict):
+                copied_part = dict(part)
+                image_url = part.get("image_url")
+                if isinstance(image_url, dict):
+                    copied_part["image_url"] = dict(image_url)
+                copied_content.append(copied_part)
+            else:
+                copied_content.append(part)
+        copied["content"] = copied_content
+
+    tool_calls = msg.get("tool_calls")
+    if isinstance(tool_calls, list) and tool_calls:
+        copied_tool_calls = []
+        for tool_call in tool_calls:
+            if isinstance(tool_call, dict):
+                copied_tool_call = dict(tool_call)
+                function = tool_call.get("function")
+                if isinstance(function, dict):
+                    copied_tool_call["function"] = dict(function)
+                copied_tool_calls.append(copied_tool_call)
+            else:
+                copied_tool_calls.append(tool_call)
+        copied["tool_calls"] = copied_tool_calls
+
+    return copied
+
+
 def apply_anthropic_cache_control(
     api_messages: List[Dict[str, Any]],
     cache_ttl: str = "5m",
@@ -92,11 +126,17 @@ def apply_anthropic_cache_control(
     messages, all at the same TTL.
 
     Returns:
-        Deep copy of messages with cache_control breakpoints injected.
+        Copy of messages with cache_control breakpoints injected. Message
+        containers are copied so later API-retry recovery can mutate the
+        returned request copy without touching canonical history.
     """
-    messages = copy.deepcopy(api_messages)
-    if not messages:
-        return messages
+    if not api_messages:
+        return []
+
+    messages = [
+        _copy_message_for_api_mutation(msg) if isinstance(msg, dict) else msg
+        for msg in api_messages
+    ]
 
     marker = _build_marker(cache_ttl)
 
diff --git a/agent/transports/chat_completions.py b/agent/transports/chat_completions.py
index ff2cdcbaee6..b20445be90a 100644
--- a/agent/transports/chat_completions.py
+++ b/agent/transports/chat_completions.py
@@ -9,7 +9,6 @@ which has provider-specific conditionals for max_tokens defaults,
 reasoning configuration, temperature handling, and extra_body assembly.
 """
 
-import copy
 from typing import Any, Dict
 
 from agent.lmstudio_reasoning import resolve_lmstudio_effort
@@ -100,7 +99,7 @@ def _is_gemini_openai_compat_base_url(base_url: Any) -> bool:
 
 
 def _model_consumes_thought_signature(model: Any) -> bool:
-    """True when the outgoing model is a Gemini family model that requires
+    """True when the outgoing model is a Gemini model that requires
     ``extra_content`` (thought_signature) to be replayed on tool calls.
 
     Gemini 3 thinking models attach ``extra_content`` to each tool call and
@@ -112,7 +111,7 @@ def _model_consumes_thought_signature(model: Any) -> bool:
     ``extra_content`` from earlier in a mixed-provider session.
     """
     m = str(model or "").lower()
-    return "gemini" in m or "gemma" in m
+    return "gemini" in m
 
 
 class ChatCompletionsTransport(ProviderTransport):
@@ -136,7 +135,7 @@ class ChatCompletionsTransport(ProviderTransport):
           ``codex_message_items`` on the message, ``call_id`` /
           ``response_item_id`` on ``tool_calls`` entries.
         - ``extra_content`` on ``tool_calls`` (Gemini thought_signature) —
-          stripped unless the outgoing ``model`` is itself Gemini-family.
+          stripped unless the outgoing ``model`` is itself Gemini.
           Gemini 3 thinking models attach it for replay, but strict providers
           (Fireworks, Mistral) reject any payload containing it with
           ``Extra inputs are not permitted, field: 'messages[N].tool_calls[M].extra_content'``.
@@ -195,27 +194,63 @@ class ChatCompletionsTransport(ProviderTransport):
         if not needs_sanitize:
             return messages
 
-        sanitized = copy.deepcopy(messages)
-        for msg in sanitized:
+        sanitized = list(messages)
+        for msg_idx, msg in enumerate(messages):
             if not isinstance(msg, dict):
                 continue
-            msg.pop("codex_reasoning_items", None)
-            msg.pop("codex_message_items", None)
-            msg.pop("tool_name", None)
-            msg.pop("timestamp", None)  # #47868 — leak into strict providers
+
+            copied_msg: dict[str, Any] | None = None
+
+            def mutable_msg() -> dict[str, Any]:
+                nonlocal copied_msg
+                if copied_msg is None:
+                    copied_msg = dict(msg)
+                    sanitized[msg_idx] = copied_msg
+                return copied_msg
+
+            if (
+                "codex_reasoning_items" in msg
+                or "codex_message_items" in msg
+                or "tool_name" in msg
+                or "timestamp" in msg  # #47868 — leak into strict providers
+            ):
+                out_msg = mutable_msg()
+                out_msg.pop("codex_reasoning_items", None)
+                out_msg.pop("codex_message_items", None)
+                out_msg.pop("tool_name", None)
+                out_msg.pop("timestamp", None)  # #47868 — leak into strict providers
+
+
             # Drop all Hermes-internal scaffolding markers (``_``-prefixed).
             # OpenAI's message schema has no ``_``-prefixed fields, so this
             # is safe and future-proofs against new markers being added.
-            for key in [k for k in msg if isinstance(k, str) and k.startswith("_")]:
-                msg.pop(key, None)
+            internal_keys = [k for k in msg if isinstance(k, str) and k.startswith("_")]
+            if internal_keys:
+                out_msg = mutable_msg()
+                for key in internal_keys:
+                    out_msg.pop(key, None)
+
             tool_calls = msg.get("tool_calls")
             if isinstance(tool_calls, list):
-                for tc in tool_calls:
+                copied_tool_calls: list[Any] | None = None
+                for tc_idx, tc in enumerate(tool_calls):
                     if isinstance(tc, dict):
-                        tc.pop("call_id", None)
-                        tc.pop("response_item_id", None)
-                        if strip_extra_content:
-                            tc.pop("extra_content", None)
+                        should_copy_tc = (
+                            "call_id" in tc
+                            or "response_item_id" in tc
+                            or (strip_extra_content and "extra_content" in tc)
+                        )
+                        if should_copy_tc:
+                            if copied_tool_calls is None:
+                                copied_tool_calls = list(tool_calls)
+                            copied_tc = dict(tc)
+                            copied_tc.pop("call_id", None)
+                            copied_tc.pop("response_item_id", None)
+                            if strip_extra_content:
+                                copied_tc.pop("extra_content", None)
+                            copied_tool_calls[tc_idx] = copied_tc
+                if copied_tool_calls is not None:
+                    mutable_msg()["tool_calls"] = copied_tool_calls
         return sanitized
 
     def convert_tools(self, tools: list[dict[str, Any]]) -> list[dict[str, Any]]:
diff --git a/plugins/model-providers/qwen-oauth/__init__.py b/plugins/model-providers/qwen-oauth/__init__.py
index a6ba29f76cb..ddf9120a65d 100644
--- a/plugins/model-providers/qwen-oauth/__init__.py
+++ b/plugins/model-providers/qwen-oauth/__init__.py
@@ -1,6 +1,4 @@
 """Qwen Portal provider profile."""
-
-import copy
 from typing import Any
 
 from providers import register_provider
@@ -10,6 +8,15 @@ from providers.base import ProviderProfile
 class QwenProfile(ProviderProfile):
     """Qwen Portal — message normalization, vl_high_resolution, metadata top-level."""
 
+    @staticmethod
+    def _copy_part_if_request_mutable(part: dict[str, Any]) -> tuple[dict[str, Any], bool]:
+        image_url = part.get("image_url")
+        if isinstance(image_url, dict):
+            copied = dict(part)
+            copied["image_url"] = dict(image_url)
+            return copied, True
+        return part, False
+
     def prepare_messages(self, messages: list[dict[str, Any]]) -> list[dict[str, Any]]:
         """Normalize content to list-of-dicts format.
 
@@ -17,37 +24,56 @@ class QwenProfile(ProviderProfile):
 
         Matches the behavior of run_agent.py:_qwen_prepare_chat_messages().
         """
-        prepared = copy.deepcopy(messages)
-        if not prepared:
-            return prepared
+        if not messages:
+            return []
 
-        for msg in prepared:
+        prepared = list(messages)
+        system_idx: int | None = None
+
+        for idx, msg in enumerate(messages):
             if not isinstance(msg, dict):
                 continue
+            if system_idx is None and msg.get("role") == "system":
+                system_idx = idx
             content = msg.get("content")
             if isinstance(content, str):
-                msg["content"] = [{"type": "text", "text": content}]
+                msg_copy = dict(msg)
+                msg_copy["content"] = [{"type": "text", "text": content}]
+                prepared[idx] = msg_copy
             elif isinstance(content, list):
                 normalized_parts = []
+                changed = False
                 for part in content:
                     if isinstance(part, str):
                         normalized_parts.append({"type": "text", "text": part})
+                        changed = True
                     elif isinstance(part, dict):
-                        normalized_parts.append(part)
-                if normalized_parts:
-                    msg["content"] = normalized_parts
+                        normalized_part, copied = self._copy_part_if_request_mutable(part)
+                        normalized_parts.append(normalized_part)
+                        changed = changed or copied
+                    else:
+                        changed = True
+                if normalized_parts and changed:
+                    msg_copy = dict(msg)
+                    msg_copy["content"] = normalized_parts
+                    prepared[idx] = msg_copy
 
         # Inject cache_control on the last part of the system message.
-        for msg in prepared:
-            if isinstance(msg, dict) and msg.get("role") == "system":
+        if system_idx is not None:
+            msg = prepared[system_idx]
+            if isinstance(msg, dict):
                 content = msg.get("content")
                 if (
                     isinstance(content, list)
                     and content
                     and isinstance(content[-1], dict)
                 ):
-                    content[-1]["cache_control"] = {"type": "ephemeral"}
-                break
+                    msg_copy = dict(msg)
+                    content_copy = list(content)
+                    content_copy[-1] = dict(content_copy[-1])
+                    content_copy[-1]["cache_control"] = {"type": "ephemeral"}
+                    msg_copy["content"] = content_copy
+                    prepared[system_idx] = msg_copy
 
         return prepared
 
diff --git a/tests/agent/test_prompt_caching.py b/tests/agent/test_prompt_caching.py
index 7a1c0495e1b..889b454d4e6 100644
--- a/tests/agent/test_prompt_caching.py
+++ b/tests/agent/test_prompt_caching.py
@@ -131,6 +131,54 @@ class TestApplyAnthropicCacheControl:
         # Original should be unmodified
         assert "cache_control" not in msgs[0].get("content", "")
 
+    def test_request_copy_protects_unmarked_nested_content(self):
+        msgs = [
+            {"role": "system", "content": "System"},
+            {"role": "user", "content": [{"type": "text", "text": "msg1"}]},
+            {"role": "assistant", "content": [{"type": "text", "text": "msg2"}]},
+            {"role": "user", "content": [{"type": "text", "text": "msg3"}]},
+            {"role": "assistant", "content": [{"type": "text", "text": "msg4"}]},
+        ]
+
+        result = apply_anthropic_cache_control(msgs)
+
+        assert result is not msgs
+        assert result[0] is not msgs[0]
+        assert result[1] is not msgs[1]
+        assert result[1]["content"] is not msgs[1]["content"]
+        assert result[1]["content"][0] is not msgs[1]["content"][0]
+        assert result[2] is not msgs[2]
+        assert result[2]["content"] is not msgs[2]["content"]
+        assert result[2]["content"][0] is not msgs[2]["content"][0]
+        assert "cache_control" not in msgs[2]["content"][0]
+        assert result[2]["content"][0]["cache_control"] == MARKER
+
+        result[1]["content"][0]["text"] = "mutated request copy"
+        assert msgs[1]["content"][0]["text"] == "msg1"
+
+    def test_request_copy_protects_nested_image_url_retry_mutation(self):
+        image_url = {"url": "data:image/png;base64,original"}
+        msgs = [
+            {"role": "system", "content": "System"},
+            {
+                "role": "user",
+                "content": [
+                    {"type": "text", "text": "see image"},
+                    {"type": "image_url", "image_url": image_url},
+                ],
+            },
+        ]
+
+        result = apply_anthropic_cache_control(msgs)
+
+        assert result[1]["content"][1] is not msgs[1]["content"][1]
+        assert result[1]["content"][1]["image_url"] is not image_url
+
+        result[1]["content"][1]["image_url"]["url"] = "data:image/png;base64,shrunk"
+        assert msgs[1]["content"][1]["image_url"]["url"] == (
+            "data:image/png;base64,original"
+        )
+
     def test_system_message_gets_marker(self):
         msgs = [
             {"role": "system", "content": "You are helpful"},
diff --git a/tests/agent/transports/test_chat_completions.py b/tests/agent/transports/test_chat_completions.py
index af24400ff51..7d1b07d2d75 100644
--- a/tests/agent/transports/test_chat_completions.py
+++ b/tests/agent/transports/test_chat_completions.py
@@ -77,13 +77,23 @@ class TestChatCompletionsBasic:
         every turn — stripping it would 400. Keep extra_content for Gemini
         targets (including aggregator slugs like google/gemini-3-pro).
         """
-        for model in ("gemini-3-pro", "google/gemini-3-pro-preview", "gemma-3-27b"):
+        for model in ("gemini-3-pro", "google/gemini-3-pro-preview"):
             msgs = self._msg_with_extra_content()
             result = transport.convert_messages(msgs, model=model)
             assert result[0]["tool_calls"][0]["extra_content"] == {
                 "google": {"thought_signature": "SIG_123"}
             }, model
 
+    def test_convert_messages_strips_extra_content_for_gemma(self, transport):
+        """Gemma can be served by Google, but it is not Gemini and rejects
+        Gemini-only request fields. Stale thought signatures from a prior
+        Gemini turn must be stripped when switching to Gemma.
+        """
+        msgs = self._msg_with_extra_content()
+        result = transport.convert_messages(msgs, model="google/gemma-4-31b-it")
+        assert "extra_content" not in result[0]["tool_calls"][0]
+        assert "extra_content" in msgs[0]["tool_calls"][0]
+
     def test_convert_messages_strips_tool_name(self, transport):
         """Internal `tool_name` (used for FTS indexing in the SQLite store) is
         not part of the OpenAI Chat Completions schema. Strict providers like
@@ -161,6 +171,78 @@ class TestChatCompletionsBasic:
         ]
         assert transport.convert_messages(msgs) is msgs
 
+    def test_convert_messages_copy_on_write_for_dirty_history(self, transport):
+        """Dirty provider metadata should not force a full-history deepcopy."""
+        clean_tool_call = {
+            "id": "call_clean",
+            "type": "function",
+            "function": {"name": "safe", "arguments": "{}"},
+        }
+        msgs = [
+            {"role": "user", "content": "hi", "metadata": {"large": ["shared"]}},
+            {
+                "role": "assistant",
+                "content": "ok",
+                "tool_calls": [
+                    clean_tool_call,
+                    {
+                        "id": "call_dirty",
+                        "call_id": "call_dirty",
+                        "response_item_id": "fc_dirty",
+                        "extra_content": {"google": {"thought_signature": "SIG"}},
+                        "type": "function",
+                        "function": {"name": "t", "arguments": "{}"},
+                    },
+                ],
+            },
+        ]
+
+        result = transport.convert_messages(msgs, model="gpt-4o")
+
+        assert result is not msgs
+        assert result[0] is msgs[0]
+        assert result[1] is not msgs[1]
+        assert result[1]["tool_calls"] is not msgs[1]["tool_calls"]
+        assert result[1]["tool_calls"][0] is clean_tool_call
+        assert result[1]["tool_calls"][1] is not msgs[1]["tool_calls"][1]
+        assert "call_id" not in result[1]["tool_calls"][1]
+        assert "response_item_id" not in result[1]["tool_calls"][1]
+        assert "extra_content" not in result[1]["tool_calls"][1]
+        assert "call_id" in msgs[1]["tool_calls"][1]
+        assert "extra_content" in msgs[1]["tool_calls"][1]
+
+    def test_same_history_survives_strict_then_gemini_model_switch(self, transport):
+        """Strict cleanup must not remove Gemini replay metadata from history."""
+        msgs = [
+            {
+                "role": "assistant",
+                "content": "ok",
+                "tool_calls": [
+                    {
+                        "id": "call_1",
+                        "call_id": "call_1",
+                        "response_item_id": "fc_1",
+                        "extra_content": {"google": {"thought_signature": "SIG_123"}},
+                        "type": "function",
+                        "function": {"name": "t", "arguments": "{}"},
+                    }
+                ],
+            }
+        ]
+
+        strict = transport.convert_messages(msgs, model="accounts/fireworks/models/llama")
+        gemini = transport.convert_messages(msgs, model="google/gemini-3-pro")
+
+        assert "extra_content" not in strict[0]["tool_calls"][0]
+        assert "call_id" not in strict[0]["tool_calls"][0]
+        assert "response_item_id" not in strict[0]["tool_calls"][0]
+        assert gemini[0]["tool_calls"][0]["extra_content"] == {
+            "google": {"thought_signature": "SIG_123"}
+        }
+        # The canonical history still has both provider-specific metadata sets.
+        assert msgs[0]["tool_calls"][0]["call_id"] == "call_1"
+        assert msgs[0]["tool_calls"][0]["extra_content"]["google"]["thought_signature"] == "SIG_123"
+
 
 class TestChatCompletionsBuildKwargs:
 
diff --git a/tests/providers/test_provider_profiles.py b/tests/providers/test_provider_profiles.py
index 3eb23403048..5f2996d8ebf 100644
--- a/tests/providers/test_provider_profiles.py
+++ b/tests/providers/test_provider_profiles.py
@@ -464,6 +464,69 @@ class TestQwenProfile:
         assert isinstance(result[1]["content"], list)
         assert result[1]["content"][0]["text"] == "hello"
 
+    def test_prepare_messages_copy_on_write(self):
+        p = get_provider_profile("qwen-oauth")
+        system_part = {"type": "text", "text": "Be helpful"}
+        msgs = [
+            {"role": "system", "content": [system_part]},
+            {"role": "assistant", "content": [{"type": "text", "text": "unchanged"}]},
+            {"role": "user", "content": ["hello"]},
+        ]
+
+        result = p.prepare_messages(msgs)
+
+        assert result is not msgs
+        assert result[0] is not msgs[0]
+        assert result[0]["content"] is not msgs[0]["content"]
+        assert result[0]["content"][0] is not system_part
+        assert result[0]["content"][0]["cache_control"] == {"type": "ephemeral"}
+        assert "cache_control" not in system_part
+        assert result[1] is msgs[1]
+        assert result[2] is not msgs[2]
+        assert result[2]["content"] == [{"type": "text", "text": "hello"}]
+        assert msgs[2]["content"] == ["hello"]
+
+    def test_prepare_messages_does_not_poison_strict_provider_history(self):
+        qwen = get_provider_profile("qwen-oauth")
+        msgs = [
+            {"role": "system", "content": [{"type": "text", "text": "Be helpful"}]},
+            {"role": "user", "content": "hello"},
+        ]
+
+        qwen_result = qwen.prepare_messages(msgs)
+
+        assert qwen_result[0]["content"][0]["cache_control"] == {"type": "ephemeral"}
+        assert "cache_control" not in msgs[0]["content"][0]
+        assert msgs[1]["content"] == "hello"
+
+    def test_prepare_messages_protects_nested_image_url_retry_mutation(self):
+        qwen = get_provider_profile("qwen-oauth")
+        image_url = {"url": "data:image/png;base64,original"}
+        msgs = [
+            {"role": "system", "content": "Be helpful"},
+            {
+                "role": "user",
+                "content": [
+                    {"type": "text", "text": "see image"},
+                    {"type": "image_url", "image_url": image_url},
+                ],
+            },
+        ]
+
+        qwen_result = qwen.prepare_messages(msgs)
+
+        assert qwen_result[1] is not msgs[1]
+        assert qwen_result[1]["content"] is not msgs[1]["content"]
+        assert qwen_result[1]["content"][1] is not msgs[1]["content"][1]
+        assert qwen_result[1]["content"][1]["image_url"] is not image_url
+
+        qwen_result[1]["content"][1]["image_url"]["url"] = (
+            "data:image/png;base64,shrunk"
+        )
+        assert msgs[1]["content"][1]["image_url"]["url"] == (
+            "data:image/png;base64,original"
+        )
+
     def test_metadata_top_level(self):
         p = get_provider_profile("qwen-oauth")
         meta = {"sessionId": "s123", "promptId": "p456"}
diff --git a/tests/run_agent/test_provider_parity.py b/tests/run_agent/test_provider_parity.py
index 8229b0f020d..df72bec8c00 100644
--- a/tests/run_agent/test_provider_parity.py
+++ b/tests/run_agent/test_provider_parity.py
@@ -297,6 +297,16 @@ class TestBuildApiKwargsOpenRouter:
         # call_id/response_item_id still stripped regardless of model
         assert "call_id" not in result["tool_calls"][0]
 
+    def test_sanitize_tool_calls_strips_extra_content_for_gemma(self, monkeypatch):
+        """Gemma is served by Google but does not consume Gemini thought signatures."""
+        agent = _make_agent(monkeypatch, "openrouter")
+        api_msg = self._api_msg_with_extra_content()
+        result = agent._sanitize_tool_calls_for_strict_api(
+            api_msg, model="google/gemma-4-31b-it"
+        )
+        assert "extra_content" not in result["tool_calls"][0]
+        assert "call_id" not in result["tool_calls"][0]
+
 
 class TestDeveloperRoleSwap:
     """GPT-5 and Codex models should get 'developer' instead of 'system' role."""