Add function processOutput to AgentDefinition and typing for an agent's output (#10447)

packages/core/src/agents/executor.ts

@@ -14,6 +14,7 @@ import type {
   FunctionCall,
   GenerateContentConfig,
   FunctionDeclaration,
+  Schema,
 } from '@google/genai';
 import { executeToolCall } from '../core/nonInteractiveToolExecutor.js';
 import { ToolRegistry } from '../tools/tool-registry.js';
@@ -36,6 +37,8 @@ import type {
 import { AgentTerminateMode } from './types.js';
 import { templateString } from './utils.js';
 import { parseThought } from '../utils/thoughtUtils.js';
+import { type z } from 'zod';
+import { zodToJsonSchema } from 'zod-to-json-schema';
 
 /** A callback function to report on agent activity. */
 export type ActivityCallback = (activity: SubagentActivityEvent) => void;
@@ -48,8 +51,8 @@ const TASK_COMPLETE_TOOL_NAME = 'complete_task';
  * This executor runs the agent in a loop, calling tools until it calls the
  * mandatory `complete_task` tool to signal completion.
  */
-export class AgentExecutor {
-  readonly definition: AgentDefinition;
+export class AgentExecutor<TOutput extends z.ZodTypeAny> {
+  readonly definition: AgentDefinition<TOutput>;
 
   private readonly agentId: string;
   private readonly toolRegistry: ToolRegistry;
@@ -67,11 +70,11 @@ export class AgentExecutor {
    * @param onActivity An optional callback to receive activity events.
    * @returns A promise that resolves to a new `AgentExecutor` instance.
    */
-  static async create(
-    definition: AgentDefinition,
+  static async create<TOutput extends z.ZodTypeAny>(
+    definition: AgentDefinition<TOutput>,
     runtimeContext: Config,
     onActivity?: ActivityCallback,
-  ): Promise<AgentExecutor> {
+  ): Promise<AgentExecutor<TOutput>> {
     // Create an isolated tool registry for this agent instance.
     const agentToolRegistry = new ToolRegistry(runtimeContext);
     const parentToolRegistry = await runtimeContext.getToolRegistry();
@@ -116,7 +119,7 @@ export class AgentExecutor {
    * instantiate the class.
    */
   private constructor(
-    definition: AgentDefinition,
+    definition: AgentDefinition<TOutput>,
     runtimeContext: Config,
     toolRegistry: ToolRegistry,
     onActivity?: ActivityCallback,
@@ -397,7 +400,36 @@ export class AgentExecutor {
         if (outputConfig) {
           const outputName = outputConfig.outputName;
           if (args[outputName] !== undefined) {
-            submittedOutput = String(args[outputName]);
+            const outputValue = args[outputName];
+            const validationResult = outputConfig.schema.safeParse(outputValue);
+
+            if (!validationResult.success) {
+              taskCompleted = false; // Validation failed, revoke completion
+              const error = `Output validation failed: ${JSON.stringify(validationResult.error.flatten())}`;
+              syncResponseParts.push({
+                functionResponse: {
+                  name: TASK_COMPLETE_TOOL_NAME,
+                  response: { error },
+                  id: callId,
+                },
+              });
+              this.emitActivity('ERROR', {
+                context: 'tool_call',
+                name: functionCall.name,
+                error,
+              });
+              continue;
+            }
+
+            const validatedOutput = validationResult.data;
+            if (this.definition.processOutput) {
+              submittedOutput = this.definition.processOutput(validatedOutput);
+            } else {
+              submittedOutput =
+                typeof outputValue === 'string'
+                  ? outputValue
+                  : JSON.stringify(outputValue, null, 2);
+            }
             syncResponseParts.push({
               functionResponse: {
                 name: TASK_COMPLETE_TOOL_NAME,
@@ -573,10 +605,14 @@ export class AgentExecutor {
     };
 
     if (outputConfig) {
-      completeTool.parameters!.properties![outputConfig.outputName] = {
-        description: outputConfig.description,
-        ...(outputConfig.schema ?? { type: Type.STRING }),
-      };
+      const jsonSchema = zodToJsonSchema(outputConfig.schema);
+      const {
+        $schema: _$schema,
+        definitions: _definitions,
+        ...schema
+      } = jsonSchema;
+      completeTool.parameters!.properties![outputConfig.outputName] =
+        schema as Schema;
       completeTool.parameters!.required!.push(outputConfig.outputName);
     }