Update close tool + add output to agent result (#1505)

# Why

Previously, agent executions could end without a structured final
response - either by hitting `maxSteps` or the LLM breaking out of its
loop without calling the `close` tool. This made it difficult to:
1. Reliably determine if a task was completed successfully
2. Extract structured data from the agent's execution

# What Changed

### Ensured Close Tool is Always Called
- Added `handleCloseToolCall` utility that forces a `close` tool call
via a separate `generateText` call when the main agent loop ends without
explicitly closing
- Integrated via new `ensureClosed` private method in
`v3AgentHandler.ts`
- Works for both `execute()` and `stream()` modes
- Triggers when `maxSteps` is reached or the LLM stops ( completes its
task)
### Added Output Schema Support (Experimental)
- Users can now pass a Zod schema to `agent.execute({ output:
z.object({...}) })` to return structured data at the end of execution
- The schema dynamically extends the close tool's input schema
- Extracted data is returned in `result.output`
- Added validation:
  - **CUA mode**: Throws `StagehandInvalidArgumentError` (not supported)
- **Non-CUA without `experimental: true`**: Throws
`ExperimentalNotConfiguredError`

### Example Usage
```typescript
const result = await agent.execute({
  instruction: "search for a shampoo on amazon and click into one of the results",
  maxSteps: 20,
  output: z.object({
    productName: z.string().describe("The name of the shampoo product"),
    price: z.string().describe("The price of the product"),
    rating: z.string().describe("The star rating of the product"),
  }),
});

console.log(result.output);
// { productName: "...", price: "$12.99", rating: "4.5 out of 5 stars" }
```

# Test Plan

- [x] Verify close tool is called when agent naturally completes (no
change in behavior)
- [x] Verify close tool is forced when `maxSteps` is reached
- [x] Verify `output` schema extracts data correctly in `execute()` mode
- [x] Verify `output` schema extracts data correctly in `stream()` mode
- [x] Verify `output` throws `StagehandInvalidArgumentError` when used
with CUA mode
- [x] Verify `output` throws `ExperimentalNotConfiguredError` when used
without `experimental: true`


<!-- This is an auto-generated description by cubic. -->
---
## Summary by cubic
Ensures every agent run ends with a structured final response and adds
optional structured output via a Zod schema. Improves reliability by
always setting completion status and final reasoning.

- **New Features**
- Always triggers a "close" tool call at the end of a run (LLM stops or
maxSteps), for both execute() and stream().
- Optional output schema: pass output: z.object({...}) to return typed
data in result.output.
- Validation: output schema is not supported in CUA (throws
StagehandInvalidArgumentError). In non-CUA, requires experimental: true
(throws ExperimentalNotConfiguredError otherwise).
- Removed "close" from the main tool list and system prompt; closing is
handled automatically post-run.

<sup>Written for commit dfb703a.
Summary will update on new commits.</sup>

<!-- End of auto-generated description by cubic. -->

---------

Co-authored-by: Nick Sweeting <github@sweeting.me>

Author: tkattkat

.changeset/quick-games-try.md

@@ -0,0 +1,5 @@
+---
+"@browserbasehq/stagehand": patch
+---
+
+Add structured output to agent result + ensure close tool is always called

packages/core/lib/v3/agent/prompts/agentSystemPrompt.ts

@@ -67,7 +67,6 @@ function buildToolsSection(
     { name: "wait", description: "Wait for a specified time" },
     { name: "navback", description: "Navigate back in browser history" },
     { name: "scroll", description: "Scroll the page x pixels up or down" },
-    { name: "close", description: "Mark the task as complete or failed" },
   ];
 
   const domTools: ToolDefinition[] = [
@@ -92,7 +91,6 @@ function buildToolsSection(
     { name: "wait", description: "Wait for a specified time" },
     { name: "navback", description: "Navigate back in browser history" },
     { name: "scroll", description: "Scroll the page x pixels up or down" },
-    { name: "close", description: "Mark the task as complete or failed" },
   ];
 
   const baseTools = isHybridMode ? hybridTools : domTools;
@@ -224,8 +222,6 @@ export function buildAgentSystemPrompt(
     <item>Always start by understanding the current page state</item>
     <item>Use the screenshot tool to verify page state when needed</item>
     <item>Use appropriate tools for each action</item>
-    <item>When the task is complete, use the "close" tool with taskComplete: true</item>
-    <item>If the task cannot be completed, use "close" with taskComplete: false</item>
   </guidelines>
   ${pageUnderstandingProtocol}
   <navigation>

packages/core/lib/v3/agent/tools/close.ts

@@ -3,7 +3,6 @@ import { actTool } from "./act";
 import { screenshotTool } from "./screenshot";
 import { waitTool } from "./wait";
 import { navBackTool } from "./navback";
-import { closeTool } from "./close";
 import { ariaTreeTool } from "./ariaTree";
 import { fillFormTool } from "./fillform";
 import { scrollTool, scrollVisionTool } from "./scroll";
@@ -87,7 +86,7 @@ export function createAgentTools(v3: V3, options?: V3AgentToolOptions) {
     ariaTree: ariaTreeTool(v3),
     click: clickTool(v3, provider),
     clickAndHold: clickAndHoldTool(v3, provider),
-    close: closeTool(),
+    //close: closeTool(),
     dragAndDrop: dragAndDropTool(v3, provider),
     extract: extractTool(v3, executionModel, options?.logger),
     fillForm: fillFormTool(v3, executionModel),
@@ -121,7 +120,6 @@ export type AgentToolTypesMap = {
   ariaTree: ReturnType<typeof ariaTreeTool>;
   click: ReturnType<typeof clickTool>;
   clickAndHold: ReturnType<typeof clickAndHoldTool>;
-  close: ReturnType<typeof closeTool>;
   dragAndDrop: ReturnType<typeof dragAndDropTool>;
   extract: ReturnType<typeof extractTool>;
   fillForm: ReturnType<typeof fillFormTool>;

packages/core/lib/v3/agent/tools/index.ts

@@ -0,0 +1,128 @@
+import { generateText, ModelMessage, LanguageModel, ToolSet } from "ai";
+import { z } from "zod";
+import { tool } from "ai";
+import { LogLine } from "../../types/public/logs";
+import { StagehandZodObject } from "../../zodCompat";
+interface CloseResult {
+  reasoning: string;
+  taskComplete: boolean;
+  messages: ModelMessage[];
+  output?: Record<string, unknown>;
+}
+
+const baseCloseSchema = z.object({
+  reasoning: z
+    .string()
+    .describe("Brief summary of what actions were taken and the outcome"),
+  taskComplete: z
+    .boolean()
+    .describe("true if the task was fully completed, false otherwise"),
+});
+
+/**
+ * Force a close tool call at the end of an agent run.
+ * This ensures we always get a structured final response,
+ * even if the main loop ended without calling close.
+ */
+export async function handleCloseToolCall(options: {
+  model: LanguageModel;
+  inputMessages: ModelMessage[];
+  instruction: string;
+  outputSchema?: StagehandZodObject;
+  logger: (message: LogLine) => void;
+}): Promise<CloseResult> {
+  const { model, inputMessages, instruction, outputSchema, logger } = options;
+
+  logger({
+    category: "agent",
+    message: "Agent calling tool: close",
+    level: 1,
+  });
+  // Merge base close schema with user-provided output schema if present
+  const closeToolSchema = outputSchema
+    ? baseCloseSchema.extend({
+        output: outputSchema.describe(
+          "The specific data the user requested from this task",
+        ),
+      })
+    : baseCloseSchema;
+
+  const outputInstructions = outputSchema
+    ? `\n\nThe user also requested the following information from this task. Provide it in the "output" field:\n${JSON.stringify(
+        Object.fromEntries(
+          Object.entries(outputSchema.shape).map(([key, value]) => [
+            key,
+            value.description || "no description",
+          ]),
+        ),
+        null,
+        2,
+      )}`
+    : "";
+
+  const systemPrompt = `You are a web automation assistant that was tasked with completing a task.
+
+The task was:
+"${instruction}"
+
+Review what was accomplished and provide your final assessment in whether the task was completed successfully. you have been provided with the history of the actions taken so far, use this to determine if the task was completed successfully.${outputInstructions}
+
+Call the "close" tool with:
+1. A brief summary of what was done
+2. Whether the task was completed successfully${outputSchema ? "\n3. The requested output data based on what you found" : ""}`;
+
+  const closeTool = tool({
+    description: outputSchema
+      ? "Complete the task with your assessment and the requested output data."
+      : "Complete the task with your final assessment.",
+    inputSchema: closeToolSchema,
+    execute: async (params) => {
+      return { success: true, ...params };
+    },
+  });
+
+  const userPrompt: ModelMessage = {
+    role: "user",
+    content: outputSchema
+      ? "Provide your final assessment and the requested output data."
+      : "Provide your final assessment.",
+  };
+
+  const result = await generateText({
+    model,
+    system: systemPrompt,
+    messages: [...inputMessages, userPrompt],
+    tools: { close: closeTool } as ToolSet,
+    toolChoice: { type: "tool", toolName: "close" },
+  });
+
+  const closeToolCall = result.toolCalls.find((tc) => tc.toolName === "close");
+  const outputMessages: ModelMessage[] = [
+    userPrompt,
+    ...(result.response?.messages || []),
+  ];
+
+  if (!closeToolCall) {
+    return {
+      reasoning: result.text || "Task execution completed",
+      taskComplete: false,
+      messages: outputMessages,
+    };
+  }
+
+  const input = closeToolCall.input as z.infer<typeof baseCloseSchema> & {
+    output?: Record<string, unknown>;
+  };
+  logger({
+    category: "agent",
+    message: `Task completed`,
+    level: 1,
+  });
+
+  return {
+    reasoning: input.reasoning,
+    taskComplete: input.taskComplete,
+    messages: outputMessages,
+    output: input.output,
+  };
+}

packages/core/lib/v3/agent/utils/handleCloseToolCall.ts

@@ -21,10 +21,10 @@ export interface AgentValidationOptions {
  * Validates agent configuration and experimental feature usage.
  *
  * This utility consolidates all validation checks for both CUA and non-CUA agent paths:
- * - Invalid argument errors for CUA (streaming, abort signal, message continuation, excludeTools are not supported)
+ * - Invalid argument errors for CUA (streaming, abort signal, message continuation, excludeTools, output schema are not supported)
  * - Experimental feature checks for integrations and tools (both CUA and non-CUA)
  * - Experimental feature checks for hybrid mode (requires experimental: true)
- * - Experimental feature checks for non-CUA only (callbacks, signal, messages, streaming, excludeTools)
+ * - Experimental feature checks for non-CUA only (callbacks, signal, messages, streaming, excludeTools, output schema)
  *
  * Throws StagehandInvalidArgumentError for invalid/unsupported configurations.
  * Throws ExperimentalNotConfiguredError if experimental features are used without experimental mode.
@@ -56,6 +56,9 @@ export function validateExperimentalFeatures(
     ) {
       unsupportedFeatures.push("excludeTools");
     }
+    if (executeOptions?.output) {
+      unsupportedFeatures.push("output schema");
+    }
 
     if (unsupportedFeatures.length > 0) {
       throw new StagehandInvalidArgumentError(
@@ -97,6 +100,9 @@ export function validateExperimentalFeatures(
     if (executeOptions.excludeTools && executeOptions.excludeTools.length > 0) {
       features.push("excludeTools");
     }
+    if (executeOptions.output) {
+      features.push("output schema");
+    }
   }
 
   if (features.length > 0) {