Working Memory (Short-term)

Agentic AI Systems advanced Building with AI 11 min read

Overview

Working memory is the ephemeral state that an AI agent holds within its current context window for immediate reasoning, tool execution, and decision-making. It is the cognitive workspace where the agent actively manipulates information — analogous to a human’s conscious thought process or a computer’s RAM. Unlike long-term memory systems that persist across sessions, working memory exists only for the duration of a single task or conversation.

For AI engineers building agentic systems, working memory is the most fundamental memory type to understand because it directly determines what the model can “see” and reason about at any given moment. Every piece of information the LLM uses — the system prompt, conversation history, retrieved documents, tool results, and intermediate reasoning steps — competes for space in this finite working memory. Managing this space effectively is the difference between an agent that reasons coherently through multi-step tasks and one that loses track of its objectives mid-execution.

The challenge is that LLMs have a fixed context window (ranging from 4K to 2M+ tokens depending on the model), and effective utilization is often far below the advertised maximum. Research shows that accuracy can drop over 30% when critical information is buried in the middle of a long context, a phenomenon known as the “lost in the middle” effect. This means working memory management is not just about fitting data in — it’s about strategically positioning and structuring information for optimal reasoning.

How It Works

The Context Window as Working Memory

An LLM’s context window functions as its working memory — the finite amount of text the model can process to generate a response. Every token in the context window consumes computational resources through the self-attention mechanism, which scales quadratically with sequence length. This creates a hard physical constraint on how much information the agent can actively reason over.

Key Concept: Advertised vs. Effective Context

A model may advertise a 128K token context window, but effective working memory — the amount of context the model can reliably use for reasoning — is significantly smaller. Frontier models often manage only a handful of active variables before reasoning quality degrades. The advertised context is like total RAM; effective working memory is like CPU cache.

Components of Working Memory

In a typical agentic system, working memory contains several distinct components that must be balanced:

The Scratchpad Pattern

One of the most effective patterns for managing working memory is the scratchpad — an explicit, structured workspace where the agent reads and writes its current state. Unlike relying on the model to implicitly “remember” information across reasoning steps, a scratchpad externalizes the agent’s internal state into a format that is both inspectable and controllable.

// Scratchpad as a tool call
const scratchpad = {
  current_objective: "Find hotels in Paris with pool amenities",
  completed_steps: [
    "Queried vector DB for 'Paris hotels'",
    "Retrieved 12 matching properties"
  ],
  pending_steps: [
    "Filter results by pool amenity",
    "Rank by user rating",
    "Format top 5 for response"
  ],
  intermediate_results: {
    total_matches: 12,
    with_pool: null  // to be filled
  },
  key_constraints: [
    "Budget: under $300/night",
    "Dates: June 15-20"
  ]
};

// Agent reads scratchpad at start of each step
// Agent writes updates after each step

Scratchpads can be implemented as tool calls that write to a file, fields in a runtime state object, or even dedicated sections within the system prompt that get updated between turns. The critical advantage is that the agent doesn’t need to hold all information in the attention mechanism simultaneously — it can selectively read and write to the scratchpad as needed.

Context Window Management Strategies

As conversations and tasks grow, working memory must be actively managed to prevent overflow and degradation:

Key Concept: Working Memory vs. Short-term Memory

These terms are often conflated but have a subtle distinction. Short-term memory stores conversation history — what was said (a transcript). Working memory stores the agent’s internal reasoning state — what it’s currently thinking about (a structured scratchpad tracking intermediate results and plan progress). In practice, both occupy the same context window, but designing them as separate concerns leads to better architectures.

Context Rot and the Middle Problem

When too much information is crammed into working memory, the model’s attention mechanism struggles to identify what’s relevant. This manifests as context rot — the agent begins to lose track of its original objectives, produce contradictory statements, or ignore critical constraints. The “lost in the middle” effect means information placed at the beginning and end of the context is reliably attended to, while information in the middle can be effectively invisible.

To combat this, engineers should place the most critical information (system prompt, current objective) at the boundaries of the context and use structured delimiters (XML tags, markdown headers) to create clear sections that help the attention mechanism identify relevant regions.

Practical Applications

Multi-Step Tool Execution

When an agent executes a chain of tool calls — querying a database, processing results, calling an API, then formatting output — each step’s result must be held in working memory. A well-designed scratchpad tracks which tools have been called, their results, and what remains to be done. Without this, agents commonly repeat tool calls they’ve already made or forget critical intermediate results.

Agentic Coding Assistants

Tools like Claude Code and GitHub Copilot manage working memory to hold the current file being edited, relevant code context, test results, and the user’s instructions simultaneously. When the context window fills up, older conversation turns are compressed or summarized to make room for new information while preserving essential context about the task.

Customer Service Agents

A support agent handling a complex ticket needs working memory to track: the customer’s original issue, account information retrieved from a CRM, previous troubleshooting steps attempted in this session, and the current action being taken. The scratchpad pattern ensures the agent doesn’t ask the customer to repeat information or re-attempt failed solutions.

Implementation with LangGraph

from langgraph.graph import StateGraph
from typing import TypedDict, List

class AgentState(TypedDict):
    messages: List[dict]        # Conversation history
    scratchpad: dict            # Structured working memory
    current_step: str           # What the agent is doing now
    tool_results: List[dict]    # Results from tool calls

def manage_working_memory(state: AgentState) -> AgentState:
    """Compress working memory when approaching limits."""
    messages = state["messages"]
    if count_tokens(messages) > MAX_CONTEXT * 0.8:
        # Summarize older messages, keep recent ones
        old_messages = messages[:-5]
        summary = summarize(old_messages)
        state["messages"] = [
            {"role": "system", "content": f"Summary: {summary}"}
        ] + messages[-5:]
    return state

graph = StateGraph(AgentState)
graph.add_node("manage_memory", manage_working_memory)
graph.add_node("reason", agent_reasoning_step)
graph.add_node("act", execute_tool)

Key Takeaways

Further Reading