Remember this: Agent state and memory with ADK

Imagine that you're a developer, and your friend is learning to code for the first time. They're struggling with some concepts, like Python dictionaries. 

But you have an idea! What if you could design an AI agent that could help your friend learn complex topics in a conversational way? What if that agent could even be personalized to your friend’s learning style, remember your friend's past performance, and adjust the learning plan in real time? With agent state and memory, all of this is possible. In this post, we'll explore how, with Agent Development Kit (ADK).

The Python Tutor agent

We'll start by designing a simple, conversational agent with ADK. This agent uses Gemini 2.5 Flash as its model, leveraging its reasoning capabilities. The agent also relies on a set of function tools that allow the agent to progress through a simple quiz on Python dictionaries.

For example, here's the start_quiz tool, which kicks off a new quiz workflow. (More on state in a minute!)

Overall, this is the quiz workflow we want our agent to take, using those tools together:

From there, we prompt our agent with system instructions to help the model reason its way through that workflow:

From here, we can define our ADK agent, and get cooking:

But how can we ask this agent to remember things that happen during a student's session, like their name, or their progress through the quiz? This is where short-term memory comes into play.

Short-term memory

In the context of an AI agent, short-term memory is what the agent can remember within one session. 

But what is a session? Think of a session like a phone call with a customer support representative—but you get a different rep every time you call. That representative can only remember what you've told them during that conversation. After you hang up the phone, all context is lost.

Short-term memory might sound like a bad thing, but it plays an important role in an AI agent. For instance, our Python Tutor agent might need to keep track of the user's quiz progress, like the number of questions completed so far. The agent probably does not need to store that progress long-term—it might just need to store their final score. 

Every user interaction with an ADK agent gets a session, and that session is managed by the ADK SessionService. Each session contains important fields, like the session ID, user ID, event history (the conversation thread), and the state.

What is session state? Think of it like the agent's scratchpad during that "phone call" with the user. Each session's state contains a list of key-value pairs, whose values are updated by the agent throughout the session.

ADK can write to session state a few different ways. One way is within a tool. We can use ADK's ToolContext to get the current session state, and create or update fields:

From there, we can instruct our Agent's model to read from the current state fields using ADK's key templating feature. All you have to do is wrap a state field in curly braces {}, and ADK will inject the state key's value into your prompt, on each model call.

Here's our Python Tutor agent updating those state fields in real-time:

By default, state fields only persist within the current session; once you start a new session, even as the same user, the values are gone. But ADK does have magic state key prefixes like user: and app: , which allow you to persist state key values either across all user sessions, or across all sessions with all users. These magic prefixes are useful if you have simple text settings you want to persist across sessions, like dark-mode=true. 

So those are the basics of ADK's short-term memory. But how does ADK store session and state data?

By default, the ADK web UI's SessionService writes session data in memory. This means that if ADK's runner crashes, or is shut down, all session data is lost. And if you're running a scaled, production-grade agent with ADK, with multiple instances of your ADK agent, you can't guarantee that user requests will always hit the same instance. This means that if request 1 goes to instance A, and request 2 goes to instance B, instance B won't have the in-memory session state stored inside instance A. 

So for production-grade agents, you should store session data outside of the agent's runtime. ADK provides two ways of doing this. The first is the DatabaseSessionService: store session data in a SQL database, like SQLLite, MySQL, or PostgreSQL. This is easy to set up - all you need is a database. Then, you can pass your database's URI into the ADK runner:

From there, you can access your SQL database and see the session and state tables.

The other option is a VertexAISessionService, where you store session data in Agent Engine. This is a good option if you're already using Agent Engine as your ADK runtime. 

Long-term memory

We've covered how ADK stores data within a session. But what if you have data you want to persist across sessions with the same user? This is where ADK's long-term memory comes in. 

An ADK agent with long-term memory is like talking to the same customer service representative every time. And that representative has access to all key information from all your past conversations.

For instance, our Python Tutor agent might want to store quiz scores over time: how has the student improved? Long-term memories could also help the agent develop a personalized learning plan for the student: what topics has the student consistently struggled with? 

ADK offers two ways to store these long-term memories. The default way is in memory, by using an InMemoryMemoryService. Here, all sessions are stored raw (with the full conversation thread), and can be retrieved by the agent in further sessions using a basic keyword search. This method is good for local development, but it has the same pitfall as the InMemorySessionService: if ADK restarts or crashes, all memories are lost forever. Another downside to this method is if you have a lot of past user sessions, and you're retrieving the session's event history in its raw form, you could overwhelm your agent's model with too much context.  

Luckily, ADK provides a way to store long-term memories persistently outside the ADK runtime, and that's with a VertexAIMemoryBankService. This memory service uses Vertex AI Memory Bank (Preview) to intelligently store and retrieve memories from past user interactions. 

Memory Bank uses the Gemini model to extract key information from session data, to store just the key memories for future use:

To store memories in Memory Bank, we implement a Callback in our Python Tutor agent to add the current session data to the bank:

Then, we instruct our agent to retrieve those memories using a vector-based similarity search, by keyword.

Now, we have an updated quiz flow that can remember the user's past quiz results:

Here's the updated Python Tutor agent in action, using Memory Bank for long-term memory retrieval:

What's happening under the hood, here? When our agent triggers memory generation in that after-agent Callback, Memory Bank uses Gemini to extract key information. So when the user starts a new quiz, triggering search_memory for past results, here's what Memory Bank hands back:

Note that while Memory Bank is part of Vertex AI Agent Engine, you don't need to run your ADK agent on Agent Engine's Runtime to use this feature. For instance, you can run your agent on Cloud Run, and just integrate Memory Bank as your long-term memory store. 

Get building! 

To recap, short- and long-term memory play a key role in AI agents, allowing the agent to remember key information within and across user sessions, resulting in a more contextually-aware and personalized user experience. ADK supports a variety of session and memory storage options, from SQL databases to Vertex AI Agent Engine. 

Check out the source code on GitHub to explore and deploy the Python Tutor agent yourself. 

And to learn more, check out these resources. 

• ADK Documentation - Session
• ADK Documentation - Memory
• Vertex AI Memory Bank documentation 
• ADK Documentation - Deploy an agent to Cloud Run 

Thanks for reading! 

Image credits: Emoji Kitchen