diff --git a/strands-py/src/strands/agent/conversation_manager/compression/pin_message.py b/strands-py/src/strands/agent/conversation_manager/compression/pin_message.py index 934a9d3c0c..ac82b14124 100644 --- a/strands-py/src/strands/agent/conversation_manager/compression/pin_message.py +++ b/strands-py/src/strands/agent/conversation_manager/compression/pin_message.py @@ -55,6 +55,24 @@ def is_pinned(messages: Messages, index: int) -> bool: return False +def count_protected_prefix(messages: Messages) -> int: + """Count the leading messages protected from eviction. + + ``is_pinned`` protects a pinned message's adjacent tool-pair partner as well, so the protected + head can extend past the explicitly pinned range. + + Args: + messages: The full messages array. + + Returns: + The number of leading messages that ``is_pinned`` protects. + """ + count = 0 + while count < len(messages) and is_pinned(messages, count): + count += 1 + return count + + def apply_pin_first(messages: Messages, count: int) -> None: """Pin the first N messages in the array permanently. diff --git a/strands-py/src/strands/agent/conversation_manager/conversation_manager.py b/strands-py/src/strands/agent/conversation_manager/conversation_manager.py index 36438c4555..fa7d45e909 100644 --- a/strands-py/src/strands/agent/conversation_manager/conversation_manager.py +++ b/strands-py/src/strands/agent/conversation_manager/conversation_manager.py @@ -78,6 +78,9 @@ def __init__(self, *, proactive_compression: Union[bool, "ProactiveCompressionCo removed_message_count: The messages that have been removed from the agents messages array. These represent messages provided by the user or LLM that have been removed, not messages included by the conversation manager through something like summarization. + pinned_head_count: The number of leading messages the manager has protected from eviction. + Restoring a session reattaches this many stored messages, so it counts every protected + message, including the tool-pair partner of a pinned message. """ # Resolve the threshold from proactive_compression parameter if proactive_compression is True: @@ -91,6 +94,7 @@ def __init__(self, *, proactive_compression: Union[bool, "ProactiveCompressionCo raise ValueError(f"compression_threshold must be between 0 (exclusive) and 1 (inclusive), got {threshold}") self.removed_message_count = 0 + self.pinned_head_count = 0 self._compression_threshold = threshold def register_hooks(self, registry: HookRegistry, **kwargs: Any) -> None: @@ -154,6 +158,7 @@ def restore_from_session(self, state: dict[str, Any]) -> list[Message] | None: if state.get("__name__") != self.__class__.__name__: raise ValueError("Invalid conversation manager state.") self.removed_message_count = state["removed_message_count"] + self.pinned_head_count = state.get("pinned_head_count", 0) return None def get_state(self) -> dict[str, Any]: @@ -161,6 +166,7 @@ def get_state(self) -> dict[str, Any]: return { "__name__": self.__class__.__name__, "removed_message_count": self.removed_message_count, + "pinned_head_count": self.pinned_head_count, } @abstractmethod diff --git a/strands-py/src/strands/agent/conversation_manager/sliding_window_conversation_manager.py b/strands-py/src/strands/agent/conversation_manager/sliding_window_conversation_manager.py index 16c3043ce7..2ff1cc4c53 100644 --- a/strands-py/src/strands/agent/conversation_manager/sliding_window_conversation_manager.py +++ b/strands-py/src/strands/agent/conversation_manager/sliding_window_conversation_manager.py @@ -11,7 +11,7 @@ from ...types.exceptions import ContextWindowOverflowException from ...types.tools import ToolResultContent from .compression.context_compression import find_valid_trim_point -from .compression.pin_message import apply_pin_first, is_pinned +from .compression.pin_message import apply_pin_first, count_protected_prefix, is_pinned from .conversation_manager import ConversationManager, ProactiveCompressionConfig logger = logging.getLogger(__name__) @@ -79,7 +79,6 @@ def __init__( self.should_truncate_results = should_truncate_results self.per_turn = per_turn self.pin_first = max(0, pin_first) if pin_first is not None else None - self._pin_first_applied = False self._model_call_count = 0 def register_hooks(self, registry: "HookRegistry", **kwargs: Any) -> None: @@ -196,9 +195,11 @@ def reduce_context(self, agent: "Agent", e: Exception | None = None, **kwargs: A messages = agent.messages # Pin first N messages permanently (only on first reduction) - if self.pin_first and not self._pin_first_applied: + if self.pin_first and self.pinned_head_count == 0: apply_pin_first(messages, self.pin_first) - self._pin_first_applied = True + # is_pinned also protects a pinned message's tool-pair partner, so the protected head can + # be longer than pin_first; restore reattaches exactly this many stored messages. + self.pinned_head_count = count_protected_prefix(messages) # window_size=0 means "remove all non-pinned messages" if self.window_size == 0: diff --git a/strands-py/src/strands/agent/conversation_manager/summarizing_conversation_manager.py b/strands-py/src/strands/agent/conversation_manager/summarizing_conversation_manager.py index e2d4b5053b..d104abd50d 100644 --- a/strands-py/src/strands/agent/conversation_manager/summarizing_conversation_manager.py +++ b/strands-py/src/strands/agent/conversation_manager/summarizing_conversation_manager.py @@ -79,7 +79,6 @@ def __init__( self.summarization_agent = summarization_agent self.summarization_system_prompt = summarization_system_prompt self.pin_first = max(0, pin_first) if pin_first is not None else None - self._pin_first_applied = False self._summary_message: Message | None = None @override @@ -98,7 +97,10 @@ def restore_from_session(self, state: dict[str, Any]) -> list[Message] | None: def get_state(self) -> dict[str, Any]: """Returns a dictionary representation of the state for the Summarizing Conversation Manager.""" - return {"summary_message": self._summary_message, **super().get_state()} + return { + "summary_message": self._summary_message, + **super().get_state(), + } def apply_management(self, agent: "Agent", **kwargs: Any) -> None: """Apply management strategy to conversation history. @@ -147,6 +149,9 @@ def reduce_context(self, agent: "Agent", e: Exception | None = None, **kwargs: A def _summarize_oldest(self, agent: "Agent") -> None: """Summarize the oldest messages and replace them with a summary. + Bookkeeping updates (removed_message_count, pinned_head_count) are applied only after + summary generation succeeds, making the operation atomic by construction. + Args: agent: The agent instance. @@ -173,9 +178,8 @@ def _summarize_oldest(self, agent: "Agent") -> None: raise ContextWindowOverflowException("Cannot summarize: insufficient messages for summarization") # Pin first N messages permanently (only on first reduction) - if self.pin_first and not self._pin_first_applied: + if self.pin_first and self.pinned_head_count == 0: apply_pin_first(agent.messages, self.pin_first) - self._pin_first_applied = True # Partition [0, messages_to_summarize_count) into pinned (preserve) and non-pinned (summarize) protected_to_preserve, to_summarize = partition_pinned(agent.messages, 0, messages_to_summarize_count) @@ -185,16 +189,17 @@ def _summarize_oldest(self, agent: "Agent") -> None: remaining_messages = agent.messages[messages_to_summarize_count:] - # Keep track of the number of messages that have been summarized thus far. - self.removed_message_count += len(to_summarize) - # If there is a summary message, don't count it in the removed_message_count. - if self._summary_message: - self.removed_message_count -= 1 + # Generate summary — if this fails, no bookkeeping is updated + summary_message = self._generate_summary(to_summarize, agent) + _ensure_tracking_id(summary_message) - # Generate summary - self._summary_message = self._generate_summary(to_summarize, agent) - # Assign tracking id to the summary message since it bypasses the append method. - _ensure_tracking_id(self._summary_message) + # All mutations below this point: summary generation succeeded, commit the changes. + removed_count = len(to_summarize) + if self._summary_message: + removed_count -= 1 + self.removed_message_count += removed_count + self.pinned_head_count = len(protected_to_preserve) + self._summary_message = summary_message # Replace summarized range with protected messages + summary + remaining agent.messages[:] = protected_to_preserve + [self._summary_message] + remaining_messages diff --git a/strands-py/src/strands/session/repository_session_manager.py b/strands-py/src/strands/session/repository_session_manager.py index 16e0a7999f..6d4b6f8d4e 100644 --- a/strands-py/src/strands/session/repository_session_manager.py +++ b/strands-py/src/strands/session/repository_session_manager.py @@ -4,6 +4,7 @@ import logging from typing import TYPE_CHECKING, Any +from ..agent.conversation_manager.compression.pin_message import pin_message from ..agent.state import AgentState from ..tools._tool_helpers import generate_missing_tool_result_content from ..types.agent import LocalAgent @@ -263,22 +264,49 @@ def initialize(self, agent: "LocalAgent", **kwargs: Any) -> None: # Restore the conversation manager to its previous state, and get the optional prepend messages prepend_messages = [] offset = 0 + pinned_head_count = 0 if isinstance(agent, Agent): prepend_messages = ( agent.conversation_manager.restore_from_session(session_agent.conversation_manager_state) or [] ) offset = agent.conversation_manager.removed_message_count + # The pinned head count is persisted in the base ConversationManager and restored above. + pinned_head_count = agent.conversation_manager.pinned_head_count + + pinned_messages: list[Message] = [] + if pinned_head_count > 0: + pinned_session_messages = self.session_repository.list_messages( + session_id=self.session_id, + agent_id=agent.agent_id, + limit=pinned_head_count, + offset=0, + ) + pinned_messages = [session_message.to_message() for session_message in pinned_session_messages] + # Re-apply pin markers to the restored head so the next compaction respects them. + for i in range(len(pinned_messages)): + pin_message(pinned_messages, i) # List the messages currently in the session, using an offset of the messages previously removed - # by the conversation manager. + # by the conversation manager. The offset must also skip past the still-live pinned head so the + # tail does not overlap with the messages reattached as ``pinned_messages`` above. session_messages = self.session_repository.list_messages( session_id=self.session_id, agent_id=agent.agent_id, - offset=offset, + offset=pinned_head_count + offset, ) self._held_records[agent.agent_id] = list(session_messages) + + # When the tail comes back empty but the store is not, seed the append cursor from the + # last stored message so appending never restarts at 0 on a non-empty store. if len(session_messages) > 0: self._latest_agent_message[agent.agent_id] = session_messages[-1] + elif pinned_head_count > 0 or offset > 0: + all_messages = self.session_repository.list_messages( + session_id=self.session_id, + agent_id=agent.agent_id, + ) + if all_messages: + self._latest_agent_message[agent.agent_id] = all_messages[-1] # Skip restoring messages when conversation is managed server-side if isinstance(agent, Agent) and agent.model.stateful: @@ -288,10 +316,13 @@ def initialize(self, agent: "LocalAgent", **kwargs: Any) -> None: self.session_id, ) else: - # Restore the agents messages array including the optional prepend messages - agent.messages = prepend_messages + [ - session_message.to_message() for session_message in session_messages - ] + # Restore the agents messages array including the optional prepend messages and the + # still-live pinned head that compaction pushed back to the front. + agent.messages = ( + pinned_messages + + prepend_messages + + [session_message.to_message() for session_message in session_messages] + ) # Fix broken session histories: https://github.com/strands-agents/harness-sdk/issues/859 agent.messages = self._fix_broken_tool_use(agent.messages) diff --git a/strands-py/tests/strands/agent/test_conversation_manager.py b/strands-py/tests/strands/agent/test_conversation_manager.py index 8a016d053a..dac53def99 100644 --- a/strands-py/tests/strands/agent/test_conversation_manager.py +++ b/strands-py/tests/strands/agent/test_conversation_manager.py @@ -1120,6 +1120,31 @@ def test_pin_first_only_applies_once(): assert "first" in [m["content"][0]["text"] for m in agent.messages] +def test_pin_first_head_count_includes_protected_tool_pair_partner(): + """``pinned_head_count`` must cover the tool-pair partner that ``is_pinned`` also protects (#4087). + + Trimming keeps a pinned assistant(toolUse) message's matching toolResult even though the result was + never explicitly pinned. A resumed session reattaches exactly ``pinned_head_count`` stored messages, + so that count has to include the partner or the restore drops the real tool result. + """ + manager = SlidingWindowConversationManager(window_size=2, should_truncate_results=False, pin_first=2) + messages = [ + {"role": "user", "content": [{"text": "use the tool"}]}, + {"role": "assistant", "content": [{"toolUse": {"toolUseId": "tu1", "name": "echo", "input": {}}}]}, + { + "role": "user", + "content": [ + {"toolResult": {"toolUseId": "tu1", "status": "success", "content": [{"text": "REAL RESULT"}]}} + ], + }, + ] + agent = _make_mock_agent(messages=messages) + manager.reduce_context(agent) + + assert manager.pinned_head_count == 3 + assert agent.messages[2]["content"][0]["toolResult"]["content"][0]["text"] == "REAL RESULT" + + def test_pinned_message_in_middle_survives_trimming(): from strands.agent.conversation_manager.compression.pin_message import pin_message diff --git a/strands-py/tests/strands/agent/test_summarizing_conversation_manager.py b/strands-py/tests/strands/agent/test_summarizing_conversation_manager.py index 4b22ab0be5..5925d7e9e4 100644 --- a/strands-py/tests/strands/agent/test_summarizing_conversation_manager.py +++ b/strands-py/tests/strands/agent/test_summarizing_conversation_manager.py @@ -948,3 +948,133 @@ def test_proactive_compression_swallows_errors(): # Should not throw — proactive compression is best-effort registry.invoke_callbacks(event) assert len(agent.messages) == 20 + + +# --------------------------------------------------------------------------- +# Regression tests for https://github.com/strands-agents/harness-sdk/issues/4027 +# +# Bug 2: a failing summarization must not mutate the bookkeeping that the +# session manager uses to restore the agent (removed_message_count, +# pinned_head_count, _summary_message). +# --------------------------------------------------------------------------- + + +def test_reduce_context_failure_rolls_back_removed_message_count(): + """A failing summary must leave removed_message_count unchanged so a subsequent restore is correct.""" + + failing_agent = Mock() + failing_agent.model = Mock() + failing_agent.model.stream = Mock( + side_effect=lambda *a, **kw: _mock_model_stream_error(RuntimeError("provider outage")) + ) + failing_agent.messages = [{"role": "user", "content": [{"text": f"msg {i}"}]} for i in range(6)] + + manager = SummarizingConversationManager( + summary_ratio=0.5, + preserve_recent_messages=1, + ) + assert manager.removed_message_count == 0 + + with patch("strands.agent.conversation_manager.summarizing_conversation_manager.logger"): + # Proactive path (e is None): swallow, but bookkeeping must be intact. + manager.reduce_context(failing_agent) + + assert manager.removed_message_count == 0 + assert manager._summary_message is None + assert manager.pinned_head_count == 0 + + +def test_reduce_context_failure_rolls_back_pinned_head_count(): + """A failing summary must leave pinned_head_count at 0 so a retry can re-apply it cleanly.""" + + failing_agent = Mock() + failing_agent.model = Mock() + failing_agent.model.stream = Mock( + side_effect=lambda *a, **kw: _mock_model_stream_error(RuntimeError("provider outage")) + ) + failing_agent.messages = [{"role": "user", "content": [{"text": f"msg {i}"}]} for i in range(6)] + + manager = SummarizingConversationManager( + summary_ratio=0.5, + preserve_recent_messages=1, + pin_first=2, + ) + assert manager.pinned_head_count == 0 + + with patch("strands.agent.conversation_manager.summarizing_conversation_manager.logger"): + manager.reduce_context(failing_agent) + + # The retry should be free to re-apply pin_first; the in-place metadata flag is idempotent + # so leaving it set is fine, but the bookkeeping must reflect that no compaction finished. + assert manager.pinned_head_count == 0 + assert manager.removed_message_count == 0 + + +def test_reduce_context_retry_after_failure_records_correct_count(): + """Retry after a failed summary records the compaction correctly. + + First attempt: reactive path raises (e is set), failure is re-raised so the agent loop + knows the overflow was not resolved. Second attempt: proactive path succeeds and commits + the bookkeeping atomically. Asserts: + 1. Failed attempt leaves all bookkeeping unchanged (atomic design). + 2. Successful retry records the exact transcript. + 3. model.stream was called twice — proves the first (failing) call was actually exercised. + """ + from strands.types.exceptions import ContextWindowOverflowException + + call_count = [0] + + def stream_with_failure(*args, **kwargs): + call_count[0] += 1 + if call_count[0] == 1: + return _mock_model_stream_error(RuntimeError("provider outage")) + return _mock_model_stream("SUMMARY") + + retry_agent = Mock() + retry_agent.model = Mock() + retry_agent.model.stream = Mock(side_effect=stream_with_failure) + retry_agent.messages = [{"role": "user", "content": [{"text": f"msg {i}"}]} for i in range(6)] + + manager = SummarizingConversationManager( + summary_ratio=0.5, + preserve_recent_messages=1, + pin_first=2, + ) + + # First call: reactive path (e is set) — failure must be re-raised so the caller knows + # the overflow was not resolved. + overflow = ContextWindowOverflowException("context overflow") + with pytest.raises(RuntimeError, match="provider outage"): + manager.reduce_context(retry_agent, e=overflow) + + # Bookkeeping must be unchanged after failure (atomic design). + assert manager.removed_message_count == 0 + assert manager.pinned_head_count == 0 + assert manager._summary_message is None + # Proves the failure path was actually exercised, not skipped. + assert retry_agent.model.stream.call_count == 1 + + # Restore messages — the failed attempt left them intact because bookkeeping is atomic. + retry_agent.messages = [{"role": "user", "content": [{"text": f"msg {i}"}]} for i in range(6)] + + # Second call: proactive path (e=None) — should succeed and commit the compaction. + manager.reduce_context(retry_agent) + + # Both model.stream calls happened: once for the failed attempt, once for the retry. + assert retry_agent.model.stream.call_count == 2 + + # With summary_ratio=0.5 and 6 messages, preserve_recent_messages=1: + # messages_to_summarize = min(3, 5) = 3 + # After partitioning: 2 pinned (0, 1), 1 to_summarize (2) + # removed_count = 1 (no prior summary, so not subtracted) + assert manager.removed_message_count == 1 + assert manager.pinned_head_count == 2 + + # Verify the exact transcript: [msg 0, msg 1, SUMMARY, msg 3, msg 4, msg 5] + texts = [m["content"][0]["text"] for m in retry_agent.messages] + assert texts[0] == "msg 0" + assert texts[1] == "msg 1" + assert "SUMMARY" in str(retry_agent.messages[2]) + assert texts[3] == "msg 3" + assert texts[4] == "msg 4" + assert texts[5] == "msg 5" diff --git a/strands-py/tests/strands/session/test_repository_session_manager.py b/strands-py/tests/strands/session/test_repository_session_manager.py index 98bfe085d4..54fde5a6b8 100644 --- a/strands-py/tests/strands/session/test_repository_session_manager.py +++ b/strands-py/tests/strands/session/test_repository_session_manager.py @@ -5,6 +5,7 @@ import pytest +from strands import tool from strands.agent.agent import Agent from strands.agent.conversation_manager.null_conversation_manager import NullConversationManager from strands.agent.conversation_manager.sliding_window_conversation_manager import SlidingWindowConversationManager @@ -187,6 +188,9 @@ def test_initialize_restores_existing_agent_with_summarizing_conversation_manage assert agent.messages[1]["role"] == "user" assert agent.messages[1]["content"][0]["text"] == "Hello" assert agent.conversation_manager.removed_message_count == 1 + # Exact transcript: summary at index 0, then the one remaining stored message. + texts = [m["content"][0].get("text", "") for m in agent.messages] + assert texts[0] == "summary" def test_append_message(session_manager): @@ -1233,6 +1237,509 @@ def tracking_read_multi_agent(session_id, multi_agent_id, **kwargs): assert read_multi_agent_calls[0] == ("existing-session", "test-multi-agent") +# --------------------------------------------------------------------------- +# Regression tests for https://github.com/strands-agents/harness-sdk/issues/4027 +# +# ``SummarizingConversationManager(pin_first=N)`` must: +# * keep the first N messages at the head of the restored transcript, and +# * not skip those pinned messages via ``list_messages(offset=removed_message_count)``. +# Both ``FileSessionManager`` and ``S3SessionManager`` go through +# ``RepositorySessionManager.initialize`` (the file-backed path is exercised here). +# --------------------------------------------------------------------------- + + +def _populate_messages(repo, agent_id, texts): + """Create SessionMessages with sequential message_ids starting at 0.""" + for i, text in enumerate(texts): + msg = SessionMessage( + message={"role": "user", "content": [ContentBlock(text=text)]}, + message_id=i, + ) + repo.create_message("test-session", agent_id, msg) + + +def test_initialize_uses_persisted_pinned_head_count_not_constructor(existing_session_manager): + """Restore must reattach the persisted head size, not constructor pin_first. + + Compaction can persist a different count than the live pin_first (constructor + changed between runs, or partition stopped mid tool-pair). Using constructor + pin_first on restore reattaches the wrong messages. + """ + conversation_manager = SummarizingConversationManager(pin_first=2) + conversation_manager.removed_message_count = 1 + conversation_manager._summary_message = {"role": "user", "content": [{"text": "summary"}]} + conversation_manager.pinned_head_count = 2 + + session_agent = SessionAgent( + agent_id="mismatch-agent", + state={}, + conversation_manager_state=conversation_manager.get_state(), + ) + existing_session_manager.session_repository.create_agent("test-session", session_agent) + + _populate_messages( + existing_session_manager.session_repository, + "mismatch-agent", + ["pinned-0", "pinned-1", "summarized", "remaining"], + ) + + # Constructor pin_first disagrees with the persisted count on purpose. + agent = Agent(agent_id="mismatch-agent", conversation_manager=SummarizingConversationManager(pin_first=5)) + existing_session_manager.initialize(agent) + + texts = [m["content"][0]["text"] for m in agent.messages] + assert texts == ["pinned-0", "pinned-1", "summary", "remaining"] + assert agent.conversation_manager.pinned_head_count == 2 + + +def test_initialize_empty_tail_does_not_restart_message_ids(existing_session_manager): + """When restore offset lands past the last stored message, append must not restart at 0.""" + conversation_manager = SummarizingConversationManager(pin_first=2) + conversation_manager.removed_message_count = 2 + conversation_manager._summary_message = {"role": "user", "content": [{"text": "summary"}]} + conversation_manager.pinned_head_count = 2 + + session_agent = SessionAgent( + agent_id="empty-tail-agent", + state={}, + conversation_manager_state=conversation_manager.get_state(), + ) + existing_session_manager.session_repository.create_agent("test-session", session_agent) + + _populate_messages( + existing_session_manager.session_repository, + "empty-tail-agent", + ["pinned-0", "pinned-1"], + ) + + agent = Agent(agent_id="empty-tail-agent", conversation_manager=SummarizingConversationManager(pin_first=2)) + existing_session_manager.initialize(agent) + + agent.messages.append({"role": "user", "content": [{"text": "NEW"}]}) + existing_session_manager.append_message(agent.messages[-1], agent) + + stored = existing_session_manager.session_repository.list_messages( + session_id="test-session", + agent_id="empty-tail-agent", + ) + texts = [m.to_message()["content"][0]["text"] for m in stored] + ids = [m.message_id for m in stored] + assert texts == ["pinned-0", "pinned-1", "NEW"] + assert ids == [0, 1, 2] + + +def test_initialize_restores_pinned_messages_after_summary_with_pin_first(existing_session_manager): + """After a compaction with pin_first, restoring the agent must reattach the pinned head. + + ``list_messages(offset=pinned_head_count + removed_message_count)`` skips past the + still-live pinned head so the tail does not overlap with the messages reattached at + the front; the resumed agent must see exactly the live transcript in order. + """ + conversation_manager = SummarizingConversationManager(pin_first=3) + conversation_manager.removed_message_count = 1 + conversation_manager._summary_message = {"role": "user", "content": [{"text": "summary of msg-3"}]} + conversation_manager.pinned_head_count = 3 + + session_agent = SessionAgent( + agent_id="pinned-agent", + state={}, + conversation_manager_state=conversation_manager.get_state(), + ) + existing_session_manager.session_repository.create_agent("test-session", session_agent) + + # 5 originally-stored messages, ids 0..4. msg-3 was the one summarized. + _populate_messages( + existing_session_manager.session_repository, + "pinned-agent", + ["pinned-0", "pinned-1", "pinned-2", "msg-3-summarized", "msg-4"], + ) + + agent = Agent(agent_id="pinned-agent", conversation_manager=SummarizingConversationManager(pin_first=3)) + existing_session_manager.initialize(agent) + + # Live transcript after compaction was: [pinned-0, pinned-1, pinned-2, summary, msg-4]. + # Restore must reproduce exactly that, in order. + texts = [m["content"][0]["text"] for m in agent.messages] + assert texts == ["pinned-0", "pinned-1", "pinned-2", "summary of msg-3", "msg-4"] + + +def test_initialize_pinned_messages_no_summary_yet(existing_session_manager): + """When pin_first is set but summarization hasn't run, restore must produce the full transcript + and must NOT invoke the pinned-head fetch path (pinned_head_count == 0). + + The gate ``if pinned_head_count > 0:`` must suppress the pinned-head list_messages call so + that restoring a pre-compaction session never issues a spurious ``limit=0`` query. If the + gate is replaced with ``if True:``, list_messages is called with ``limit=0`` which is + detectable via the call record. + """ + conversation_manager = SummarizingConversationManager(pin_first=2) + # No compaction has happened, so pinned_head_count == 0 in the persisted state. + + session_agent = SessionAgent( + agent_id="fresh-pin-agent", + state={}, + conversation_manager_state=conversation_manager.get_state(), + ) + existing_session_manager.session_repository.create_agent("test-session", session_agent) + + _populate_messages( + existing_session_manager.session_repository, + "fresh-pin-agent", + ["keep-0", "keep-1", "keep-2"], + ) + + # Spy on list_messages to detect any call that passes limit=0, which would indicate the + # gate fired when it should not have (pinned_head_count == 0). + real_list = existing_session_manager.session_repository.list_messages + limit_zero_calls = [] + + def spy_list_messages(session_id, agent_id, limit=None, offset=0): + if limit == 0: + limit_zero_calls.append((session_id, agent_id, limit, offset)) + return real_list(session_id, agent_id, limit=limit, offset=offset) + + existing_session_manager.session_repository.list_messages = spy_list_messages + + agent = Agent(agent_id="fresh-pin-agent", conversation_manager=SummarizingConversationManager(pin_first=2)) + existing_session_manager.initialize(agent) + + texts = [m["content"][0]["text"] for m in agent.messages] + assert texts == ["keep-0", "keep-1", "keep-2"] + # The gate must have suppressed the fetch — no limit=0 call is acceptable. + assert limit_zero_calls == [], "gate was bypassed: list_messages(limit=0) was called when pinned_head_count == 0" + + +def test_initialize_restores_from_legacy_state_missing_pinned_head_count(existing_session_manager): + """Restore a session persisted before pinned_head_count existed. + + Legacy state has no ``pinned_head_count`` key. The restore path must default to 0, + preserving existing behaviour: no pinned head is reattached, and the offset is + calculated from ``removed_message_count`` alone. + """ + # Simulate a legacy state dict: no pinned_head_count key at all. + legacy_state = { + "__name__": "SummarizingConversationManager", + "removed_message_count": 1, + "summary_message": {"role": "user", "content": [{"text": "summary"}]}, + } + + session_agent = SessionAgent( + agent_id="legacy-agent", + state={}, + conversation_manager_state=legacy_state, + ) + existing_session_manager.session_repository.create_agent("test-session", session_agent) + + _populate_messages( + existing_session_manager.session_repository, + "legacy-agent", + ["a", "b", "c"], + ) + + agent = Agent(agent_id="legacy-agent", conversation_manager=SummarizingConversationManager()) + existing_session_manager.initialize(agent) + + # Legacy behaviour: offset=removed_message_count (1), so we get [summary, b, c]. + texts = [m["content"][0]["text"] for m in agent.messages] + assert texts == ["summary", "b", "c"] + assert agent.conversation_manager.pinned_head_count == 0 + + +def test_restored_pinned_messages_have_pin_markers(existing_session_manager): + """Pinned messages restored from session must have their pin markers re-applied. + + Without re-applying pin markers, the next compaction would summarize away the + pinned messages that were supposed to be protected. + """ + conversation_manager = SummarizingConversationManager(pin_first=2) + conversation_manager.removed_message_count = 1 + conversation_manager._summary_message = {"role": "user", "content": [{"text": "summary"}]} + conversation_manager.pinned_head_count = 2 + + session_agent = SessionAgent( + agent_id="pin-marker-agent", + state={}, + conversation_manager_state=conversation_manager.get_state(), + ) + existing_session_manager.session_repository.create_agent("test-session", session_agent) + + _populate_messages( + existing_session_manager.session_repository, + "pin-marker-agent", + ["pinned-0", "pinned-1", "summarized", "remaining"], + ) + + agent = Agent(agent_id="pin-marker-agent", conversation_manager=SummarizingConversationManager(pin_first=2)) + existing_session_manager.initialize(agent) + + # Verify the pinned messages have their markers + assert agent.messages[0].get("metadata", {}).get("custom", {}).get("pinned") is True + assert agent.messages[1].get("metadata", {}).get("custom", {}).get("pinned") is True + # The summary message (at index 2) should NOT have the pinned marker + assert agent.messages[2].get("metadata", {}).get("custom", {}).get("pinned") is not True + + +# --------------------------------------------------------------------------- +# End-to-end File-backed repro from issue #4027. +# These tests verify the fix by actually running compaction and restore with +# a FileSessionManager, simulating a process restart. +# --------------------------------------------------------------------------- + + +class StubModel: + """A model that reports token usage for proactive compression triggers. + + Args: + fail_summarisation: Raise instead of answering when the summariser calls. + """ + + def __init__(self, *, fail_summarisation: bool = False, context_window: int = 200): + self.fail_summarisation = fail_summarisation + self.context_window = context_window + self.answers = 0 + self.summarisation_calls = 0 + self.stateful = False + self._utilization_limit_warned = False + + def update_config(self, **model_config): + pass + + def get_config(self): + return {"context_window_limit": self.context_window} + + @property + def context_window_limit(self): + return self.context_window + + def estimate_utilization(self, input_tokens): + if self.context_window is None: + return 0.0 + return input_tokens / self.context_window + + async def stream(self, messages, tool_specs=None, system_prompt=None, **kwargs): + """Answer one turn, or one summarisation request.""" + agent_prompt = "Agent System Prompt" + if system_prompt != agent_prompt: + self.summarisation_calls += 1 + if self.fail_summarisation: + raise RuntimeError("the provider is unavailable") + text = "SUMMARY" + else: + self.answers += 1 + text = f"answer {self.answers}" + + input_tokens = len(str(messages)) // 4 + output_tokens = max(1, len(text) // 4) + yield {"messageStart": {"role": "assistant"}} + yield {"contentBlockDelta": {"delta": {"text": text}}} + yield {"contentBlockStop": {}} + yield {"messageStop": {"stopReason": "end_turn"}} + yield { + "metadata": { + "usage": { + "inputTokens": input_tokens, + "outputTokens": output_tokens, + "totalTokens": input_tokens + output_tokens, + }, + "metrics": {"latencyMs": 0}, + } + } + + async def count_tokens(self, messages, tool_specs=None, system_prompt=None): + return len(str(messages)) // 4 + + +def _build_file_agent(tmp_path, model, pin_first=2, compression_threshold=0.5): + """Build an agent with FileSessionManager for end-to-end testing.""" + from strands.session.file_session_manager import FileSessionManager + + return Agent( + model=model, + system_prompt="Agent System Prompt", + conversation_manager=SummarizingConversationManager( + pin_first=pin_first, + proactive_compression={"compression_threshold": compression_threshold}, + ), + session_manager=FileSessionManager(session_id="repro", storage_dir=str(tmp_path)), + callback_handler=None, + ) + + +def _texts(agent): + """Extract the text from each message in the agent's transcript.""" + return [message["content"][0].get("text", "") for message in agent.messages] + + +def _converse_until_compaction(agent, model, max_turns=30): + """Hold a conversation until the SDK compacts it, return the turn count.""" + for turn in range(1, max_turns + 1): + agent(f"question number {turn}") + if model.summarisation_calls: + return turn + raise RuntimeError(f"{max_turns} turns produced no summarisation call") + + +@tool +def echo(text: str) -> str: + """Return the supplied text unchanged.""" + return text + + +class ToolCallingModel: + """A model stub that calls ``echo`` on the first turn and answers with text on every later turn.""" + + def __init__(self, context_window: int = 200): + self.context_window = context_window + self.stateful = False + self._utilization_limit_warned = False + self.calls = 0 + + def update_config(self, **model_config): + pass + + def get_config(self): + return {"context_window_limit": self.context_window} + + @property + def context_window_limit(self): + return self.context_window + + def estimate_utilization(self, input_tokens): + return input_tokens / self.context_window + + async def stream(self, messages, tool_specs=None, system_prompt=None, **kwargs): + """Emit a single ``echo`` tool call, then plain text once the call is in the transcript.""" + self.calls += 1 + called_tool = any("toolUse" in block for message in messages for block in message.get("content", [])) + + yield {"messageStart": {"role": "assistant"}} + if not called_tool: + yield {"contentBlockStart": {"start": {"toolUse": {"toolUseId": "tu1", "name": "echo"}}}} + yield {"contentBlockDelta": {"delta": {"toolUse": {"input": '{"text": "REAL RESULT"}'}}}} + yield {"contentBlockStop": {}} + yield {"messageStop": {"stopReason": "tool_use"}} + else: + yield {"contentBlockDelta": {"delta": {"text": f"answer {self.calls}"}}} + yield {"contentBlockStop": {}} + yield {"messageStop": {"stopReason": "end_turn"}} + yield { + "metadata": { + "usage": {"inputTokens": 1, "outputTokens": 1, "totalTokens": 2}, + "metrics": {"latencyMs": 0}, + } + } + + +def _build_sliding_file_agent(tmp_path, model, pin_first=2, window_size=4): + """Build a sliding-window agent backed by a FileSessionManager for end-to-end testing.""" + return Agent( + model=model, + system_prompt="Agent System Prompt", + tools=[echo], + conversation_manager=SlidingWindowConversationManager( + pin_first=pin_first, window_size=window_size, should_truncate_results=False + ), + session_manager=FileSessionManager(session_id="repro", storage_dir=str(tmp_path)), + callback_handler=None, + ) + + +def _transcript(agent): + """Extract role + content blocks for each message, ignoring per-message metadata such as pin markers.""" + return [(message["role"], copy.deepcopy(message["content"])) for message in agent.messages] + + +def test_file_session_sliding_window_restore_matches_in_memory(tmp_path): + """End-to-end: a resumed sliding-window session matches the live transcript after a pin_first trim.""" + agent = _build_sliding_file_agent(tmp_path, StubModel(), window_size=2) + agent("first question") + agent("second question") + agent("third question") + in_memory = _transcript(agent) + assert agent.conversation_manager.removed_message_count > 0, "no trimming happened, restore is not exercised" + + restored_agent = _build_sliding_file_agent(tmp_path, StubModel(), window_size=2) + restored = _transcript(restored_agent) + + assert restored == in_memory, ( + f"a resumed sliding-window session is not the conversation the agent had.\n" + f" in memory ({len(in_memory)}): {in_memory}\n" + f" restored ({len(restored)}): {restored}" + ) + assert restored_agent.conversation_manager.pinned_head_count == 2 + + +def test_file_session_sliding_window_restore_keeps_pinned_tool_pair(tmp_path): + """End-to-end: the restored head covers the pinned tool pair, so the real tool result survives (#4087). + + ``pin_first=2`` pins the opening user message and the assistant toolUse, and trimming also keeps the + toolResult partner. The persisted head count has to cover all three; otherwise a resumed session + reattaches a toolUse with no result and the session manager fabricates an interrupted-tool error. + """ + agent = _build_sliding_file_agent(tmp_path, ToolCallingModel()) + agent("use the tool") + agent("second question") + agent("third question") + in_memory = _transcript(agent) + assert agent.conversation_manager.pinned_head_count == 3 + assert agent.conversation_manager.removed_message_count > 0, "no trimming happened, restore is not exercised" + + restored_agent = _build_sliding_file_agent(tmp_path, ToolCallingModel()) + restored = _transcript(restored_agent) + + assert restored == in_memory, ( + f"a resumed sliding-window session lost the protected tool pair.\n" + f" in memory ({len(in_memory)}): {in_memory}\n" + f" restored ({len(restored)}): {restored}" + ) + flattened = str(restored) + assert "REAL RESULT" in flattened + assert "Tool was interrupted." not in flattened + + +def test_file_session_resumed_conversation_matches_in_memory(tmp_path): + """End-to-end: after compaction, rebuilding the agent on the same session restores exactly what it had. + + This is the first half of issue #4027: the offset was short by exactly ``pin_first``, + so ``list_messages(offset=…)`` started ``pin_first`` messages too early. + """ + model = StubModel() + agent = _build_file_agent(tmp_path, model) + _converse_until_compaction(agent, model) + in_memory = _texts(agent) + + # Simulate a process restart: fresh FileSessionManager over the same directory + restored_agent = _build_file_agent(tmp_path, StubModel()) + restored = _texts(restored_agent) + + assert restored == in_memory, ( + f"a resumed session is not the conversation the agent had.\n" + f" in memory ({len(in_memory)}): {in_memory}\n" + f" restored ({len(restored)}): {restored}" + ) + + +def test_file_session_failed_summarisation_keeps_messages(tmp_path): + """End-to-end: if summarisation fails, the messages that would have been summarised are preserved. + + This is the second half of issue #4027: mutations happened before summary generation, + so a failed compaction left ``removed_message_count`` incremented with nothing removed. + """ + model = StubModel(fail_summarisation=True) + agent = _build_file_agent(tmp_path, model) + _converse_until_compaction(agent, model) + in_memory = _texts(agent) + + restored_agent = _build_file_agent(tmp_path, StubModel()) + restored = _texts(restored_agent) + + assert restored == in_memory, ( + f"a resumed session is missing messages that no summary replaced.\n" + f" in memory ({len(in_memory)}): {in_memory}\n" + f" restored ({len(restored)}): {restored}" + ) + + def test_fix_broken_tool_use_removes_orphaned_tool_result_mid_conversation(session_manager): """Orphaned toolResult after a message that carries no toolUse is removed (issue #3651).""" messages = [