mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-10-26 13:16:24 +00:00
fix(models): keep Claude prompt-cache markers out of thread messages (#5892)
* fix(models): keep Claude prompt-cache markers out of thread messages ClaudeChatModel._apply_prompt_caching set cache_control on payload blocks in place. langchain-anthropic forwards Claude-native blocks (image or document blocks with a source, search results) by reference, so the marker landed in the thread's own message and was checkpointed with it. Each later request carried those stale markers on top of its own four, and from the third turn of a thread with such blocks the API rejected every request (400, "A maximum of 4 blocks with cache_control may be provided"). Breakpoints are now placed on copies, and markers already in the request are dropped first, so threads whose checkpoints already carry them recover without a migration. The OAuth path's _strip_cache_control popped markers from the same shared objects; it now copies too, and prompt caching reuses it. * fix(models): drop stored Claude cache markers with caching off too The strip only ran inside _apply_prompt_caching, so with enable_prompt_caching: false a thread whose checkpoint already carried five or more markers kept failing: exactly the switch an operator would flip in response to these 400s. _get_request_payload now strips on the non-caching path as well, so recovery no longer depends on the flag. The stored-checkpoint recovery is now pinned through _get_request_payload with real messages, for caching on and off. The docstring and the AGENTS.md note now say which objects are really shared (blocks, list-form system blocks, a reused binding's tool dicts), that the message-dict and list copies are defensive, and that _apply_prompt_caching must strip first. --------- Co-authored-by: Totoro-qaq <279883115+Totoro-qaq@users.noreply.github.com>
This commit is contained in:
parent
e2c5027b0c
commit
4028b63647
@ -63,6 +63,12 @@ the same policy first so run metadata reports the effective values. Design note:
|
||||
- `ClaudeChatModel.model_post_init` calls `load_claude_code_credential()` for every instance, and `create_chat_model` builds fresh instances per run (lead agent, title, summarization, subagents)
|
||||
- `$CLAUDE_CODE_OAUTH_TOKEN_FILE_DESCRIPTOR` is a one-shot handoff: a pipe returns EOF and a file keeps its advanced offset. `_read_secret_from_file_descriptor` therefore caches a non-empty secret per `(env_var, fd)` under a lock held across the read. Do not drop the cache or the lock — later instances would get no credential, and the Anthropic SDK raises `TypeError: Could not resolve authentication method` before sending. Empty reads and `OSError` are not cached. The key is the descriptor number on purpose — a closed handoff keeps serving its token, and a secret placed on a recycled number in-process is not re-read unless the cache is cleared. The cache is per process, so a new process (e.g. a uvicorn `--reload` worker) cannot recover a drained descriptor. Pinned by `tests/test_credential_loader.py`, including a two-instance `ClaudeChatModel` test
|
||||
|
||||
### Claude Prompt Caching (`packages/harness/deerflow/models/claude_provider.py`)
|
||||
|
||||
- The request payload shares objects with the caller: langchain-anthropic forwards Claude-native blocks (an image or document with a `source`, search results) and list-form system blocks by reference, and a reused tool binding passes its own tool dicts (the lead agent re-binds per call, so its tool dicts are fresh). Writing `cache_control` in place checkpointed the markers with the thread's messages, and the stale ones pushed every later request past the 4-breakpoint limit
|
||||
- Every request goes through `_strip_cache_control`, with caching on or off, so markers stored by older checkpoints never reach the API. It copies a marked block without its marker and replaces the system, message, content and tool lists and every message dict with copies; langchain-anthropic already builds fresh message dicts and content lists, so that part is defensive
|
||||
- `_apply_prompt_caching` must call `_strip_cache_control` first: it then replaces slots in those payload-owned lists with marked copies and writes `msg["content"]` on copied message dicts, placing at most four breakpoints. Pinned by `tests/test_claude_provider_prompt_caching.py`
|
||||
|
||||
### vLLM Provider (`packages/harness/deerflow/models/vllm_provider.py`)
|
||||
|
||||
- `VllmChatModel` subclasses `langchain_openai:ChatOpenAI` for vLLM 0.19.0 OpenAI-compatible endpoints
|
||||
|
||||
@ -146,6 +146,11 @@ class ClaudeChatModel(ChatAnthropic):
|
||||
|
||||
if self.enable_prompt_caching:
|
||||
self._apply_prompt_caching(payload)
|
||||
else:
|
||||
# Checkpoints written before breakpoints were placed on copies can
|
||||
# carry markers on old blocks; a thread with more than four of them
|
||||
# must not stay broken just because caching was switched off.
|
||||
self._strip_cache_control(payload)
|
||||
|
||||
if self.auto_thinking_budget:
|
||||
self._apply_thinking_budget(payload)
|
||||
@ -197,28 +202,38 @@ class ClaudeChatModel(ChatAnthropic):
|
||||
placed on the *last* eligible blocks because later breakpoints cover a
|
||||
larger prefix and yield better cache hit rates.
|
||||
|
||||
The budget covers the whole request, so markers already present are
|
||||
dropped first. Checkpoints written before breakpoints were placed on
|
||||
copies can still carry markers on old blocks, and those would otherwise
|
||||
push every later request of the thread past the limit. A marked block
|
||||
is always a copy; ``_strip_cache_control`` explains why.
|
||||
|
||||
The system prompt is expected to be fully static (no per-user memory or
|
||||
current date). Dynamic context is injected per-turn via
|
||||
DynamicContextMiddleware as a <system-reminder> in the first HumanMessage.
|
||||
"""
|
||||
MAX_CACHE_BREAKPOINTS = 4
|
||||
|
||||
# Collect candidate blocks in document order:
|
||||
# Must run first: stripping leaves every list and message dict below
|
||||
# owned by this payload, which is what makes replacing a candidate
|
||||
# slot (and writing msg["content"]) safe for any payload.
|
||||
self._strip_cache_control(payload)
|
||||
|
||||
# Collect candidate slots (list, index) in document order:
|
||||
# 1. system text blocks
|
||||
# 2. content blocks of the last prompt_cache_size messages
|
||||
# 3. the last tool definition
|
||||
candidates: list[dict] = []
|
||||
candidates: list[tuple[list, int]] = []
|
||||
|
||||
# 1. System blocks
|
||||
system = payload.get("system")
|
||||
if system and isinstance(system, list):
|
||||
for block in system:
|
||||
for index, block in enumerate(system):
|
||||
if isinstance(block, dict) and block.get("type") == "text":
|
||||
candidates.append(block)
|
||||
candidates.append((system, index))
|
||||
elif system and isinstance(system, str):
|
||||
new_block: dict = {"type": "text", "text": system}
|
||||
payload["system"] = [new_block]
|
||||
candidates.append(new_block)
|
||||
payload["system"] = [{"type": "text", "text": system}]
|
||||
candidates.append((payload["system"], 0))
|
||||
|
||||
# 2. Recent message blocks
|
||||
messages = payload.get("messages", [])
|
||||
@ -229,23 +244,22 @@ class ClaudeChatModel(ChatAnthropic):
|
||||
continue
|
||||
content = msg.get("content")
|
||||
if isinstance(content, list):
|
||||
for block in content:
|
||||
for index, block in enumerate(content):
|
||||
if isinstance(block, dict):
|
||||
candidates.append(block)
|
||||
candidates.append((content, index))
|
||||
elif isinstance(content, str) and content:
|
||||
new_block = {"type": "text", "text": content}
|
||||
msg["content"] = [new_block]
|
||||
candidates.append(new_block)
|
||||
msg["content"] = [{"type": "text", "text": content}]
|
||||
candidates.append((msg["content"], 0))
|
||||
|
||||
# 3. Last tool definition
|
||||
tools = payload.get("tools", [])
|
||||
if tools and isinstance(tools[-1], dict):
|
||||
candidates.append(tools[-1])
|
||||
candidates.append((tools, len(tools) - 1))
|
||||
|
||||
# Apply cache_control only to the last MAX_CACHE_BREAKPOINTS candidates
|
||||
# to stay within the API limit.
|
||||
for block in candidates[-MAX_CACHE_BREAKPOINTS:]:
|
||||
block["cache_control"] = {"type": "ephemeral"}
|
||||
for container, index in candidates[-MAX_CACHE_BREAKPOINTS:]:
|
||||
container[index] = {**container[index], "cache_control": {"type": "ephemeral"}}
|
||||
|
||||
def _apply_thinking_budget(self, payload: dict) -> None:
|
||||
"""Auto-allocate thinking budget (80% of max_tokens)."""
|
||||
@ -262,26 +276,46 @@ class ClaudeChatModel(ChatAnthropic):
|
||||
|
||||
@staticmethod
|
||||
def _strip_cache_control(payload: dict) -> None:
|
||||
"""Remove cache_control markers before OAuth requests reach Anthropic."""
|
||||
for section in ("system", "messages"):
|
||||
items = payload.get(section)
|
||||
if not isinstance(items, list):
|
||||
continue
|
||||
for item in items:
|
||||
if not isinstance(item, dict):
|
||||
continue
|
||||
item.pop("cache_control", None)
|
||||
content = item.get("content")
|
||||
if isinstance(content, list):
|
||||
for block in content:
|
||||
if isinstance(block, dict):
|
||||
block.pop("cache_control", None)
|
||||
"""Remove cache_control markers without writing into the payload's objects.
|
||||
|
||||
Every request is stripped: prompt caching strips before placing its own
|
||||
breakpoints, a request without caching is stripped in
|
||||
``_get_request_payload``, and OAuth requests again before they reach
|
||||
Anthropic.
|
||||
|
||||
The payload shares objects with the caller. langchain-anthropic passes
|
||||
Claude-native blocks (an image or document with a ``source``, a search
|
||||
result) and list-form system blocks through by reference, and a reused
|
||||
tool binding passes its own tool dicts. Changing them in place changes
|
||||
the thread's messages, and a marker written there is checkpointed with
|
||||
them. So a block that carries a marker is copied without it, and the
|
||||
system, message, content and tool lists and every message dict are
|
||||
replaced with copies. langchain-anthropic already builds fresh message
|
||||
dicts and content lists, so those copies are defensive; they are what
|
||||
lets ``_apply_prompt_caching`` write into them for any payload.
|
||||
"""
|
||||
|
||||
def without_markers(items: list) -> list:
|
||||
return [{key: value for key, value in item.items() if key != "cache_control"} if isinstance(item, dict) and "cache_control" in item else item for item in items]
|
||||
|
||||
system = payload.get("system")
|
||||
if isinstance(system, list):
|
||||
payload["system"] = without_markers(system)
|
||||
|
||||
messages = payload.get("messages")
|
||||
if isinstance(messages, list):
|
||||
copied: list = []
|
||||
for message in messages:
|
||||
if isinstance(message, dict):
|
||||
message = {key: value for key, value in message.items() if key != "cache_control"}
|
||||
if isinstance(message.get("content"), list):
|
||||
message["content"] = without_markers(message["content"])
|
||||
copied.append(message)
|
||||
payload["messages"] = copied
|
||||
|
||||
tools = payload.get("tools")
|
||||
if isinstance(tools, list):
|
||||
for tool in tools:
|
||||
if isinstance(tool, dict):
|
||||
tool.pop("cache_control", None)
|
||||
payload["tools"] = without_markers(tools)
|
||||
|
||||
def _create(self, payload: dict) -> Any:
|
||||
if self._is_oauth:
|
||||
|
||||
@ -2,12 +2,14 @@
|
||||
|
||||
Validates that the function never places more than 4 cache_control breakpoints
|
||||
(the hard limit enforced by the Anthropic API and AWS Bedrock) regardless of
|
||||
how many system blocks, message content blocks, or tool definitions are present.
|
||||
how many system blocks, message content blocks, or tool definitions are present,
|
||||
and that it never writes into blocks the caller still owns.
|
||||
"""
|
||||
|
||||
from unittest import mock
|
||||
|
||||
import pytest
|
||||
from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
|
||||
|
||||
from deerflow.models.claude_provider import ClaudeChatModel
|
||||
|
||||
@ -247,3 +249,99 @@ def test_old_messages_outside_cache_window_not_cached(model):
|
||||
# Only the last message should be within the cache window
|
||||
assert "cache_control" not in payload["messages"][0]["content"][0]
|
||||
assert payload["messages"][1]["content"][0].get("cache_control") == {"type": "ephemeral"}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Caller-owned blocks
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _native_image() -> dict:
|
||||
"""A Claude-native image block, which langchain-anthropic forwards by reference."""
|
||||
return {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": "iVBORw0KGgo="}}
|
||||
|
||||
|
||||
def test_breakpoints_are_placed_on_copies_of_caller_blocks(model):
|
||||
system_block = {"type": "text", "text": "sys"}
|
||||
image = _native_image()
|
||||
tool = {"name": "bash"}
|
||||
payload: dict = {
|
||||
"system": [system_block],
|
||||
"messages": [{"role": "user", "content": [image]}],
|
||||
"tools": [tool],
|
||||
}
|
||||
|
||||
model._apply_prompt_caching(payload)
|
||||
|
||||
assert _count_cache_control(payload) == 3
|
||||
assert payload["messages"][0]["content"][0] == {**_native_image(), "cache_control": {"type": "ephemeral"}}
|
||||
assert "cache_control" not in system_block
|
||||
assert "cache_control" not in image
|
||||
assert "cache_control" not in tool
|
||||
|
||||
|
||||
def test_repeated_requests_leave_thread_messages_unmarked(model):
|
||||
"""Regression: a marker written into a message's own block was checkpointed.
|
||||
|
||||
Every later request then carried the stale markers of earlier turns on top
|
||||
of its own four, and the API rejected the thread from the third turn on.
|
||||
"""
|
||||
history: list = [SystemMessage(content="You are helpful.")]
|
||||
for turn in range(1, 6):
|
||||
history.append(HumanMessage(content=[_native_image()], id=f"human-{turn}"))
|
||||
payload = model._get_request_payload(history)
|
||||
assert _count_cache_control(payload) <= 4, f"turn {turn}"
|
||||
history.append(AIMessage(content=f"answer {turn}", id=f"ai-{turn}"))
|
||||
|
||||
assert all(message.content == [_native_image()] for message in history if isinstance(message, HumanMessage))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("enable_prompt_caching", [True, False])
|
||||
def test_markers_stored_in_the_thread_stop_reaching_the_api(model, enable_prompt_caching):
|
||||
"""Threads checkpointed before the fix carry a marker on every such block.
|
||||
|
||||
Recovery must not depend on caching being on: switching it off is the
|
||||
obvious reaction to the 400s, and five stored markers still exceed the limit.
|
||||
"""
|
||||
model.enable_prompt_caching = enable_prompt_caching
|
||||
stored = {**_native_image(), "cache_control": {"type": "ephemeral"}}
|
||||
history: list = [SystemMessage(content="You are helpful.")]
|
||||
for turn in range(1, 6):
|
||||
history.append(HumanMessage(content=[dict(stored)], id=f"human-{turn}"))
|
||||
payload = model._get_request_payload(history)
|
||||
assert _count_cache_control(payload) <= (4 if enable_prompt_caching else 0), f"turn {turn}"
|
||||
history.append(AIMessage(content=f"answer {turn}", id=f"ai-{turn}"))
|
||||
|
||||
|
||||
def test_markers_already_in_the_request_are_replaced(model):
|
||||
"""Checkpoints written by older versions can carry markers on old blocks."""
|
||||
stale = [{**_native_image(), "cache_control": {"type": "ephemeral"}} for _ in range(5)]
|
||||
payload: dict = {
|
||||
"system": [{"type": "text", "text": "sys", "cache_control": {"type": "ephemeral"}}],
|
||||
"messages": [{"role": "user", "content": [block]} for block in stale],
|
||||
"tools": [{"name": "bash", "cache_control": {"type": "ephemeral"}}],
|
||||
}
|
||||
|
||||
model._apply_prompt_caching(payload)
|
||||
|
||||
# Candidates: the system block, the last three messages, the tool — the
|
||||
# system block falls outside the last four.
|
||||
assert _count_cache_control(payload) == 4
|
||||
assert [index for index, message in enumerate(payload["messages"]) if "cache_control" in message["content"][0]] == [2, 3, 4]
|
||||
assert "cache_control" not in payload["system"][0]
|
||||
assert payload["tools"][0]["cache_control"] == {"type": "ephemeral"}
|
||||
assert all("cache_control" in block for block in stale)
|
||||
|
||||
|
||||
def test_strip_cache_control_copies_instead_of_popping():
|
||||
system_block = {"type": "text", "text": "sys", "cache_control": {"type": "ephemeral"}}
|
||||
image = {**_native_image(), "cache_control": {"type": "ephemeral"}}
|
||||
tool = {"name": "bash", "cache_control": {"type": "ephemeral"}}
|
||||
message = {"role": "user", "content": [image], "cache_control": {"type": "ephemeral"}}
|
||||
payload: dict = {"system": [system_block], "messages": [message], "tools": [tool]}
|
||||
|
||||
ClaudeChatModel._strip_cache_control(payload)
|
||||
|
||||
assert _count_cache_control(payload) == 0
|
||||
assert payload["messages"][0] == {"role": "user", "content": [_native_image()]}
|
||||
assert all("cache_control" in item for item in (system_block, image, tool, message))
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user