fix(google-genai): route streaming chunks to GeminiPassthroughLogging…

[awais786]

…Handler so success_callbacks fire

streaming_iterator was tagging Google GenAI chunks as VERTEX_AI, causing
_route_streaming_logging_to_handler to use the wrong parser. Result was
None so async_complete_streaming_response was never set and all success
callbacks were silently skipped for /models/{model}:streamGenerateContent.

• Add EndpointType.GOOGLE_GENAI enum value
• Fix streaming_iterator to use GOOGLE_GENAI instead of VERTEX_AI
• Add GOOGLE_GENAI branch in streaming_handler routing to GeminiPassthroughLoggingHandler
• Add regression tests (7 tests, all passing)

Fixes #24097

Relevant issues

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:

Type

🆕 New Feature
🐛 Bug Fix
🧹 Refactoring
📖 Documentation
🚄 Infrastructure
✅ Test

Changes

[vercel]

The latest updates on your projects. Learn more about Vercel for GitHub.

Project	Deployment	Actions	Updated (UTC)
litellm	Ready	Preview, Comment	Mar 19, 2026 11:52am

[codspeed-hq]

Merging this PR will not alter performance

✅ 16 untouched benchmarks

---

_{Comparing awais786:fix/google-genai-success-callbacks (de5c810) with main (e5baa22)}

[greptile-apps]

Greptile Summary

This PR fixes a silent callback regression for /models/{model}:streamGenerateContent (Google GenAI pass-through streaming). The root cause was that BaseGoogleGenAIGenerateContentStreamingIterator was tagging collected chunks as EndpointType.VERTEX_AI, which caused _route_streaming_logging_to_handler to fall through with no matching branch — leaving async_complete_streaming_response unset and all success_callbacks silently skipped.

Key changes:

• Adds EndpointType.GOOGLE_GENAI = "google-genai" to the enum.
• Fixes streaming_iterator.py to pass endpoint_type=EndpointType.GOOGLE_GENAI, the correct url_route, and model=self.model explicitly.
• Adds a GOOGLE_GENAI routing branch in _route_streaming_logging_to_handler that dispatches to GeminiPassthroughLoggingHandler._handle_logging_gemini_collected_chunks.
• Correctly avoids pre-setting async_complete_streaming_response before calling async_success_handler — doing so would have triggered the early-return guard in litellm_logging.py:2492 and silently skipped every callback loop anyway.
• Adds a well-designed integration test using a real LiteLLMLoggingObj and a spy CustomLogger to verify the end-to-end callback dispatch path.

Remaining gaps:

• The synchronous GoogleGenAIGenerateContentStreamingIterator.__next__ still never calls _handle_async_streaming_logging on StopIteration, so sync callers of the Google GenAI streaming API remain silently broken after this fix.
• chunk_processor does not forward model when creating the logging task, leaving a latent fallback-to-URL-parsing gap if GOOGLE_GENAI is ever routed through that path.
• Several test-quality issues (fragile inspect.getsource assertions, sys.path manipulation, raw_bytes type mismatch, dead-code assertion messages) were noted in previous review threads and have not all been addressed.

Confidence Score: 3/5

• The async streaming fix is logically sound, but the sync iterator path remains broken and test-quality issues from previous review rounds are not fully addressed.
• The core async fix is correct — the new GOOGLE_GENAI enum and routing branch work together properly, and the avoidance of pre-setting async_complete_streaming_response is the right call. However, the synchronous GoogleGenAIGenerateContentStreamingIterator.next still never triggers logging callbacks (a real functional gap for sync callers), and chunk_processor does not pass model for GOOGLE_GENAI endpoint types. Combined with multiple unresolved test-quality issues flagged in prior review threads, this needs at least one more iteration before merging.
• Pay close attention to litellm/google_genai/streaming_iterator.py (sync iterator logging gap) and litellm/proxy/pass_through_endpoints/streaming_handler.py (chunk_processor missing model parameter).

Important Files Changed

Filename	Overview
litellm/google_genai/streaming_iterator.py	Correctly switches endpoint_type to GOOGLE_GENAI and url_route to the real streaming path; now also passes model=self.model explicitly. However, the sync iterator (GoogleGenAIGenerateContentStreamingIterator.next) still never calls _handle_async_streaming_logging on StopIteration, meaning synchronous callers remain silently broken.
litellm/proxy/pass_through_endpoints/streaming_handler.py	Adds GOOGLE_GENAI branch routing to GeminiPassthroughLoggingHandler and correctly avoids pre-setting async_complete_streaming_response (which would have triggered an early-return guard). chunk_processor does not forward model for the new GOOGLE_GENAI endpoint type, leaving a latent fallback-to-URL-parsing gap.
litellm/types/passthrough_endpoints/pass_through_endpoints.py	Minimal, correct addition of GOOGLE_GENAI = "google-genai" to the EndpointType enum; no issues.
litellm/litellm_core_utils/litellm_logging.py	Formatting-only changes (parenthesis removal, line wrapping); no logic modifications.
litellm/proxy/management_endpoints/team_endpoints.py	Formatting-only change (wrapping a long type annotation); no logic modifications.
tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_google_genai_success_callbacks.py	The integration test (TestStreamingHandlerGoogleGenAIRouting) is well-designed — it uses a real LiteLLMLoggingObj and spy logger to verify that callbacks actually fire through the full async_success_handler path. Several test-quality issues have been flagged in previous review threads (dead-code assertion messages, fragile sys.path insertion, raw_bytes type mismatch, inspect-based source tests).

Sequence Diagram

sequenceDiagram
    participant SI as streaming_iterator.py<br/>(Async/SyncIterator)
    participant SH as streaming_handler.py<br/>(_route_streaming_logging_to_handler)
    participant GH as GeminiPassthroughLoggingHandler
    participant LL as litellm_logging.py<br/>(async_success_handler)
    participant CB as success_callbacks

    Note over SI: StopAsyncIteration raised
    SI->>SH: _route_streaming_logging_to_handler(<br/>endpoint_type=GOOGLE_GENAI,<br/>url_route=/models/{model}:streamGenerateContent,<br/>model=self.model)
    SH->>GH: _handle_logging_gemini_collected_chunks(...)
    GH-->>SH: {result: ModelResponse, kwargs: {response_cost, model}}
    Note over SH: Does NOT pre-set<br/>async_complete_streaming_response
    SH->>LL: async_success_handler(result, **kwargs)
    Note over LL: Guard at line 2492:<br/>async_complete_streaming_response<br/>NOT in model_call_details → continues
    Note over LL: call_type == "pass_through_endpoint"<br/>→ sets async_complete_streaming_response
    LL->>CB: async_log_success_event(kwargs, response_obj, ...)
    Note over SI,CB: ❌ Sync iterator (GoogleGenAIGenerateContentStreamingIterator)<br/>never reaches this flow — StopIteration raised without logging

Loading

Comments Outside Diff (2)

1. litellm/google_genai/streaming_iterator.py, line 96-104 (link)

   Sync iterator never fires logging callbacks

   GoogleGenAIGenerateContentStreamingIterator.__next__ raises StopIteration without calling _handle_async_streaming_logging, so any synchronous caller of this iterator will silently skip all success callbacks — the same root bug that this PR fixes for the async path still exists in the sync path.

   The only fix applied is in AsyncGoogleGenAIGenerateContentStreamingIterator.__anext__ (line 158), which correctly awaits _handle_async_streaming_logging() on StopAsyncIteration. The sync __next__ has no equivalent.

   Since _handle_async_streaming_logging is async, it cannot be await-ed directly from __next__. A possible approach is to schedule logging via asyncio.get_event_loop().run_until_complete(...) if a running loop is available, or introduce a sync logging helper. At a minimum, this should be documented as a known limitation of the sync iterator.

       def __next__(self):
           try:
               chunk = next(self.stream_iterator)
               self.collected_chunks.append(chunk)
               return chunk
           except StopIteration:
               # TODO: _handle_async_streaming_logging is async and cannot be
               # awaited here — sync callers do NOT fire success callbacks.
               raise StopIteration

2. litellm/proxy/pass_through_endpoints/streaming_handler.py, line 83-94 (link)

   chunk_processor never passes model for GOOGLE_GENAI

   chunk_processor creates the logging task without forwarding model to _route_streaming_logging_to_handler. For the new GOOGLE_GENAI branch inside _route_streaming_logging_to_handler, the handler falls back to extract_model_from_url(url_route).

   If endpoint_type == EndpointType.GOOGLE_GENAI is ever routed through chunk_processor (rather than the streaming_iterator.py path), model extraction silently falls back to URL parsing and will return "unknown" for any URL that doesn't contain /models/{model}. This could silently compute an incorrect cost.

   The VERTEX_AI branch inside _route_streaming_logging_to_handler already receives model because VertexPassthroughLoggingHandler._handle_logging_vertex_collected_chunks accepts it; parity should be maintained for the new GOOGLE_GENAI branch.

_{Last reviewed commit: "fix(google-genai): r..."}

[greptile-apps]

Custom assertion messages are dead code

The pattern mock_gemini.assert_called_once(), ("message string") creates a tuple expression — it does not attach the string as the failure message to the assertion.

If assert_called_once() fails, it raises AssertionError with the default mock message (e.g., "Expected 'mock' to be called once. Called 0 times."), not the custom string. The string after the comma is evaluated and discarded.

The same bug appears on line 168 (assert_not_called()).

To attach a custom failure message, use pytest's assert statement:

Suggested change

	mock_gemini.assert_called_once(), (
	"GeminiPassthroughLoggingHandler was NOT called for GOOGLE_GENAI endpoint — "
	"callbacks would be silently skipped (issue #24097 not fixed)"
	)
	assert mock_gemini.call_count == 1, (
	"GeminiPassthroughLoggingHandler was NOT called for GOOGLE_GENAI endpoint — "
	"callbacks would be silently skipped (issue #24097 not fixed)"
	)

[greptile-apps]

Developer local machine path should be removed

The docstring contains a personal filesystem path and environment activation command specific to the author's machine. This will be confusing and misleading to other contributors.

Suggested change

	cd /Users/awais.qureshi/Documents/devstack/forks/litellm
	pyenv activate wagtail-chat-env
	python -m pytest tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_google_genai_success_callbacks.py -v
	Run:
	python -m pytest tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_google_genai_success_callbacks.py -v

[greptile-apps]

raw_bytes type mismatch

_route_streaming_logging_to_handler expects raw_bytes: List[bytes], but b"data: {}\n" is a plain bytes object, not a list. This doesn't cause a test failure here because _convert_raw_bytes_to_str_lines is mocked, but it misrepresents the correct call signature and could confuse future readers.

Suggested change

	raw_bytes=b"data: {}\n",
	raw_bytes=[b"data: {}\n"],

[greptile-apps]

Same raw_bytes type mismatch

Same issue as line 121 in the other test — b"data: {}\n" is bytes, not List[bytes].

Suggested change

	model="gemini-1.5-pro",
	raw_bytes=[b"data: {}\n"],

[greptile-apps]

Same dead-code assertion message pattern

mock_gemini.assert_not_called(), ("message") has the same issue as line 125-128 — the custom message string is unreachable. Use:

Suggested change

	mock_gemini.assert_not_called(), (
	"GeminiPassthroughLoggingHandler was called for VERTEX_AI endpoint — "
	"routing regression detected"
	)
	assert mock_gemini.call_count == 0, (
	"GeminiPassthroughLoggingHandler was called for VERTEX_AI endpoint — "
	"routing regression detected"
	)

[greptile-apps]

async_complete_streaming_response now fires for ALL passthrough endpoint types

The placement of this new block — after the if/elif chain — means async_complete_streaming_response is now set for every endpoint type: Anthropic, Vertex AI, OpenAI, and Google GenAI. Before this PR, it was never set, so function-based success_callbacks were silently skipped for all passthrough streaming endpoints. Now they will suddenly fire for Anthropic, Vertex AI, and OpenAI passthrough callers that previously never received them.

This is a broader behavioral change than the PR description suggests and could surprise existing users of those other passthrough endpoints. Per the project's backwards-compatibility policy, unintended behaviour changes should be gated behind a feature flag unless they are clearly a universally correct fix.

If this is intentional (fixing the same silent-skip bug for all providers), it should be called out explicitly in the PR description and ideally covered by tests for each of the other endpoint types. If it's only meant to fix Google GenAI, the block should be placed inside the elif endpoint_type == EndpointType.GOOGLE_GENAI: branch.

Rule Used: What: avoid backwards-incompatible changes without... (source)

[greptile-apps]

Hardcoded url_route won't match streamGenerateContent logging checks

The url_route is hardcoded as "/v1/generateContent", but the actual endpoint being called for streaming is /models/{model}:streamGenerateContent. This discrepancy flows into _build_complete_streaming_response, which checks:

if "generateContent" in url_route or "streamGenerateContent" in url_route:

Because "generateContent" is a substring of "streamGenerateContent", the current hardcoded value happens to pass the check. However, it also means extract_model_from_url("/v1/generateContent") would return "unknown" if the model kwarg were ever missing — the URL lacks the /models/{model} segment that the regex expects. The bug is latent today because model=self.model is passed explicitly, but the hardcoded URL does not reflect reality and could silently break model-extraction if the call signature changes.

Consider either passing the real URL or using a constant that includes the /models/{model} segment to make the intent clear.

[greptile-apps]

Fragile sys.path insertion uses a working-directory-relative path

os.path.abspath("../../..") is resolved relative to the current working directory at the time the test runs, not relative to the test file's location. When pytest is run from the repository root (the typical CI invocation), this expands to three levels above the repo root — an entirely wrong path that adds nothing useful.

If a path insertion is genuinely needed, use the file's own location as the anchor:

Suggested change

	sys.path.insert(0, os.path.abspath("../../.."))
	sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../..")))

That said, the project's standard test infrastructure (pytest.ini / pyproject.toml with pythonpath or an installed package) should already make litellm importable without manual sys.path manipulation. Removing this line entirely may be the cleanest fix.

[greptile-apps]

Source-inspection tests are fragile

inspect.getsource parses the raw source text of the class and checks for string literals. This means the assertion:

assert "EndpointType.VERTEX_AI" not in src

would falsely fail if anyone adds a docstring or comment to BaseGoogleGenAIGenerateContentStreamingIterator that mentions the old type — e.g. # Previously tagged as EndpointType.VERTEX_AI. Comments and docstrings are included in getsource() output.

A more robust approach is to directly test the runtime behaviour rather than parsing source text:

def test_uses_google_genai_not_vertex_ai(self):
    """streaming_iterator must tag chunks as GOOGLE_GENAI, not VERTEX_AI."""
    from litellm.types.passthrough_endpoints.pass_through_endpoints import EndpointType
    from unittest.mock import MagicMock, patch, AsyncMock

    captured = {}

    async def capture_call(**kwargs):
        captured.update(kwargs)

    with patch(
        "litellm.proxy.pass_through_endpoints.streaming_handler"
        ".PassThroughStreamingHandler._route_streaming_logging_to_handler",
        side_effect=capture_call,
    ):
        import asyncio, litellm.google_genai.streaming_iterator as si
        # ... drive one iteration through an AsyncIterator
        # then assert captured["endpoint_type"] == EndpointType.GOOGLE_GENAI

This tests the actual runtime value that ends up in the call, making it immune to comment or docstring changes.

[greptile-apps]

Missing model kwarg leaves extraction to URL parsing

_route_streaming_logging_to_handler accepts an optional model parameter that is forwarded directly to _handle_logging_gemini_collected_chunks. The iterator already has self.model available, but it is not passed here — instead the handler falls back to parsing the model name out of the URL string via extract_model_from_url.

If self.model is ever None, the formatted URL becomes /models/None:streamGenerateContent, and extract_model_from_url will return the literal string "None". That string is then fed into litellm.completion_cost(model="None"), which will silently return an incorrect cost (or raise).

Passing model=self.model explicitly would be more direct and avoids this edge case.

Suggested change

	PassThroughStreamingHandler._route_streaming_logging_to_handler(
	litellm_logging_obj=self.litellm_logging_obj,
	passthrough_success_handler_obj=GLOBAL_PASS_THROUGH_SUCCESS_HANDLER_OBJ,
	url_route="/v1/generateContent",
	url_route=f"/models/{self.model}:streamGenerateContent",
	request_body=self.request_body or {},
	endpoint_type=EndpointType.VERTEX_AI,
	endpoint_type=EndpointType.GOOGLE_GENAI,
	start_time=self.start_time,
	raw_bytes=self.collected_chunks,
	end_time=end_time,
	PassThroughStreamingHandler._route_streaming_logging_to_handler(
	litellm_logging_obj=self.litellm_logging_obj,
	passthrough_success_handler_obj=GLOBAL_PASS_THROUGH_SUCCESS_HANDLER_OBJ,
	url_route=f"/models/{self.model}:streamGenerateContent",
	request_body=self.request_body or {},
	endpoint_type=EndpointType.GOOGLE_GENAI,
	start_time=self.start_time,
	raw_bytes=self.collected_chunks,
	end_time=end_time,
	model=self.model,

[greptile-apps]

Pre-setting async_complete_streaming_response triggers unconditional early return, silently skipping all callbacks

async_success_handler has an early-exit guard at line 2486–2487 of litellm_logging.py:

if "async_complete_streaming_response" in self.model_call_details:
    return  # break out of this.

By setting async_complete_streaming_response on model_call_details before calling async_success_handler, this block causes the handler to return immediately — before the callback loops at lines 2571+ (CustomLogger, function-based, DynamoDB, OpenMeter, etc.) are ever reached.

Effect: All callbacks are silently skipped for every passthrough provider — Anthropic, Vertex AI, OpenAI, and Google GenAI alike. This is worse than the original bug (#24097), which only affected Google GenAI.

The prior working path for pass-through endpoints was: async_success_handler set async_complete_streaming_response internally (within the elif self.call_type == "pass_through_endpoint" branch at line 2543), and then execution continued to the callback loop below. By pre-setting the key before the call, that path is bypassed entirely.

The fix should set the key inside async_success_handler (or remove the pre-set), not before calling it.

[greptile-apps]

Mock hides the regression this test is designed to catch

async_success_handler is replaced with a bare AsyncMock, so the test never exercises the real guard in litellm_logging.py:

# litellm_logging.py ~line 2492
if "async_complete_streaming_response" in self.model_call_details:
    return  # break out of this.

Because streaming_handler.py pre-sets async_complete_streaming_response on model_call_details before calling async_success_handler, the real handler hits this guard and returns immediately — skipping every callback loop. The test then asserts the pre-set key as a positive assertion (lines 128–133), effectively validating the broken behaviour as correct.

A test that truly verifies "success callbacks fire" must either:

1. Not mock async_success_handler (let the real one run), or
2. Use side_effect to verify that the callbacks list is iterated.

As written, the test passes even when no callback ever executes, defeating its stated purpose.

Rule Used: # Code Review Rule: Mock Test Integrity

What:... (source)

[dkindlund]

Hi @awais786 — heads up that I've opened #25960 as a revival of this PR's approach, since the underlying issue (#24097) is still affecting users.

The new PR preserves your structural fix (the EndpointType.GOOGLE_GENAI enum, streaming_iterator routing change, and streaming_handler branch) and addresses the outstanding Greptile feedback that was open when this PR was closed:

• Removed the inspect.getsource source-text tests in favor of runtime-behavior tests
• Removed the sys.path manipulation hack
• Fixed the dead-code assertion-message-as-tuple pattern
• Fixed raw_bytes type mismatch in stub args
• Iterator now passes model explicitly + uses the real streamGenerateContent URL
• Added an end-to-end test that uses a real Logging instance + real CustomLogger subclass (the original mocked async_success_handler itself, which would have masked the regression)

Also closed one functional gap I noticed while reviewing: the sync iterator's next never invoked the logging route on StopIteration, so sync callers were silently broken even with your async-side fix. The new PR mirrors the async path on the sync side.

You're credited as the original author via Co-Authored-By trailer. Happy to defer to you if you'd rather pick this back up — just let me know.