3. Adding State

This section assumes you’ve completed the workflow version in the previous step.

We want to add now a state object that persists across handlers to:

1. Track the conversation history
2. Track the tool responses

Below are the changes needed to do these changes. For the changes, we’re using a new file named 3-adding-state.ts.

1. Add state middleware and a state type

Add the import of the stateful middleware and define a state object to track conversation and tool progress:

import { createStatefulMiddleware } from "@llamaindex/workflow-core/middleware/state";

type AgentWorkflowState = {
  expectedToolCount: number;
  messages: InputMessage[];
  toolResponses: Array<ToolResponseMessage>;
};

Then wrap your workflow instance using withState() to add state support:

const stateful = createStatefulMiddleware((state: AgentWorkflowState) => state);
const workflow = stateful.withState(createWorkflow());

2. Use state inside the user input handler

Now we can simplify the initial userInputEvent handler by storing the messages in the state and removing the logic of collecting tool responses (since later we will put the tool responses into the context state):

workflow.handle([userInputEvent], async (context, { data }) => {
  const { sendEvent, state } = context;
  const { messages } = data;

  const response = await llm(messages, tools);
  state.messages = [...messages, response];

  if (response.tool_calls && response.tool_calls.length > 0) {
    state.expectedToolCount = response.tool_calls.length;
    for (const toolCall of response.tool_calls) {
      if (toolCall.type !== "function") {
        throw new Error("Unsupported tool call type");
      }
      sendEvent(toolCallEvent.with({ toolCall }));
    }
  } else {
    sendEvent(finalResponseEvent.with(response.content || ""));
  }
});

3. Add a handler to deal with tool responses

Instead of collecting all tool responses locally, we add a new handler dealing with tool responses and storing them in the state:

workflow.handle([toolResponseEvent], async (context, { data }) => {
  const { sendEvent, state } = context;

  state.toolResponses.push(data);

  if (state.toolResponses.length === state.expectedToolCount) {
    const finalMessages = [...state.messages, ...state.toolResponses];
    sendEvent(userInputEvent.with({ messages: finalMessages }));
  }
});

Once all tool responses have arrived, we continue with a userInputEvent containing all the tool response messages to resume the LLM interaction.

4. Initialize the context with default state

As we’re adding a state object, we need to provide initial values when creating the workflow context:

const { stream, sendEvent } = workflow.createContext({
  expectedToolCount: 0,
  messages: [],
  toolResponses: [],
});

What stayed the same

• The event definitions (userInputEvent, toolCallEvent, toolResponseEvent, finalResponseEvent) remain the same.
• The llm() and callTool() functions are the same as before.
• The handler that executes tool calls and sends back a toolResponseEvent remains unchanged.

Why this helps

• Centralized state: Conversation history and pending tool responses live in one place.
• Simpler control flow: Handlers are smaller; no ad-hoc streaming aggregation.
• Extensible: Easy to add more state (e.g., retries, metrics) without changing event shapes.

If you want to see the complete stateful version in one place, check demo/express/3-adding-state.ts.