Skip to content

feat: implement runtime schema validation for AgentTool#456

Closed
AmaadMartin wants to merge 1 commit into
google:mainfrom
AmaadMartin:feat/agent-tool-schema-validation
Closed

feat: implement runtime schema validation for AgentTool#456
AmaadMartin wants to merge 1 commit into
google:mainfrom
AmaadMartin:feat/agent-tool-schema-validation

Conversation

@AmaadMartin

Copy link
Copy Markdown
Collaborator

Description

This PR implements runtime input and output schema validation for agents wrapped as tools using AgentTool.

Key Changes

  • On-Demand Schema Conversion: Modified LlmAgent to store original schemas (inputSchema and outputSchema as LlmAgentSchema which can be Zod or GenAI Schema) directly, rather than eagerly converting them to GenAI Schema in the constructor. Conversion is now performed on-demand when needed:
    • In AgentTool._getDeclaration() for tool declaration.
    • In BasicLlmRequestProcessor for passing to model config.
  • Runtime Validation in AgentTool:
    • Input Validation: Before executing the sub-agent, AgentTool.runAsync validates the input arguments against the agent's inputSchema (if it's a Zod object). Throws an error on validation failure.
    • Output Validation: After executing the sub-agent, AgentTool.runAsync parses the output as JSON and validates it against the agent's outputSchema (if it's a Zod object). Throws an error on validation failure.
  • Improved Helper: Updated zodObjectToSchema helper to accept and return Schema as-is if it is not a Zod object, simplifying call sites.

Testing

  • Added unit tests in core/test/tools/agent_tool_test.ts covering:
    • Input validation success and failure.
    • Output validation success and failure.
    • Bypassing validation for raw GenAI schemas.
  • Added integration tests in tests/integration/tools/agent_tool_test.ts to verify:
    • Schema validation in full runner flow.
    • Error propagation (validation errors wrapped in tool response instead of crashing).
  • Verified E2E behavior using a manual scratch script on Vertex AI.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant