Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@
)

from .config import PredictStateMapping
from .serialization import serialize_tool_args

logger = logging.getLogger(__name__)

Expand Down Expand Up @@ -241,7 +242,7 @@ async def _execute_proxy_tool(self, args: Dict[str, Any], tool_context: Any) ->
logger.debug(f"Emitted TOOL_CALL_START for {tool_call_id}")

# Emit TOOL_CALL_ARGS event
args_json = json.dumps(args)
args_json = serialize_tool_args(args)
args_event = ToolCallArgsEvent(
type=EventType.TOOL_CALL_ARGS,
tool_call_id=tool_call_id,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@
from google.adk.events import Event as ADKEvent

from .config import PredictStateMapping, normalize_predict_state
from .serialization import serialize_tool_args

import logging
logger = logging.getLogger(__name__)
Expand Down Expand Up @@ -761,9 +762,7 @@ async def translate_lro_function_calls(self,adk_event: ADKEvent)-> AsyncGenerato
parent_message_id=None
)
if hasattr(fc, 'args') and fc.args:
# Convert args to string (JSON format)
import json
args_str = json.dumps(fc.args) if isinstance(fc.args, dict) else str(fc.args)
args_str = serialize_tool_args(fc.args)
yield ToolCallArgsEvent(
type=EventType.TOOL_CALL_ARGS,
tool_call_id=fc.id,
Expand Down Expand Up @@ -839,8 +838,7 @@ async def _translate_function_calls(

# Emit TOOL_CALL_ARGS if we have arguments
if hasattr(func_call, 'args') and func_call.args:
# Convert args to string (JSON format)
args_str = json.dumps(func_call.args) if isinstance(func_call.args, dict) else str(func_call.args)
args_str = serialize_tool_args(func_call.args)

yield ToolCallArgsEvent(
type=EventType.TOOL_CALL_ARGS,
Expand Down Expand Up @@ -1198,7 +1196,7 @@ def _translate_function_calls_to_tool_calls(function_calls: List[Any]) -> List[T
type="function",
function=FunctionCall(
name=fc.name,
arguments=json.dumps(fc.args) if hasattr(fc, 'args') and fc.args else "{}"
arguments=serialize_tool_args(fc.args) if hasattr(fc, 'args') and fc.args else "{}"
)
)
tool_calls.append(tool_call)
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,30 @@
"""Shared JSON serialization helpers for tool-call arguments.

Standard ``json.dumps`` fails when args dicts contain Pydantic models or
Python ``Enum`` values (e.g. ``SecuritySchemeType``). The helper here uses
Pydantic's ``TypeAdapter`` which knows how to serialize those types.
"""

from __future__ import annotations

from typing import Any

from pydantic import TypeAdapter

_dict_adapter: TypeAdapter[dict[str, Any]] = TypeAdapter(dict[str, Any])


def serialize_tool_args(args: Any) -> str:
"""Serialize tool-call *args* to a JSON string.

Handles dicts that may contain Pydantic models, Enums, or other
non-stdlib-serializable values by delegating to Pydantic's
``TypeAdapter.dump_json``.

Returns:
A JSON-encoded string. For non-dict values the result is
``str(args)``.
"""
if isinstance(args, dict):
return _dict_adapter.dump_json(args).decode()
return str(args)
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,8 @@
from google.adk.events import Event as ADKEvent
from google.genai import types

from ..serialization import serialize_tool_args

logger = logging.getLogger(__name__)

def _get_text_value(item: Union[dict, TextInputContent]) -> Optional[str]:
Expand Down Expand Up @@ -234,7 +236,7 @@ def convert_adk_event_to_ag_ui_message(event: ADKEvent) -> Optional[Message]:
type="function",
function=FunctionCall(
name=part.function_call.name,
arguments=json.dumps(part.function_call.args) if hasattr(part.function_call, 'args') else "{}"
arguments=serialize_tool_args(part.function_call.args) if hasattr(part.function_call, 'args') else "{}"
)
))

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -766,7 +766,7 @@ async def test_translate_function_calls_basic(self, translator, mock_adk_event):
assert events[0].tool_call_id == "call_123"
assert events[0].tool_call_name == "test_function"
assert events[1].tool_call_id == "call_123"
assert events[1].delta == '{"param1": "value1"}'
assert json.loads(events[1].delta) == {"param1": "value1"}
assert events[2].tool_call_id == "call_123"

@pytest.mark.asyncio
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
"""

import asyncio
import json
import os
import uuid
import warnings
Expand Down Expand Up @@ -276,6 +277,85 @@ async def test_lro_emitted_ids_cleared_on_reset(self, translator):
translator.reset()
assert translator.lro_emitted_ids_by_name == {}

@pytest.mark.asyncio
async def test_lro_adk_request_credential_oauth2(self, translator):
"""Regression (#1331): adk_request_credential with OAuth2 AuthConfig must serialize.

ADK emits a long-running function call named ``adk_request_credential``
whose args dict contains an ``AuthConfig`` Pydantic model. The model
in turn nests ``OAuth2`` which has a ``type_: SecuritySchemeType`` enum
field. Before the fix, ``json.dumps`` raised:

TypeError: Object of type SecuritySchemeType is not JSON serializable
"""
from fastapi.openapi.models import OAuthFlowAuthorizationCode
from google.adk.auth.auth_schemes import OAuth2, OAuthFlows, SecuritySchemeType
from google.adk.auth import AuthConfig
from google.adk.auth.auth_credential import (
AuthCredential,
AuthCredentialTypes,
OAuth2Auth,
)

auth_scheme = OAuth2(
flows=OAuthFlows(
authorizationCode=OAuthFlowAuthorizationCode(
authorizationUrl="https://accounts.google.com/o/oauth2/auth",
tokenUrl="https://oauth2.googleapis.com/token",
scopes={
"https://www.googleapis.com/auth/calendar": "Calendar access",
},
),
),
)
raw_credential = AuthCredential(
auth_type=AuthCredentialTypes.OAUTH2,
oauth2=OAuth2Auth(
client_id="123456.apps.googleusercontent.com",
client_secret="GOCSPX-secret",
),
)
auth_config = AuthConfig(
auth_scheme=auth_scheme,
raw_auth_credential=raw_credential,
)

fc = MagicMock()
fc.id = "adk-cred-123"
fc.name = "adk_request_credential"
fc.args = {
"function_call_id": "adk-cred-123",
"auth_config": auth_config,
}
part = MagicMock()
part.function_call = fc
part.text = None
evt = MagicMock()
evt.content = MagicMock()
evt.content.parts = [part]
evt.long_running_tool_ids = ["adk-cred-123"]

events = []
async for e in translator.translate_lro_function_calls(evt):
events.append(e)

assert len(events) == 3
assert events[0].type == EventType.TOOL_CALL_START
assert events[0].tool_call_name == "adk_request_credential"

args_event = events[1]
assert args_event.type == EventType.TOOL_CALL_ARGS
parsed = json.loads(args_event.delta)
assert parsed["function_call_id"] == "adk-cred-123"

ac = parsed["auth_config"]
assert ac["auth_scheme"]["type_"] == "oauth2"
assert ac["raw_auth_credential"]["auth_type"] == "oauth2"
assert ac["raw_auth_credential"]["oauth2"]["client_id"] == "123456.apps.googleusercontent.com"
auth_code_flow = ac["auth_scheme"]["flows"]["authorizationCode"]
assert auth_code_flow["authorizationUrl"] == "https://accounts.google.com/o/oauth2/auth"
assert auth_code_flow["tokenUrl"] == "https://oauth2.googleapis.com/token"

@pytest.mark.asyncio
async def test_parallel_same_name_lro_calls_all_emitted(self, translator):
"""Multiple parallel LRO calls to the same tool should all be emitted (issue #1334)."""
Expand Down
73 changes: 73 additions & 0 deletions integrations/adk-middleware/python/tests/test_serialization.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,73 @@
"""Tests for the shared serialize_tool_args helper.

Covers plain dicts, dicts containing Python Enums (the SecuritySchemeType
scenario), dicts containing Pydantic models, non-dict values, and edge cases.
"""

import enum
import json

from pydantic import BaseModel

from ag_ui_adk.serialization import serialize_tool_args


class FakeSecuritySchemeType(enum.Enum):
oauth2 = "oauth2"
apiKey = "apiKey"


class NestedModel(BaseModel):
url: str
scheme_type: FakeSecuritySchemeType


class TestSerializeToolArgs:

def test_plain_dict(self):
args = {"city": "Seattle", "units": "metric"}
result = serialize_tool_args(args)
assert json.loads(result) == args

def test_dict_with_enum_value(self):
"""Regression (#1331): SecuritySchemeType-like enums must not raise TypeError."""
args = {
"auth_type": FakeSecuritySchemeType.oauth2,
"scopes": ["read", "write"],
}
result = serialize_tool_args(args)
parsed = json.loads(result)
assert parsed["auth_type"] == "oauth2"
assert parsed["scopes"] == ["read", "write"]

def test_dict_with_pydantic_model_value(self):
args = {
"endpoint": NestedModel(
url="https://example.com",
scheme_type=FakeSecuritySchemeType.apiKey,
)
}
result = serialize_tool_args(args)
parsed = json.loads(result)
assert parsed["endpoint"]["url"] == "https://example.com"
assert parsed["endpoint"]["scheme_type"] == "apiKey"

def test_dict_with_nested_enum(self):
args = {
"config": {
"type": FakeSecuritySchemeType.oauth2,
"enabled": True,
}
}
result = serialize_tool_args(args)
parsed = json.loads(result)
assert parsed["config"]["type"] == "oauth2"

def test_string_args_passthrough(self):
assert serialize_tool_args("raw_string") == "raw_string"

def test_non_dict_non_string(self):
assert serialize_tool_args(42) == "42"

def test_empty_dict(self):
assert serialize_tool_args({}) == "{}"
Original file line number Diff line number Diff line change
Expand Up @@ -399,7 +399,7 @@ def test_convert_assistant_event_with_function_call(self):
assert tool_call.id == "call_123"
assert tool_call.type == "function"
assert tool_call.function.name == "get_weather"
assert tool_call.function.arguments == '{"location": "Boston"}'
assert json.loads(tool_call.function.arguments) == {"location": "Boston"}

def test_convert_assistant_event_with_text_and_function_call(self):
"""Test converting assistant event with both text and function call."""
Expand Down
Loading