From 6f6212089ee6b011e0d985b384e2e1033a87a525 Mon Sep 17 00:00:00 2001 From: Nishanth Date: Tue, 27 Oct 2026 14:23:17 +0530 Subject: [PATCH] fix(models): reject non-object Apigee tool-call arguments Validate completed Chat Completions arguments before assigning them to FunctionCall.args. Drop invalid calls individually and report a structured MALFORMED_FUNCTION_CALL error while preserving valid siblings, text, and metadata. Preserve malformed-JSON exceptions and partial streaming, with provider MAX_TOKENS and SAFETY reasons retaining precedence. Fixes #7740 AI assistance: OpenAI Codex generated the implementation and regression tests and performed automated review and validation. --- src/google/adk/models/apigee_llm.py | 59 +++- .../models/test_completions_http_client.py | 326 ++++++++++++++++++ 2 files changed, 377 insertions(+), 8 deletions(-) diff --git a/src/google/adk/models/apigee_llm.py b/src/google/adk/models/apigee_llm.py index 477c82698d2..21702cf364d 100644 --- a/src/google/adk/models/apigee_llm.py +++ b/src/google/adk/models/apigee_llm.py @@ -1015,6 +1015,7 @@ def __init__(self) -> None: self.content_parts = '' self.tool_call_parts: dict[int, types.Part] = {} self._tool_call_deltas: dict[int, list[str]] = {} + self._invalid_tool_call_indices: set[int] = set() self._tool_call_trackers: dict[int, streaming_utils._JsonPathTracker] = {} self.role = '' self.streaming_complete = False @@ -1061,6 +1062,16 @@ def process_response(self, response: dict[str, Any]) -> LlmResponse: content=types.Content(role=role, parts=parts), usage_metadata=usage_metadata, finish_reason=self._map_finish_reason(choice.get('finish_reason')), + error_code=( + self._map_finish_reason(choice.get('finish_reason')) + if self._invalid_tool_call_indices + else None + ), + error_message=( + 'Tool call arguments must be a JSON object.' + if self._invalid_tool_call_indices + else None + ), logprobs_result=logprobs_result, model_version=response.get('model'), custom_metadata=custom_metadata, @@ -1135,6 +1146,16 @@ def process_chunk( parts=self._get_content_parts(), ), finish_reason=self._map_finish_reason(finish_reason), + error_code=( + self._map_finish_reason(finish_reason) + if self._invalid_tool_call_indices + else None + ), + error_message=( + 'Tool call arguments must be a JSON object.' + if self._invalid_tool_call_indices + else None + ), custom_metadata=self.custom_metadata, model_version=self.model, usage_metadata=usage_metadata, @@ -1144,6 +1165,11 @@ def process_chunk( return def _map_finish_reason(self, reason: str | None) -> types.FinishReason: + if self._invalid_tool_call_indices and reason not in ( + 'length', + 'content_filter', + ): + return types.FinishReason.MALFORMED_FUNCTION_CALL if reason == 'stop': return types.FinishReason.STOP if reason == 'length': @@ -1277,12 +1303,14 @@ def _get_content_parts(self) -> list[types.Part]: parts.append(types.Part.from_text(text=self.content_parts)) sorted_indices = sorted(self.tool_call_parts.keys()) for index in sorted_indices: + if index in self._invalid_tool_call_indices: + continue part = self.tool_call_parts[index] if part.function_call and not part.function_call.args: accumulated_args = ''.join(self._tool_call_deltas.get(index, [])) if accumulated_args: try: - part.function_call.args = _json_utils.safe_json_loads( + args = _json_utils.safe_json_loads( accumulated_args, context=f'tool call arguments: {accumulated_args}', ) @@ -1300,17 +1328,31 @@ def _get_content_parts(self) -> list[types.Part]: if isinstance(parsed_delta, dict): merged_args.update(parsed_delta) else: - logger.warning('Parsed delta is not a dict: %s', parsed_delta) + self._validate_tool_call_args(index, parsed_delta) except ValueError: logger.warning( 'Failed to parse individual delta as JSON: %s', delta ) if not merged_args: raise - part.function_call.args = merged_args + if index in self._invalid_tool_call_indices: + continue + args = merged_args + # Validate outside the decoder's ValueError fallback: a complete + # non-object must not be salvaged into different arguments. + if not self._validate_tool_call_args(index, args): + continue + part.function_call.args = args parts.append(part) return parts + def _validate_tool_call_args(self, index: int, args: object) -> bool: + """Rejects non-object arguments without changing valid sibling calls.""" + if not isinstance(args, dict): + self._invalid_tool_call_indices.add(index) + return False + return True + def _upsert_tool_call( self, tool_call: dict[str, Any], is_streaming: bool = False ) -> types.Part: @@ -1365,11 +1407,12 @@ def _upsert_tool_call( args = _json_utils.safe_json_loads( args_delta, context=f'tool call arguments: {args_delta}' ) - chunk_function_call.args = args - if not function_call.args: - function_call.args = dict(args) - else: - function_call.args.update(args) + if self._validate_tool_call_args(index, args): + chunk_function_call.args = args + if not function_call.args: + function_call.args = args.copy() + else: + function_call.args.update(args) func_name = func.get('name') if func_name: diff --git a/tests/unittests/models/test_completions_http_client.py b/tests/unittests/models/test_completions_http_client.py index 4a0a09b2522..c2f00b8ac6c 100644 --- a/tests/unittests/models/test_completions_http_client.py +++ b/tests/unittests/models/test_completions_http_client.py @@ -14,19 +14,345 @@ import asyncio import atexit +import copy import gc import json from unittest import mock from unittest.mock import AsyncMock import weakref +from google.adk.agents.llm_agent import Agent from google.adk.models.apigee_llm import ChatCompletionsResponseHandler from google.adk.models.apigee_llm import CompletionsHTTPClient from google.adk.models.llm_request import LlmRequest +from google.adk.models.llm_response import LlmResponse from google.genai import types import httpx import pytest +from .. import testing_utils + + +def _completed_argument_responses( + arguments: list[str], streaming: bool, finish: str, *, legacy: bool = False +) -> list[LlmResponse]: + handler = ChatCompletionsResponseHandler() + calls = [ + { + 'index': index, + 'id': f'call_{index}', + 'type': 'function', + 'function': {'name': f'tool_{index}', 'arguments': value}, + } + for index, value in enumerate(arguments) + ] + message = {'role': 'assistant', 'content': 'Looking that up.'} + if legacy: + message['function_call'] = calls[0]['function'] + else: + message['tool_calls'] = calls + payload = { + 'model': 'test-model', + 'id': 'response-id', + 'usage': {'prompt_tokens': 3, 'completion_tokens': 4, 'total_tokens': 7}, + } + original = copy.deepcopy(message) + if streaming: + responses = list( + handler.process_chunk({ + 'model': 'test-model', + 'id': 'response-id', + 'choices': [{'delta': message}], + }) + ) + responses.extend( + handler.process_chunk({ + **payload, + 'choices': [{'delta': {}, 'finish_reason': finish}], + }) + ) + assert all(response.partial for response in responses[:-1]) + for response in responses[:-1]: + for part in response.content.parts: + if part.function_call: + assert part.function_call.args is None + assert part.function_call.will_continue + else: + responses = [ + handler.process_response({ + **payload, + 'choices': [{'message': message, 'finish_reason': finish}], + }) + ] + assert message == original + return responses + + +@pytest.mark.parametrize('streaming', [False, True]) +@pytest.mark.parametrize('finish', ['tool_calls', 'stop', 'length']) +@pytest.mark.parametrize( + 'arguments', + [ + '[1, 2]', + '[]', + '[["a", 1]]', + '"x"', + '""', + '5', + '0', + 'false', + 'true', + 'null', + ], +) +def test_completed_non_object_arguments_rejected( + arguments: str, streaming: bool, finish: str +) -> None: + """Completed non-object calls report errors without losing response metadata.""" + response = _completed_argument_responses([arguments], streaming, finish)[-1] + expected = ( + types.FinishReason.MAX_TOKENS + if finish == 'length' + else types.FinishReason.MALFORMED_FUNCTION_CALL + ) + assert response.finish_reason == expected + assert response.error_code == expected + assert 'JSON object' in response.error_message + assert [p.text for p in response.content.parts] == ['Looking that up.'] + assert response.usage_metadata.total_token_count == 7 + assert response.model_version == 'test-model' + assert response.custom_metadata['id'] == 'response-id' + + +@pytest.mark.parametrize('streaming', [False, True]) +@pytest.mark.parametrize('finish', ['tool_calls', 'stop']) +@pytest.mark.parametrize( + 'arguments', ['{}', '{"a": 1}', '{"a": {"b": [1, null]}}'] +) +def test_completed_object_arguments_preserved( + arguments: str, streaming: bool, finish: str +) -> None: + """Nested objects and empty objects remain valid function arguments.""" + response = _completed_argument_responses([arguments], streaming, finish)[-1] + assert response.error_code is None + assert response.finish_reason == types.FinishReason.STOP + assert response.content.parts[1].function_call.args == json.loads(arguments) + + +@pytest.mark.parametrize('streaming', [False, True]) +@pytest.mark.parametrize('finish', ['tool_calls', 'stop']) +@pytest.mark.parametrize('invalid_index', [0, 1]) +def test_non_object_arguments_preserve_valid_sibling( + streaming: bool, finish: str, invalid_index: int +) -> None: + """Only the invalid call is removed regardless of its position.""" + arguments = ['{"a": 1}', '{"a": 1}'] + arguments[invalid_index] = 'null' + response = _completed_argument_responses(arguments, streaming, finish)[-1] + assert response.error_code == types.FinishReason.MALFORMED_FUNCTION_CALL + calls = [p.function_call for p in response.content.parts if p.function_call] + assert len(calls) == 1 + assert calls[0].id == f'call_{1 - invalid_index}' + assert calls[0].args == {'a': 1} + + +@pytest.mark.parametrize('arguments', ['null', '[]', '""']) +def test_non_object_legacy_function_call_rejected(arguments: str) -> None: + """Legacy non-streaming calls use the same object validation.""" + response = _completed_argument_responses( + [arguments], False, 'stop', legacy=True + )[-1] + assert response.error_code == types.FinishReason.MALFORMED_FUNCTION_CALL + assert not any(p.function_call for p in response.content.parts) + + +@pytest.mark.parametrize('streaming', [False, True]) +def test_empty_argument_buffer_preserves_absent_args(streaming: bool) -> None: + """An absent argument buffer remains distinct from a JSON empty string.""" + response = _completed_argument_responses([''], streaming, 'stop')[-1] + assert response.error_code is None + assert response.content.parts[1].function_call.args is None + + +@pytest.mark.parametrize('streaming', [False, True]) +def test_malformed_argument_json_still_raises(streaming: bool) -> None: + """Syntax errors keep the established exception contract.""" + with pytest.raises(ValueError, match='Invalid JSON'): + _completed_argument_responses(['{"a":'], streaming, 'stop') + + +@pytest.mark.parametrize( + 'fragments', + [ + ['{"a": 1}', '[1, 2]'], + ['[1, 2]', '{"a": 1}'], + ], +) +def test_non_object_fallback_fragment_is_not_silently_dropped( + fragments: list[str], +) -> None: + """Fallback merging cannot conceal a non-object argument fragment.""" + handler = ChatCompletionsResponseHandler() + for fragment in fragments: + list( + handler.process_chunk({ + 'choices': [{ + 'delta': { + 'tool_calls': [{ + 'index': 0, + 'id': 'call_0', + 'function': {'name': 'tool', 'arguments': fragment}, + }] + } + }] + }) + ) + response = list( + handler.process_chunk( + {'choices': [{'delta': {}, 'finish_reason': 'stop'}]} + ) + )[-1] + assert response.error_code == types.FinishReason.MALFORMED_FUNCTION_CALL + assert not any(p.function_call for p in response.content.parts) + + +@pytest.mark.parametrize('streaming', [False, True]) +async def test_valid_sibling_executes_with_malformed_response( + streaming: bool, +) -> None: + """The default Runner executes the retained call, never the rejected one.""" + response = _completed_argument_responses( + ['null', '{"a": 1}'], streaming, 'stop' + )[-1] + invoked = [] + + def tool_0() -> None: + invoked.append('invalid') + + def tool_1(a: int) -> int: + invoked.append(a) + return a + 1 + + model = testing_utils.MockModel.create( + responses=[ + response, + LlmResponse( + content=types.Content( + role='model', parts=[types.Part(text='done')] + ) + ), + ] + ) + agent = Agent(name='agent', model=model, tools=[tool_0, tool_1]) + runner = testing_utils.TestInMemoryRunner(agent=agent) + events = await runner.run_async_with_new_session('test') + + assert invoked == [1] + assert any( + event.error_code == types.FinishReason.MALFORMED_FUNCTION_CALL + for event in events + ) + results = [ + p.function_response + for event in events + if event.content + for p in event.content.parts + if p.function_response + ] + assert len(results) == 1 + assert results[0].name == 'tool_1' + assert results[0].response == {'result': 2} + + +@pytest.mark.parametrize('arguments', ['{"a": 1}', '[1, 2]']) +def test_fragmented_arguments_validate_only_on_completion( + arguments: str, +) -> None: + """Incomplete fragments are hints; only completed calls are validated.""" + handler = ChatCompletionsResponseHandler() + for fragment in (arguments[:3], arguments[3:]): + responses = list( + handler.process_chunk({ + 'choices': [{ + 'delta': { + 'tool_calls': [{ + 'index': 2, + 'id': 'call_2', + 'function': {'name': 'tool', 'arguments': fragment}, + }] + } + }] + }) + ) + assert all(r.partial and r.error_code is None for r in responses) + assert responses[0].content.parts[0].function_call.args is None + final = list( + handler.process_chunk( + {'choices': [{'delta': {}, 'finish_reason': 'stop'}]} + ) + )[-1] + if arguments.startswith('{'): + assert final.error_code is None + assert final.content.parts[0].function_call.args == {'a': 1} + else: + assert final.error_code == types.FinishReason.MALFORMED_FUNCTION_CALL + assert final.content.parts == [] + + +@pytest.mark.parametrize( + 'fragments', + [ + ['{"a":', 'null'], + ['null', '{"a":'], + ['null', 'null'], + ], +) +def test_unsalvageable_fallback_json_still_raises(fragments: list[str]) -> None: + """Non-object fragments do not hide an existing malformed-JSON exception.""" + handler = ChatCompletionsResponseHandler() + for fragment in fragments: + list( + handler.process_chunk({ + 'choices': [{ + 'delta': { + 'tool_calls': [{ + 'index': 0, + 'function': {'name': 'tool', 'arguments': fragment}, + }] + } + }] + }) + ) + with pytest.raises(ValueError, match='Invalid JSON'): + list( + handler.process_chunk( + {'choices': [{'delta': {}, 'finish_reason': 'stop'}]} + ) + ) + + +def test_non_object_arguments_in_finish_bearing_chunk() -> None: + """A finish signal on the argument chunk still rejects non-object args.""" + handler = ChatCompletionsResponseHandler() + responses = list( + handler.process_chunk({ + 'choices': [{ + 'delta': { + 'tool_calls': [{ + 'index': 0, + 'id': 'bad', + 'function': {'name': 'tool', 'arguments': 'null'}, + }] + }, + 'finish_reason': 'tool_calls', + }] + }) + ) + assert responses[0].partial + assert responses[0].error_code is None + assert responses[-1].error_code == types.FinishReason.MALFORMED_FUNCTION_CALL + assert responses[-1].content.parts == [] + @pytest.fixture def client():