Back to Copilotkit

State streaming

showcase/shell-docs/src/content/docs/integrations/mastra/shared-state/predictive-state-updates.mdx

1.63.25.0 KB
Original Source

import RunAndConnect from "@/snippets/integrations/mastra/run-and-connect.mdx" import { IframeSwitcher } from "@/components/content"

<IframeSwitcher id="mastra-state-streaming-example" exampleUrl="https://feature-viewer.copilotkit.ai/mastra/feature/shared_state?sidebar=false&chatDefaultOpen=false" codeUrl="https://feature-viewer.copilotkit.ai/mastra/feature/shared_state?view=code&sidebar=false&codeLayout=tabs" exampleLabel="Demo" codeLabel="Code" height="700px" />

<Callout type="info"> This example demonstrates streaming shared state in the [CopilotKit Feature Viewer](https://feature-viewer.copilotkit.ai/mastra/feature/shared_state). </Callout>

What is this?

An agent's working memory usually surfaces to the UI once, when a run finishes. But a single run can take many seconds and write a lot of content — a drafted document, a plan, a summary — that the user would rather watch appear as it is generated.

Agent-native applications reflect what the agent is doing to the end-user as continuously as possible. CopilotKit enables this for Mastra by streaming working-memory updates to the frontend token-by-token, so agent.state grows live during the run instead of jumping to the final value at the end.

When should I use this?

Use this when you want to give the user real-time feedback about what your agent is producing, specifically to:

  • Keep users engaged by avoiding long loading indicators
  • Build trust by showing the work as it is written
  • Enable agent steering — letting users react before the run completes

How it works

Mastra injects a built-in updateWorkingMemory tool whenever working memory is enabled. When the agent calls it during a run, CopilotKit's Mastra integration intercepts the tool call's streaming arguments, parses the growing payload, and emits an initial STATE_SNAPSHOT followed by incremental STATE_DELTA updates against the working-memory schema. The frontend's useAgent subscription re-renders on every delta.

<Callout type="info" title="Important"> This guide assumes you are embedding your Mastra agent inside Copilot Runtime, like so.
ts
const runtime = new CopilotRuntime({
  agents: MastraAgent.getLocalAgents({ mastra }),
});
</Callout>

Implementation

<Steps> <Step> ### Run and connect your agent <RunAndConnect components={props.components} /> </Step> <Step> ### Define the streamed state Provide a working-memory schema whose field holds the content you want to stream. Here `document` is a single string that the agent will fill in progressively.
```ts title="mastra/agents/streaming-agent.ts"
import { openai } from "@ai-sdk/openai";
import { Agent } from "@mastra/core/agent";
import { LibSQLStore } from "@mastra/libsql";
import { Memory } from "@mastra/memory";
import { z } from "zod";

// 1. Define the streamed state schema
export const StreamingAgentState = z.object({
  document: z.string().default(""),
});

// 2. Create the agent with working memory enabled
export const streamingAgent = new Agent({
  name: "Streaming Agent",
  model: openai("gpt-4o"),
  // The prompt drives the model to write the full document straight into
  // working memory via the built-in updateWorkingMemory tool, rather than
  // pasting it into a chat message.
  instructions: `You are a collaborative writing assistant. Whenever the user asks you to write, draft, or revise anything, call the \`updateWorkingMemory\` tool with the FULL content under the \`document\` field. Never paste the document into a chat message — it belongs in shared state, and the UI renders it live as you stream it.`,
  memory: new Memory({
    storage: new LibSQLStore({ id: "mastra-storage", url: ":memory:" }),
    options: {
      workingMemory: {
        enabled: true, // [!code highlight]
        schema: StreamingAgentState, // [!code highlight]
      },
    },
  }),
});
```
</Step> <Step> ### Observe the stream Subscribe to the agent state with `useAgent`. `agent.state.document` updates on every delta as the agent writes, so the panel fills in live.
```tsx title="ui/app/page.tsx"
import { useAgent, UseAgentUpdate } from "@copilotkit/react-core"; // [!code highlight]

const YourMainContent = () => {
  // [!code highlight:4]
  const { agent } = useAgent({
    agentId: "streamingAgent",
    updates: [UseAgentUpdate.OnStateChanged, UseAgentUpdate.OnRunStatusChanged],
  });

  const document = (agent.state as { document?: string })?.document ?? "";

  return (
    <div>
      <h1>Document {agent.isRunning && <span>· LIVE</span>}</h1>
      <pre>{document || "Ask the agent to write something…"}</pre>
    </div>
  );
};
```
</Step> <Step> ### Give it a try! Ask the agent to write a poem, draft an email, or explain a topic. The document panel fills in token-by-token as the agent streams, and the run-end working-memory snapshot reconciles any drift. </Step> </Steps>