Forward cache_control from message-level placement in ChatDatabricks

[adamgurary]

Context

This issue is one finding from a hands-on evaluation (a Critical User Journey run) of building, tracing, and improving an agent that runs off Databricks. The agent is a LangGraph rebuild of a customer-notes workflow. It runs on a laptop and calls the Databricks Foundation Model API endpoint databricks-claude-opus-4-8. Its MLflow traces are written to Unity Catalog.

Full writeup:

• Summary: https://databricks.atlassian.net/wiki/spaces/~7120200d4c393f21a14d67bba42ffb3e92a69f/pages/6623789187
• Walkthrough: https://databricks.atlassian.net/wiki/spaces/~7120200d4c393f21a14d67bba42ffb3e92a69f/pages/6624084274

Environment

• MLflow 3.14.0, LangGraph 1.2.7, LangChain 1.3.11, databricks-langchain 0.17.0, databricks-sdk 0.94.0, openai 2.24.0, Python 3.12.
• Databricks workspace: logfood (adb-2548836972759138). SQL warehouse 79415a9aea7a3a7e.
• MLflow experiment: customer-notes-agent-uc (id 2824204872263693).
• UC trace location: catalog home_adam_gurary, schema personal_agents, table prefix customer_notes.
• Per run: roughly 1.6M input tokens across about 14 to 16 LLM calls.

ChatDatabricks drops cache_control unless it is inside a typed content block

Summary

The LangChain ChatDatabricks client (from databricks-langchain) forwards an Anthropic cache_control breakpoint only when it sits 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. ChatDatabricks CAN cache. The breakpoint just has to be in the one shape the client preserves, and nothing tells the user that.

Steps to reproduce

Run the attached repro (repro_cache_control_and_autolog.py, Bug 1 section) against a Databricks-served Claude endpoint. It sends the same large static system block three ways through the same ChatDatabricks client and reads the warm-call cache count each time.

Actual behavior

Warm-call cache_read by placement (from the repro):

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

Isolation proof that ChatDatabricks actually writes the cache, not just that the flag reaches the wire: a never-seen block sent ONLY through ChatDatabricks in the content-block form, then read once via the OpenAI client, reported cache_read_input_tokens > 0.

Root cause

databricks_langchain.chat_models._convert_message_to_dict (v0.17.0) builds the outbound message as {"content": message.content} and copies nothing else:

def _convert_message_to_dict(message: BaseMessage) -> dict:
    message_dict: Dict[str, Any] = {"content": message.content}
    ...
    elif isinstance(message, SystemMessage):
        return {"role": "system", **message_dict}

So cache_control survives only when it is nested inside a content block that lives in message.content. A message-level sibling key and additional_kwargs sit outside content and are dropped, with no warning.

The fix in this PR

_convert_message_to_dict now forwards a message-level cache_control from additional_kwargs onto the outbound message dict, so caching activates without nesting the flag inside a content block. 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. The full test_chat_models.py unit suite passes (108 passed, 2 skipped).

This is a distinct bug from issue #446 (ChatDatabricks: UsageMetadata inflates input_tokens when prompt caching is used), which concerns token-count accounting once caching is active. This PR is about the breakpoint being dropped so caching never activates for a message-level placement.

Evidence and references

Full repro attached in the CUJ writeup (mlflow-caching-repro-2026-07-23.zip): README Bug 1, repro-notes.md "Bug 1 mechanism", and the FORM matrix in expected-output.txt. Warm-call cache_read was 0 for the message-level and additional_kwargs placements and 8803 for the typed content block, through the same client. This is the corrected version of CUJ step 35, which first read this as "ChatDatabricks cannot cache at all." See the walkthrough linked above.