How to maintain state in a LlamaIndex agent

Agent sessions often need to carry facts, tool results, or workflow counters beyond one prompt. A LlamaIndex agent keeps that session data in a Context object so a later run can continue from values saved during an earlier run.

The Context belongs to the workflow that created it. JsonSerializer converts the object into JSON-compatible data, while Context.from_dict() rebuilds it for another workflow.run(…, ctx=restored_ctx) call using the same workflow definition.

A deterministic FunctionCallingLLM test double exercises both agent turns without hosted model credentials. Its first tool call saves the name, and its post-restoration tool call reads that value from state before producing the later agent response.

Steps to maintain LlamaIndex agent state with Context:

Install the LlamaIndex core package in the active Python environment.
```
$ python3 -m pip install --upgrade llama-index-core
```
A project virtual environment prevents this package from changing a shared system Python environment.
Related: How to install LlamaIndex with pip

Create llamaindex-agent-state.py with the foundational state-tool section.

llamaindex-agent-state.py

import asyncio
import json
from typing import Any
 
from llama_index.core.agent.workflow import AgentWorkflow
from llama_index.core.base.llms.types import (
    ChatMessage,
    ChatResponse,
    CompletionResponse,
    LLMMetadata,
    MessageRole,
)
from llama_index.core.llms.function_calling import FunctionCallingLLM
from llama_index.core.llms.llm import ToolSelection
from llama_index.core.workflow import Context, JsonSerializer
 
TOOL_CALL_LOG: list[str] = []
 
 
async def set_name(ctx: Context, name: str) -> str:
    """Save the user's preferred name in agent state."""
    async with ctx.store.edit_state() as ctx_state:
        ctx_state["state"]["name"] = name
    TOOL_CALL_LOG.append(f"set_name(name={name})")
    return name
 
 
async def read_name(ctx: Context) -> str:
    """Read the user's preferred name from agent state."""
    state = await ctx.store.get("state")
    TOOL_CALL_LOG.append("read_name()")
    return state["name"]

The workflow stores initial_state under the Context store's state key. edit_state() keeps the name change inside the Context that will be serialized.

Add the deterministic model metadata and tool-call parser below read_name().

class StateAwareMockLLM(FunctionCallingLLM):
    @property
    def metadata(self) -> LLMMetadata:
        return LLMMetadata(
            is_chat_model=True,
            is_function_calling_model=True,
            model_name="state-aware-mock-llm",
        )
 
    def _prepare_chat_with_tools(
        self,
        tools: list[Any],
        user_msg: str | ChatMessage | None = None,
        chat_history: list[ChatMessage] | None = None,
        **kwargs: Any,
    ) -> dict[str, Any]:
        messages = list(chat_history or [])
        if isinstance(user_msg, str):
            messages.append(ChatMessage(role=MessageRole.USER, content=user_msg))
        elif isinstance(user_msg, ChatMessage):
            messages.append(user_msg)
        return {"messages": messages}
 
    def get_tool_calls_from_response(
        self,
        response: ChatResponse,
        error_on_no_tool_call: bool = True,
        **kwargs: Any,
    ) -> list[ToolSelection]:
        calls = response.message.additional_kwargs.get("tool_calls", [])
        if not calls and error_on_no_tool_call:
            raise ValueError("No tool calls returned")
        return calls

Add the mock model's state-tool selection below get_tool_calls_from_response().

    def chat(self, messages: list[ChatMessage], **kwargs: Any) -> ChatResponse:
        if messages and messages[-1].role == MessageRole.TOOL:
            return ChatResponse(
                message=ChatMessage(
                    role=MessageRole.ASSISTANT,
                    content=f"The saved name is {messages[-1].content}.",
                )
            )
 
        latest_user = next(
            str(message.content)
            for message in reversed(messages)
            if message.role == MessageRole.USER
        )
        if "What name did I ask" not in latest_user:
            tool_call = ToolSelection(
                tool_id="call_set_name",
                tool_name="set_name",
                tool_kwargs={"name": "Laurie"},
            )
        else:
            tool_call = ToolSelection(
                tool_id="call_read_name",
                tool_name="read_name",
                tool_kwargs={},
            )
 
        return ChatResponse(
            message=ChatMessage(
                role=MessageRole.ASSISTANT,
                content="",
                additional_kwargs={"tool_calls": [tool_call]},
            )
        )

The second response uses the tool message content rather than a hard-coded name. A provider-backed application can replace StateAwareMockLLM after this local state path works.

Complete the model adapters below chat().

    async def achat(
        self, messages: list[ChatMessage], **kwargs: Any
    ) -> ChatResponse:
        return self.chat(messages, **kwargs)
 
    def stream_chat(self, messages: list[ChatMessage], **kwargs: Any):
        yield self.chat(messages, **kwargs)
 
    async def astream_chat(self, messages: list[ChatMessage], **kwargs: Any):
        async def response_stream():
            yield self.chat(messages, **kwargs)
 
        return response_stream()
 
    def complete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
        return CompletionResponse(text="")
 
    async def acomplete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
        return self.complete(prompt, **kwargs)
 
    def stream_complete(self, prompt: str, **kwargs: Any):
        yield CompletionResponse(text="")
 
    async def astream_complete(self, prompt: str, **kwargs: Any):
        yield CompletionResponse(text="")

Add the workflow factory below StateAwareMockLLM.

def build_workflow() -> AgentWorkflow:
    return AgentWorkflow.from_tools_or_functions(
        [set_name, read_name],
        llm=StateAwareMockLLM(),
        system_prompt="Use the state tools to save or retrieve the user's name.",
        initial_state={"name": "unset"},
    )

Start main() with the first agent turn below build_workflow().

async def main() -> None:
    workflow = build_workflow()
    ctx = Context(workflow)
 
    saved_response = await workflow.run(
        user_msg="Remember that my name is Laurie.", ctx=ctx
    )
    print(f"saved_agent={saved_response}")

The first agent turn calls set_name() with the active Context instead of changing the state directly outside the agent loop.

Append the Context serialization block inside main().
```
    ctx_dict = ctx.to_dict(serializer=JsonSerializer())
    with open("agent-state.json", "w", encoding="utf-8") as state_file:
        json.dump(ctx_dict, state_file)
```
JsonSerializer accepts only JSON-serializable values in state. API clients, model handles, open files, and indexes belong outside the saved state.

Append the Context restoration block after the JSON write.

    with open("agent-state.json", encoding="utf-8") as state_file:
        saved_ctx_dict = json.load(state_file)
 
    restored_ctx = Context.from_dict(
        workflow,
        saved_ctx_dict,
        serializer=JsonSerializer(),
    )

The restored dictionary must use the workflow definition that created the original Context.

Append the post-restoration agent turn below Context.from_dict().

    restored_response = await workflow.run(
        user_msg="What name did I ask you to remember?", ctx=restored_ctx
    )
    print(f"restored_agent={restored_response}")
    print(f"tool_calls={TOOL_CALL_LOG}")

The second workflow.run() receives restored_ctx. Its read_name() tool call must return the value saved by the earlier agent turn.

Finish the program with the asynchronous entry point.

if __name__ == "__main__":
    asyncio.run(main())

Inspect the assembled llamaindex-agent-state.py file for every constructed section.

llamaindex-agent-state.py

import asyncio
import json
from typing import Any
 
from llama_index.core.agent.workflow import AgentWorkflow
from llama_index.core.base.llms.types import (
    ChatMessage,
    ChatResponse,
    CompletionResponse,
    LLMMetadata,
    MessageRole,
)
from llama_index.core.llms.function_calling import FunctionCallingLLM
from llama_index.core.llms.llm import ToolSelection
from llama_index.core.workflow import Context, JsonSerializer
 
TOOL_CALL_LOG: list[str] = []
 
 
async def set_name(ctx: Context, name: str) -> str:
    """Save the user's preferred name in agent state."""
    async with ctx.store.edit_state() as ctx_state:
        ctx_state["state"]["name"] = name
    TOOL_CALL_LOG.append(f"set_name(name={name})")
    return name
 
 
async def read_name(ctx: Context) -> str:
    """Read the user's preferred name from agent state."""
    state = await ctx.store.get("state")
    TOOL_CALL_LOG.append("read_name()")
    return state["name"]
 
 
class StateAwareMockLLM(FunctionCallingLLM):
    @property
    def metadata(self) -> LLMMetadata:
        return LLMMetadata(
            is_chat_model=True,
            is_function_calling_model=True,
            model_name="state-aware-mock-llm",
        )
 
    def _prepare_chat_with_tools(
        self,
        tools: list[Any],
        user_msg: str | ChatMessage | None = None,
        chat_history: list[ChatMessage] | None = None,
        **kwargs: Any,
    ) -> dict[str, Any]:
        messages = list(chat_history or [])
        if isinstance(user_msg, str):
            messages.append(ChatMessage(role=MessageRole.USER, content=user_msg))
        elif isinstance(user_msg, ChatMessage):
            messages.append(user_msg)
        return {"messages": messages}
 
    def get_tool_calls_from_response(
        self,
        response: ChatResponse,
        error_on_no_tool_call: bool = True,
        **kwargs: Any,
    ) -> list[ToolSelection]:
        calls = response.message.additional_kwargs.get("tool_calls", [])
        if not calls and error_on_no_tool_call:
            raise ValueError("No tool calls returned")
        return calls
 
    def chat(self, messages: list[ChatMessage], **kwargs: Any) -> ChatResponse:
        if messages and messages[-1].role == MessageRole.TOOL:
            return ChatResponse(
                message=ChatMessage(
                    role=MessageRole.ASSISTANT,
                    content=f"The saved name is {messages[-1].content}.",
                )
            )
 
        latest_user = next(
            str(message.content)
            for message in reversed(messages)
            if message.role == MessageRole.USER
        )
        if "What name did I ask" not in latest_user:
            tool_call = ToolSelection(
                tool_id="call_set_name",
                tool_name="set_name",
                tool_kwargs={"name": "Laurie"},
            )
        else:
            tool_call = ToolSelection(
                tool_id="call_read_name",
                tool_name="read_name",
                tool_kwargs={},
            )
 
        return ChatResponse(
            message=ChatMessage(
                role=MessageRole.ASSISTANT,
                content="",
                additional_kwargs={"tool_calls": [tool_call]},
            )
        )
 
    async def achat(
        self, messages: list[ChatMessage], **kwargs: Any
    ) -> ChatResponse:
        return self.chat(messages, **kwargs)
 
    def stream_chat(self, messages: list[ChatMessage], **kwargs: Any):
        yield self.chat(messages, **kwargs)
 
    async def astream_chat(self, messages: list[ChatMessage], **kwargs: Any):
        async def response_stream():
            yield self.chat(messages, **kwargs)
 
        return response_stream()
 
    def complete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
        return CompletionResponse(text="")
 
    async def acomplete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
        return self.complete(prompt, **kwargs)
 
    def stream_complete(self, prompt: str, **kwargs: Any):
        yield CompletionResponse(text="")
 
    async def astream_complete(self, prompt: str, **kwargs: Any):
        yield CompletionResponse(text="")
 
 
def build_workflow() -> AgentWorkflow:
    return AgentWorkflow.from_tools_or_functions(
        [set_name, read_name],
        llm=StateAwareMockLLM(),
        system_prompt="Use the state tools to save or retrieve the user's name.",
        initial_state={"name": "unset"},
    )
 
 
async def main() -> None:
    workflow = build_workflow()
    ctx = Context(workflow)
 
    saved_response = await workflow.run(
        user_msg="Remember that my name is Laurie.", ctx=ctx
    )
    print(f"saved_agent={saved_response}")
 
    ctx_dict = ctx.to_dict(serializer=JsonSerializer())
    with open("agent-state.json", "w", encoding="utf-8") as state_file:
        json.dump(ctx_dict, state_file)
 
    with open("agent-state.json", encoding="utf-8") as state_file:
        saved_ctx_dict = json.load(state_file)
 
    restored_ctx = Context.from_dict(
        workflow,
        saved_ctx_dict,
        serializer=JsonSerializer(),
    )
    restored_response = await workflow.run(
        user_msg="What name did I ask you to remember?", ctx=restored_ctx
    )
    print(f"restored_agent={restored_response}")
    print(f"tool_calls={TOOL_CALL_LOG}")
 
 
if __name__ == "__main__":
    asyncio.run(main())

Run the completed state round trip.
```
$ python3 llamaindex-agent-state.py
saved_agent=The saved name is Laurie.
restored_agent=The saved name is Laurie.
tool_calls=['set_name(name=Laurie)', 'read_name()']
```
The restored response must come after read_name() in the tool log. A missing call or an unset response means the later agent turn did not consume the saved Context value.

Author: Mohd Shakir Zakaria
Mohd Shakir Zakaria is a cloud architect with deep roots in software development and open-source advocacy. Certified in AWS, Red Hat, VMware, ITIL, and Linux, he specializes in designing and managing robust cloud and on-premises infrastructures.