From f1e48e7dad214a7cf890973f31794fc5d1c9ea88 Mon Sep 17 00:00:00 2001 From: Andrew White Date: Fri, 24 Jul 2026 12:22:55 -0500 Subject: [PATCH] fix: accept 0 as a valid token_count in ContextManager.add_message add_message annotated token_count as int but defaulted to None, and the guard treated 0 as unset. This meant callers could not explicitly set the running context usage to 0. Annotate token_count as int | None and check so 0 is honored. Signed-off-by: Andrew White --- agent/context_manager/manager.py | 4 ++-- tests/unit/test_context_manager.py | 16 ++++++++++++++++ 2 files changed, 18 insertions(+), 2 deletions(-) create mode 100644 tests/unit/test_context_manager.py diff --git a/agent/context_manager/manager.py b/agent/context_manager/manager.py index 85e96af0f..1c7a54263 100644 --- a/agent/context_manager/manager.py +++ b/agent/context_manager/manager.py @@ -257,9 +257,9 @@ def _load_system_prompt( f"Tools={len(tool_specs)}]" ) - def add_message(self, message: Message, token_count: int = None) -> None: + def add_message(self, message: Message, token_count: int | None = None) -> None: """Add a message to the history""" - if token_count: + if token_count is not None: self.running_context_usage = token_count self.items.append(message) if self.on_message_added: diff --git a/tests/unit/test_context_manager.py b/tests/unit/test_context_manager.py new file mode 100644 index 000000000..237fcb68d --- /dev/null +++ b/tests/unit/test_context_manager.py @@ -0,0 +1,16 @@ +"""Unit tests for the context manager.""" + +from agent.context_manager.manager import ContextManager + + +def test_add_message_accepts_zero_token_count(): + """A token_count of 0 should update running_context_usage, not be treated as missing.""" + manager = ContextManager(tool_specs=[]) + manager.running_context_usage = 100 + + class _Message: + role = "user" + content = "hi" + + manager.add_message(_Message(), token_count=0) + assert manager.running_context_usage == 0