Agent Memory and State Management: Short-Term/Long-Term Memory and Context Compression

This post is Part 2 of the 'AI Agent Engineering' series. Moving beyond simple conversation history passing, we dive deep into state schema design, short-term/long-term memory architectures, and context window compression techniques required to reliably execute complex multi-step tasks.

1. Agent State Schema Design: Pydantic v2 & TypedDict

Large Language Models (LLMs) themselves are merely stateless functions. Given the same prompt and options, they produce independent token outputs each time and cannot remember previous calls or tool execution results on their own.

Therefore, the core of an AI Agent runtime lies in: "How to consistently define, maintain immutability for, and pass the State inside the agent Graph?"

While the LLM itself is stateless, the agent runtime is a stateful state transition system. At each state transition, it deterministically merges the previous state and node outputs through a Reducer function.

Three Core Principles of State Schema Design

1. Single Source of Truth: The State object passing through agent nodes must be the sole repository containing all information about the task in progress (conversation history, tool results, sub-task success status, user metadata).
2. Explicit Transition Rules (Reducer Pattern): When a node updates the state, it must explicitly define whether to overwrite, append to a list, or merge existing values.
3. Strict Runtime Type Validation: Type validation based on Pydantic v2 or TypedDict is essential to prevent invalid data from being passed between nodes.

TypedDict vs Pydantic v2 Comparison & Real-World Implementation

• TypedDict: A Python built-in type hinting mechanism with near-zero overhead. Advantageous for high-speed data transfer between nodes and Reducer type annotations within agent graphs (e.g., LangGraph).
• Pydantic v2: Provides robust runtime data validation powered by C++ engine (pydantic-core), custom validators, and data serialization/deserialization features (JSON Schema, Redis/RDBMS Persistence). Crucial for validating agent state at API boundaries or external gateways.

Here is a production-grade agent state schema example combining TypedDict and Pydantic v2:

Copy1from typing import Annotated, Any, Dict, List, Literal, Optional, TypedDict2from pydantic import BaseModel, Field, ConfigDict3 4class ChatMessage(BaseModel):5    """Single message schema"""6    role: Literal["system", "user", "assistant", "tool"]7    content: str8    name: Optional[str] = None9    tool_call_id: Optional[str] = None10    metadata: Dict[str, Any] = Field(default_factory=dict)11 12def add_messages(left: List[ChatMessage], right: List[ChatMessage]) -> List[ChatMessage]:13    """Reducer function to safely append message lists"""14    return left + right15 16# 1. TypedDict-based Graph State Schema (for graph runtimes such as LangGraph)17class GraphAgentState(TypedDict):18    session_id: str19    user_id: str20    messages: Annotated[List[ChatMessage], add_messages]21    current_step: int22    max_steps: int23    scratchpad: Dict[str, Any]24    is_completed: bool25 26# 2. Pydantic v2-based Runtime API Boundary Validation Model27class AgentStateModel(BaseModel):28    model_config = ConfigDict(arbitrary_types_allowed=True)29 30    session_id: str = Field(description="Unique session identifier for conversation & execution")31    user_id: str = Field(description="User ID")32    messages: List[ChatMessage] = Field(33        default_factory=list, description="Buffer of conversation history and tool execution messages"34    )35    current_step: int = Field(default=0, ge=0, description="Current loop step count")36    max_steps: int = Field(default=15, gt=0, description="Maximum allowed loop steps")37    scratchpad: Dict[str, Any] = Field(38        default_factory=dict, description="Scratchpad for temporary data sharing between nodes"39    )40    is_completed: bool = Field(default=False, description="Whether the task is finished")

2. Short-Term Memory vs Long-Term Memory Architecture

Just as the human brain is divided into working memory, episodic memory, and semantic memory, production AI agents feature a memory hierarchy structured by role and volatility.

Memory Hierarchy State Diagram

Conversation context exceeding thresholds in the short-term memory buffer undergoes dynamic summarization and branches off to episodic memory (Vector DB) and semantic memory (Fact Base / RDBMS) for permanent storage.

Detailed Memory Hierarchy Comparison Matrix (Working vs Episodic vs Semantic vs Procedural)

Applying Experience Replay (Episodic Memory Recall) allows the agent to search for past successful execution traces that overcame similar errors when executing complex data transformations or API calls, injecting them as few-shot examples. This reduces tool compatibility error rates by more than 40% compared to zero-shot prompting.

Long-Term Memory Integration Mechanisms (RAG & Experience Replay)

Long-term memory goes beyond simple RAG (Retrieval-Augmented Generation) document retrieval. Advanced agent engineering utilizes "episodic memory recall" techniques:

1. Similarity Search: Upon receiving user input, the system searches the Vector Store for execution traces of the most similar past tasks.
2. In-Context Few-Shot Injection: Retrieved past success cases are inserted as few-shot examples into the system prompt, significantly reducing hallucinations and boosting tool call accuracy.

3. Context Window Limits & Context Compression Strategies

Although LLM context windows have expanded dramatically to 128k or 1M+ tokens, blindly retaining massive context remains a major bottleneck in production.

Indiscriminately preserving huge context windows increases API costs and causes the "Lost in the Middle" phenomenon—where key information in the middle of the prompt is forgotten—greatly degrading the agent's intent comprehension performance.

Three Major Issues with Retaining Massive Context

1. Surging Cost and Latency: API costs and Time-to-First-Token (TTFT) increase linearly or quadratically in proportion to the input token count.
2. Lost in the Middle: LLM attention mechanisms frequently miss critical instructions or tool execution results positioned in the middle of long prompts.
3. Context Drift: As conversations lengthen, initial system instructions and user goals become buried and distorted under accumulated tool-call payload messages.

Context Compression Pipeline Sequence Flow

This sequence illustrates message pruning, token estimation, dynamic summarization, and long-term memory flushing executed when an agent runtime receives a user query to assemble the final LLM prompt window.

Comparison of 3 Core Context Compression Techniques

4. Hands-on Implementation: Python-Based Window Memory & State Compression Pipeline

Below is the architecture and production-grade implementation of AgentMemoryManager built with Pydantic v2 and Python asyncio. It handles token budget-based sliding windows, automatic tool output truncation, LLM-driven summarization, and memory flushing.

Memory Manager Pipeline Flowchart

In production environments, when implementing estimate_tokens, it is recommended to bind the tiktoken library to accurately parse BPE token counts for your specific model rather than using character-count approximations.

Copy1import asyncio2from typing import List, Dict, Any, Optional, Literal3from pydantic import BaseModel, Field4 5class ChatMessage(BaseModel):6    """Message data model"""7    role: Literal["system", "user", "assistant", "tool"] = Field(..., description="Sender role")8    content: str = Field(..., description="Message body")9    name: Optional[str] = Field(None, description="Tool name or author")10    tool_call_id: Optional[str] = Field(None, description="Tool call identifier")11 12class MemoryConfig(BaseModel):13    """Memory manager configuration model"""14    max_token_budget: int = Field(default=4000, description="Maximum allowed token budget")15    sliding_window_size: int = Field(default=10, description="Number of recent messages to preserve in sliding window")16    tool_pruning_threshold: int = Field(default=150, description="Character threshold to truncate older tool outputs")17    summary_trigger_ratio: float = Field(default=0.8, description="Token budget ratio to trigger summarization")18 19class AgentMemoryManager:20    """21    Memory manager responsible for preventing token budget overflow22    and managing context compression.23    """24    def __init__(self, config: MemoryConfig) -> None:25        self.config: MemoryConfig = config26        self.summary_buffer: str = ""27 28    def estimate_tokens(self, text: str) -> int:29        """Estimate token count (In production, tiktoken is recommended; uses character approximation here)"""30        if not text:31            return 032        return max(1, len(text) // 3)33 34    def prune_tool_outputs(self, messages: List[ChatMessage]) -> List[ChatMessage]:35        """Truncate large historical tool execution results to preserve token budget"""36        pruned_messages: List[ChatMessage] = []37        for i, msg in enumerate(messages):38            # Preserve tool outputs in the latest 3 messages; prune older ones39            if msg.role == "tool" and len(messages) - i > 3:40                if len(msg.content) > self.config.tool_pruning_threshold:41                    truncated_content = (42                        msg.content[: self.config.tool_pruning_threshold]43                        + f"\n... [Tool Output Truncated: Total Length {len(msg.content)} chars]"44                    )45                    pruned_messages.append(46                        ChatMessage(47                            role=msg.role,48                            content=truncated_content,49                            name=msg.name,50                            tool_call_id=msg.tool_call_id,51                        )52                    )53                    continue54            pruned_messages.append(msg)55        return pruned_messages56 57    async def summarize_old_messages(self, messages_to_summarize: List[ChatMessage]) -> str:58        """Summarize older conversation history using an LLM (mock async LLM call)"""59        print(f"[MemoryManager] Summarizing {len(messages_to_summarize)} old messages...")60        await asyncio.sleep(0.05)  # Simulate LLM API latency61 62        extracted_facts = [63            f"- {m.role}: {m.content[:50]}..."64            for m in messages_to_summarize65            if m.role in ["user", "assistant"]66        ]67        return "Previous Conversation Key Summary:\n" + "\n".join(extracted_facts)68 69    async def prepare_context_window(70        self, system_prompt: ChatMessage, conversation_history: List[ChatMessage]71    ) -> List[ChatMessage]:72        """Pipeline to assemble and compress the final context window sent to the LLM"""73        # Step 1: Tool Output Pruning74        pruned_history = self.prune_tool_outputs(conversation_history)75 76        # Step 2: Total Token Calculation77        total_tokens = self.estimate_tokens(system_prompt.content) + self.estimate_tokens(self.summary_buffer)78        for msg in pruned_history:79            total_tokens += self.estimate_tokens(msg.content)80 81        trigger_limit = int(self.config.max_token_budget * self.config.summary_trigger_ratio)82 83        # Step 3: Summarization & Windowing if token budget exceeded84        if total_tokens > trigger_limit and len(pruned_history) > self.config.sliding_window_size:85            cutoff_index = len(pruned_history) - self.config.sliding_window_size86            old_messages = pruned_history[:cutoff_index]87            recent_messages = pruned_history[cutoff_index:]88 89            # Summarize old messages & update summary buffer90            new_summary = await self.summarize_old_messages(old_messages)91            self.summary_buffer = f"{self.summary_buffer}\n{new_summary}".strip()92 93            pruned_history = recent_messages94 95        # Step 4: Assemble final message list (System Prompt -> Summary Message -> Recent Window)96        final_messages: List[ChatMessage] = [system_prompt]97 98        if self.summary_buffer:99            final_messages.append(100                ChatMessage(101                    role="system",102                    content=f"[Persistent Memory Summary]\n{self.summary_buffer}",103                )104            )105 106        final_messages.extend(pruned_history)107        return final_messages108 109# --- Usage Example and Test Run ---110async def main() -> None:111    # Verify compression behavior with 100 token budget and sliding window size of 3112    config = MemoryConfig(max_token_budget=100, sliding_window_size=3, tool_pruning_threshold=150)113    memory_mgr = AgentMemoryManager(config)114 115    sys_prompt = ChatMessage(role="system", content="You are a skilled Python agent.")116 117    history = [118        ChatMessage(role="user", content="How do I configure database connections?"),119        ChatMessage(role="assistant", content="Preparing SQLAlchemy configuration for PostgreSQL connection."),120        ChatMessage(role="tool", content="{" * 100 + "DB_CONNECTION_SUCCESS_FULL_LOG" + "}" * 100),121        ChatMessage(role="user", content="How many connection pool sizes should I set?"),122        ChatMessage(role="assistant", content="We recommend 10 by default."),123        ChatMessage(role="user", content="Show me an async asyncpg configuration example as well."),124    ]125 126    compressed_context = await memory_mgr.prepare_context_window(sys_prompt, history)127 128    print("\n=== Compressed Context Window Passed to Final LLM ===")129    for idx, msg in enumerate(compressed_context):130        print(f"[{idx}] {msg.role.upper()}: {msg.content[:80]}...")131 132if __name__ == "__main__":133    asyncio.run(main())

5. Defensive Engineering Checklist for Production Agents

To prevent incidents caused by inadequate memory and state management in production, agents must adhere to the following checklist and risk mitigation framework.

In production agents, defensive programming measures against state concurrency, personally identifiable information (PII) leakage, and infinite-loop token consumption are imperative.

Defensive Engineering Risk & Response Matrix

1. Preventing Infinite Loop Token Leaks

• Maximum Step Limit: Explicitly set max_steps to block infinite retries in specific nodes.
• Duplicate Call Detection: If identical tool calls and results repeat consecutively 3 times or more, halt execution and isolate to a higher-level exception handler.

2. Blocking Memory Contamination & Hallucination

• Separate User Claims from Facts: Never directly persist unverified user claims (e.g., "Our revenue grew 10x") into long-term episodic memory. Only transfer facts validated through system verification processes.
• Preserve System Prompt Anchor: When applying sliding windows, ensure the system role persona and safety guardrails are always pinned at array index 0.

3. Concurrency & Distributed State Management

• State Atomicity: In multi-agent systems or distributed environments, use Redis Locks or database transactions to prevent state overwrite collisions (race conditions) on the same session_id.
• PII Masking: Before persisting conversation history into long-term memory (Vector DB), mask personal identifiers such as social security numbers, API keys, and email addresses using regex or PII detector services.

6. Conclusion

In this post, we explored state schema design, memory hierarchy architectures, and context compression techniques that underpin the persistence and reasoning accuracy of AI agents. To summarize key takeaways:

1. State must ensure immutability and runtime stability using Pydantic v2, TypedDict, and the Reducer pattern.
2. Short-Term Memory token bloat must be managed through Tool Output Pruning, Sliding Window, and LLM Summarization pipelines.
3. Long-Term Memory should extend beyond simple context retrieval to episodic memory that recalls an agent's past successful execution traces.

In the upcoming [AI Agent Engineering Episode 3] post, we will cover "Tool Calling & Safe Execution Environment (Sandbox Security & Failure Recovery)" to inspect how agents can invoke internal and external tools accurately and safely.