feat: Add strict input schema validation for LlmAgent workflow nodes

In preparation for natively supporting task-mode agents as Workflow nodes, we need to enforce input boundaries at the wrapper level. Currently, bad `node_input` bypasses validation and generates invalid JSON strings for the LLM.

Co-authored-by: Shangjie Chen <deanchen@google.com>
PiperOrigin-RevId: 944894496

Author: DeanChensj

src/google/adk/utils/_schema_utils.py

@@ -125,9 +125,81 @@ def validate_schema(schema: SchemaType, json_text: str) -> Any:
   elif is_list_of_basemodel(schema):
     # For list[BaseModel], use TypeAdapter to validate
     type_adapter = TypeAdapter(schema)
-    validated = type_adapter.validate_json(json_text)
+    validated: list[Any] = type_adapter.validate_json(json_text)
     return [item.model_dump(exclude_none=True) for item in validated]
   else:
     # For other schema types (list[str], dict, Schema, etc.),
-    # just parse JSON without pydantic validation
     return json.loads(json_text)
+
+
+def validate_node_data(
+    schema: Optional[SchemaType],
+    data: Any,
+    *,
+    preserve_content: bool = False,
+) -> Any:
+  """Validates and sanitizes node input or output data against a schema."""
+  if data is None or schema is None:
+    return data
+
+  if isinstance(schema, (dict, types.Schema)):
+    return data
+
+  def _to_serializable(val: Any) -> Any:
+    if isinstance(val, BaseModel):
+      return val.model_dump(exclude_none=True)
+    if isinstance(val, list):
+      return [_to_serializable(item) for item in val]
+    if isinstance(val, dict):
+      return {k: _to_serializable(v) for k, v in val.items()}
+    return val
+
+  def _validate_python_object(val: Any) -> Any:
+    validated: Any = TypeAdapter(schema).validate_python(val)
+    return _to_serializable(validated)
+
+  # If schema expects Content, do not unwrap
+  if isinstance(schema, type) and issubclass(schema, types.Content):
+    return _validate_python_object(data)
+  if schema is types.Content:
+    return _validate_python_object(data)
+
+  if isinstance(data, types.Content):
+    # Extract text part
+    text_parts = [p.text for p in data.parts if p.text] if data.parts else []
+    text_str = "".join(text_parts)
+
+    # Validate the text
+    if schema is str:
+      validated_payload = text_str
+    else:
+      # Try to parse text as JSON first
+      try:
+        parsed_json = json.loads(text_str)
+        validated_payload = _validate_python_object(parsed_json)
+      except json.JSONDecodeError:
+        # Fallback to validate raw string
+        validated_payload = _validate_python_object(text_str)
+
+    if not preserve_content:
+      return validated_payload
+
+    # Re-wrap in Content
+    new_parts = [p for p in data.parts if not p.text] if data.parts else []
+    new_parts.append(
+        types.Part(
+            text=json.dumps(validated_payload)
+            if not isinstance(validated_payload, str)
+            else validated_payload
+        )
+    )
+    return types.Content(role=data.role, parts=new_parts)
+
+  # If data is a string (but not wrapped in Content)
+  if isinstance(data, str):
+    if schema is str:
+      return data
+    return _validate_python_object(data)
+
+  # For any other Python object (dict, BaseModel instance, etc.)
+  return _validate_python_object(data)

src/google/adk/workflow/_base_node.py

@@ -19,16 +19,13 @@
 from typing import final
 from typing import TYPE_CHECKING
 
-from google.genai import types
 from pydantic import BaseModel
 from pydantic import ConfigDict
 from pydantic import Field
 from pydantic import field_validator
-from pydantic import TypeAdapter
-from pydantic import ValidationError
 
 from ..utils._schema_utils import SchemaType
-from ..utils.content_utils import extract_text_from_content
+from ..utils._schema_utils import validate_node_data
 from ._retry_config import RetryConfig
 
 if TYPE_CHECKING:
@@ -123,64 +120,16 @@ def _validate_name(cls, v: str) -> str:
   """
 
   def _validate_schema(self, data: Any, schema: Any) -> Any:
-    """Validates data against a schema using ``TypeAdapter``.
-
-    Handles BaseModel, list[BaseModel], primitive types, unions, and
-    generic aliases.  Descriptive schemas (``types.Schema``, raw
-    ``dict``) are skipped.  Any BaseModel instances in the validated
-    result are converted to dicts via ``model_dump()`` to keep
-    ``Event.output`` JSON-serializable.
-    """
-    if data is None or schema is None:
-      return data
-
-    if isinstance(schema, (dict, types.Schema)):
-      return data
-
-    validated = TypeAdapter(schema).validate_python(data)
-    return self._to_serializable(validated)
+    """Validates data against a schema using validate_node_data helper."""
+    return validate_node_data(schema, data)
 
   def _validate_input_data(self, data: Any) -> Any:
     """Validates data against input_schema if set."""
-    if self.input_schema and isinstance(data, types.Content):
-      # Extract text from Content (e.g. user input from START node).
-      text = extract_text_from_content(data)
-      if self.input_schema is str:
-        return text
-      # If schema is defined, try to parse the text as JSON.
-      try:
-        return TypeAdapter(self.input_schema).validate_json(text)
-      except ValidationError:
-        # Fallback to validate_python if it's a raw string matching the schema.
-        try:
-          return TypeAdapter(self.input_schema).validate_python(text)
-        except ValidationError:
-          pass
-    return self._validate_schema(data, self.input_schema)
+    return validate_node_data(self.input_schema, data, preserve_content=False)
 
   def _validate_output_data(self, data: Any) -> Any:
     """Validates data against output_schema if set."""
-    if not self.output_schema:
-      return data
-
-    # 1. Try standard validation first
-    try:
-      return self._validate_schema(data, self.output_schema)
-    except ValidationError as e:
-      # 2. If failed, try to parse JSON ONLY if it's Content
-      if isinstance(data, types.Content):
-        text = extract_text_from_content(data)
-        if self.output_schema is str:
-          return text
-        if text.strip():
-          try:
-            validated = TypeAdapter(self.output_schema).validate_json(text)
-            return self._to_serializable(validated)
-          except ValidationError:
-            pass
-
-      # 3. If not Content or parsing failed, re-raise original error
-      raise e
+    return validate_node_data(self.output_schema, data, preserve_content=False)
 
   @staticmethod
   def _to_serializable(data: Any) -> Any:

src/google/adk/workflow/_dynamic_node_scheduler.py

@@ -170,10 +170,9 @@ async def __call__(
       try:
         node_input = node._validate_input_data(node_input)
       except ValidationError as e:
-        raise ValueError(
-            'Runtime schema validation failed for dynamic node'
-            f" '{node_name or node.name}'. Input does not match"
-            f' input_schema: {e}'
+        raise ValidationError.from_exception_data(
+            title=f"dynamic node '{node_name or node.name}'",
+            line_errors=e.errors(),  # type: ignore[arg-type]
         ) from e
 
     logger.debug('node %s schedule start.', node_path)

tests/unittests/utils/test_schema_utils.py

@@ -17,8 +17,12 @@
 from google.adk.utils._schema_utils import get_list_inner_type
 from google.adk.utils._schema_utils import is_basemodel_schema
 from google.adk.utils._schema_utils import is_list_of_basemodel
+from google.adk.utils._schema_utils import validate_node_data
 from google.adk.utils._schema_utils import validate_schema
+from google.genai import types
 from pydantic import BaseModel
+from pydantic import ValidationError
+import pytest
 
 
 class SampleModel(BaseModel):
@@ -144,3 +148,63 @@ def test_dict_schema(self):
     json_text = '{"key1": 1, "key2": 2}'
     result = validate_schema(dict[str, int], json_text)
     assert result == {'key1': 1, 'key2': 2}
+
+
+class TestValidateNodeData:
+  """Tests for validate_node_data function."""
+
+  def test_none_schema_or_data_returns_data(self):
+    """Bypasses validation if schema or data is None."""
+    assert validate_node_data(None, 'some_data') == 'some_data'
+    assert validate_node_data(SampleModel, None) is None
+
+  def test_dict_or_types_schema_returns_data(self):
+    """Bypasses validation if schema is dict or types.Schema."""
+    assert validate_node_data({'key': int}, 'some_data') == 'some_data'
+    # Mock types.Schema
+    schema = types.Schema(type=types.Type.STRING)
+    assert validate_node_data(schema, 'some_data') == 'some_data'
+
+  def test_content_schema_returns_data(self):
+    """Bypasses validation if target schema is types.Content or subclass."""
+    result = validate_node_data(
+        types.Content, types.Content(role='user', parts=[])
+    )
+    assert result == {'role': 'user', 'parts': []}
+
+  def test_plain_basemodel_schema_validates_raw_dict(self):
+    """Validates raw dict data against BaseModel schema."""
+    result = validate_node_data(SampleModel, {'name': 'test', 'value': 42})
+    assert result == {'name': 'test', 'value': 42}
+
+  def test_content_data_and_preserve_content(self):
+    """Validates wrapped content and wraps result back into Content."""
+    data = types.Content(
+        role='user',
+        parts=[types.Part(text='{"name": "test", "value": 42}')],
+    )
+    result = validate_node_data(SampleModel, data, preserve_content=True)
+    assert isinstance(result, types.Content)
+    assert result.role == 'user'
+    assert len(result.parts) == 1
+    assert result.parts[0].text == '{"name": "test", "value": 42}'
+
+  def test_content_data_no_preserve_content(self):
+    """Validates wrapped content and returns unwrapped dictionary."""
+    data = types.Content(
+        role='user',
+        parts=[types.Part(text='{"name": "test", "value": 42}')],
+    )
+    result = validate_node_data(SampleModel, data, preserve_content=False)
+    assert isinstance(result, dict)
+    assert result == {'name': 'test', 'value': 42}
+
+  def test_raw_json_string_validated_against_basemodel_schema(self):
+    """Raw JSON string fails validation against BaseModel schema (not auto-parsed)."""
+    with pytest.raises(ValidationError):
+      validate_node_data(SampleModel, '{"name": "test", "value": 42}')
+
+  def test_raw_string_not_parsed_with_str_schema(self):
+    """Bypasses JSON parsing if schema is str."""
+    result = validate_node_data(str, 'hello')
+    assert result == 'hello'

tests/unittests/workflow/test_dynamic_node_scheduler.py

@@ -32,6 +32,7 @@
 from google.adk.workflow._node_status import NodeStatus
 from google.adk.workflow._workflow import _LoopState
 from pydantic import BaseModel
+from pydantic import ValidationError
 import pytest
 
 # --- Fixtures ---
@@ -639,10 +640,7 @@ async def test_runtime_schema_validation_raises():
 
   node = BaseNode(name='child', input_schema=_ModelA)
 
-  with pytest.raises(
-      ValueError,
-      match=r"Runtime schema validation failed for dynamic node 'child'",
-  ):
+  with pytest.raises(ValidationError):
     await scheduler(
         ctx,
         node,

tests/unittests/workflow/test_llm_agent_as_node.py

@@ -1361,3 +1361,67 @@ async def test_three_layer_llm_agent_transfer_round_trip(
       if p.text
   ]
   assert any('Welcome back to root!' in t for t in content_texts3)
+
+
+@pytest.mark.asyncio
+async def test_workflow_node_with_valid_input_schema_completes_successfully(
+    request: pytest.FixtureRequest,
+):
+  """A valid node_input payload successfully passes schema validation."""
+
+  # Arrange
+  class InputSchema(BaseModel):
+    required_field: str
+
+  agent = LlmAgent(
+      name='schema_agent',
+      model='test_model',
+      input_schema=InputSchema,
+      instruction='Just say hi',
+      mode='single_turn',
+  )
+  wrapper = build_node(agent)
+  wf = Workflow(
+      name='test_workflow',
+      edges=[('START', wrapper)],
+  )
+  runner = _new_workflow_runner(wf, request.function.__name__)
+  agent_clone = next(n for n in wf.graph.nodes if n.name == wrapper.name)
+
+  # Act
+  with _mock_agent_run(agent_clone, content_text='hi'):
+    events = await runner.run_async('{"required_field": "hello"}')
+
+  # Assert
+  assert len(events) > 0
+
+
+@pytest.mark.asyncio
+async def test_workflow_node_with_invalid_input_schema_raises_validation_error(
+    request: pytest.FixtureRequest,
+):
+  """An invalid node_input payload raises a pydantic ValidationError."""
+
+  # Arrange
+  class InputSchema(BaseModel):
+    required_field: str
+
+  agent = LlmAgent(
+      name='schema_agent',
+      model='test_model',
+      input_schema=InputSchema,
+      instruction='Just say hi',
+      mode='single_turn',
+  )
+  wrapper = build_node(agent)
+  wf = Workflow(
+      name='test_workflow',
+      edges=[('START', wrapper)],
+  )
+  runner = _new_workflow_runner(wf, request.function.__name__)
+  agent_clone = next(n for n in wf.graph.nodes if n.name == wrapper.name)
+
+  # Act / Assert
+  with _mock_agent_run(agent_clone, content_text='hi'):
+    with pytest.raises(ValidationError):
+      await runner.run_async('{"wrong_field": "hello"}')

tests/unittests/workflow/test_workflow_failures.py

@@ -1082,13 +1082,13 @@ async def test_fail_fast_preserves_completed_siblings(
   async def succeeding_node(ctx: Context):
     nonlocal node_success_started, node_success_completed
     node_success_started = True
-    await asyncio.sleep(0.1)
+    await asyncio.sleep(0)
     node_success_completed = True
     return 'success_output'
 
   @node()
   async def failing_node(ctx: Context):
-    await asyncio.sleep(0.1)
+    await asyncio.sleep(0)
     raise ValueError('Fail')
 
   wf = Workflow(

tests/unittests/workflow/test_workflow_schema.py

@@ -631,10 +631,7 @@ async def _run_impl(
   msg = types.Content(parts=[types.Part(text='hello')], role='user')
 
   # We expect it to raise ValidationError
-  with pytest.raises(
-      ValueError,
-      match=r"Runtime schema validation failed for dynamic node 'node'",
-  ):
+  with pytest.raises(ValidationError):
     async for event in runner.run_async(
         user_id='u', session_id=session.id, new_message=msg
     ):