Get started with Vercel Workflow Development Kit

import type { NextConfig } from 'next';
import { withWorkflow } from 'workflow/next';

const nextConfig: NextConfig = {};

export default withWorkflow(nextConfig);

import { getWorkflowMetadata } from 'workflow';
import type { InvocationData } from '@ably/ai-transport';
import { failRun, openRun, runInference, runTool, type TurnIds } from './steps';

export async function chatWorkflow(invocation: InvocationData): Promise<void> {
  'use workflow';
  const { workflowRunId } = getWorkflowMetadata();

  let ids: TurnIds | undefined;
  try {
    ids = await openRun(invocation, workflowRunId);
    let outcome = await runInference(invocation, ids);

    while (outcome.kind === 'server-tools') {
      for (const toolCall of outcome.serverToolCalls) {
        await runTool(invocation, ids, toolCall);
      }
      outcome = await runInference(invocation, ids);
    }
  } catch (error) {
    // Retries exhausted: if a Run was opened, end it in error so observers unstick.
    if (ids) {
      const message = error instanceof Error ? error.message : String(error);
      try {
        await failRun(invocation, ids, message);
      } catch {
        /* best-effort */
      }
    }
    throw error;
  }
}

import { z } from 'zod';
import type { Tool } from 'ai';

export const tools: Record<string, Tool> = {
  getStockPrice: {
    description: 'Get the current stock price for a ticker symbol.',
    inputSchema: z.object({
      symbol: z.string().describe('The ticker symbol, for example "AAPL"'),
    }),
    execute: async ({ symbol }: { symbol: string }) => {
      // Intentionally flaky: throws on an odd price so you can watch the retry re-roll it.
      const priceUSD = Math.round(50 + Math.random() * 500);
      if (priceUSD % 2 !== 0) {
        throw new Error(`stock price service returned an odd price (${priceUSD}), retry me`);
      }
      return { symbol, priceUSD };
    },
  },
};

import * as Ably from 'ably';
import { convertToModelMessages, stepCountIs, streamText } from 'ai';
import { getStepMetadata } from 'workflow';
import { anthropic } from '@ai-sdk/anthropic';
import { Invocation, type InvocationData } from '@ably/ai-transport';
import {
  createAgentSession,
  pendingToolCalls,
  stripToolExecutes,
  vercelRunOutcome,
} from '@ably/ai-transport/vercel';
import { tools } from './tools';

export interface TurnIds {
  runId: string;
  invocationId: string;
  triggerEventId: string;
}

export interface ToolCallInfo {
  toolCallId: string;
  toolName: string;
  input: unknown;
}

// server-tools is the only non-terminal kind; the other kinds are already published.
export type InferenceOutcome =
  | { kind: 'complete' }
  | { kind: 'suspend' }
  | { kind: 'cancelled' }
  | { kind: 'error'; errorMessage: string }
  | { kind: 'server-tools'; serverToolCalls: ToolCallInfo[] };

type AgentSession = ReturnType<typeof createAgentSession>;
type AgentRun = ReturnType<AgentSession['adoptRun']> | ReturnType<AgentSession['createRun']>;

// A fresh session per step. detach (not end) leaves any open Run on the channel
// for the next step, or for a retry to adopt.
async function withAgentSession<T>(
  channelName: string,
  body: (session: AgentSession) => Promise<T>,
): Promise<T> {
  const client = new Ably.Realtime({ key: process.env.ABLY_API_KEY! });
  const session = createAgentSession({ client, channelName });
  try {
    await session.connect();
    return await body(session);
  } finally {
    try {
      await session.detach();
    } catch {
      /* best-effort */
    }
    client.close();
  }
}

// Open the Run and return its ids. No model call here; the first inference is its own step.
export async function openRun(invocationData: InvocationData, workflowRunId: string): Promise<TurnIds> {
  'use step';
  const invocation = Invocation.fromJSON(invocationData);
  const invocationId = `inv:${workflowRunId}`;
  return withAgentSession(invocation.sessionName, async (session) => {
    // Pin the ids to the replay-stable workflow run id so a retry re-enters the same Run.
    const run = session.createRun(invocation, { runId: `run:${workflowRunId}`, invocationId });

    // Drain history so the trigger folds in, then open (or resume) the Run.
    while (run.view.hasOlder()) await run.view.loadOlder();
    await run.start();

    return { runId: run.runId, invocationId, triggerEventId: invocation.inputEventId };
  });
}

// Adopt the open Run, run one model call as a Step, and publish its terminal.
export async function runInference(invocationData: InvocationData, ids: TurnIds): Promise<InferenceOutcome> {
  'use step';
  const invocation = Invocation.fromJSON(invocationData);
  return withAgentSession(invocation.sessionName, async (session) => {
    const { stepId } = getStepMetadata();
    const run = session.adoptRun(ids);
    await run.load({ timeoutMs: 15_000 });
    while (run.view.hasOlder()) await run.view.loadOlder();

    const step = run.createStep({ stepId });
    await step.start();

    const conversation = run.view.getMessages().map((entry) => entry.message);
    const result = streamText({
      model: anthropic('claude-sonnet-4-20250514'),
      messages: await convertToModelMessages(conversation),
      tools: stripToolExecutes(tools),
      abortSignal: run.abortSignal,
      // The workflow drives the loop; this call runs one step only.
      stopWhen: stepCountIs(1),
    });

    const pipeResult = await step.pipe(result.toUIMessageStream());
    const outcome = await vercelRunOutcome(pipeResult, result.finishReason);
    await step.end(outcome.reason === 'error' ? { reason: 'failed' } : undefined);

    let inference: InferenceOutcome;
    if (outcome.reason === 'error') {
      inference = { kind: 'error', errorMessage: outcome.error.message };
    } else if (outcome.reason === 'cancelled') {
      inference = { kind: 'cancelled' };
    } else if (outcome.reason === 'complete') {
      inference = { kind: 'complete' };
    } else {
      // Server tools (execute in the registry) run as steps; anything else suspends.
      const serverToolCalls = pendingToolCalls(run.messages)
        .filter((call) => typeof tools[call.toolName]?.execute === 'function')
        .map((call) => ({ toolCallId: call.toolCallId, toolName: call.toolName, input: call.input }));
      inference = serverToolCalls.length > 0 ? { kind: 'server-tools', serverToolCalls } : { kind: 'suspend' };
    }

    await publishTerminal(run, inference);
    return inference;
  });
}

// Execute one server tool and publish its result as a Step. A throw retries under the same stepId.
export async function runTool(invocationData: InvocationData, ids: TurnIds, toolCall: ToolCallInfo): Promise<void> {
  'use step';
  const invocation = Invocation.fromJSON(invocationData);
  await withAgentSession(invocation.sessionName, async (session) => {
    const run = session.adoptRun(ids);
    await run.load({ timeoutMs: 15_000 });

    const step = run.createStep({ stepId: getStepMetadata().stepId });
    await step.start();
    // Look the tool up by the model-provided name and narrow it to the callable shape.
    const tool = tools[toolCall.toolName] as { execute?: (input: unknown) => Promise<unknown> };
    if (!tool?.execute) throw new Error(`tool '${toolCall.toolName}' has no execute`);
    const output = await tool.execute(toolCall.input);
    await step.send({ type: 'tool-output-available', toolCallId: toolCall.toolCallId, output });
    await step.end();
  });
}

// The workflow's failure terminal, run from its catch when a step exhausts retries.
// A load() rejection means the Run is already gone; nothing to clean up.
export async function failRun(invocationData: InvocationData, ids: TurnIds, errorMessage: string): Promise<void> {
  'use step';
  const invocation = Invocation.fromJSON(invocationData);
  await withAgentSession(invocation.sessionName, async (session) => {
    const run = session.adoptRun(ids);
    try {
      await run.load({ timeoutMs: 15_000 });
    } catch {
      return;
    }
    await run.end({ reason: 'error', error: new Ably.ErrorInfo(errorMessage, 104000, 500) });
  });
}

// Cleanup is best-effort and must not cascade: one attempt, no retries.
failRun.maxRetries = 0;

// Publish the lifecycle event the outcome implies; server-tools publishes nothing.
async function publishTerminal(run: AgentRun, outcome: InferenceOutcome): Promise<void> {
  switch (outcome.kind) {
    case 'server-tools':
      return;
    case 'suspend':
      await run.suspend();
      return;
    case 'error':
      await run.end({ reason: 'error', error: new Ably.ErrorInfo(outcome.errorMessage, 104000, 500) });
      return;
    default:
      await run.end({ reason: outcome.kind });
  }
}

import { start } from 'workflow/api';
import type { InvocationData } from '@ably/ai-transport';
import { chatWorkflow } from '../../workflows/turn';

export async function POST(req: Request): Promise<Response> {
  const invocation = (await req.json()) as InvocationData;
  const run = await start(chatWorkflow, [invocation]);
  return Response.json({ workflowRunId: run.runId });
}