diff --git a/litellm/integrations/anthropic_cache_control_hook.py b/litellm/integrations/anthropic_cache_control_hook.py index 545b0f40018..3519240dda9 100644 --- a/litellm/integrations/anthropic_cache_control_hook.py +++ b/litellm/integrations/anthropic_cache_control_hook.py @@ -755,6 +755,7 @@ class AnthropicCacheControlHook(CustomPromptManagement): def record_gateway_injection( request_kwargs: Mapping[str, object], added: int, + injected_for_every_deployment: bool = False, ) -> None: """Name the deployment whose payload the gateway, not the client, put breakpoints on. @@ -771,7 +772,16 @@ class AnthropicCacheControlHook(CustomPromptManagement): A pass that runs before a deployment is chosen, which is what the proxy does for prompt templates, injects into the payload every leg goes on to send, so it marks - the request for all of them rather than for one. + the request for all of them rather than for one. Such a pass says so with + ``injected_for_every_deployment`` instead of relying on the shape of + ``request_kwargs``: the router's prompt-management factory stamps a provisional + deployment's ``model_info`` into kwargs before the prompt pass runs, and billing + the request through any other deployment would silently drop the credit. An + every-deployment mark, once written, also never narrows: a later per-leg stamp + (the Bedrock converse tool_config one included) describes one leg of a payload + every leg sends, so narrowing to it would uncredit whichever leg gets billed + after a failover. Both losses are fail-closed under-crediting, which is why the + guard only protects the sentinel and per-leg marks still overwrite each other. Only what this pass actually placed counts. A ``tool_config`` point is placed by the Bedrock converse transform, and only when the request carries tools, so the @@ -801,13 +811,19 @@ class AnthropicCacheControlHook(CustomPromptManagement): ), None, ) - if bucket is not None: - model_info: Final = request_kwargs.get("model_info") - bucket[GATEWAY_INJECTED_CACHE_METADATA_KEY] = ( - model_info.get("id", GATEWAY_INJECTED_FOR_EVERY_DEPLOYMENT) - if isinstance(model_info, dict) - else GATEWAY_INJECTED_FOR_EVERY_DEPLOYMENT - ) + if bucket is None: + return + if bucket.get(GATEWAY_INJECTED_CACHE_METADATA_KEY) == GATEWAY_INJECTED_FOR_EVERY_DEPLOYMENT: + return + if injected_for_every_deployment: + bucket[GATEWAY_INJECTED_CACHE_METADATA_KEY] = GATEWAY_INJECTED_FOR_EVERY_DEPLOYMENT + return + model_info: Final = request_kwargs.get("model_info") + bucket[GATEWAY_INJECTED_CACHE_METADATA_KEY] = ( + model_info.get("id", GATEWAY_INJECTED_FOR_EVERY_DEPLOYMENT) + if isinstance(model_info, dict) + else GATEWAY_INJECTED_FOR_EVERY_DEPLOYMENT + ) @staticmethod def maybe_inject_cache_control( diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 9a6fb11f978..f94e86b4460 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -901,6 +901,7 @@ class Logging(LiteLLMLoggingBaseClass): prompt_label: str | None = None, prompt_version: int | None = None, request_kwargs: dict[str, object] | None = None, # mutable-ok: marker stamped into live request kwargs + injected_for_every_deployment: bool = False, ) -> tuple[str, list[AllMessageValues], dict]: from litellm.integrations.anthropic_cache_control_hook import AnthropicCacheControlHook @@ -933,6 +934,7 @@ class Logging(LiteLLMLoggingBaseClass): AnthropicCacheControlHook.record_gateway_injection( request_kwargs, AnthropicCacheControlHook.count_request_cache_breakpoints(messages) - breakpoints_before, + injected_for_every_deployment=injected_for_every_deployment, ) self.messages = messages return model, messages, non_default_params @@ -950,6 +952,7 @@ class Logging(LiteLLMLoggingBaseClass): prompt_label: str | None = None, prompt_version: int | None = None, request_kwargs: dict[str, object] | None = None, # mutable-ok: marker stamped into live request kwargs + injected_for_every_deployment: bool = False, ) -> tuple[str, list[AllMessageValues], dict]: from litellm.integrations.anthropic_cache_control_hook import AnthropicCacheControlHook @@ -985,6 +988,7 @@ class Logging(LiteLLMLoggingBaseClass): AnthropicCacheControlHook.record_gateway_injection( request_kwargs, AnthropicCacheControlHook.count_request_cache_breakpoints(messages) - breakpoints_before, + injected_for_every_deployment=injected_for_every_deployment, ) self.messages = messages return model, messages, non_default_params diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 051d36c4d0f..cab2bd6d9db 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1524,6 +1524,7 @@ class ProxyLogging: prompt_label=data.pop("prompt_label", None) or {}, prompt_version=data.pop("prompt_version", None) or {}, request_kwargs=data, + injected_for_every_deployment=True, ) data.update(optional_params) diff --git a/litellm/router.py b/litellm/router.py index 6e4405ebfef..462d5414456 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -4006,6 +4006,7 @@ class Router: prompt_variables=prompt_variables, prompt_label=prompt_label, request_kwargs=kwargs, + injected_for_every_deployment=True, ) # Filter out prompt management specific parameters from data before merging diff --git a/tests/e2e/test_junit_properties.py b/tests/e2e/test_junit_properties.py index c0596177cc1..f7d1f70c5ec 100644 --- a/tests/e2e/test_junit_properties.py +++ b/tests/e2e/test_junit_properties.py @@ -115,7 +115,7 @@ class TestResultProperties: ("logging/test_x.py", 40, "TestFoo.test_bar"), (FakeMarker("covers", "LOG-1", "LOG-2"),), ) - assert result_properties(item) == ( + assert result_properties(item) == ( # pyright: ignore[reportArgumentType] # duck-typed Item double ("package", "logging"), ("covers", "LOG-1,LOG-2"), ("source", "tests/e2e/logging/test_x.py:41"), @@ -125,8 +125,8 @@ class TestResultProperties: """Collection can run the hook more than once; a second pass must not double the entries in the report.""" item = FakeItem("logging/test_x.py::test_bar", ("logging/test_x.py", 40, "test_bar")) - attach_result_properties(item) - attach_result_properties(item) + attach_result_properties(item) # pyright: ignore[reportArgumentType] # duck-typed Item double + attach_result_properties(item) # pyright: ignore[reportArgumentType] # duck-typed Item double assert [name for name, _ in item.user_properties] == ["package", "covers", "source"] diff --git a/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py b/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py index e995cbae782..de8b654987b 100644 --- a/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py +++ b/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py @@ -2858,6 +2858,27 @@ class TestRecordGatewayInjection: AnthropicCacheControlHook.record_gateway_injection(kwargs, 0) assert kwargs["litellm_metadata"][self.KEY] == self.DEPLOYMENT + def test_an_every_deployment_mark_survives_a_later_per_deployment_stamp(self): + """A per-leg stamp like the Bedrock converse tool_config one describes one leg of + a payload every leg sends, so narrowing an every-deployment mark to that leg's + deployment would uncredit whichever leg gets billed after a failover.""" + kwargs: dict = {"litellm_metadata": {self.KEY: ""}, "model_info": {"id": self.DEPLOYMENT}} + AnthropicCacheControlHook.record_gateway_injection(kwargs, 1) + assert kwargs["litellm_metadata"][self.KEY] == "" + + def test_a_pre_choice_pass_stamps_the_sentinel_over_a_provisional_deployment(self): + """The router's prompt-management factory stamps a provisional deployment's + model_info into kwargs before the prompt pass runs, and any other deployment can + end up billed, so the pass declares every-deployment scope explicitly.""" + kwargs: dict = {"litellm_metadata": {}, "model_info": {"id": self.DEPLOYMENT}} + AnthropicCacheControlHook.record_gateway_injection(kwargs, 1, injected_for_every_deployment=True) + assert kwargs["litellm_metadata"][self.KEY] == "" + + def test_a_per_deployment_mark_still_follows_the_latest_leg(self): + kwargs: dict = {"litellm_metadata": {self.KEY: "dep-old"}, "model_info": {"id": self.DEPLOYMENT}} + AnthropicCacheControlHook.record_gateway_injection(kwargs, 1) + assert kwargs["litellm_metadata"][self.KEY] == self.DEPLOYMENT + def test_v1_messages_auto_injection_stamps_the_marker(self, monkeypatch): monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) kwargs: dict = {"litellm_metadata": {}, "model_info": {"id": self.DEPLOYMENT}} diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index 366f61ded49..f1de7390b5b 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -6002,7 +6002,9 @@ async def test_prompt_hook_injection_marker_recorded_for_every_surface(logging_o """The savings gate reads litellm_gateway_injected_cache from the request's metadata bucket. Recording lives in the shared prompt-hook wrappers, so chat, /v1/responses, router prompt deployments, and proxy prompt templates all mark - injected requests the same way; a hook that injects nothing leaves no marker.""" + injected requests the same way; a hook that injects nothing leaves no marker. + A pass that runs before deployment choice declares it and gets the every-deployment + sentinel, which a later per-deployment pass never narrows.""" from litellm.integrations.custom_prompt_management import CustomPromptManagement class _InjectingHook(CustomPromptManagement): @@ -6086,6 +6088,28 @@ async def test_prompt_hook_injection_marker_recorded_for_every_surface(logging_o ) assert "litellm_gateway_injected_cache" not in untouched["metadata"] + pre_choice = {"metadata": {}, "model_info": {"id": "dep-of-this-attempt"}} + logging_obj.get_chat_completion_prompt( + model="claude-sonnet-5", + messages=[{"role": "user", "content": "hi"}], + non_default_params={}, + prompt_variables=None, + prompt_management_logger=_InjectingHook(), + request_kwargs=pre_choice, + injected_for_every_deployment=True, + ) + assert pre_choice["metadata"]["litellm_gateway_injected_cache"] == "" + + await logging_obj.async_get_chat_completion_prompt( + model="claude-sonnet-5", + messages=[{"role": "user", "content": "a fresh turn"}], + non_default_params={}, + prompt_variables=None, + prompt_management_logger=_InjectingHook(), + request_kwargs=pre_choice, + ) + assert pre_choice["metadata"]["litellm_gateway_injected_cache"] == "" + def test_get_standard_logging_object_payload_reads_overhead_from_logging_obj_for_dict_results(logging_obj): """LIT-5466: /v1/messages returns a plain dict with no _hidden_params, so the overhead diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 84f6344be35..058ed5bb3a7 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -11730,3 +11730,55 @@ class TestPreRoutingTierDrivesFallbacks: response = await router.acompletion(model="smart-router", messages=[{"role": "user", "content": "hi"}]) assert response.choices[0].message.content == "from backup-b" + + +@pytest.mark.asyncio +async def test_prompt_management_factory_marks_injection_for_every_deployment(monkeypatch): + """The factory stamps a provisional deployment's model_info into kwargs before the + prompt pass runs, then routes on the returned model, so any deployment can end up + billed. An injection recorded there must carry the every-deployment sentinel, never + the provisional deployment's id, or a differently-billed deployment loses the credit.""" + import time + + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging + + router = litellm.Router( + model_list=[ + { + "model_name": "cached-claude", + "litellm_params": { + "model": "anthropic_cache_control_hook/claude-sonnet-5", + "prompt_id": "cache-points", + }, + "model_info": {"id": "provisional-dep"}, + } + ] + ) + captured: dict = {} + + async def _capture_acompletion(**kwargs): + captured.update(kwargs) + return litellm.ModelResponse() + + monkeypatch.setattr(litellm, "acompletion", _capture_acompletion) + logging_obj = LiteLLMLogging( + model="cached-claude", + messages=[{"role": "user", "content": "hi"}], + stream=False, + call_type="acompletion", + start_time=time.time(), + litellm_call_id="lit-6445", + function_id="f", + ) + await router.acompletion( + model="cached-claude", + messages=[ + {"role": "system", "content": "a static system prompt"}, + {"role": "user", "content": "hi"}, + ], + cache_control_injection_points=[{"location": "message", "role": "system"}], + litellm_logging_obj=logging_obj, + ) + bucket = captured.get("litellm_metadata") or captured["metadata"] + assert captured["model_info"]["id"] == "provisional-dep" + assert bucket["litellm_gateway_injected_cache"] == ""