Resumable AI SDK Chat Streaming with Vercel Workflow - Demo

ai-sdk-workflow-resumable-example.md

Resumable AI SDK Chat Streaming with Vercel Workflow

Demo link: https://firstdomain.world

This gist demonstrates how to implement resumable AI SDK chat streaming using Vercel Workflow, allowing chat sessions to survive network interruptions and function timeouts.

Key Features

• Resumable streaming: Chat continues even after network interruptions
• Persistent state: Chat history saved in localStorage
• Workflow durability: Long-running operations survive function timeouts
• Error recovery: Automatic reconnection with stored workflow run IDs

Architecture

Files Structure

• chat.tsx - Main chat component with WorkflowChatTransport
• api-chat-route.ts - Initial chat API route (/api/chat/route.ts)
• api-chat-resume-route.ts - Resume API route (/api/chat/[id]/stream/route.ts)
• chat-workflow.ts - Workflow implementation with durable steps
• chat-schema.ts - TypeScript types for messages
• chat-input.tsx - Chat input component

Key Components

1. WorkflowChatTransport: Handles resumable streaming with workflow run ID storage
2. Workflow Functions: Durable orchestration that survives timeouts
3. Step Functions: Individual operations with full Node.js access
4. Persistence Layer: localStorage for demo (use database in production)

How Resumability Works

1. Initial Request: Client sends message to /api/chat
2. Workflow Creation: Server starts workflow and returns run ID in x-workflow-run-id header
3. ID Storage: Client stores workflow run ID in localStorage
4. Interruption Handling: On network errors, client reconnects using stored run ID
5. Resume Stream: Client calls /api/chat/{workflowRunId}/stream with startIndex
6. State Recovery: Workflow resumes from last checkpoint

Code Highlights

Client-side Resumption

transport: new WorkflowChatTransport({
  prepareReconnectToStreamRequest: ({ id, api, ...rest }) => {
    const workflowRunId = localStorage.getItem('active-workflow-run-id');
    return {
      ...rest,
      api: `/api/chat/${encodeURIComponent(workflowRunId)}/stream`,
    };
  },
})

Server-side Durability

export async function chat(messages: UIMessage[]) {
  'use workflow';
  
  const writable = getWorkflowWritableStream<UIMessageChunk>();
  
  // Durable steps that survive function timeouts
  for (let i = 0; i < MAX_STEPS; i++) {
    const result = await streamTextStep(i, currMessages, writable);
    // State automatically persisted between steps
  }
}

Usage

1. Set up Vercel Workflow SDK in your Next.js project
2. Copy these files to appropriate locations:
   • chat.tsx → app/chat.tsx
   • api-chat-route.ts → app/api/chat/route.ts
   • api-chat-resume-route.ts → app/api/chat/[id]/stream/route.ts
   • chat-workflow.ts → workflows/chat.ts
3. Install dependencies: @vercel/workflow-ai, @ai-sdk/react, ai
4. Configure your AI provider in the workflow

Benefits

• Resilient UX: Users don't lose progress during network issues
• Scalable: Handles long-running AI operations without function timeouts
• Developer-friendly: Standard React hooks with workflow durability
• Production-ready: Built on Vercel's durable execution primitives

Perfect for AI applications requiring reliable, long-running conversations!

api-chat-resume-route.ts

import { getWorkflowReadableStream } from '@vercel/workflow-core/runtime';
import { createUIMessageStreamResponse, type UIMessageChunk } from 'ai';

// Uncomment to simulate a long running Vercel Function timing
// out due to a long running agent. The client-side will
// automatically reconnect to the stream.
//export const maxDuration = 5;

export async function GET(
  request: Request,
  { params }: { params: Promise<{ id: string }> }
) {
  const { id } = await params;
  const { searchParams } = new URL(request.url);
  const startIndexParam = searchParams.get('startIndex');
  const startIndex =
    startIndexParam !== null ? parseInt(startIndexParam, 10) : undefined;

  return createUIMessageStreamResponse({
    stream: getWorkflowReadableStream<UIMessageChunk>(id, { startIndex }),
  });
}

api-chat-route.ts

import { getWorkflowReadableStream } from '@vercel/workflow-core/runtime';
import {
  createUIMessageStreamResponse,
  type UIMessage,
  type UIMessageChunk,
} from 'ai';
import { chat } from '@/workflows/chat';

// Uncomment to simulate a long running Vercel Function timing
// out due to a long running agent. The client-side will
// automatically reconnect to the stream.
//export const maxDuration = 8;

export async function POST(req: Request) {
  const { messages }: { messages: UIMessage[] } = await req.json();

  const workflowHandle = await chat(messages);

  // TODO: To avoid the cast, use the `start()` function instead of
  // calling the `chat()` function directly?
  const workflowRunId = (workflowHandle as any).runId;

  return createUIMessageStreamResponse({
    stream: getWorkflowReadableStream<UIMessageChunk>(workflowRunId),
    headers: {
      // The workflow run ID is stored into `localStorage` on the client side,
      // which influences the `resume` flag in the `useChat` hook.
      'x-workflow-run-id': workflowRunId,
    },
  });
}

chat-input.tsx

import { useState } from 'react';

export default function ChatInput({
  status,
  onSubmit,
  onNewChat,
  inputRef,
}: {
  status: string;
  onSubmit: (text: string) => void;
  onNewChat: () => void;
  inputRef: React.RefObject<HTMLInputElement | null>;
}) {
  const [text, setText] = useState('');

  return (
    <form
      onSubmit={(e) => {
        e.preventDefault();
        if (text.trim() === '') return;
        onSubmit(text);
        setText('');
      }}
    >
      <div className="flex flex-row fixed bottom-0 w-full max-w-2xl mb-8 gap-2">
        <input
          ref={inputRef}
          className="p-3 border border-gray-300 rounded-lg shadow-xl w-full bg-[#fff]"
          placeholder="Ask me about flights, airports, or bookings..."
          disabled={status !== 'ready'}
          value={text}
          onChange={(e) => setText(e.target.value)}
        />
        <button
          className="px-3 py-1 text-sm transition-colors bg-gray-200 rounded-md hover:bg-gray-300"
          type="button"
          onClick={() => {
            onNewChat();
          }}
        >
          New Chat
        </button>
      </div>
    </form>
  );
}

chat-schema.ts

import type { UIDataTypes, UIMessage } from 'ai';
import { z } from 'zod/v4';

export const myMessageMetadataSchema = z.object({
  createdAt: z.number(),
});

export type MyMessageMetadata = z.infer<typeof myMessageMetadataSchema>;

export type MyUIMessage = UIMessage<MyMessageMetadata, UIDataTypes>;

chat-workflow.ts

import { getWorkflowWritableStream } from '@vercel/workflow-core';
import {
  convertToModelMessages,
  type FinishReason,
  generateId,
  type ModelMessage,
  stepCountIs,
  streamText,
  type UIMessage,
  type UIMessageChunk,
} from 'ai';
import { FLIGHT_ASSISTANT_PROMPT, flightBookingTools } from './chat-tools';

const MAX_STEPS = 10;

/** A Stream Text Step */
export async function streamTextStep(
  step: number,
  messages: ModelMessage[],
  writable: WritableStream<UIMessageChunk>
) {
  'use step';

  // Send start data message
  const writer = writable.getWriter();
  writer.write({
    id: generateId(),
    type: 'data-workflow',
    data: {
      message: `Workflow step "streamTextStep" started (#${step})`,
    },
  });

  // Mimic a random network error
  if (Math.random() < 0.3) {
    await new Promise((resolve) => setTimeout(resolve, 2000));
    writer.write({
      id: generateId(),
      type: 'data-workflow',
      data: {
        message: `Workflow step "streamTextStep" errored (#${step})`,
        type: 'error',
      },
    });
    throw new Error('Error connecting to LLM');
  }

  // Make the LLM request
  console.log('Sending request to LLM');
  const result = streamText({
    model: 'bedrock/claude-4-sonnet-20250514-v1',
    messages,
    system: FLIGHT_ASSISTANT_PROMPT,
    // We'll handle the back and forth ourselves
    stopWhen: stepCountIs(1),
    tools: flightBookingTools,
    headers: {
      'anthropic-beta': 'interleaved-thinking-2025-05-14',
    },
    providerOptions: {
      anthropic: {
        thinking: { type: 'enabled', budgetTokens: 16000 },
      },
    },
  });

  // Pipe the stream to the client
  const reader = result
    // We send these chunks outside the loop
    .toUIMessageStream({ sendStart: false, sendFinish: false })
    .getReader();
  try {
    while (true) {
      const { done, value } = await reader.read();
      if (done) break;
      await writer.write(value);
    }
  } finally {
    writer.write({
      id: generateId(),
      type: 'data-workflow',
      data: {
        message: `Workflow step "streamTextStep" completed (#${step})`,
      },
    });
    reader.releaseLock();
    writer.releaseLock();
  }

  // Return the values back to the workflow
  const finishReason = await result.finishReason;

  // Workflow will retry errors
  if (finishReason === 'error') {
    writer.write({
      id: generateId(),
      type: 'data-workflow',
      data: {
        message: `Workflow step "streamTextStep" errored (#${step})`,
        type: 'error',
      },
    });
    throw new Error('LLM error from streamTextStep');
  }

  return {
    messages: (await result.response).messages,
    finishReason,
  };
}

export async function startStream(writable: WritableStream<UIMessageChunk>) {
  'use step';
  const writer = writable.getWriter();

  writer.write({
    type: 'start',
    messageMetadata: {
      createdAt: Date.now(),
      messageId: generateId(),
    },
  });

  writer.write({
    id: generateId(),
    type: 'data-workflow',
    data: { message: 'Starting workflow stream' },
  });

  writer.releaseLock();
}

export async function endStream(writable: WritableStream<UIMessageChunk>) {
  'use step';
  const writer = writable.getWriter();

  console.log('Closing workflow stream');

  writer.write({
    id: generateId(),
    type: 'data-workflow',
    data: {
      message: 'Closing workflow stream',
    },
  });

  writer.write({
    type: 'finish',
  });

  writer.close();
  writer.releaseLock();
}

/**
 * The main chat workflow
 */
export async function chat(messages: UIMessage[]) {
  'use workflow';

  console.log('Starting workflow');

  const writable = getWorkflowWritableStream<UIMessageChunk>();

  // Write the "start" message to the client
  await startStream(writable);

  const currMessages: ModelMessage[] = convertToModelMessages(messages);
  let finishReason: FinishReason = 'unknown';

  // Run `streamText` in a loop while we have tool calls
  for (let i = 0; i < MAX_STEPS; i++) {
    console.log(`Running step ${i + 1}`);

    const result = await streamTextStep(i, currMessages, writable);

    currMessages.push(...result.messages);
    finishReason = result.finishReason;

    if (finishReason !== 'tool-calls') {
      break;
    }
  }

  // Write the "finish" message to the client
  await endStream(writable);

  console.log('Finished workflow');
}

chat.tsx

'use client';

import { useChat } from '@ai-sdk/react';
import { WorkflowChatTransport } from '@vercel/workflow-ai';
import { useEffect, useMemo, useRef } from 'react';
import type { MyUIMessage } from '@/util/chat-schema';
import ChatInput from './chat-input';
import Message from './message';

export default function ChatComponent() {
  const inputRef = useRef<HTMLInputElement>(null);

  const activeWorkflowRunId = useMemo(() => {
    if (typeof window === 'undefined') return;
    return localStorage.getItem('active-workflow-run-id') ?? undefined;
  }, []);

  const chat = useChat<MyUIMessage>({
    resume: !!activeWorkflowRunId,
    onError(error) {
      console.error('onError', error);
    },
    onFinish(data) {
      console.log('onFinish', data);

      // Update the chat history in `localStorage` to include the latest bot message
      console.log('Saving chat history to localStorage', data.messages);
      localStorage.setItem('chat-history', JSON.stringify(data.messages));

      requestAnimationFrame(() => {
        inputRef.current?.focus();
      });
    },

    transport: new WorkflowChatTransport({
      onChatSendMessage: (response, options) => {
        console.log('onChatSendMessage', response, options);

        // Update the chat history in `localStorage` to include the latest user message
        localStorage.setItem('chat-history', JSON.stringify(options.messages));

        // We'll store the workflow run ID in `localStorage` to allow the client
        // to resume the chat session after a page refresh or network interruption
        const workflowRunId = response.headers.get('x-workflow-run-id');
        if (!workflowRunId) {
          throw new Error(
            'Workflow run ID not found in "x-workflow-run-id" response header'
          );
        }
        localStorage.setItem('active-workflow-run-id', workflowRunId);
      },
      onChatEnd: ({ chatId, chunkIndex }) => {
        console.log('onChatEnd', chatId, chunkIndex);

        // Once the chat stream ends, we can remove the workflow run ID from `localStorage`
        localStorage.removeItem('active-workflow-run-id');
      },
      // Configure reconnection to use the stored workflow run ID
      prepareReconnectToStreamRequest: ({ id, api, ...rest }) => {
        console.log('prepareReconnectToStreamRequest', id);
        const workflowRunId = localStorage.getItem('active-workflow-run-id');
        if (!workflowRunId) {
          throw new Error('No active workflow run ID found');
        }
        // Use the workflow run ID instead of the chat ID for reconnection
        return {
          ...rest,
          api: `/api/chat/${encodeURIComponent(workflowRunId)}/stream`,
        };
      },
      // Optional: Configure error handling for reconnection attempts
      maxConsecutiveErrors: 5,
    }),
  });

  // Load chat history from `localStorage`. In a real-world application,
  // this would likely be done on the server side and loaded from a database,
  // but for the purposes of this demo, we'll load it from `localStorage`.
  useEffect(() => {
    const chatHistory = localStorage.getItem('chat-history');
    if (!chatHistory) return;
    chat.setMessages(JSON.parse(chatHistory) as MyUIMessage[]);
  }, [chat.setMessages]);

  // Activate the input field
  useEffect(() => {
    inputRef.current?.focus();
  }, []);

  return (
    <div className="flex flex-col w-full max-w-2xl pt-12 pb-24 mx-auto stretch">
      <div className="mb-8 text-center">
        <h1 className="text-3xl font-bold mb-2">✈️ Flight Booking Agent</h1>
        <p className="text-gray-600">Book a flight using workflows</p>
      </div>

      {chat.messages.length === 0 && (
        <div className="mb-8 p-6 bg-blue-50 rounded-lg">
          <h2 className="text-lg font-semibold mb-3">
            How can I help you today?
          </h2>
          <p className="text-gray-700 mb-4">I can assist you with:</p>
          <ul className="space-y-2 text-gray-600">
            <li>🔍 Search for flights between cities</li>
            <li>📊 Check real-time flight status</li>
            <li>🛫 Get airport information</li>
            <li>🎫 Book flights for you</li>
            <li>🧳 Check baggage allowances</li>
          </ul>
          <p className="mt-4 text-sm text-gray-500">
            Try asking: "Find me flights from San Francisco to Los Angeles" or
            "What's the status of flight UA123?"
          </p>
        </div>
      )}

      {chat.messages.map((message) => (
        <Message
          key={message.id}
          message={message}
          regenerate={chat.regenerate}
          sendMessage={chat.sendMessage}
          addToolResult={chat.addToolResult}
          status={chat.status}
        />
      ))}
      <ChatInput
        status={chat.status}
        onSubmit={(text: string) => {
          chat.sendMessage({ text, metadata: { createdAt: Date.now() } });
        }}
        onNewChat={async () => {
          await chat.stop();
          localStorage.removeItem('active-workflow-run-id');
          localStorage.removeItem('chat-history');
          chat.setMessages([]);
        }}
        inputRef={inputRef}
      />
    </div>
  );
}