From b55f03740cb1e56ac6b8250f2fa81af2c032be59 Mon Sep 17 00:00:00 2001 From: yucheng Date: Thu, 29 Oct 2026 04:33:31 -0700 Subject: [PATCH] fix(guardrails): keep every response step listed on generateContent The generateContent adapter logs its nested completion, spend row included, before the response hooks run, so dropping never-ran steps afterwards left the header and the trace naming fewer steps than the spend row. The route now keeps main's list. The integration fixtures give each video its own id and accept any non-hit cache_hit value. --- litellm/proxy/utils.py | 5 +++-- .../observability/test_guardrail_pipeline_applied_header.py | 5 ++--- .../proxy/utils/proxy_logging/test_guardrail_pipeline.py | 3 +++ 3 files changed, 8 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 732c72c54ae..d80a3f4cbd6 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -744,11 +744,12 @@ def _steps_that_ran(result: PipelineExecutionResult) -> frozenset[str]: def _whole_response_log_waits_for_its_pipelines(data: Mapping[str, object], call_type: str) -> bool: """ The success log is parked until the response hooks finish, and the headers are built - after them, except on the LLM pass-through route. + after them. The LLM pass-through route builds its headers first, and the generateContent + adapter logs its nested completion before the response hooks run. """ logging_obj: Final = data.get("litellm_logging_obj") return ( - call_type != CallTypes.allm_passthrough_route.value + call_type not in (CallTypes.allm_passthrough_route.value, CallTypes.agenerate_content.value) and isinstance(logging_obj, Logging) and logging_obj.enqueue_deferred_logging is not None ) diff --git a/tests/integration/observability/test_guardrail_pipeline_applied_header.py b/tests/integration/observability/test_guardrail_pipeline_applied_header.py index b62e38a1a37..bebef167f92 100644 --- a/tests/integration/observability/test_guardrail_pipeline_applied_header.py +++ b/tests/integration/observability/test_guardrail_pipeline_applied_header.py @@ -90,7 +90,6 @@ _FINE_TUNING_JOB: Final[dict[str, JsonValue]] = { "validation_file": None, } _VIDEO: Final = { - "id": "video_pipeline", "object": "video", "model": "sora-2", "status": "queued", @@ -240,7 +239,7 @@ def _completion(request: Request) -> Reply: if request.target in ("/v1/images/generations", "/v1/images/edits"): return Reply(body=json.dumps({"created": 1, "data": [{"b64_json": "aW1n", "revised_prompt": None}]}).encode()) if request.target == "/v1/videos": - return Reply(body=json.dumps(_VIDEO).encode()) + return Reply(body=json.dumps({**_VIDEO, "id": f"video_{uuid.uuid4().hex}"}).encode()) if request.target == "/v1/fine_tuning/jobs": return Reply(body=json.dumps(_FINE_TUNING_JOB).encode()) if request.target.split("?", 1)[0].endswith(":generateContent"): @@ -823,7 +822,7 @@ def test_a_cache_hit_names_every_response_step_that_ran_in_the_spend_log_and_the expected: Final = [_RESPONSE_FIRST, _RESPONSE_SECOND] assert [first.status_code, cached.status_code] == [200, 200], (first.text, cached.text) assert _steps_that_ran(rig) == (_RESPONSE_FIRST, _RESPONSE_SECOND) * 2 - assert [row["cache_hit"] for row in rows] == ["False", "True"], rows + assert [row["cache_hit"] == "True" for row in rows] == [False, True], rows assert [sorted(response.headers.get(_HEADER, "").split(",")) for response in (first, cached)] == [expected] * 2 assert [sorted(object_value(row["metadata"]).get("applied_guardrails") or []) for row in rows] == [expected] * 2 assert [sorted(ast.literal_eval(str(span["attributes"]["metadata.applied_guardrails"]))) for span in spans] == [ diff --git a/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py index a0ba2723443..99a06f4e141 100644 --- a/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -1366,6 +1366,9 @@ async def test_pre_call_hook_names_a_step_that_ran_in_the_litellm_metadata_a_lat pytest.param( True, "allm_passthrough_route", "metadata", ["first", "second"], id="pass-through-headers-already-out" ), + pytest.param( + True, "agenerate_content", "metadata", ["first", "second"], id="generate-content-adapter-logged-its-call" + ), ], ) async def test_post_call_pipeline_drops_the_steps_that_did_not_run_only_while_nothing_has_read_them(