AI Agents

Streaming Updates from Tools

After building a durable AI agent, we already get UI message chunks for displaying tool invocations and return values. However, for long-running steps, we may want to show progress updates, or stream step output to the user while it's being generated.

Workflow DevKit enables this by letting step functions write custom chunks to the same stream the agent uses. These chunks appear as data parts in your messages, which you can render however you like.

As an example, we'll extend out Flight Booking Agent to use emit more granular progress updates while searching for flights.

Define Your Data Part Type

First, define a TypeScript type for your custom data part. This ensures type safety across your tool and client code:

schemas/chat.ts
export interface FoundFlightDataPart {
  type: "data-found-flight"; 
  id: string;
  data: {
    flightNumber: string;
    from: string;
    to: string;
  };
}

The type field must be a string starting with data- followed by your custom identifier. The id field should match the toolCallId so the client can associate the data with the correct tool invocation. Learn more about data parts in the AI SDK documentation.

Emit Updates from Your Tool

Use getWritable() inside a step function to get a handle to the stream. This is the same stream that the LLM and other tools calls are writing to, so we can inject out own data packets directly.

workflows/chat/steps/tools.ts
import { getWritable } from "workflow"; 
import type { UIMessageChunk } from "ai";

export async function searchFlights(
  { from, to, date }: { from: string; to: string; date: string },
  { toolCallId }: { toolCallId: string } 
) {
  "use step";

  const writable = getWritable<UIMessageChunk>(); 
  const writer = writable.getWriter(); 

  // ... existing logic to generate flights ...

  for (const flight of generatedFlights) { 

    // Simulate the time it takes to find each flight
    await new Promise((resolve) => setTimeout(resolve, 1000)); 

    await writer.write({ 
      id: `${toolCallId}-${flight.flightNumber}`, 
      type: "data-found-flight", 
      data: flight, 
    }); 
  } 

  writer.releaseLock(); 

  return {
    message: `Found ${generatedFlights.length} flights from ${from} to ${to} on ${date}`,
    flights: generatedFlights.sort((a, b) => a.price - b.price), // Sort by price
  };
}

Key points:

  • Call getWritable<UIMessageChunk>() to get the stream
  • Use getWriter() to acquire a writer
  • Write objects with type, id, and data fields
  • Always call releaseLock() when done writing (learn more about streaming)

Handle Data Parts in the Client

Update your chat component to detect and render the custom data parts. Data parts are stored in the message's parts array alongside text and tool invocation parts:

app/page.tsx
{message.parts.map((part, partIndex) => {
  // Render text parts
  if (part.type === "text") {
    return (
      <Response key={`${message.id}-text-${partIndex}`}>
        {part.text}
      </Response>
    );
  }

  // Render streaming flight data parts
  if (part.type === "data-found-flight") { 
    const flight = part.data as { 
      flightNumber: string; 
      airline: string; 
      from: string; 
      to: string; 
    }; 
    return ( 
      <div key={`${part.id}-${flight.flightNumber}`} className="p-3 bg-muted rounded-md">
        <div className="font-medium">{flight.airline} - {flight.flightNumber}</div>
        <div className="text-muted-foreground">{flight.from} → {flight.to}</div>
      </div>
    ); 
  } 

  // ... other rendering logic ...
})}

The pattern is:

  1. Data parts have a type field starting with data-
  2. Match the type to your custom identifier (e.g., data-found-flight)
  3. Use the data part's payload to display progress or intermediate results

Now, when you run the agent to search for flights, you'll see the flight results pop up one after another. This will be most useful if you have tool calls that take minutes to complete, and you need to show granular progress updates to the user.