Writing agent state

Run and connect your agent#

Start your AG2 backend and connect your CopilotKit frontend to the AG-UI /chat endpoint.

Define the Agent State#

Create your AG2 backend with ContextVariables and emit StateSnapshotEvent whenever state changes:

from typing import Annotated

from ag_ui.core import EventType, StateSnapshotEvent
from fastapi import FastAPI, Header
from fastapi.responses import StreamingResponse
from autogen import ContextVariables, ConversableAgent, LLMConfig
from autogen.ag_ui import AGUIStream, RunAgentInput

def read_state(context: ContextVariables) -> dict:
    return context.get("agent_state", {"language": "english"})

def write_state(context: ContextVariables, state: dict) -> StateSnapshotEvent:
    context["agent_state"] = state
    return StateSnapshotEvent(type=EventType.STATE_SNAPSHOT, snapshot=state)

agent = ConversableAgent(
    name="assistant",
    system_message=(
        "You are a helpful assistant for tracking language. "
        "Always respond in the current language."
    ),
    llm_config=LLMConfig({"model": "gpt-5.4-mini"}),
)

@agent.register_for_llm(description="Update the language in shared state.")
def set_language(
    context: ContextVariables,
    language: Annotated[str, "language such as english or spanish"],
) -> StateSnapshotEvent:
    return write_state(context, {"language": language.lower()})

agent.register_for_execution(name="set_language")(set_language)

stream = AGUIStream(agent)
app = FastAPI()

@app.post("/chat")
async def run_agent(
    message: RunAgentInput,
accept: str | None = Header(None),
):
    return StreamingResponse(
        stream.dispatch(message, accept=accept),
        media_type=accept or "text/event-stream",
    )

Call setState function from the useAgent hook#

useAgent returns an agent object with a setState function that you can use to update the agent state. Calling this will update the agent state and trigger a rerender of anything that depends on the agent state.

// Define the agent state type, should match the actual state of your agent
type AgentState = {
language: "english" | "spanish";
}

// Example usage in a pseudo React component
function YourMainContent() {
  const { agent } = useAgent({ // [!code highlight]
    agentId: "my_agent", // MUST match the agent name in CopilotRuntime
    initialState: { language: "english" }  // optionally provide an initial state
  });

  // ...

  const toggleLanguage = () => {
    agent.setState({ language: agent.state.language === "english" ? "spanish" : "english" }); // [!code highlight]
  };

  // ...

  return (
    // style excluded for brevity
    <div>
      <h1>Your main content</h1>
      {/* [!code highlight:2] */}
      <p>Language: {agent.state.language}</p>
      <button onClick={toggleLanguage}>Toggle Language</button>
    </div>
  );
}

Give it a try#

You can now use agent.setState to update the agent state and agent.state to read it. Try toggling the language button and talking to your agent. You'll see the language change to match the agent's state.