Contextual Memory for Single Request Lifecycles

Introduction

In our previous lesson, we established a robust foundation for our orchestrator by implementing a SharedState class. This "blackboard" system, complete with a transactional update mechanism, ensures data integrity even when multiple agents operate in parallel. This shared state is, in essence, a primitive form of memory for the duration of a single, complex request.

Today, we will formalize and expand upon this concept as we tackle our first learning outcome in the "Conversation and State Management" module: implementing short-term memory for maintaining context within a single request lifecycle. We will explore what short-term memory means for an AI agent, differentiate between passive conversation history and active working memory, and implement a sophisticated memory system using a dedicated tool that integrates seamlessly with our existing SharedState architecture.

What is Short-Term Memory for an AI Agent?

For an agent to perform any non-trivial task, it needs to remember what has been said and what it has learned. Within the scope of a single user request, this short-term memory typically takes two forms:

1. Conversation History: This is the most basic form of memory. It is a running log of the entire interaction so far—user messages, assistant replies, and tool outputs. It provides immediate conversational context, allowing the agent to understand follow-up questions and references.
2. Working Memory: This is a more structured and dynamic "scratchpad." As an agent processes information, it can synthesize facts, infer user preferences, and store intermediate conclusions. This is distinct from the raw conversation log; it's a curated set of knowledge the agent builds and maintains to inform its decisions.

A robust agentic system needs both. The conversation history provides the "what was said," while the working memory provides the "what is known."

Basic Implementation: Passing Conversation History

The simplest way to implement short-term memory is to manage the conversation history. With each new user prompt, the entire sequence of previous messages is sent back to the LLM. The Vercel AI SDK is designed around this pattern, using an array of CoreMessage objects to represent the conversation.

Let's review this fundamental pattern, which you may have encountered in your prior work with the Vercel AI SDK.

This pattern is effective and straightforward. The messages array passed to generateText or streamText is the conversation history.

import { CoreMessage, streamText } from 'ai';
import { openai } from '@ai-sdk/openai';

// Client-side state would hold this array
const messages: CoreMessage[] = [
  { role: 'user', content: 'What is the capital of France?' },
  { role: 'assistant', content: 'The capital of France is Paris.' },
  { role: 'user', content: 'What is its population?' }
];

async function getNextResponse() {
  const { textStream } = await streamText({
    model: openai('gpt-4-turbo'),
    messages: messages, // The entire history is passed
  });
  
  // As the response streams in, you would update the UI
  // and append the final assistant message to the `messages` array.
}

While simple, this approach has a clear scaling problem: the message history grows with every turn, consuming more tokens and increasing costs. We will address this with truncation and summarization strategies in a future lesson. For now, let's focus on a more advanced way to manage the working memory.

Advanced Implementation: A Custom Memory Tool

Instead of just passively receiving history, a truly intelligent agent should be able to actively manage its own memory. We can empower it to do so by providing a memory tool. This allows the agent to decide for itself when to store a new fact, update an existing one, or search its memory for relevant information.

This approach transforms memory from a simple log into an active component of the agent's reasoning process. The Vercel AI SDK documentation provides an excellent guide for building exactly this.

Let's implement this pattern, integrating it with the SharedState class we built in the last lesson. Our SharedState instance will serve as the in-memory backend for our tool—the "filesystem abstraction" mentioned in the guide.

First, we define the tool using a Zod schema, just as described in the documentation. This tool will operate on our SharedState.

// src/tools/memory-tool.ts

import { tool } from 'ai';
import { z } from 'zod';
import { SharedState } from '../shared-state';

// The schema for the tool's input
const memoryInputSchema = z.object({
  command: z
    .enum(['read', 'write', 'search'])
    .describe('The memory action: read a key, write to a key, or search values.'),
  key: z.string().describe('The key in memory to read from or write to.'),
  value: z
    .string()
    .optional()
    .describe('The value to write. Required for the "write" command.'),
  query: z
    .string()
    .optional()
    .describe('A search query. Required for the "search" command.')
});

// The memory tool definition
export const memoryTool = tool({
  description: `Access and modify a key-value memory store. Use this to remember user preferences or facts across tasks.
- To store information: use the 'write' command with a descriptive key and the value to store.
- To retrieve information: use the 'read' command with the exact key.
- To find information when you don't know the key: use the 'search' command with a query.`,
  inputSchema: memoryInputSchema,
  execute: async ({ command, key, value, query }, { experimental_context }) => {
    // We expect the SharedState instance to be passed in the context
    const { sharedState } = experimental_context as { sharedState: SharedState };

    if (!sharedState) {
      throw new Error("SharedState was not provided to the memory tool's execution context.");
    }

    try {
      switch (command) {
        case 'write':
          if (value === undefined) {
            return { success: false, message: 'A "value" is required for the "write" command.' };
          }
          // Use our transactional write to prevent race conditions
          await sharedState.transact(key, () => value);
          return { success: true, message: `Wrote to memory with key: ${key}` };

        case 'read':
          const readValue = sharedState.read(key);
          return readValue !== undefined
            ? { success: true, key, value: readValue }
            : { success: false, message: `No value found for key: ${key}` };

        case 'search':
            if (query === undefined) {
                return { success: false, message: 'A "query" is required for the "search" command.' };
            }
            const allState = sharedState.toObject();
            const searchResults: Record<string, any> = {};
            for (const [k, v] of Object.entries(allState)) {
                if (typeof v === 'string' && v.includes(query)) {
                    searchResults[k] = v;
                }
            }
            return { success: true, results: searchResults };

        default:
          throw new Error(`Unknown memory command: ${command}`);
      }
    } catch (error) {
      return { success: false, message: (error as Error).message };
    }
  },
});

This tool is now a powerful extension for our agents. An agent can, for example, be prompted: "My favorite color is blue. Please remember this for our conversation." A well-instructed agent would recognize the instruction to "remember" and call memoryTool({ command: 'write', key: 'user_favorite_color', value: 'blue' }).

Injecting Core Memory with prepareCall

A common use for working memory is to store "core" facts that should always be available to the agent, such as the user's name, preferences, or high-level goals. Constantly forcing the agent to call the read command for this core information is inefficient.

The Vercel AI SDK provides an elegant solution: the prepareCall hook. This function runs before every LLM call within an agent's reasoning loop, allowing you to dynamically modify the prompt or other settings. We can use it to fetch core memory and inject it directly into the agent's system prompt.

Let's adapt this pattern. We'll define an agent that uses our memoryTool and has a prepareCall hook to inject information stored under a special core_memory key in our SharedState.

// src/agents/memory-aware-agent.ts

import { ToolLoopAgent } from 'ai';
import { openai } from '@ai-sdk/openai';
import { memoryTool } from '../tools/memory-tool';
import { SharedState } from '../shared-state';

const today = new Date().toISOString().slice(0, 10);

export const memoryAwareAgent = new ToolLoopAgent({
  model: openai('gpt-4-turbo'),
  tools: {
    memory: memoryTool,
  },
  prepareCall: async (settings) => {
    // The sharedState object is the single source of truth
    const { sharedState } = settings.experimental_context as { sharedState: SharedState };

    const coreMemoryContent = sharedState.read('core_memory') || 'No core memory set.';

    return {
      ...settings,
      // We inject the fetched core memory into the system prompt
      instructions: `You are a helpful assistant. Today's date is ${today}.
      
      Core Memory Snapshot:
      ---
      ${coreMemoryContent}
      ---
      
      You have a 'memory' tool to read and write information. Use it to remember facts.
      If a user asks you to remember something, use the 'write' command.`,
    };
  },
});

With this setup, if an agent uses the memoryTool to update the core_memory key, the prepareCall hook ensures that the very next thought process of the agent will see the updated information in its system prompt. This creates an immediate feedback loop, making the agent highly adaptive.

Conclusion

In this lesson, we have moved beyond the basic state management of our previous session to implement a formal short-term memory system for our agents. You learned to distinguish between passive conversation history and active working memory, and you have implemented a sophisticated, agent-controlled mechanism for managing the latter.

Key Takeaways:

• Short-term memory within a single request consists of conversation history (the raw message log) and working memory (a structured scratchpad for synthesized facts).
• The most basic memory implementation involves passing the entire CoreMessage history with each API call, a pattern native to the Vercel AI SDK.
• A more advanced and powerful approach is to provide the agent with a dedicated memory tool, allowing it to consciously decide what to remember and recall.
• We integrated this tool with our existing SharedState class, which serves as the in-memory backend, demonstrating how to build upon our core architecture.
• The prepareCall hook is a crucial feature for efficiently managing memory, enabling us to inject "core" facts into the agent's context before each reasoning step.

Our agents are now not only capable of executing tasks but also of learning and adapting within the scope of a single conversation. However, this memory is still ephemeral—it is lost once the request is complete. In our next lesson, we will tackle this limitation by exploring conversation history storage and retrieval for multi-turn interactions, allowing our agents to remember users and context across different sessions.