feat: context window management (Feature 5)

BinglanLi · BinglanLi · commit e583e4319a37 · 2026-04-14T14:20:26.000-07:00
Adds a configurable sliding window to prevent context overflow during
long single-agent runs. Truncation is applied to both LLM call sites
in NodeExecutor (generate and execute_self_critic) while preserving
the full message history in state for checkpointing.

- BaseAgentConfig: max_context_messages field (None = disabled, int &gt;= 2)
  with BASE_AGENT_MAX_CONTEXT_MESSAGES env var and validation
- BaseAgent: max_context_messages kwarg wired from config
- NodeExecutor._truncate_messages(): keeps first msg (user task) + last N-1 msgs
- 21 unit tests covering logic, config validation, and node wiring
diff --git a/BaseAgent/base_agent.py b/BaseAgent/base_agent.py
@@ -65,6 +65,7 @@ def __init__(
         checkpoint_db_path: str | None = None,
         require_approval: str | None = None,
         skills_directory: str | None = None,
+        max_context_messages: int | None = None,
         spec: AgentSpec | None = None,
     ):
         """
@@ -81,6 +82,10 @@ def __init__(
                 "always" (default) — interrupt before every code block;
                 "never" — never interrupt.
             skills_directory: Path to a directory of SKILL.md files to load on startup.
+            max_context_messages: Sliding window size for context management. None (default)
+                disables truncation. Integer >= 2 limits the number of messages passed to
+                the LLM per call, always preserving the first message (user task) plus the
+                most recent messages. The full history is retained in state for checkpointing.
             spec: Optional agent identity and persona.  When provided, ``spec.name``,
                 ``spec.role``, ``spec.tool_names``, ``spec.skill_names``,
                 ``spec.llm``, ``spec.source``, and ``spec.temperature`` override
@@ -109,6 +114,7 @@ def __init__(
         self.checkpoint_db_path = checkpoint_db_path if checkpoint_db_path is not None else default_config.checkpoint_db_path
         self.require_approval = require_approval if require_approval is not None else default_config.require_approval
         self.skills_directory = skills_directory if skills_directory is not None else default_config.skills_directory
+        self.max_context_messages = max_context_messages if max_context_messages is not None else default_config.max_context_messages
         self.thread_id: str | None = None
         self._interrupted: bool = False
         self._run_config: dict | None = None
diff --git a/BaseAgent/config.py b/BaseAgent/config.py
@@ -59,6 +59,9 @@ class BaseAgentConfig:
     # Skills
     skills_directory: str | None = None  # directory of SKILL.md files to load on startup
 
+    # Context window management
+    max_context_messages: int | None = None  # None = disabled; int >= 2 = sliding window size
+
     def __post_init__(self):
         """Load any environment variable overrides if they exist."""
         # Check for environment variable overrides (optional)
@@ -90,6 +93,16 @@ def __post_init__(self):
             self.require_approval = val
         if os.getenv("BASE_AGENT_SKILLS_DIRECTORY"):
             self.skills_directory = os.getenv("BASE_AGENT_SKILLS_DIRECTORY")
+        if os.getenv("BASE_AGENT_MAX_CONTEXT_MESSAGES"):
+            try:
+                val = int(os.getenv("BASE_AGENT_MAX_CONTEXT_MESSAGES"))
+            except ValueError:
+                raise ValueError("BASE_AGENT_MAX_CONTEXT_MESSAGES must be an integer")
+            if val < 2:
+                raise ValueError("BASE_AGENT_MAX_CONTEXT_MESSAGES must be >= 2")
+            self.max_context_messages = val
+        if self.max_context_messages is not None and self.max_context_messages < 2:
+            raise ValueError("max_context_messages must be None or >= 2")
 
     def to_dict(self) -> dict:
         """Convert config to dictionary for easy access."""
@@ -105,6 +118,7 @@ def to_dict(self) -> dict:
             "checkpoint_db_path": self.checkpoint_db_path,
             "require_approval": self.require_approval,
             "skills_directory": self.skills_directory,
+            "max_context_messages": self.max_context_messages,
         }
 
 
diff --git a/BaseAgent/nodes.py b/BaseAgent/nodes.py
@@ -28,6 +28,28 @@ class NodeExecutor:
     def __init__(self, agent: "BaseAgent"):
         self.agent = agent
 
+    # ------------------------------------------------------------------
+    # Context window management
+    # ------------------------------------------------------------------
+
+    def _truncate_messages(self, messages: list) -> list:
+        """Sliding window: keep first message + last (N-1) messages.
+
+        The first message in state["input"] is the initial user HumanMessage
+        (the task). It is always preserved so the LLM retains its objective.
+        The system prompt is prepended separately in generate() and is not
+        part of state["input"].
+
+        When max_context_messages is None (default), all messages are passed
+        through unchanged.
+        """
+        if not messages:
+            return messages
+        max_msgs = self.agent.max_context_messages
+        if max_msgs is None or len(messages) <= max_msgs:
+            return messages
+        return [messages[0]] + messages[-(max_msgs - 1):]
+
     # ------------------------------------------------------------------
     # Core nodes
     # ------------------------------------------------------------------
@@ -44,8 +66,8 @@ def generate(self, state: "AgentState") -> "AgentState":
                 "cache_control": {"type": "ephemeral"}
             }
 
-        # Add the system prompt to the input to LLM
-        input = [system_message] + state["input"]
+        # Add the system prompt to the input to LLM; apply sliding window if configured
+        input = [system_message] + self._truncate_messages(state["input"])
         output = agent.llm.invoke(input)
 
         usage_metrics = extract_usage_metrics(agent.source, output, model=getattr(agent.llm, "model_name", None))
@@ -191,8 +213,10 @@ def execute_self_critic(self, state: "AgentState", test_time_scale_round: int) -
         """LLM feedback generation for self-critic mode."""
         agent = self.agent
         if agent.critic_count < test_time_scale_round:
-            # Generate feedback based on message history
-            input = state["input"]
+            # Generate feedback based on message history; apply sliding window if configured.
+            # Note: no system message is prepended here — first element of state["input"]
+            # is the user HumanMessage (the task).
+            input = self._truncate_messages(state["input"])
             feedback_prompt = get_feedback_prompt(agent.user_task)
             feedback = agent.llm.invoke(input + [HumanMessage(content=feedback_prompt)])
 
diff --git a/BaseAgent/tests/test_context_window.py b/BaseAgent/tests/test_context_window.py
@@ -0,0 +1,263 @@
+"""Unit tests for Feature 5: Context Window Management.
+
+Tests cover:
+- _truncate_messages() logic (edge cases, boundary conditions)
+- BaseAgentConfig validation for max_context_messages
+- Wiring: generate() and execute_self_critic() pass truncated input to llm.invoke()
+"""
+
+from __future__ import annotations
+
+import os
+from unittest.mock import MagicMock, patch
+
+import pytest
+from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
+
+from BaseAgent.config import BaseAgentConfig
+from BaseAgent.nodes import NodeExecutor
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def make_agent(max_context_messages=None):
+    """Minimal mock agent for NodeExecutor tests."""
+    agent = MagicMock()
+    agent.source = "Anthropic"
+    agent.system_prompt = "You are a helpful assistant."
+    agent.use_tool_retriever = False
+    agent.timeout_seconds = 30
+    agent.critic_count = 0
+    agent.user_task = "test task"
+    agent.max_context_messages = max_context_messages
+    agent.llm.invoke.return_value = MagicMock(
+        content="<solution>answer</solution>", usage_metadata=None
+    )
+    return agent
+
+
+def make_messages(n: int) -> list:
+    """Return a list of n alternating HumanMessage / AIMessage objects."""
+    msgs = []
+    for i in range(n):
+        if i % 2 == 0:
+            msgs.append(HumanMessage(content=f"human {i}"))
+        else:
+            msgs.append(AIMessage(content=f"ai {i}"))
+    return msgs
+
+
+def make_state(messages=None):
+    return {
+        "input": messages if messages is not None else [HumanMessage(content="task")],
+        "next_step": None,
+        "pending_code": None,
+        "pending_language": None,
+    }
+
+
+# ---------------------------------------------------------------------------
+# _truncate_messages unit tests
+# ---------------------------------------------------------------------------
+
+
+class TestTruncateMessages:
+    def test_empty_list(self):
+        executor = NodeExecutor(make_agent())
+        assert executor._truncate_messages([]) == []
+
+    def test_disabled_by_default(self):
+        """max_context_messages=None passes all messages through unchanged."""
+        executor = NodeExecutor(make_agent(max_context_messages=None))
+        msgs = make_messages(20)
+        result = executor._truncate_messages(msgs)
+        assert result is msgs  # same object, no copy
+
+    def test_no_truncation_below_limit(self):
+        """Fewer messages than the limit: no truncation."""
+        executor = NodeExecutor(make_agent(max_context_messages=10))
+        msgs = make_messages(5)
+        result = executor._truncate_messages(msgs)
+        assert result == msgs
+
+    def test_no_truncation_at_exact_limit(self):
+        """Exactly at the limit: no truncation (boundary condition)."""
+        executor = NodeExecutor(make_agent(max_context_messages=5))
+        msgs = make_messages(5)
+        result = executor._truncate_messages(msgs)
+        assert result == msgs
+
+    def test_truncation_above_limit(self):
+        """10 messages, limit 5: keeps first + last 4."""
+        executor = NodeExecutor(make_agent(max_context_messages=5))
+        msgs = make_messages(10)
+        result = executor._truncate_messages(msgs)
+        assert len(result) == 5
+        assert result[0] is msgs[0]   # first message preserved
+        assert result[1:] == msgs[-4:]  # last 4 messages
+
+    def test_first_message_always_preserved(self):
+        """The initial user task (index 0) is always the first element."""
+        executor = NodeExecutor(make_agent(max_context_messages=3))
+        msgs = make_messages(20)
+        result = executor._truncate_messages(msgs)
+        assert result[0] is msgs[0]
+
+    def test_state_input_not_mutated_by_generate(self):
+        """state["input"] must not be modified during generate()."""
+        agent = make_agent(max_context_messages=3)
+        executor = NodeExecutor(agent)
+        msgs = make_messages(10)
+        state = make_state(list(msgs))  # copy so we can compare
+        original_len = len(state["input"])
+        with patch("BaseAgent.nodes.extract_usage_metrics", return_value=None):
+            executor.generate(state)
+        # state["input"] grows by 1 (the AI response is appended), but is not truncated
+        assert len(state["input"]) == original_len + 1
+
+
+# ---------------------------------------------------------------------------
+# BaseAgentConfig validation tests
+# ---------------------------------------------------------------------------
+
+
+class TestConfigValidation:
+    def test_default_is_none(self):
+        config = BaseAgentConfig()
+        assert config.max_context_messages is None
+
+    def test_valid_value(self):
+        config = BaseAgentConfig(max_context_messages=10)
+        assert config.max_context_messages == 10
+
+    def test_value_2_is_minimum(self):
+        config = BaseAgentConfig(max_context_messages=2)
+        assert config.max_context_messages == 2
+
+    def test_value_1_raises(self):
+        with pytest.raises(ValueError, match="max_context_messages"):
+            BaseAgentConfig(max_context_messages=1)
+
+    def test_value_0_raises(self):
+        with pytest.raises(ValueError, match="max_context_messages"):
+            BaseAgentConfig(max_context_messages=0)
+
+    def test_negative_raises(self):
+        with pytest.raises(ValueError, match="max_context_messages"):
+            BaseAgentConfig(max_context_messages=-5)
+
+    def test_env_var_valid(self, monkeypatch):
+        monkeypatch.setenv("BASE_AGENT_MAX_CONTEXT_MESSAGES", "20")
+        config = BaseAgentConfig()
+        assert config.max_context_messages == 20
+        monkeypatch.delenv("BASE_AGENT_MAX_CONTEXT_MESSAGES", raising=False)
+
+    def test_env_var_too_small_raises(self, monkeypatch):
+        monkeypatch.setenv("BASE_AGENT_MAX_CONTEXT_MESSAGES", "1")
+        with pytest.raises(ValueError, match="BASE_AGENT_MAX_CONTEXT_MESSAGES"):
+            BaseAgentConfig()
+        monkeypatch.delenv("BASE_AGENT_MAX_CONTEXT_MESSAGES", raising=False)
+
+    def test_env_var_non_integer_raises(self, monkeypatch):
+        monkeypatch.setenv("BASE_AGENT_MAX_CONTEXT_MESSAGES", "abc")
+        with pytest.raises(ValueError, match="BASE_AGENT_MAX_CONTEXT_MESSAGES"):
+            BaseAgentConfig()
+        monkeypatch.delenv("BASE_AGENT_MAX_CONTEXT_MESSAGES", raising=False)
+
+    def test_to_dict_includes_field(self):
+        config = BaseAgentConfig(max_context_messages=15)
+        d = config.to_dict()
+        assert "max_context_messages" in d
+        assert d["max_context_messages"] == 15
+
+
+# ---------------------------------------------------------------------------
+# Wiring tests: generate() passes truncated input to llm.invoke()
+# ---------------------------------------------------------------------------
+
+
+class TestGenerateNodeWiring:
+    def test_generate_truncates_llm_input(self):
+        """generate() must pass a truncated list to llm.invoke() when limit is set."""
+        agent = make_agent(max_context_messages=3)
+        executor = NodeExecutor(agent)
+
+        msgs = make_messages(10)
+        state = make_state(list(msgs))
+
+        with patch("BaseAgent.nodes.extract_usage_metrics", return_value=None):
+            executor.generate(state)
+
+        call_args = agent.llm.invoke.call_args[0][0]
+        # call_args = [SystemMessage] + truncated(state["input"])
+        # truncated = [msgs[0]] + msgs[-2:] = 3 messages
+        # total = 1 (system) + 3 = 4
+        assert len(call_args) == 4
+        assert isinstance(call_args[0], SystemMessage)
+        assert call_args[1] is msgs[0]  # first message preserved
+
+    def test_generate_no_truncation_when_disabled(self):
+        """generate() passes all messages when max_context_messages=None."""
+        agent = make_agent(max_context_messages=None)
+        executor = NodeExecutor(agent)
+
+        msgs = make_messages(10)
+        state = make_state(list(msgs))
+
+        with patch("BaseAgent.nodes.extract_usage_metrics", return_value=None):
+            executor.generate(state)
+
+        call_args = agent.llm.invoke.call_args[0][0]
+        # 1 system + 10 messages = 11
+        assert len(call_args) == 11
+
+
+# ---------------------------------------------------------------------------
+# Wiring tests: execute_self_critic() passes truncated input to llm.invoke()
+# ---------------------------------------------------------------------------
+
+
+class TestSelfCriticWiring:
+    def test_self_critic_truncates_llm_input(self):
+        """execute_self_critic() must pass a truncated list to llm.invoke()."""
+        agent = make_agent(max_context_messages=3)
+        # Set critic_count=0 and test_time_scale_round=1 to exercise the branch
+        agent.critic_count = 0
+        agent.llm.invoke.return_value = MagicMock(content="feedback text", usage_metadata=None)
+
+        executor = NodeExecutor(agent)
+        msgs = make_messages(10)
+        state = make_state(list(msgs))
+        state["next_step"] = "generate"
+
+        with patch("BaseAgent.nodes.extract_usage_metrics", return_value=None):
+            with patch("BaseAgent.nodes.get_feedback_prompt", return_value="give feedback"):
+                executor.execute_self_critic(state, test_time_scale_round=1)
+
+        call_args = agent.llm.invoke.call_args[0][0]
+        # truncated = [msgs[0]] + msgs[-2:] = 3 messages
+        # + 1 feedback HumanMessage appended after truncation = 4 total
+        assert len(call_args) == 4
+        assert call_args[0] is msgs[0]  # first message preserved
+        assert isinstance(call_args[-1], HumanMessage)  # feedback appended last
+
+    def test_self_critic_no_truncation_when_disabled(self):
+        """execute_self_critic() passes all messages when max_context_messages=None."""
+        agent = make_agent(max_context_messages=None)
+        agent.critic_count = 0
+        agent.llm.invoke.return_value = MagicMock(content="feedback", usage_metadata=None)
+
+        executor = NodeExecutor(agent)
+        msgs = make_messages(10)
+        state = make_state(list(msgs))
+
+        with patch("BaseAgent.nodes.extract_usage_metrics", return_value=None):
+            with patch("BaseAgent.nodes.get_feedback_prompt", return_value="give feedback"):
+                executor.execute_self_critic(state, test_time_scale_round=1)
+
+        call_args = agent.llm.invoke.call_args[0][0]
+        # 10 messages + 1 feedback HumanMessage = 11
+        assert len(call_args) == 11
diff --git a/BaseAgent/tests/test_interrupt_resume.py b/BaseAgent/tests/test_interrupt_resume.py
@@ -25,6 +25,7 @@ def make_agent(require_approval="never", source="Anthropic"):
     agent.critic_count = 0
     agent.user_task = "test task"
     agent.require_approval = require_approval
+    agent.max_context_messages = None
     resp = MagicMock()
     resp.content = "<solution>answer</solution>"
     resp.usage_metadata = None
diff --git a/BaseAgent/tests/test_nodes.py b/BaseAgent/tests/test_nodes.py
@@ -19,6 +19,7 @@ def make_agent(source="Anthropic", use_tool_retriever=False):
     agent.timeout_seconds = 30
     agent.critic_count = 0
     agent.user_task = "test task"
+    agent.max_context_messages = None
     # llm.invoke returns a mock response
     agent.llm.invoke.return_value = MagicMock(content="<solution>answer</solution>", usage_metadata=None)
     return agent
diff --git a/Multi_Agent_Future_Roadmap.md b/Multi_Agent_Future_Roadmap.md
diff --git a/examples/05_custom_configuration.py b/examples/05_custom_configuration.py
