In-Context Memory Management

8 min read Module 4 of 10 Topic 10 of 30

What you'll learn

  • Understand the context window as a finite resource and its impact on agent cost and reliability
  • Implement message trimming strategies to keep token usage bounded
  • Use running summarization to preserve information from trimmed context
  • Design a scratchpad pattern for multi-step agent state accumulation
Building this at your company? For enterprise and company teams taking this to production: book a 30-minute session with our AI engineers for architecture guidance, code review, and a rollout plan for your use case.
Book a Team Session

The Context Window Is Your Most Precious Resource

Every LLM has a finite context window, the number of tokens it can process in a single call. For production agents, this isn’t an abstract concern: it’s the primary constraint that shapes architecture.

Here’s the math for a typical multi-step agent task:

System prompt:           ~500 tokens
Initial user message:    ~100 tokens
Tool schemas:            ~300 tokens
× 15 iterations:
  - Each LLM response:   ~200 tokens
  - Each tool result:    ~500 tokens average

Total after 15 steps:    ~(500+100+300) + 15×700 = ~11,400 tokens

For a task that runs 30-50 steps (common for complex research or coding tasks), you’re looking at 20,000-40,000 tokens. At $3/M tokens (GPT-5.6 Sol), that’s $0.06-$0.12 per run. At scale (10,000 runs/day), that’s $600-$1,200/day just for input tokens.

Context management is engineering, not optimization.


The Five Zones of Context

graph TB
    Z1["SYSTEM PROMPT + Tool Schemas\nAlways include, the contract\nNever trimmed"]
    Z2["COMPRESSED HISTORY\nOld messages, summarized\nKeep condensed, rotate out"]
    Z3["RECENT MESSAGES last N turns\nFull detail preserved\nSliding window"]
    Z4["SCRATCHPAD / Accumulated Facts\nWorking memory, curated\nAgent-managed"]
    Z5["RESPONSE BUFFER\n~2K tokens headroom\nSpace for the next reply"]
    Z1 --- Z2 --- Z3 --- Z4 --- Z5
    style Z1 fill:#EEF0F7,stroke:#6366F1,color:#0F172A
    style Z2 fill:#EEF0F7,stroke:#818CF8,color:#0F172A
    style Z3 fill:#fff7ed,stroke:#f59e0b,color:#0F172A
    style Z4 fill:#fff7ed,stroke:#f59e0b,color:#0F172A
    style Z5 fill:#f0fdf9,stroke:#0D9488,color:#0F172A

Message Trimming with LangChain

trim_messages reduces the token footprint of older messages while always preserving the system prompt and the most recent turns where context is densest.

from langchain_core.messages import trim_messages, SystemMessage

def trim_conversation_history(
    messages: list[BaseMessage],
    max_tokens: int = 8000,
    keep_last_n: int = 10,
) -> list[BaseMessage]:
    """Trim messages while preserving the system prompt and recent context."""
    
    # System messages hold the agent's instructions: always keep them regardless of token budget
    system_msgs = [m for m in messages if isinstance(m, SystemMessage)]
    other_msgs = [m for m in messages if not isinstance(m, SystemMessage)]
    
    # Sliding window: always keep the last N messages in full (highest information density)
    recent = other_msgs[-keep_last_n:]
    older = other_msgs[:-keep_last_n]
    
    if not older:
        # nothing old enough to trim: return everything as-is
        return system_msgs + recent
    
    # Trim older messages by token count
    trimmed = trim_messages(
        older,
        max_tokens=max_tokens - 2000,  # reserve 2K for system + recent messages
        token_counter=len,  # replace with tiktoken for accuracy, len() counts chars, not tokens
        strategy="last",    # keep the most recent of the older messages, drop the oldest first
        include_system=False,
    )
    
    return system_msgs + trimmed + recent

# Use in a LangGraph node
def call_llm_with_trimming(state: AgentState) -> dict:
    # trim before each LLM call so token count stays bounded regardless of run length
    trimmed_messages = trim_conversation_history(
        state["messages"],
        max_tokens=12000,
        keep_last_n=8,
    )
    
    # the model sees the trimmed list: the full history still lives in state["messages"]
    response = llm_with_tools.invoke(trimmed_messages)
    return {"messages": [response]}

Running Summarization

When trimming would lose important information, summarize it instead:

class SummarizingState(TypedDict):
    messages: Annotated[list[BaseMessage], add_messages]
    summary: str  # running summary of trimmed history, starts empty, grows over time
    message_count: int

SUMMARIZE_EVERY_N = 10  # summarize every N messages to bound context growth

def maybe_summarize(state: SummarizingState) -> dict:
    """Summarize old messages when the conversation gets long."""
    
    # below the threshold: no action needed
    if len(state["messages"]) < SUMMARIZE_EVERY_N:
        return {}
    
    # keep the 4 most recent messages in full: summarize everything older
    messages_to_summarize = state["messages"][:-4]
    recent_messages = state["messages"][-4:]
    
    # Build summarization prompt
    existing_summary = state.get("summary", "")
    # format messages as "type: content" lines for the summarization LLM
    context = "\n".join([f"{m.type}: {m.content}" for m in messages_to_summarize])
    
    summary_prompt = f"""
Update this summary of a conversation by incorporating the new messages.

Existing summary:
{existing_summary or "None yet."}

New messages to incorporate:
{context}

Provide an updated, concise summary that preserves all important facts, decisions, and findings.
"""
    
    summary_response = llm.invoke([HumanMessage(content=summary_prompt)])
    
    return {
        # replace the full message list with only recent messages: older ones are now in the summary
        "messages": recent_messages,
        "summary": summary_response.content,
        "message_count": len(recent_messages),
    }

def call_llm_with_summary(state: SummarizingState) -> dict:
    """Inject the running summary as context for the model."""
    messages = list(state["messages"])
    
    if state.get("summary"):
        # prepend the compressed history as a SystemMessage so the model treats it as background context
        summary_msg = SystemMessage(
            content=f"CONVERSATION HISTORY SUMMARY:\n{state['summary']}\n\n"
                    f"The above is a summary of earlier conversation. Current messages follow:"
        )
        messages = [summary_msg] + messages
    
    response = llm_with_tools.invoke(messages)
    return {"messages": [response]}

The Scratchpad Pattern

A scratchpad is a dedicated state field for the agent to accumulate curated information separate from raw message history:

class AgentStateWithScratchpad(TypedDict):
    messages: Annotated[list[BaseMessage], add_messages]
    scratchpad: str  # agent's self-maintained working notes, separate from message history
    task: str

def update_scratchpad(state: AgentStateWithScratchpad) -> dict:
    """Let the model update its own scratchpad after each tool result."""
    # state["messages"][-1] is the most recent ToolMessage after a tool call completes
    last_tool_result = state["messages"][-1].content
    
    update_prompt = f"""
You are maintaining a research scratchpad for the task: {state['task']}

Current scratchpad:
{state['scratchpad'] or '(empty)'}

New information from the last tool call:
{last_tool_result}

Update the scratchpad with any new facts, insights, or progress notes.
Keep it concise (max 500 words) and focused on what's needed for the task.
Discard information that is no longer relevant.
"""
    
    result = llm.invoke([HumanMessage(content=update_prompt)])
    # returns only the scratchpad field: messages are unchanged by this node
    return {"scratchpad": result.content}

def call_llm_with_scratchpad(state: AgentStateWithScratchpad) -> dict:
    """Give the model access to its scratchpad in the system context."""
    
    # Build context-efficient message list: scratchpad in system prompt, only last 6 raw messages
    # This means the model's curated notes are always present without bloating the message list
    messages = [
        SystemMessage(content=f"""Task: {state['task']}

Your current notes:
{state['scratchpad'] or 'No notes yet.'}

Use the tools to complete the task. Update your notes as you discover information."""),
    ] + state["messages"][-6:]  # only last 6 messages in full, older content lives in scratchpad
    
    response = llm_with_tools.invoke(messages)
    return {"messages": [response]}

Token Counting with Tiktoken

For accurate token budgeting, count tokens before constructing your context:

import tiktoken

# load the tokenizer for the specific model: token counts vary between model families
encoder = tiktoken.encoding_for_model("gpt-5.6-sol")

def count_tokens(messages: list[BaseMessage]) -> int:
    total = 0
    for msg in messages:
        # handle both string content and structured content (list of dicts for multimodal)
        content = msg.content if isinstance(msg.content, str) else str(msg.content)
        # +4 per message accounts for the role header and separator tokens the API adds
        total += len(encoder.encode(content)) + 4  # +4 for message overhead
    return total

def build_context_within_budget(
    system_prompt: str,
    messages: list[BaseMessage],
    max_tokens: int = 100_000,
    response_budget: int = 4_000,  # reserve tokens for the model's output
) -> list[BaseMessage]:
    """Build a message list that stays within token budget."""
    # subtract the response budget: we need headroom for the model to reply
    available = max_tokens - response_budget
    system_msg = SystemMessage(content=system_prompt)
    # always include the system prompt first, then subtract its cost from the budget
    available -= count_tokens([system_msg])
    
    # Add messages from most recent backward until budget is exhausted
    # greedy from the end preserves recency: oldest messages are dropped first
    selected = []
    for msg in reversed(messages):
        msg_tokens = count_tokens([msg])
        if available - msg_tokens < 500:  # keep 500 token buffer to avoid off-by-one overflows
            break
        selected.insert(0, msg)
        available -= msg_tokens
    
    return [system_msg] + selected

The combination of trimming, summarization, scratchpad, and accurate token counting gives you precise control over what the model sees at each step: keeping agents fast, cheap, and reliable even on long-running tasks.

Knowledge Check

3 questions to test your understanding

1 Why does context window exhaustion cause agent failures even when the model has a large context (e.g. 200K tokens)?

2 What is the role of a 'scratchpad' in an agent's state?

3 When should you summarize old messages rather than simply trimming them?

Discussion

Questions and notes from learners on this topic

Loading discussion…

Go further with expert guidance

Ready to build production AI?
Talk to our R&D team.

These courses give you the foundation. Our embedded AI teams take you from prototype to production in 30–90 days, with your team, your codebase, your goals. Book a free strategy call to see how we can accelerate your AI initiative.

30 minutes · No obligation · Expert AI engineers, not sales reps

AI Architecture Review

Audit your current stack and identify high-impact improvements

Project Review

Get expert feedback on your AI implementation and codebase

Team Mentoring

Upskill your engineers with hands-on AI coaching sessions

AI Strategy

Define your AI roadmap, prioritization, and implementation plan