fix: allow andAgent schema to accept Output.* specs (arrays, choi…

.changeset/tall-dryers-run.md

@@ -0,0 +1,32 @@
+---
+"@voltagent/core": patch
+---
+
+fix: allow `andAgent` schema to accept `Output.*` specs (arrays, choices, json, text)
+
+`andAgent` now supports the same output flexibility as `agent.generateText`, so you can return non-object
+structures from workflow steps.
+
+Usage:
+
+```ts
+import { Output } from "ai";
+import { z } from "zod";
+import { Agent, createWorkflowChain } from "@voltagent/core";
+
+const agent = new Agent({
+  name: "Tagger",
+  model: "openai/gpt-4o-mini",
+  instructions: "Return tags only.",
+});
+
+const workflow = createWorkflowChain({
+  id: "tag-workflow",
+  input: z.object({ topic: z.string() }),
+}).andAgent(async ({ data }) => `List tags for ${data.topic}`, agent, {
+  schema: Output.array({ element: z.string() }),
+});
+
+const result = await workflow.run({ topic: "workflows" });
+// result: string[]
+```

packages/core/src/workflow/chain.spec-d.ts

@@ -1,3 +1,4 @@
+import { Output } from "ai";
 import { describe, expectTypeOf, it } from "vitest";
 import { z } from "zod";
 import type { Agent } from "../agent/agent";
@@ -443,6 +444,27 @@ describe("workflow chain - type inference", () => {
       expectTypeOf(workflow).not.toBeNever();
     });
 
+    it("should infer andAgent output from output specs", () => {
+      const workflow = createWorkflowChain({
+        id: "test-agent-output",
+        name: "Test Agent Output",
+        input: z.object({ prompt: z.string() }),
+        result: z.array(z.string()),
+      })
+        .andAgent(async ({ data }) => `List: ${data.prompt}`, mockAgent, {
+          schema: Output.array({ element: z.string() }),
+        })
+        .andThen({
+          id: "ensure-array",
+          execute: async ({ data }) => {
+            expectTypeOf(data).toEqualTypeOf<string[]>();
+            return data;
+          },
+        });
+
+      expectTypeOf(workflow).not.toBeNever();
+    });
+
     it("should allow mapping andAgent output into existing data", () => {
       const workflow = createWorkflowChain({
         id: "test-agent-map",

packages/core/src/workflow/chain.ts

@@ -3,7 +3,7 @@ import type { Logger } from "@voltagent/internal";
 import type { DangerouslyAllowAny } from "@voltagent/internal/types";
 import type { UIMessage } from "ai";
 import type { z } from "zod";
-import type { Agent, BaseGenerationOptions } from "../agent/agent";
+import type { Agent } from "../agent/agent";
 import { createWorkflow } from "./core";
 import type {
   InternalAnyWorkflowStep,
@@ -42,6 +42,7 @@ import {
   andWhen,
   andWorkflow,
 } from "./steps";
+import type { AgentConfig, AgentOutputSchema, InferAgentOutput } from "./steps/and-agent";
 import type { InternalWorkflow } from "./steps/types";
 import type {
   Workflow,
@@ -55,23 +56,7 @@ import type {
   WorkflowStreamWriter,
 } from "./types";
 
-/**
- * Agent configuration for the chain
- */
-export type AgentConfig<
-  SCHEMA extends z.ZodTypeAny,
-  INPUT_SCHEMA extends InternalBaseWorkflowInputSchema,
-  CURRENT_DATA,
-> = BaseGenerationOptions & {
-  schema:
-    | SCHEMA
-    | ((
-        context: Omit<
-          WorkflowExecuteContext<WorkflowInput<INPUT_SCHEMA>, CURRENT_DATA, any, any>,
-          "suspend" | "writer"
-        >,
-      ) => SCHEMA | Promise<SCHEMA>);
-};
+export type { AgentConfig } from "./steps/and-agent";
 
 /**
  * A workflow chain that provides a fluent API for building workflows
@@ -164,7 +149,7 @@ export class WorkflowChain<
    * @param map - Optional mapper to shape or merge the agent output with existing data
    * @returns A workflow step that executes the agent with the task
    */
-  andAgent<SCHEMA extends z.ZodTypeAny>(
+  andAgent<SCHEMA extends AgentOutputSchema>(
     task:
       | string
       | UIMessage[]
@@ -177,10 +162,16 @@ export class WorkflowChain<
           any
         >,
     agent: Agent,
-    config: AgentConfig<SCHEMA, INPUT_SCHEMA, CURRENT_DATA>,
-  ): WorkflowChain<INPUT_SCHEMA, RESULT_SCHEMA, z.infer<SCHEMA>, SUSPEND_SCHEMA, RESUME_SCHEMA>;
+    config: AgentConfig<SCHEMA, WorkflowInput<INPUT_SCHEMA>, CURRENT_DATA>,
+  ): WorkflowChain<
+    INPUT_SCHEMA,
+    RESULT_SCHEMA,
+    InferAgentOutput<SCHEMA>,
+    SUSPEND_SCHEMA,
+    RESUME_SCHEMA
+  >;
 
-  andAgent<SCHEMA extends z.ZodTypeAny, NEW_DATA>(
+  andAgent<SCHEMA extends AgentOutputSchema, NEW_DATA>(
     task:
       | string
       | UIMessage[]
@@ -193,9 +184,9 @@ export class WorkflowChain<
           any
         >,
     agent: Agent,
-    config: AgentConfig<SCHEMA, INPUT_SCHEMA, CURRENT_DATA>,
+    config: AgentConfig<SCHEMA, WorkflowInput<INPUT_SCHEMA>, CURRENT_DATA>,
     map: (
-      output: z.infer<SCHEMA>,
+      output: InferAgentOutput<SCHEMA>,
       context: WorkflowExecuteContext<WorkflowInput<INPUT_SCHEMA>, CURRENT_DATA, any, any>,
     ) => Promise<NEW_DATA> | NEW_DATA,
   ): WorkflowChain<INPUT_SCHEMA, RESULT_SCHEMA, NEW_DATA, SUSPEND_SCHEMA, RESUME_SCHEMA>;
@@ -213,7 +204,7 @@ export class WorkflowChain<
           any
         >,
     agent: Agent,
-    config: AgentConfig<z.ZodTypeAny, INPUT_SCHEMA, CURRENT_DATA>,
+    config: AgentConfig<AgentOutputSchema, WorkflowInput<INPUT_SCHEMA>, CURRENT_DATA>,
     map?: (
       output: unknown,
       context: WorkflowExecuteContext<WorkflowInput<INPUT_SCHEMA>, CURRENT_DATA, any, any>,

packages/core/src/workflow/steps/and-agent.spec.ts

@@ -1,4 +1,5 @@
 import { safeStringify } from "@voltagent/internal";
+import { Output } from "ai";
 import { describe, expect, it } from "vitest";
 import { z } from "zod";
 import {
@@ -49,4 +50,32 @@ describe("andAgent", () => {
       emailType: { type: "support", priority: "high" },
     });
   });
+
+  it("supports output specs for non-object schemas", async () => {
+    const agent = createTestAgent({
+      model: createMockLanguageModel({
+        doGenerate: {
+          ...defaultMockResponse,
+          content: [
+            {
+              type: "text" as const,
+              text: safeStringify({ elements: ["alpha", "beta"] }),
+            },
+          ],
+        },
+      }),
+    });
+
+    const step = andAgent(({ data }) => `List tags for ${data.topic}`, agent, {
+      schema: Output.array({ element: z.string() }),
+    });
+
+    const result = await step.execute(
+      createMockWorkflowExecuteContext({
+        data: { topic: "workflow" },
+      }),
+    );
+
+    expect(result).toEqual(["alpha", "beta"]);
+  });
 });

packages/core/src/workflow/steps/and-agent.ts

@@ -1,24 +1,47 @@
 import type { ModelMessage } from "@ai-sdk/provider-utils";
-import { Output, type UIMessage } from "ai";
+import { type InferGenerateOutput, Output, type UIMessage } from "ai";
 import type { z } from "zod";
 import type { Agent, BaseGenerationOptions } from "../../agent/agent";
 import { convertUsage } from "../../utils/usage-converter";
 import type { InternalWorkflowFunc, WorkflowExecuteContext } from "../internal/types";
 import type { WorkflowStepAgent } from "./types";
 
-export type AgentConfig<SCHEMA extends z.ZodTypeAny, INPUT, DATA> = BaseGenerationOptions & {
+type OutputSpec = Output.Output<unknown, unknown>;
+
+export type AgentOutputSchema = OutputSpec | z.ZodTypeAny;
+
+export type InferAgentOutput<SCHEMA extends AgentOutputSchema> = SCHEMA extends OutputSpec
+  ? InferGenerateOutput<SCHEMA>
+  : SCHEMA extends z.ZodTypeAny
+    ? z.infer<SCHEMA>
+    : never;
+
+export type AgentConfig<SCHEMA extends AgentOutputSchema, INPUT, DATA> = Omit<
+  BaseGenerationOptions,
+  "output"
+> & {
   schema:
     | SCHEMA
     | ((
         context: Omit<WorkflowExecuteContext<INPUT, DATA, any, any>, "suspend" | "writer">,
       ) => SCHEMA | Promise<SCHEMA>);
 };
 
-type AgentResultMapper<INPUT, DATA, SCHEMA extends z.ZodTypeAny, RESULT> = (
-  output: z.infer<SCHEMA>,
+type AgentResultMapper<INPUT, DATA, SCHEMA extends AgentOutputSchema, RESULT> = (
+  output: InferAgentOutput<SCHEMA>,
   context: WorkflowExecuteContext<INPUT, DATA, any, any>,
 ) => Promise<RESULT> | RESULT;
 
+const isOutputSpec = (value: unknown): value is OutputSpec => {
+  if (!value || typeof value !== "object") return false;
+  const candidate = value as OutputSpec;
+  return (
+    typeof candidate.parseCompleteOutput === "function" &&
+    typeof candidate.parsePartialOutput === "function" &&
+    "responseFormat" in candidate
+  );
+};
+
 /**
  * Creates an agent step for a workflow
  *
@@ -39,11 +62,16 @@ type AgentResultMapper<INPUT, DATA, SCHEMA extends z.ZodTypeAny, RESULT> = (
  *
  * @param task - The task (prompt) to execute for the agent, can be a string or a function that returns a string
  * @param agent - The agent to execute the task using `generateText`
- * @param config - The config for the agent (schema) `generateText` call
+ * @param config - The config for the agent (schema/output) `generateText` call
  * @param map - Optional mapper to shape or merge the agent output with existing data
  * @returns A workflow step that executes the agent with the task
  */
-export function andAgent<INPUT, DATA, SCHEMA extends z.ZodTypeAny, RESULT = z.infer<SCHEMA>>(
+export function andAgent<
+  INPUT,
+  DATA,
+  SCHEMA extends AgentOutputSchema,
+  RESULT = InferAgentOutput<SCHEMA>,
+>(
   task:
     | UIMessage[]
     | ModelMessage[]
@@ -64,10 +92,11 @@ export function andAgent<INPUT, DATA, SCHEMA extends z.ZodTypeAny, RESULT = z.in
       const { schema, ...restConfig } = config;
       const finalTask = typeof task === "function" ? await task(context) : task;
       const finalSchema = typeof schema === "function" ? await schema(context) : schema;
+      const output = isOutputSpec(finalSchema)
+        ? finalSchema
+        : Output.object({ schema: finalSchema });
 
-      const output = Output.object({ schema: finalSchema });
-
-      const mapOutput = async (outputValue: z.infer<SCHEMA>) => {
+      const mapOutput = async (outputValue: InferAgentOutput<SCHEMA>) => {
         if (map) {
           return (await map(outputValue, context)) as RESULT;
         }
@@ -99,7 +128,7 @@ export function andAgent<INPUT, DATA, SCHEMA extends z.ZodTypeAny, RESULT = z.in
           }
           state.usage.totalTokens += convertedUsage?.totalTokens || 0;
         }
-        return mapOutput(result.output as z.infer<SCHEMA>);
+        return mapOutput(result.output as InferAgentOutput<SCHEMA>);
       }
 
       // Step start event removed - now handled by OpenTelemetry spans
@@ -131,7 +160,7 @@ export function andAgent<INPUT, DATA, SCHEMA extends z.ZodTypeAny, RESULT = z.in
           state.usage.totalTokens += convertedUsage?.totalTokens || 0;
         }
 
-        return mapOutput(result.output as z.infer<SCHEMA>);
+        return mapOutput(result.output as InferAgentOutput<SCHEMA>);
       } catch (error) {
         // Check if this is a suspension, not an error
         if (

website/docs/workflows/steps/and-agent.md

@@ -33,20 +33,20 @@ const result = await workflow.run({ text: "I love this!" });
 
 ## How It Works
 
-`andAgent` = AI prompt + structured output schema:
+`andAgent` = AI prompt + structured output schema (Zod or `Output.*`):
 
 ```typescript
 .andAgent(
   prompt,     // What to ask the AI
   agent,      // Which AI to use
-  { schema }, // What shape the answer should be
+  { schema }, // Zod schema or ai-sdk Output.* spec
   map?        // Optional: merge/shape output with existing data
 )
 ```
 
 If you pass a function for `prompt`, it must return a Promise (use `async`) and resolve to a string, `UIMessage[]`, or `ModelMessage[]`.
 
-**Important:** `andAgent` uses `generateText` with `Output.object` under the hood, which means:
+**Important:** `andAgent` uses `generateText`. If you pass a Zod schema, it is wrapped with `Output.object`. If you pass an `Output.*` spec, it is used directly. This means:
 
 - ✅ You get **structured, typed responses** based on your schema
 - ✅ The agent **can use tools** during this step
@@ -90,6 +90,14 @@ By default, the step result replaces the workflow data with the agent output. If
   { schema: z.object({ type: z.enum(["support", "sales", "spam"]) }) },
   (output, { data }) => ({ ...data, emailType: output })
 )
+
+// Use Output.* for non-object outputs (arrays, choices, json, text)
+// (Requires: import { Output } from "ai";)
+.andAgent(
+  async ({ data }) => `List tags for: ${data.topic}`,
+  agent,
+  { schema: Output.array({ element: z.string() }) }
+)
 ```
 
 ## Middleware and Retries
@@ -201,6 +209,24 @@ Retry behavior:
 )
 ```
 
+### Arrays and Choices
+
+Requires `import { Output } from "ai";`
+
+```typescript
+.andAgent(
+  async ({ data }) => `Give 3 tags for: ${data.topic}`,
+  agent,
+  { schema: Output.array({ element: z.string() }) }
+)
+
+.andAgent(
+  async ({ data }) => `Pick a category for: ${data.title}`,
+  agent,
+  { schema: Output.choice({ options: ["news", "blog", "doc"] }) }
+)
+```
+
 ### Merge Output With Existing Data
 
 ```typescript