perf: avoid broad message prep deepcopies

(cherry picked from commit 030746c560)
This commit is contained in:
pedrommaiaa 2026-06-04 20:25:44 -04:00 committed by kshitij
parent 4ed910c689
commit 724ab9098d
7 changed files with 341 additions and 37 deletions

View file

@ -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)

View file

@ -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]]:

View file

@ -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

View file

@ -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"},

View file

@ -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:

View file

@ -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"}

View file

@ -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."""