Skip to content

Forward cache_control from message-level placement in ChatDatabricks - #456

Open
adamgurary wants to merge 1 commit into
databricks:mainfrom
adamgurary:cuj-chatdatabricks-cache-control
Open

Forward cache_control from message-level placement in ChatDatabricks#456
adamgurary wants to merge 1 commit into
databricks:mainfrom
adamgurary:cuj-chatdatabricks-cache-control

Conversation

@adamgurary

@adamgurary adamgurary commented Jul 23, 2026

Copy link
Copy Markdown

Context

This fix came out of a hands-on end-to-end evaluation of building, tracing, and improving a LangGraph agent on Databricks.

Problem: ChatDatabricks drops cache_control unless it is inside a typed content block

ChatDatabricks forwards an Anthropic cache_control breakpoint only when it is nested inside a typed content block. A message-level key or an additional_kwargs placement is dropped silently, with no error, so prompt caching never activates for those placements.

Steps to reproduce

Send the same large static system block three ways through the same ChatDatabricks client and check warm-call cache_read per placement:

message-level sibling key {"role","content":<str>,"cache_control":...} : cache_read = 0       -> DROPPED
SystemMessage(additional_kwargs={"cache_control":...})                 : cache_read = 0       -> DROPPED
typed content block [{"type":"text","text":...,"cache_control":...}]   : cache_read = 8803    -> FORWARDED

Root cause

_convert_message_to_dict (v0.17.0) builds the outbound message as {"content": message.content} and copies nothing else from additional_kwargs. So cache_control survives only when it is nested inside a content block that lives in message.content.

Fix

_convert_message_to_dict now forwards a message-level cache_control from additional_kwargs onto the outbound message dict. The forwarding fires only when cache_control is present, so the no-breakpoint path is unchanged.

if (cache_control := message.additional_kwargs.get("cache_control")) is not None:
    message_dict["cache_control"] = cache_control

Two unit tests cover the present and absent cases.

This is a distinct bug from #446 (token-count inflation when caching is active). This PR addresses the breakpoint being dropped so caching never activates.

_convert_message_to_dict built the outbound message from message.content
only, so an Anthropic cache_control breakpoint reached the endpoint solely
when nested inside a typed content block. A message-level breakpoint (stored
by LangChain in additional_kwargs) was dropped silently, so prompt caching
never activated for that placement.

Forward cache_control from additional_kwargs onto the outbound message dict.
Add unit tests for the present and absent cases.

Co-authored-by: Isaac
Signed-off-by: Adam Gurary <guraryadam@gmail.com>
@adamgurary
adamgurary force-pushed the cuj-chatdatabricks-cache-control branch from 670d455 to d9d5216 Compare July 30, 2026 23:32
@adamgurary

Copy link
Copy Markdown
Author

Hey @fanzeyi and @aravind-segu, this PR fixes a silent drop of message-level cache_control in _convert_message_to_dict. Two unit tests added. Would appreciate a review when you have a moment.

@adamgurary

Copy link
Copy Markdown
Author

One thing on my side is blocked: the test workflow is waiting on a maintainer to click "Approve and run workflows" (first PR I've opened here, so the runs are sitting at action_required). Once that runs, CI should be green. The two new unit tests cover the present and absent cache_control cases and pass locally with 108 passed / 2 skipped.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant