CopilotKit

State Streaming

Stream partial agent state updates to the UI while a tool call is still running.


"""LangGraph agent backing the State Streaming demo.Demonstrates per-token state-delta streaming. The agent writes a long`document` string into shared agent state via a `write_document` tool;`StateStreamingMiddleware(StateItem(...))` tells CopilotKit to forward*every token* of the tool's `content` argument directly into the`document` state key as it is generated. The UI (useAgent) sees`state.document` grow token-by-token, without waiting for the tool callto finish.This is the canonical per-token state-streaming pattern:docs.copilotkit.ai/integrations/langgraph/shared-state/predictive-state-updates"""import uuidfrom langchain.agents import AgentState as BaseAgentState, create_agentfrom langchain.tools import ToolRuntime, toolfrom langchain_core.messages import ToolMessagefrom langchain_openai import ChatOpenAIfrom langgraph.types import Commandfrom copilotkit import (    CopilotKitMiddleware,    StateItem,    StateStreamingMiddleware,)class AgentState(BaseAgentState):    """Shared state. `document` is streamed token-by-token."""    document: str@tooldef write_document(document: str, runtime: ToolRuntime) -> Command:    """Write a document for the user.    Always call this tool when the user asks you to write or draft    something of any length (an essay, poem, email, summary, etc.).    The `document` argument is streamed *per token* into shared agent    state under the `document` key, so the UI can render it as it is    generated.    """    return Command(        update={            "document": document,            "messages": [                ToolMessage(                    content="Document written to shared state.",                    name="write_document",                    id=str(uuid.uuid4()),                    tool_call_id=runtime.tool_call_id,                )            ],        }    )graph = create_agent(    model=ChatOpenAI(model="gpt-5.4"),    tools=[write_document],    middleware=[        CopilotKitMiddleware(),        # Forward every token of write_document's `document` argument        # straight into state["document"] while the tool call is still        # streaming. Without this, `document` would only update once        # the tool call completes.        #        # NOTE: the frontend `usePredictStateSubscription` hook indexes        # the (partial-JSON-parsed) tool args by `state_key`, so the        # tool's argument name MUST match `state_key` ("document") for        # per-token deltas to land in `state.document`.        StateStreamingMiddleware(            StateItem(                state_key="document",                tool="write_document",                tool_argument="document",            )        ),    ],    state_schema=AgentState,    system_prompt=(        "You are a collaborative writing assistant. Whenever the user asks "        "you to write, draft, or revise any piece of text, ALWAYS call the "        "`write_document` tool with the full content as a single string in "        "the `document` argument. Never paste the document into a chat "        "message directly — the document belongs in shared state and the "        "UI renders it live as you type."    ),)

What is this?#

By default, agent state only updates between LangGraph node transitions, so a long-running tool call (writing a full document, drafting an email) appears to the UI as one big burst at the end. For agent-native apps, that feels broken: users expect to watch the output materialise.

State streaming forwards the value of a specific tool argument straight into an agent state key as the argument is being generated. The UI, subscribed via useAgent, re-renders every token.

When should I use this?#

Use state streaming whenever a tool's output is long-form text or a growing structured value and you want the user to see it assemble in real time. Common shapes:

  • A collaborative writing agent that emits a document
  • A research agent that accumulates a list of findings
  • A planning agent that builds up a step-by-step plan

Without streaming, the user stares at a spinner. With streaming, they see the answer grow token-by-token.

The backend: one streaming state mapping#

Install the LangGraph Python SDK

uv add copilotkit
poetry add copilotkit
pip install copilotkit --extra-index-url https://copilotkit.gateway.scarf.sh/simple/
conda install copilotkit -c copilotkit-channel

Add state streaming middleware

StateStreamingMiddleware maps a tool argument to an agent state key and forwards partial tool-call arguments as state updates while the model is still generating.

shared_state_streaming.py
import uuid

from langchain.agents import AgentState as BaseAgentState, create_agent
from langchain.tools import ToolRuntime, tool
from langchain_core.messages import ToolMessage
from langchain_openai import ChatOpenAI
from langgraph.types import Command

from copilotkit import (
    CopilotKitMiddleware,
    StateItem,
    StateStreamingMiddleware,
)


class AgentState(BaseAgentState):
    """Shared state. `document` is streamed token-by-token."""

    document: str


@tool
def write_document(document: str, runtime: ToolRuntime) -> Command:
    """Write a document for the user.

    Always call this tool when the user asks you to write or draft
    something of any length (an essay, poem, email, summary, etc.).
    The `document` argument is streamed *per token* into shared agent
    state under the `document` key, so the UI can render it as it is
    generated.
    """
    return Command(
        update={
            "document": document,
            "messages": [
                ToolMessage(
                    content="Document written to shared state.",
                    name="write_document",
                    id=str(uuid.uuid4()),
                    tool_call_id=runtime.tool_call_id,
                )
            ],
        }
    )


graph = create_agent(
    model=ChatOpenAI(model="gpt-5.4"),
    tools=[write_document],
    middleware=[
        CopilotKitMiddleware(),
        # Forward every token of write_document's `document` argument
        # straight into state["document"] while the tool call is still
        # streaming. Without this, `document` would only update once
        # the tool call completes.
        #
        # NOTE: the frontend `usePredictStateSubscription` hook indexes
        # the (partial-JSON-parsed) tool args by `state_key`, so the
        # tool's argument name MUST match `state_key` ("document") for
        # per-token deltas to land in `state.document`.
        StateStreamingMiddleware(
            StateItem(
                state_key="document",
                tool="write_document",
                tool_argument="document",
            )
        ),
    ],
    state_schema=AgentState,
    system_prompt=(
        "You are a collaborative writing assistant. Whenever the user asks "
        "you to write, draft, or revise any piece of text, ALWAYS call the "
        "`write_document` tool with the full content as a single string in "
        "the `document` argument. Never paste the document into a chat "
        "message directly — the document belongs in shared state and the "
        "UI renders it live as you type."
    ),
)

The backend pattern is always the same: map one streaming tool argument to one shared-state key. In Python prebuilt agents, that is StateStreamingMiddleware with one or more StateItem(...) entries. TypeScript graphs use copilotkitCustomizeConfig with an emitIntermediateState mapping for the same shape. When the LLM streams that argument, CopilotKit writes every partial value into shared state before the tool even finishes executing.

shared_state_streaming.py
import uuidfrom langchain.agents import AgentState as BaseAgentState, create_agentfrom langchain.tools import ToolRuntime, toolfrom langchain_core.messages import ToolMessagefrom langchain_openai import ChatOpenAIfrom langgraph.types import Commandfrom copilotkit import (    CopilotKitMiddleware,    StateItem,    StateStreamingMiddleware,)class AgentState(BaseAgentState):    """Shared state. `document` is streamed token-by-token."""    document: str@tooldef write_document(document: str, runtime: ToolRuntime) -> Command:    """Write a document for the user.    Always call this tool when the user asks you to write or draft    something of any length (an essay, poem, email, summary, etc.).    The `document` argument is streamed *per token* into shared agent    state under the `document` key, so the UI can render it as it is    generated.    """    return Command(        update={            "document": document,            "messages": [                ToolMessage(                    content="Document written to shared state.",                    name="write_document",                    id=str(uuid.uuid4()),                    tool_call_id=runtime.tool_call_id,                )            ],        }    )graph = create_agent(    model=ChatOpenAI(model="gpt-5.4"),    tools=[write_document],    middleware=[        CopilotKitMiddleware(),        # Forward every token of write_document's `document` argument        # straight into state["document"] while the tool call is still        # streaming. Without this, `document` would only update once        # the tool call completes.        #        # NOTE: the frontend `usePredictStateSubscription` hook indexes        # the (partial-JSON-parsed) tool args by `state_key`, so the        # tool's argument name MUST match `state_key` ("document") for        # per-token deltas to land in `state.document`.        StateStreamingMiddleware(            StateItem(                state_key="document",                tool="write_document",                tool_argument="document",            )        ),    ],    state_schema=AgentState,    system_prompt=(        "You are a collaborative writing assistant. Whenever the user asks "        "you to write, draft, or revise any piece of text, ALWAYS call the "        "`write_document` tool with the full content as a single string in "        "the `document` argument. Never paste the document into a chat "        "message directly — the document belongs in shared state and the "        "UI renders it live as you type."    ),)

A few things to note:

  • The state_key must exist on your AgentState schema (document: str in this demo).
  • The tool and tool_argument name the exact LLM-facing tool and argument to forward.
  • When the tool call completes, its final return value is written to the same key, so the streamed partial eventually becomes the authoritative final value.

The frontend: useAgent + OnStateChanged#

The UI side is identical to any other shared-state subscription: useAgent with OnStateChanged gives you a reactive agent.state. Add OnRunStatusChanged if you want a "LIVE" / "done" indicator.

page.tsx
  // Subscribe to BOTH state changes and run-status changes. The former  // drives the per-token document rerender; the latter toggles the  // "LIVE" badge when the agent starts / stops.  const { agent } = useAgent({    agentId: "shared-state-streaming",    updates: [UseAgentUpdate.OnStateChanged, UseAgentUpdate.OnRunStatusChanged],  });

From there, agent.state.document is just a string that grows on every token, and agent.isRunning tells you whether to show a streaming indicator.