fix: add PEP 604 union syntax in function tool parameters

- Support PEP 604 unions in direct parameter parser.
- Fix a bug where collapsing simple unions (e.g., `Optional[list[T]]`) lost
  nested schema properties (like `items` or `properties`) by replacing the
  parent schema with the collapsed inner schema instead of just copying its
  type.
- Update `test_required_fields_set_in_json_schema_fallback` to use
  `tuple[str, ...]` to ensure the fallback path remains tested.

Change-Id: Idc1cb55e265ba888c03aa923f60d2d4b3d1ae131

Author: sbunchgoogle

src/google/adk/tools/_function_parameter_parse_util.py

@@ -247,7 +247,7 @@ def _parse_schema_from_parameter(
     _raise_if_schema_unsupported(variant, schema)
     return schema
   if (
-      get_origin(param.annotation) is Union
+      get_origin(param.annotation) in (Union, typing_types.UnionType)
       # only parse simple UnionType, example int | str | float | bool
       # complex types.UnionType will be invoked in raise branch
       and all(
@@ -276,8 +276,10 @@ def _parse_schema_from_parameter(
         schema.any_of.append(schema_in_any_of)
         unique_types.add(schema_in_any_of.model_dump_json(exclude_none=True))
     if len(schema.any_of) == 1:  # param: list | None -> Array
-      schema.type = schema.any_of[0].type
-      schema.any_of = None
+      collapsed = schema.any_of[0]
+      if schema.nullable:
+        collapsed.nullable = True
+      schema = collapsed
     if (
         param.default is not inspect.Parameter.empty
         and param.default is not None
@@ -287,8 +289,10 @@ def _parse_schema_from_parameter(
       schema.default = param.default
     _raise_if_schema_unsupported(variant, schema)
     return schema
-  if isinstance(param.annotation, _GenericAlias) or isinstance(
-      param.annotation, typing_types.GenericAlias
+  if (
+      isinstance(param.annotation, _GenericAlias)
+      or isinstance(param.annotation, typing_types.GenericAlias)
+      or isinstance(param.annotation, typing_types.UnionType)
   ):
     origin = get_origin(param.annotation)
     args = get_args(param.annotation)
@@ -330,7 +334,7 @@ def _parse_schema_from_parameter(
         schema.default = param.default
       _raise_if_schema_unsupported(variant, schema)
       return schema
-    if origin is Union:
+    if origin in (Union, typing_types.UnionType):
       schema.any_of = []
       schema.type = types.Type.OBJECT
       unique_types = set()
@@ -365,8 +369,10 @@ def _parse_schema_from_parameter(
           schema.any_of.append(schema_in_any_of)
           unique_types.add(schema_in_any_of.model_dump_json(exclude_none=True))
       if len(schema.any_of) == 1:  # param: Union[List, None] -> Array
-        schema.type = schema.any_of[0].type
-        schema.any_of = None
+        collapsed = schema.any_of[0]
+        if schema.nullable:
+          collapsed.nullable = True
+        schema = collapsed
       if (
           param.default is not None
           and param.default is not inspect.Parameter.empty

tests/unittests/tools/test_from_function_with_options.py

@@ -324,7 +324,7 @@ async def test_function(param: str) -> AsyncGenerator[Dict[str, str], None]:
 def test_required_fields_set_in_json_schema_fallback():
   """Test that required fields are populated when the json_schema fallback path is used.
 
-  When a parameter has a complex union type (e.g. list[str] | None) that
+  When a parameter has a complex type (e.g. tuple[str, ...] | None) that
   _parse_schema_from_parameter can't handle, from_function_with_options falls
   back to the parameters_json_schema branch. This test verifies that the
   required fields are correctly populated in that fallback branch.
@@ -333,7 +333,7 @@ def test_required_fields_set_in_json_schema_fallback():
   def complex_tool(
       query: str,
       mode: str = 'default',
-      tags: list[str] | None = None,
+      tags: tuple[str, ...] | None = None,
   ) -> str:
     """A tool where one param has a complex union type."""
     return query

tests/unittests/tools/test_set_model_response_tool.py

@@ -15,10 +15,13 @@
 """Tests for SetModelResponseTool."""
 
 import inspect
+from typing import Optional
 
 from google.adk.agents.invocation_context import InvocationContext
 from google.adk.agents.llm_agent import LlmAgent
 from google.adk.agents.run_config import RunConfig
+from google.adk.features._feature_registry import FeatureName
+from google.adk.features._feature_registry import temporary_feature_override
 from google.adk.sessions.in_memory_session_service import InMemorySessionService
 from google.adk.tools.set_model_response_tool import SetModelResponseTool
 from google.adk.tools.tool_context import ToolContext
@@ -467,3 +470,67 @@ async def test_run_async_dict_schema():
   assert result is not None
   assert isinstance(result, dict)
   assert result == {'a': 1, 'b': 2, 'c': 3}
+
+
+class SubSchema(BaseModel):
+
+  field1: str = Field(description='Field 1')
+  field2: int = Field(description='Field 2')
+
+
+class ConsolidatedOptionalSchema(BaseModel):
+
+  nested: Optional[SubSchema] = Field(default=None, description='Nested model')
+  nested_list: Optional[list[SubSchema]] = Field(
+      default=None, description='Nested list of models'
+  )
+  pep604_nested: SubSchema | None = Field(
+      default=None, description='PEP 604 optional nested model'
+  )
+  pep604_raw_list: list | None = Field(default=None, description='Raw list')
+
+
+def test_get_declaration_optional_fields():
+  """Test that tool declaration preserves properties for various optional fields."""
+  with temporary_feature_override(FeatureName.JSON_SCHEMA_FOR_FUNC_DECL, False):
+    tool = SetModelResponseTool(ConsolidatedOptionalSchema)
+
+    declaration = tool._get_declaration()
+
+  assert declaration is not None
+  assert declaration.name == 'set_model_response'
+  params_schema = declaration.parameters
+  assert params_schema is not None
+  assert params_schema.type == 'OBJECT'
+
+  # 1. Optional[SubSchema]
+  assert 'nested' in params_schema.properties
+  nested_schema = params_schema.properties['nested']
+  assert nested_schema.type == 'OBJECT'
+  assert nested_schema.properties is not None
+  assert nested_schema.properties['field1'].type == 'STRING'
+  assert nested_schema.properties['field2'].type == 'INTEGER'
+
+  # 2. Optional[list[SubSchema]]
+  assert 'nested_list' in params_schema.properties
+  nested_list_schema = params_schema.properties['nested_list']
+  assert nested_list_schema.type == 'ARRAY'
+  assert nested_list_schema.items is not None
+  items_schema = nested_list_schema.items
+  assert items_schema.type == 'OBJECT'
+  assert items_schema.properties is not None
+  assert items_schema.properties['field1'].type == 'STRING'
+  assert items_schema.properties['field2'].type == 'INTEGER'
+
+  # 3. SubSchema | None (PEP 604)
+  assert 'pep604_nested' in params_schema.properties
+  pep604_nested_schema = params_schema.properties['pep604_nested']
+  assert pep604_nested_schema.type == 'OBJECT'
+  assert pep604_nested_schema.properties is not None
+  assert pep604_nested_schema.properties['field1'].type == 'STRING'
+  assert pep604_nested_schema.properties['field2'].type == 'INTEGER'
+
+  # 4. list | None (PEP 604)
+  assert 'pep604_raw_list' in params_schema.properties
+  pep604_raw_list_schema = params_schema.properties['pep604_raw_list']
+  assert pep604_raw_list_schema.type == 'ARRAY'