fix(task): support raw dict and types.Schema output_schema in FinishTaskTool (#7935)

src/google/adk/agents/llm/task/_finish_task_tool.py

@@ -16,6 +16,7 @@
 
 from __future__ import annotations
 
+import copy
 from typing import Any
 from typing import Optional
 from typing import TYPE_CHECKING
@@ -92,12 +93,14 @@ def __init__(
     self.output_schema: SchemaType = (
         output_schema if output_schema is not None else _DefaultTaskOutput
     )
-    self._adapter: TypeAdapter[Any] = TypeAdapter(self.output_schema)
-    raw_schema = self._adapter.json_schema()
+    if isinstance(self.output_schema, (dict, types.Schema)):
+      self._adapter: Optional[TypeAdapter[Any]] = None
+    else:
+      self._adapter = TypeAdapter(self.output_schema)
     # FunctionDeclaration parameters must be a JSON object schema.
     # If the schema is already an object (e.g. BaseModel), use it directly.
     # Otherwise wrap it in an object with a single key.
-    self._wrapper_key: str | None = get_output_wrapper_key(raw_schema)
+    self._wrapper_key: str | None = get_output_wrapper_key(self.output_schema)
 
     description = (
         'Signal that this agent has completed its delegated task. Call this'
@@ -114,7 +117,10 @@ def __init__(
   @override
   def _get_declaration(self) -> Optional[types.FunctionDeclaration]:
     """Get the function declaration for this tool."""
-    raw_schema = self._adapter.json_schema()
+    if self._adapter is not None:
+      raw_schema = self._adapter.json_schema()
+    else:
+      raw_schema = copy.deepcopy(schema_to_json_schema(self.output_schema))
     if self._wrapper_key:
       # Extract $defs to the root level so $ref pointers remain valid
       # after wrapping the schema inside an object property.
@@ -186,8 +192,11 @@ async def run_async(
     """
     try:
       raw_value = args.get(self._wrapper_key) if self._wrapper_key else args
-      validated = self._adapter.validate_python(raw_value)
-      validated_output = self._adapter.dump_python(validated, mode='json')
+      if self._adapter is not None:
+        validated = self._adapter.validate_python(raw_value)
+        validated_output = self._adapter.dump_python(validated, mode='json')
+      else:
+        validated_output = raw_value
     except ValidationError as e:
       return {
           'error': (

src/google/adk/utils/_schema_utils.py

@@ -130,6 +130,10 @@ def schema_to_json_schema(schema: SchemaType) -> dict[str, Any]:
   """
   if isinstance(schema, dict):
     return schema
+  if isinstance(schema, types.Schema):
+    raw_schema = schema.model_dump(mode="json", exclude_none=True)
+    lowercase_schema_types(raw_schema)
+    return raw_schema
   return TypeAdapter(schema).json_schema()
 
 

tests/unittests/agents/llm/task/test_finish_task_tool.py

@@ -19,8 +19,10 @@
 from typing import Any
 from unittest import mock
 
+from google.adk.agents import LlmAgent
 from google.adk.agents.llm.task._finish_task_tool import FINISH_TASK_TOOL_NAME
 from google.adk.agents.llm.task._finish_task_tool import FinishTaskTool
+from google.genai import types
 from pydantic import BaseModel
 import pytest
 
@@ -72,6 +74,28 @@ def test_init_with_output_schema(self):
     assert 'Signal that this agent has completed' in tool.description
     assert 'output data' in tool.description
 
+  def test_init_with_raw_dict_output_schema(self):
+    """Test FinishTaskTool initialization with a raw dict output schema."""
+    schema = {"type": "object", "properties": {"answer": {"type": "string"}}}
+    agent = _make_task_agent(output_schema=schema)
+    tool = FinishTaskTool(task_agent=agent)
+    assert tool.name == FINISH_TASK_TOOL_NAME
+    assert tool.output_schema == schema
+    assert tool._adapter is None
+    assert tool._wrapper_key is None
+
+  def test_init_with_types_schema_output_schema(self):
+    """Test FinishTaskTool initialization with a types.Schema output schema."""
+    schema = types.Schema(
+        type="OBJECT", properties={"answer": types.Schema(type="STRING")}
+    )
+    agent = _make_task_agent(output_schema=schema)
+    tool = FinishTaskTool(task_agent=agent)
+    assert tool.name == FINISH_TASK_TOOL_NAME
+    assert tool.output_schema == schema
+    assert tool._adapter is None
+    assert tool._wrapper_key is None
+
   def test_get_declaration_without_output_schema(self):
     """Test function declaration generation without output schema."""
     agent = _make_task_agent()
@@ -97,6 +121,39 @@ def test_get_declaration_with_output_schema(self):
     assert 'result' in schema.get('properties', {})
     assert 'count' in schema.get('properties', {})
 
+  def test_get_declaration_with_raw_dict_output_schema(self):
+    """Test function declaration generation with raw dict output schema."""
+    schema = {
+        'type': 'object',
+        'properties': {'answer': {'type': 'string'}},
+        'required': ['answer'],
+    }
+    agent = _make_task_agent(output_schema=schema)
+    tool = FinishTaskTool(task_agent=agent)
+    declaration = tool._get_declaration()
+
+    assert declaration is not None
+    assert declaration.name == FINISH_TASK_TOOL_NAME
+    assert declaration.parameters_json_schema == schema
+    # Verify original dict was not mutated
+    assert tool.output_schema == schema
+
+  def test_get_declaration_with_types_schema_output_schema(self):
+    """Test function declaration generation with types.Schema output schema."""
+    schema = types.Schema(
+        type='OBJECT', properties={'answer': types.Schema(type='STRING')}
+    )
+    agent = _make_task_agent(output_schema=schema)
+    tool = FinishTaskTool(task_agent=agent)
+    declaration = tool._get_declaration()
+
+    assert declaration is not None
+    assert declaration.name == FINISH_TASK_TOOL_NAME
+    assert declaration.parameters_json_schema == {
+        'type': 'object',
+        'properties': {'answer': {'type': 'string'}},
+    }
+
   @pytest.mark.asyncio
   async def test_run_async_returns_confirmation(self):
     """Test that run_async returns a confirmation message."""
@@ -245,6 +302,16 @@ class TestFinishTaskToolAllSchemaTypes:
       [
           (SampleOutputSchema, None),
           (dict[str, Any], None),
+          ({'type': 'object', 'properties': {'x': {'type': 'string'}}}, None),
+          ({'type': 'string'}, 'result'),
+          (
+              types.Schema(
+                  type='OBJECT',
+                  properties={'x': types.Schema(type='STRING')},
+              ),
+              None,
+          ),
+          (types.Schema(type='STRING'), 'result'),
           (str, 'result'),
           (int, 'result'),
           (bool, 'result'),
@@ -256,6 +323,10 @@ class TestFinishTaskToolAllSchemaTypes:
       ids=[
           'BaseModel',
           'dict',
+          'raw_dict_object',
+          'raw_dict_string',
+          'types_schema_object',
+          'types_schema_string',
           'str',
           'int',
           'bool',
@@ -305,6 +376,21 @@ def test_get_declaration_wrapped_schema(self, output_schema):
               {'result': 'done', 'count': 5},
           ),
           (dict[str, Any], {'key': 'value'}, {'key': 'value'}),
+          (
+              {'type': 'object', 'properties': {'answer': {'type': 'string'}}},
+              {'answer': 'hello'},
+              {'answer': 'hello'},
+          ),
+          ({'type': 'string'}, {'result': 'hello'}, 'hello'),
+          (
+              types.Schema(
+                  type='OBJECT',
+                  properties={'answer': types.Schema(type='STRING')},
+              ),
+              {'answer': 'hello'},
+              {'answer': 'hello'},
+          ),
+          (types.Schema(type='STRING'), {'result': 'hello'}, 'hello'),
           (str, {'result': 'hello'}, 'hello'),
           (int, {'result': 42}, 42),
           (bool, {'result': True}, True),
@@ -321,6 +407,10 @@ def test_get_declaration_wrapped_schema(self, output_schema):
       ids=[
           'BaseModel',
           'dict',
+          'raw_dict_object',
+          'raw_dict_string',
+          'types_schema_object',
+          'types_schema_string',
           'str',
           'int',
           'bool',
@@ -361,3 +451,37 @@ def test_get_declaration_list_basemodel_defs_at_root(self):
     assert 'SampleOutputSchema' in schema['$defs']
     # The nested result schema should not contain $defs
     assert '$defs' not in schema['properties']['result']
+
+
+class TestFinishTaskToolLlmAgentIntegration:
+  """Integration tests constructing LlmAgent with mode='task'."""
+
+  def test_task_agent_with_raw_dict_output_schema(self):
+    schema = {'type': 'object', 'properties': {'answer': {'type': 'string'}}}
+    agent = LlmAgent(
+        name='task_worker',
+        model='gemini-2.5-flash',
+        output_schema=schema,
+        mode='task',
+    )
+    assert agent.output_schema == schema
+    finish_tools = [t for t in agent.tools if isinstance(t, FinishTaskTool)]
+    assert len(finish_tools) == 1
+    assert finish_tools[0]._adapter is None
+    assert finish_tools[0]._wrapper_key is None
+
+  def test_task_agent_with_types_schema_output_schema(self):
+    schema = types.Schema(
+        type='OBJECT', properties={'answer': types.Schema(type='STRING')}
+    )
+    agent = LlmAgent(
+        name='task_worker',
+        model='gemini-2.5-flash',
+        output_schema=schema,
+        mode='task',
+    )
+    assert agent.output_schema == schema
+    finish_tools = [t for t in agent.tools if isinstance(t, FinishTaskTool)]
+    assert len(finish_tools) == 1
+    assert finish_tools[0]._adapter is None
+    assert finish_tools[0]._wrapper_key is None

tests/unittests/utils/test_schema_utils.py

@@ -466,6 +466,17 @@ def test_list_of_basemodel_schema_becomes_an_array_of_objects(self):
     ref = result["items"]["$ref"].rsplit("/", 1)[-1]
     assert result["$defs"][ref]["properties"]["name"]["type"] == "string"
 
+  def test_types_schema_is_converted_to_json_schema(self):
+    schema = types.Schema(
+        type="OBJECT",
+        properties={"answer": types.Schema(type="STRING")},
+    )
+    result = schema_to_json_schema(schema)
+    assert result == {
+        "type": "object",
+        "properties": {"answer": {"type": "string"}},
+    }
+
 
 class TestLowercaseSchemaTypes:
   """Tests for lowercase_schema_types function."""