fix(vertex_ai): Bake tool_choice into Gemini CachedContent body to prevent silent drop

[adriangomez24]

Relevant issues

Fixes:

• [Bug]: tool_choice silently dropped on cached Gemini requests — affects both gemini/ and vertex_ai/ providers #29088

This PR is inspired by this PR which fixes the issue but has merge conflicts: #25659

Linear ticket

Pre-Submission checklist

Please complete all items before asking a LiteLLM maintainer to review your PR

• I have Added testing in the tests/test_litellm/ directory, Adding at least 1 test is a hard requirement - see details
• My PR passes all unit tests on make test-unit
• My PR's scope is as isolated as possible, it only solves 1 specific problem
• I have requested a Greptile review by commenting @greptileai and received a Confidence Score of at least 4/5 before requesting a maintainer review

Delays in PR merge?

If you're seeing a delay in your PR being merged, ping the LiteLLM Team on Slack (#pr-review).

CI (LiteLLM team)

CI status guideline:

• 50-55 passing tests: main is stable with minor issues.
• 45-49 passing tests: acceptable but needs attention
• <= 40 passing tests: unstable; be careful with your merges and assess the risk.

• Branch creation CI run
  Link:

• CI run for the last commit
  Link:

• Merge / cherry-pick CI run
  Links:

Screenshots / Proof of Fix

Screen.Recording.2026-05-28.at.12.46.12.PM.mov

Screen.Recording.2026-05-28.at.12.48.30.PM.mov

Save as repro.py, then run with GEMINI_API_KEY=... python repro.py:

import litellm
litellm.modify_params = True   # engage the #26077 guard
MODEL = "gemini/gemini-2.5-flash"
LONG_SYSTEM = (
    "You are a deterministic test fixture. Always respond with one short sentence. "
    "Never volunteer tool calls unless explicitly asked. " * 200
)
TOOL = {
    "type": "function",
    "function": {
        "name": "log_test_event",
        "description": "Internal logging tool.",
        "parameters": {
            "type": "object",
            "properties": {"event_name": {"type": "string"}},
            "required": ["event_name"],
        },
    },
}
def run(label, cache_control, tool_choice):
    sys_block = {"type": "text", "text": LONG_SYSTEM}
    if cache_control:
        sys_block["cache_control"] = {"type": "ephemeral"}
    kwargs = {
        "model": MODEL,
        "messages": [
            {"role": "system", "content": [sys_block]},
            {"role": "user", "content": "Hi, how are you?"},
        ],
        "tools": [TOOL],
    }
    if tool_choice:
        kwargs["tool_choice"] = tool_choice
    resp = litellm.completion(**kwargs)
    msg = resp.choices[0].message
    print(f"{label}: tool_calls={len(msg.tool_calls or [])}, content={(msg.content or '')[:60]!r}")
run("C (no cache, required)", cache_control=False, tool_choice="required")
run("B (cache, no choice)",    cache_control=True,  tool_choice=None)
run("A (cache, required)",     cache_control=True,  tool_choice="required")

Before this fix:

C (no cache, required): tool_calls=1, content=''
B (cache, no choice):   tool_calls=0, content='I am functioning correctly.'
A (cache, required):    tool_calls=0, content='I am functioning correctly.'   ← BUG: tool_choice dropped

A vs C is the smoking gun: identical tool_choice="required", but adding the cache_control marker silently drops the value. A and B return identical content, proving the request reaching Gemini was the same in both cases — tool_choice never made it on the wire.
After this fix:

C (no cache, required): tool_calls=1, content=''
B (cache, no choice):   tool_calls=0, content='I am functioning correctly.'
A (cache, required):    tool_calls=1, content=''                              ← FIX: tool call now made via log_test_event

usage.cached_tokens remains non-zero on B and A, confirming the cache still engages — the fix preserves all caching benefits while restoring tool_choice semantics on cached requests.

Type

🐛 Bug Fix

Changes

LiteLLM's Gemini context-caching path correctly pops tools from optional_params and bakes them into the CachedContent body but does the same for tool_choice for neither the sync nor async path. The value stays in optional_params, gets translated to toolConfig by _transform_request_body, and either:

1. Pre-1.84.0 / modify_params=False: ends up on the follow-up generateContent call alongside cachedContent → Vertex 400 INVALID_ARGUMENT "Tool config, tools and system instruction should not be set in the request when using cached content."
2. 1.84.0+ with modify_params=True (fix(vertex_ai): omit system_instruction/tools/toolConfig when cachedContent set #26077): defensively stripped from the generate body, but never written to the cache either. The net effect is that tool_choice is silently dropped on every cached request. Gemini receives no function_calling_config and falls back to its implicit default regardless of what the caller asked for.

This PR closes that asymmetry by treating tool_choice symmetrically with tools at cache-creation time:

• Pop tool_choice from optional_params in both check_and_create_cache and async_check_and_create_cache.
• Include tool_choice in the local cache key so two requests with different tool_choice values (e.g. "auto" vs "required") produce distinct cache entries — without this, the second request silently reuses the first cache's toolConfig.
• Bake tool_choice into the CachedContent body as toolConfig so cache hits inherit the caller's tool_choice semantics. The fix(vertex_ai): omit system_instruction/tools/toolConfig when cachedContent set #26077 generate-side guard then continues to work as designed (omits the field from the follow-up generate body because it's already in the cache).

Applied to both the sync and async code paths.

For testing, in addition to the above script I added 8 new mocked tests in
-tests/test_litellm/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py, all parametrized across gemini, vertex_ai, and vertex_ai_beta providers:

• test_check_and_create_cache_tool_choice_popped_from_optional_params — sync, with cached messages: tool_choice is popped.
• test_check_and_create_cache_tool_choice_not_popped_when_no_cached_messages — sync, no cached messages: tool_choice is not popped (early-return).
• test_async_check_and_create_cache_tool_choice_popped_from_optional_params — async equivalent of the first.
• test_async_check_and_create_cache_tool_choice_not_popped_when_no_cached_messages — async equivalent of the second.
• test_check_and_create_cache_tool_choice_in_request_body — sync end-to-end: toolConfig appears in the cache-creation HTTP POST body and the cache key includes tool_choice.
• test_async_check_and_create_cache_tool_choice_in_request_body — async equivalent.
• test_check_and_create_cache_omits_tool_config_when_tool_choice_unset — when caller doesn't pass tool_choice, no toolConfig is written to the cache body.
• test_check_and_create_cache_tool_choice_function_pin — tool_choice as a function-pin dict ({"type": "function", "function": {"name": "..."}}) survives the cache body intact.
• test_check_and_create_cache_distinct_tool_choices_use_distinct_keys — two requests with different tool_choice values produce different cache keys (regression guard for the cache-key-collision scenario flagged in fix(gemini): pop tool_choice from optional_params when creating cached content #25659's review).
  Also updated two existing mock_cache_obj.get_cache_key.assert_called_once_with assertions to include tool_choice=None in the expected signature (the local cache key now always has it as a positional named argument, even when the caller didn't pass one).

Adding tool_choice to local_cache_obj.get_cache_key(...) changes the hash for every cached entry — pre-fix entries will miss once on the first request after deploy, then resume normal cache-read behavior. This is intentional:

• Pre-fix CachedContent objects were created without toolConfig baked in. Reusing them post-fix would silently reproduce the bug this PR is fixing.
• Cost impact is bounded: one cache-creation call per unique (messages, tools, tool_choice, model) tuple, the first time it's seen after deploy. No recurring overhead.
• Stale pre-fix entries expire on their own via Vertex's CachedContent TTL (default 5 min). No cleanup pass needed.

In short: a one-time re-cache on deploy is the correct migration. Reusing pre-fix entries would perpetuate the bug.

[CLAassistant]

Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you all sign our Contributor License Agreement before we can accept your contribution.
2 out of 3 committers have signed the CLA.

✅ adriangomez24
✅ yuneng-berri
❌ shin-berri
_{You have signed the CLA already but the status is still pending? Let us recheck it.}

[adriangomez24]

@greptileai

[greptile-apps]

Greptile Summary

This PR fixes a silent tool_choice drop that occurred when Vertex AI / Gemini context caching was active. Previously, tool_choice sat in optional_params and was quietly excluded from both the CachedContent creation body and the downstream generate request body (because Vertex rejects toolConfig alongside cachedContent). The fix mirrors the existing tools treatment: pop tool_choice from optional_params, bake it into the CachedContent request as toolConfig, and incorporate it into the local cache key so that distinct tool_choice values produce distinct cache entries.

• Both the synchronous check_and_create_cache and async async_check_and_create_cache paths receive the same change, keeping them in sync.
• Nine new mock-only tests cover the pop behaviour, early-return preservation, end-to-end body injection, absence when unset, function-pin dict pass-through, and cache-key differentiation across tool_choice values.

Confidence Score: 4/5

The production fix is safe to merge; the only gap is that the new end-to-end tests mock the wrong input type for tool_choice.

The implementation correctly pops tool_choice after the early-return guards and assigns the already-converted Gemini ToolConfig object to toolConfig in the CachedContent request body — consistent with how tools is handled and with the existing transformation.py path. Both sync and async paths are updated identically. The new tests are thorough in coverage but inject raw OpenAI strings ("required") as tool_choice, whereas by the time context caching runs in production map_openai_params has already converted those strings to Gemini ToolConfig dicts. The assertion toolConfig == "required" would never be true in a real request, so that specific path isn't exercised with production-realistic data.

The test file's end-to-end body-injection tests use OpenAI-format tool_choice strings instead of the Gemini ToolConfig dicts that actually arrive at check_and_create_cache in production.

Important Files Changed

Filename	Overview
litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py	Adds tool_choice extraction and inclusion as toolConfig in the CachedContent creation body, mirroring the existing tools handling, for both sync and async paths. The change is internally consistent and correctly uses the Gemini-format ToolConfig object that map_openai_params has already placed in optional_params.
tests/test_litellm/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py	Adds nine new tests (sync and async) covering tool_choice pop, early-return preservation, end-to-end body injection, absence when unset, function-pin dicts, and distinct cache keys. Updates two existing get_cache_key assertion call sites to include tool_choice=None. Tests use raw OpenAI strings ("required") rather than the Gemini ToolConfig dicts that would actually arrive at runtime.

_{Reviews (1): Last reviewed commit: "fix(vertex_ai): bake tool_choice into Ge..." | Re-trigger Greptile}

[greptile-apps]

The tests inject "required" (a raw OpenAI string) directly into optional_params["tool_choice"], but at the point check_and_create_cache runs in production, map_openai_params has already converted that string to a Gemini ToolConfig dict such as {"functionCallingConfig": {"mode": "ANY"}}. Because the code just passes the value through unchanged, the assertions like toolConfig == "required" will never be true in production. Using the actual Gemini-format value in the test would catch a regression where the wrong format reaches the API, and would also serve as clearer documentation of what the function really receives.

Suggested change

	optional_params = self.sample_optional_params.copy()
	optional_params["tool_choice"] = "required"
	self.context_caching.check_and_create_cache(
	messages=self.sample_messages,
	optional_params=optional_params,
	api_key="test_key",
	api_base=None,
	model="gemini-1.5-pro",
	client=self.mock_client,
	timeout=30.0,
	logging_obj=self.mock_logging,
	custom_llm_provider=custom_llm_provider,
	vertex_project="test_project",
	vertex_location="test_location",
	vertex_auth_header="vertext_test_token",
	)
	self.mock_client.post.assert_called_once()
	call_args = self.mock_client.post.call_args
	assert call_args.kwargs["json"]["tools"] == self.sample_tools
	assert call_args.kwargs["json"]["toolConfig"] == "required"
	mock_cache_obj.get_cache_key.assert_called_once_with(
	messages=cached_messages,
	tools=self.sample_tools,
	tool_choice="required",
	model="gemini-1.5-pro",
	)
	# In production optional_params["tool_choice"] is a Gemini ToolConfig
	# dict, not a raw OpenAI string, because map_openai_params calls
	# map_tool_choice_values before sync_transform_request_body is reached.
	gemini_tool_choice = {"functionCallingConfig": {"mode": "ANY"}}
	optional_params = self.sample_optional_params.copy()
	optional_params["tool_choice"] = gemini_tool_choice
	self.context_caching.check_and_create_cache(
	messages=self.sample_messages,
	optional_params=optional_params,
	api_key="test_key",
	api_base=None,
	model="gemini-1.5-pro",
	client=self.mock_client,
	timeout=30.0,
	logging_obj=self.mock_logging,
	custom_llm_provider=custom_llm_provider,
	vertex_project="test_project",
	vertex_location="test_location",
	vertex_auth_header="vertext_test_token",
	)
	self.mock_client.post.assert_called_once()
	call_args = self.mock_client.post.call_args
	assert call_args.kwargs["json"]["tools"] == self.sample_tools
	assert call_args.kwargs["json"]["toolConfig"] == gemini_tool_choice
	mock_cache_obj.get_cache_key.assert_called_once_with(
	messages=cached_messages,
	tools=self.sample_tools,
	tool_choice=gemini_tool_choice,
	model="gemini-1.5-pro",
	)

[adriangomez24]

Fixed in following commit.

[greptile-apps]

Greptile Summary

This PR fixes a silent data-loss bug where tool_choice passed by callers was dropped when Vertex AI context caching was active, because only tools was baked into the CachedContent request body while tool_choice was left in optional_params until the main request transform silently discarded it (since the Gemini transform drops toolConfig when cachedContent is present).

• tool_choice is now popped from optional_params alongside tools in both the sync and async check_and_create_cache paths, and is included in the Gemini toolConfig field of the cache-creation POST body when non-None.
• The cache key generation is extended with the new tool_choice kwarg so that requests that differ only in their tool-calling mode produce separate cache entries.
• Eight new unit tests cover the pop-behavior, request body plumbing, early-return preservation, and key parameterisation for all three providers (gemini, vertex_ai, vertex_ai_beta).

Confidence Score: 4/5

Safe to merge; the change is narrowly scoped to the context-caching path and preserves existing behaviour when tool_choice is absent.

The production fix is correct and symmetric across sync/async paths. The only findings are in the tests: they seed tool_choice with pre-translation OpenAI strings rather than the Gemini-translated ToolConfig dicts that actually flow through this code in production, and one test asserts forwarding of inputs to get_cache_key rather than distinctness of the resulting hash. Neither issue affects runtime correctness.

Test file — the new test cases use OpenAI-format tool_choice values rather than the Gemini-translated format that reaches this code in real requests.

Important Files Changed

Filename	Overview
litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py	Pops tool_choice from optional_params and includes it in both the cache key and the CachedContent request body (toolConfig); logic is symmetric across sync and async paths.
tests/test_litellm/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py	Adds 8 new test cases covering pop-behavior, request body inclusion, and cache-key parameterisation; existing assertion updated to pass tool_choice=None.

_{Reviews (2): Last reviewed commit: "fix(vertex_ai): bake tool_choice into Ge..." | Re-trigger Greptile}

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[greptile-apps]

Test uses OpenAI-format values, not Gemini-translated ones

By the time check_and_create_cache is called in production, tool_choice has already been translated by VertexGeminiConfig.map_tool_choice_values into a Gemini ToolConfig dict (e.g. "required" → {"functionCallingConfig": {"mode": "ANY"}}). The tests inject raw OpenAI strings like "required" and the function-pin dict in OpenAI format ({"type": "function", "function": {"name": "…"}}). As a result the assertions assert call_args.kwargs["json"]["toolConfig"] == "required" and == function_pin pass, but they only prove the plumbing routes a value through unchanged — they don't validate that the Gemini-translated value is accepted by the CachedContent API. Consider adding at least one test that seeds optional_params["tool_choice"] with the already-translated Gemini format to match the real production code path.

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

[adriangomez24]

Fixed in following commit

[greptile-apps]

Test name claims distinct keys but only verifies distinct inputs

test_check_and_create_cache_distinct_tool_choices_use_distinct_keys mocks get_cache_key to always return "test_cache_key" regardless of inputs. The assertions confirm that get_cache_key was called with "auto" and "required" respectively, but they never verify that the real cache-key implementation produces two different hash strings. The test name implies end-to-end key isolation, which it doesn't actually prove. Either rename the test to reflect what it truly covers (that the correct tool_choice argument is forwarded to get_cache_key), or remove the mock and let the real Cache.get_cache_key run to confirm distinct outputs.

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

[adriangomez24]

Fixed in following commit

[hclsys]

looks right — both paths covered, and the raw assign to toolConfig is fine because map_openai_params already ran map_tool_choice_values and wrote the gemini ToolConfig shape back to optional_params["tool_choice"], so it's mapped by the time you pop it here (not the openai "required" string). test value matches that.

one thing worth a line in the PR: adding tool_choice to get_cache_key changes the key, so existing cached entries miss once. that's correct (a different tool_choice shouldn't reuse a cache) but it does mean a one-time re-cache on deploy — intended, right?

[adriangomez24]

@hclsys Thanks for the review! Just updated the PR comment to address your comment!

[krrish-berri-2]

@adriangomez24 — could you add a screenshot or short video showing that this change works as expected? It really helps reviewers verify the fix quickly. Thanks!

[adriangomez24]

Hey @krrish-berri-2 thanks for taking the time to add feedback! Added two videos running the script in the proof of fix test running real gemini 2.5 calls.

The first video is pinned to litellm 1.84.0 from PyPI. Condition A (cache_control: ephemeral. + tool_choice="required") returns tool_calls = 0 and a text response showing tool_choice is silently dropped on the cached path even though the caller forced it.

In the second video, it's based off the changes in this PR, and it shows that a tool use is actually triggered.

Both videos confirm that cache was still engaged from the usage.cached_tokens log.

[hclsys]

thanks for crediting the RCA, and good — both sync + async check_and_create_cache are covered, which was my main worry on #29088.

checked the one thing that looked risky: whether tool_choice reaches here in raw OpenAI shape ("auto", {"type":"function",...}) — it doesn't. map_tool_choice_values already converts it to the Gemini ToolConfig and writes it back to optional_params["tool_choice"] during param mapping (vertex_and_google_ai_studio_gemini.py:1243-1250), so by the time you pop it, cached_content_request_body["toolConfig"] = tool_choice is the right shape. correct.

one minor test gap: map_tool_choice_values returns a ToolConfig pydantic model, but the new tests set optional_params["tool_choice"] = {"functionCallingConfig": {...}} (plain dict). worth one case that pops the actual mapped ToolConfig object so the cached body exercises the same serialization the non-cached path does — otherwise a model-vs-dict mismatch in the request body wouldn't be caught.

[adriangomez24]

Hey @hclsys, thanks for the feedback! I added a test covers this case in the last commit.