diff --git a/python/packages/gemini/agent_framework_gemini/_chat_client.py b/python/packages/gemini/agent_framework_gemini/_chat_client.py index 515bab2f3f4..6c43c2a6cfd 100644 --- a/python/packages/gemini/agent_framework_gemini/_chat_client.py +++ b/python/packages/gemini/agent_framework_gemini/_chat_client.py @@ -19,6 +19,7 @@ ChatResponse, ChatResponseUpdate, Content, + FinishReason, FinishReasonLiteral, FunctionInvocationConfiguration, FunctionInvocationLayer, @@ -1260,18 +1261,20 @@ def _parse_usage(self, usage: types.GenerateContentResponseUsageMetadata | None) details["reasoning_output_token_count"] = v return details or None - def _map_finish_reason(self, reason: str | None) -> FinishReasonLiteral | None: + def _map_finish_reason(self, reason: str | None) -> FinishReasonLiteral | FinishReason | None: """Map a Gemini finish reason string to the framework's FinishReasonLiteral. Args: reason: The finish reason name from the Gemini API (e.g. ``"STOP"``), or None. Returns: - The corresponding ``FinishReasonLiteral``, or None if the reason is absent or unmapped. + The corresponding ``FinishReasonLiteral`` for known reasons, the raw reason string for + values not yet mapped (so callers still see that the response terminated abnormally + instead of losing it), or None if the reason is absent or ``FINISH_REASON_UNSPECIFIED``. """ - if not reason: + if not reason or reason == "FINISH_REASON_UNSPECIFIED": return None - return _FINISH_REASON_MAP.get(reason) + return _FINISH_REASON_MAP.get(reason, FinishReason(reason)) # endregion diff --git a/python/packages/gemini/tests/test_gemini_client.py b/python/packages/gemini/tests/test_gemini_client.py index e23ae0a8c6e..68e39e7282e 100644 --- a/python/packages/gemini/tests/test_gemini_client.py +++ b/python/packages/gemini/tests/test_gemini_client.py @@ -463,19 +463,39 @@ async def test_get_response_no_usage_when_metadata_absent() -> None: @pytest.mark.parametrize( ("gemini_reason", "expected"), [ + # Currently-mapped values: must keep mapping to exactly what they map to today. ("STOP", "stop"), ("MAX_TOKENS", "length"), ("SAFETY", "content_filter"), ("RECITATION", "content_filter"), + ("LANGUAGE", "content_filter"), ("BLOCKLIST", "content_filter"), ("PROHIBITED_CONTENT", "content_filter"), ("SPII", "content_filter"), + ("IMAGE_SAFETY", "content_filter"), + ("IMAGE_PROHIBITED_CONTENT", "content_filter"), + ("IMAGE_RECITATION", "content_filter"), ("MALFORMED_FUNCTION_CALL", "tool_calls"), - ("OTHER", None), + ("UNEXPECTED_TOOL_CALL", "tool_calls"), + # Real google-genai FinishReason values with no entry in _FINISH_REASON_MAP: must now + # pass through as the raw string instead of being silently dropped to None. + ("OTHER", "OTHER"), + ("TOO_MANY_TOOL_CALLS", "TOO_MANY_TOOL_CALLS"), + ("NO_IMAGE", "NO_IMAGE"), + ("IMAGE_OTHER", "IMAGE_OTHER"), + # Absent / unspecified: must still yield None. + (None, None), + ("FINISH_REASON_UNSPECIFIED", None), ], ) -async def test_finish_reason_mapping(gemini_reason: str, expected: str | None) -> None: - """Maps Gemini finish reason strings to the correct FinishReasonLiteral values.""" +async def test_finish_reason_mapping(gemini_reason: str | None, expected: str | None) -> None: + """Maps Gemini finish reason strings to the correct FinishReasonLiteral values. + + Unmapped-but-real provider values (e.g. TOO_MANY_TOOL_CALLS) must pass through as the raw + string rather than vanishing to None, mirroring the fallback PR #7105 added for the other + chat clients (bedrock, claude, core, github_copilot, ollama, openai) but not gemini. + FINISH_REASON_UNSPECIFIED and an absent reason remain the legitimate None cases. + """ client, mock = _make_gemini_client() mock.aio.models.generate_content = AsyncMock( return_value=_make_response([_make_part(text="Hi")], finish_reason=gemini_reason) @@ -486,6 +506,37 @@ async def test_finish_reason_mapping(gemini_reason: str, expected: str | None) - assert response.finish_reason == expected +async def test_unmapped_finish_reason_still_attaches_usage_on_streamed_final_chunk() -> None: + """An unmapped provider finish reason must not suppress the final chunk's usage/token accounting. + + Before this fix, `_process_chunk` attached usage only when `_map_finish_reason` returned a + non-None value. Since TOO_MANY_TOOL_CALLS is a real google-genai FinishReason absent from + `_FINISH_REASON_MAP`, it mapped to None, which silently dropped both the finish reason and + the whole turn's usage/token accounting for the final streamed chunk. + """ + client, mock = _make_gemini_client() + chunks = [ + _make_response([_make_part(text="Hello ")], finish_reason=None, prompt_tokens=None, output_tokens=None), + _make_response( + [_make_part(text="world!")], + finish_reason="TOO_MANY_TOOL_CALLS", + prompt_tokens=10, + output_tokens=5, + ), + ] + mock.aio.models.generate_content_stream = AsyncMock(return_value=_async_iter(chunks)) + + stream = client.get_response( + messages=[Message(role="user", contents=[Content.from_text("Hi")])], + stream=True, + ) + + updates = [update async for update in stream] + + assert updates[-1].finish_reason == "TOO_MANY_TOOL_CALLS" + assert any(c.type == "usage" for c in updates[-1].contents) + + # message conversion