From b5277601c0c3b9ff0f321373946c5027e30e855a Mon Sep 17 00:00:00 2001 From: strandly-the-agent Date: Wed, 28 Oct 2026 16:00:49 +0000 Subject: [PATCH 1/2] fix(context-manager): never stash an offload preview under its original's key After a session restore onto a host whose stash no longer holds the originals, the one-time backfill stored every restored block, including already-truncated tool-result previews, under the original's reference. retrieve_context then served the preview as the original. Backfill now skips blocks that are offload markers, and retrieve_context refuses to serve a stashed entry that is itself a marker, which also covers stores already polluted by earlier releases. Fixes #5620 --- .../_context_manager/retrieval_tool.py | 8 +++ .../src/strands/_context_manager/stash.py | 22 ++++++- .../_context_manager/test_context_manager.py | 26 ++++++++ .../_context_manager/test_retrieval_tool.py | 10 ++++ .../strands/_context_manager/test_stash.py | 59 ++++++++++++++++++- 5 files changed, 123 insertions(+), 2 deletions(-) diff --git a/strands-py/src/strands/_context_manager/retrieval_tool.py b/strands-py/src/strands/_context_manager/retrieval_tool.py index 7904077d47..ba9d52bc94 100644 --- a/strands-py/src/strands/_context_manager/retrieval_tool.py +++ b/strands-py/src/strands/_context_manager/retrieval_tool.py @@ -14,6 +14,7 @@ from ..types.tools import ToolResult, ToolResultContent, ToolSpec from ..vended_plugins.context_offloader.search import _search_content from .methods.truncate import CHARS_PER_TOKEN +from .stash import _is_offload_marker if TYPE_CHECKING: from ..types.content import Message @@ -72,6 +73,13 @@ async def _invoke(tool_use: dict[str, Any], **kwargs: Any) -> ToolResult: text = f"Error: reference not found: {reference}" return ToolResult(toolUseId=tool_use_id, status="error", content=[ToolResultContent(text=text)]) + stashed_text = _extract_text(result) + if stashed_text is not None and _is_offload_marker(stashed_text): + text = ( + f"Error: original content for {reference} is no longer available; only its offloaded preview remains." + ) + return ToolResult(toolUseId=tool_use_id, status="error", content=[ToolResultContent(text=text)]) + if pattern is None and line_range is None: if isinstance(result, dict): media = _restore_media(result) diff --git a/strands-py/src/strands/_context_manager/stash.py b/strands-py/src/strands/_context_manager/stash.py index e689c7a3f7..1330822070 100644 --- a/strands-py/src/strands/_context_manager/stash.py +++ b/strands-py/src/strands/_context_manager/stash.py @@ -9,6 +9,7 @@ import base64 import json import logging +import re from typing import TYPE_CHECKING, Any from ..storage.storage import _NAMESPACED, _NamespacedStorage @@ -54,6 +55,21 @@ def _format_stash_refs(refs: list[str]) -> str: return f" [refs: {', '.join(refs)}]" +_STASH_REFS_SUFFIX = re.compile(r"\[refs?: [^\]]*_\d+\]$") + + +def _is_offload_marker(text: str) -> bool: + """True when ``text`` is an offload replacement (preview, summary, drop marker) ending in a stash ref token.""" + return _STASH_REFS_SUFFIX.search(text.rstrip()) is not None + + +def _block_is_offload_marker(block: ContentBlock) -> bool: + """True when a block carries a stash ref token, i.e. it stands in for an original rather than being one.""" + if "toolResult" in block: + return any(_is_offload_marker(item["text"]) for item in block["toolResult"]["content"] if "text" in item) + return "text" in block and _is_offload_marker(block["text"]) + + def _resolve_stash_namespace( storage: Storage, session_id: str, agent_id: str, *, custom_stash_namespace: bool ) -> Storage: @@ -107,9 +123,13 @@ async def store_message( ) -> None: """Eagerly persist all stashable blocks from a message. - With ``keep_existing``, entries already stored under a block's key are left untouched. + With ``keep_existing`` (backfill of restored history), entries already stored under a block's key are + left untouched, and blocks that are themselves offload markers are skipped: they stand in for an + original and must not be stored under its key. """ for block_index, block in enumerate(message["content"]): + if keep_existing and _block_is_offload_marker(block): + continue if "toolResult" in block: tool_result = block["toolResult"] if skip_tool_use_ids and tool_result["toolUseId"] in skip_tool_use_ids: diff --git a/strands-py/tests/strands/_context_manager/test_context_manager.py b/strands-py/tests/strands/_context_manager/test_context_manager.py index 2573ad5c07..dc669c5e59 100644 --- a/strands-py/tests/strands/_context_manager/test_context_manager.py +++ b/strands-py/tests/strands/_context_manager/test_context_manager.py @@ -413,6 +413,32 @@ async def test_backfill_does_not_overwrite_existing_stash_entries(self, mock_age assert await cm.stash.retrieve("pre-tu-1_0") == {"text": "original full result"} assert await cm.stash.retrieve("seed_0") == {"text": "seed message"} + @pytest.mark.asyncio + async def test_backfill_does_not_stash_previews_when_original_is_gone(self, mock_agent): + """Restored onto an empty stash, a truncated preview must not be stored under its original's key.""" + mock_agent.session_id = "test-session" + mock_agent.storage = None + preview = "[Truncated: 1 block, ~7,500 tokens]\n\nrow-1\nrow-2\n\n[Stashed] [ref: pre-tu-1_0]" + mock_agent.messages = [ + Message(role="user", content=[ContentBlock(text="seed message")], tracking_id="seed"), + Message( + role="user", + content=[ + ContentBlock( + toolResult=ToolResult(toolUseId="pre-tu-1", status="success", content=[{"text": preview}]) + ) + ], + ), + ] + cm = ContextManager(strategies=[], stash=True) + cm.init_agent(mock_agent) + + event = BeforeModelCallEvent(agent=mock_agent, projected_input_tokens=100) + await mock_agent.hooks.invoke_callbacks_async(event) + + assert await cm.stash.retrieve("pre-tu-1_0") is None + assert await cm.stash.retrieve("seed_0") == {"text": "seed message"} + @pytest.mark.asyncio async def test_backfill_runs_only_once(self): agent = unittest.mock.MagicMock() diff --git a/strands-py/tests/strands/_context_manager/test_retrieval_tool.py b/strands-py/tests/strands/_context_manager/test_retrieval_tool.py index 635a9eda52..2d861c6775 100644 --- a/strands-py/tests/strands/_context_manager/test_retrieval_tool.py +++ b/strands-py/tests/strands/_context_manager/test_retrieval_tool.py @@ -72,6 +72,16 @@ async def test_returns_error_for_unknown_reference(self, stash): assert result["status"] == "error" assert "not found" in result["content"][0]["text"] + @pytest.mark.asyncio + async def test_refuses_to_serve_a_stashed_preview_as_the_original(self, stash): + """A stash polluted with an offload preview (pre-fix backfill) must not pass it off as the original.""" + ref = await _store_text(stash, "[Truncated: 1 block, ~7,500 tokens]\n\nrow-1\n\n[Stashed] [ref: tool-1_0]") + tool = _create_retrieval_tool(stash) + for inputs in ({"reference": ref}, {"reference": ref, "pattern": "row-4321"}): + result = await tool._tool_func({"toolUseId": "t1", "input": inputs}) + assert result["status"] == "error" + assert "no longer available" in result["content"][0]["text"] + @pytest.mark.asyncio async def test_retrieves_json_content(self, stash): data = json.dumps({"json": {"key": "value"}}).encode("utf-8") diff --git a/strands-py/tests/strands/_context_manager/test_stash.py b/strands-py/tests/strands/_context_manager/test_stash.py index ad2d03df5a..268c61bd65 100644 --- a/strands-py/tests/strands/_context_manager/test_stash.py +++ b/strands-py/tests/strands/_context_manager/test_stash.py @@ -5,7 +5,7 @@ import pytest -from strands._context_manager.stash import Stash, _BytesEncoder, _format_stash_refs +from strands._context_manager.stash import Stash, _BytesEncoder, _format_stash_refs, _is_offload_marker from strands.storage.in_memory_storage import InMemoryStorage from strands.types.content import ContentBlock, Message from strands.types.tools import ToolResult @@ -233,6 +233,63 @@ def test_multiple_refs(self): assert result == " [refs: tu-1_0, tu-1_1]" +class TestIsOffloadMarker: + """Tests for the _is_offload_marker helper.""" + + @pytest.mark.parametrize( + "text", + [ + "[Truncated: 1 block, ~7,500 tokens]\n\nhead\n\n[... 100 chars elided ...]\n\n[Stashed] [ref: tu-1_0]", + "[Dropped] [refs: tu-1_0, tu-1_1]", + "[Offloaded: ~500 tokens] [ref: track-1_2]", + "[Summarized: ~900 tokens]\n\nsummary [ref: tooluse_abc_0]\n", + ], + ) + def test_matches_offload_replacements(self, text): + assert _is_offload_marker(text) + + @pytest.mark.parametrize( + "text", + [ + "plain result", + "see [ref: RFC-2119] for details", + "[ref: tu-1_0] appears at the start, not the end", + "", + ], + ) + def test_ignores_originals(self, text): + assert not _is_offload_marker(text) + + +class TestStoreMessageSkipsOffloadMarkers: + """Backfill must never write a preview under the key of the original it replaced.""" + + def _preview_message(self): + preview = "[Truncated: 1 block, ~7,500 tokens]\n\nhead\n\n[Stashed] [ref: tu-1_0]" + return Message( + role="user", + content=[ + ContentBlock(toolResult=ToolResult(toolUseId="tu-1", status="success", content=[{"text": preview}])), + ContentBlock(text="[Dropped] [ref: track-1_1]"), + ContentBlock(text="an ordinary block"), + ], + tracking_id="track-1", + ) + + @pytest.mark.asyncio + async def test_keep_existing_skips_marker_blocks_on_empty_stash(self, stash): + await stash.store_message(self._preview_message(), keep_existing=True) + assert await stash.retrieve("tu-1_0") is None + assert await stash.retrieve("track-1_1") is None + assert await stash.retrieve("track-1_2") == {"text": "an ordinary block"} + + @pytest.mark.asyncio + async def test_live_path_still_stores_marker_lookalikes(self, stash): + await stash.store_message(self._preview_message()) + assert await stash.retrieve("tu-1_0") is not None + assert await stash.retrieve("track-1_1") == {"text": "[Dropped] [ref: track-1_1]"} + + class TestBytesEncoder: """Tests for _BytesEncoder.""" From 9021ae4acaf35fac75848fc64ad4a5bc429bb29e Mon Sep 17 00:00:00 2001 From: strandly-the-agent Date: Wed, 28 Oct 2026 16:21:26 +0000 Subject: [PATCH 2/2] fix(context-manager): only treat a marker as a preview when it names its own key Review follow-up: bound the ref-token regex (no backtracking across '[') and require the marker to reference the key it would be stored under, so an original that merely cites another stash entry is still stored and served. --- .../_context_manager/retrieval_tool.py | 2 +- .../src/strands/_context_manager/stash.py | 26 +++++++++++-------- .../_context_manager/test_retrieval_tool.py | 7 +++++ .../strands/_context_manager/test_stash.py | 23 ++++++++++++---- 4 files changed, 41 insertions(+), 17 deletions(-) diff --git a/strands-py/src/strands/_context_manager/retrieval_tool.py b/strands-py/src/strands/_context_manager/retrieval_tool.py index ba9d52bc94..fc569a09d7 100644 --- a/strands-py/src/strands/_context_manager/retrieval_tool.py +++ b/strands-py/src/strands/_context_manager/retrieval_tool.py @@ -74,7 +74,7 @@ async def _invoke(tool_use: dict[str, Any], **kwargs: Any) -> ToolResult: return ToolResult(toolUseId=tool_use_id, status="error", content=[ToolResultContent(text=text)]) stashed_text = _extract_text(result) - if stashed_text is not None and _is_offload_marker(stashed_text): + if stashed_text is not None and _is_offload_marker(stashed_text, [reference]): text = ( f"Error: original content for {reference} is no longer available; only its offloaded preview remains." ) diff --git a/strands-py/src/strands/_context_manager/stash.py b/strands-py/src/strands/_context_manager/stash.py index 1330822070..dff970e862 100644 --- a/strands-py/src/strands/_context_manager/stash.py +++ b/strands-py/src/strands/_context_manager/stash.py @@ -10,6 +10,7 @@ import json import logging import re +from collections.abc import Iterable from typing import TYPE_CHECKING, Any from ..storage.storage import _NAMESPACED, _NamespacedStorage @@ -55,19 +56,22 @@ def _format_stash_refs(refs: list[str]) -> str: return f" [refs: {', '.join(refs)}]" -_STASH_REFS_SUFFIX = re.compile(r"\[refs?: [^\]]*_\d+\]$") +_STASH_REFS_SUFFIX = re.compile(r"\[refs?: ([^\[\]]*)\]$") -def _is_offload_marker(text: str) -> bool: - """True when ``text`` is an offload replacement (preview, summary, drop marker) ending in a stash ref token.""" - return _STASH_REFS_SUFFIX.search(text.rstrip()) is not None +def _is_offload_marker(text: str, references: Iterable[str]) -> bool: + """True when ``text`` is an offload replacement (preview, summary, drop marker) for one of ``references``.""" + match = _STASH_REFS_SUFFIX.search(text.rstrip()) + return match is not None and not set(match.group(1).split(", ")).isdisjoint(references) -def _block_is_offload_marker(block: ContentBlock) -> bool: - """True when a block carries a stash ref token, i.e. it stands in for an original rather than being one.""" +def _block_is_offload_marker(block: ContentBlock, references: list[str]) -> bool: + """True when a block stands in for one of ``references`` rather than being the original stored there.""" if "toolResult" in block: - return any(_is_offload_marker(item["text"]) for item in block["toolResult"]["content"] if "text" in item) - return "text" in block and _is_offload_marker(block["text"]) + texts = [item["text"] for item in block["toolResult"]["content"] if "text" in item] + else: + texts = [block["text"]] if "text" in block else [] + return any(_is_offload_marker(text, references) for text in texts) def _resolve_stash_namespace( @@ -124,11 +128,11 @@ async def store_message( """Eagerly persist all stashable blocks from a message. With ``keep_existing`` (backfill of restored history), entries already stored under a block's key are - left untouched, and blocks that are themselves offload markers are skipped: they stand in for an - original and must not be stored under its key. + left untouched, and a block that is itself the offload marker for its own key is skipped: it stands in + for the original and must not be stored in its place. """ for block_index, block in enumerate(message["content"]): - if keep_existing and _block_is_offload_marker(block): + if keep_existing and _block_is_offload_marker(block, self.refs_for(block, message, block_index)): continue if "toolResult" in block: tool_result = block["toolResult"] diff --git a/strands-py/tests/strands/_context_manager/test_retrieval_tool.py b/strands-py/tests/strands/_context_manager/test_retrieval_tool.py index 2d861c6775..29f0f241f0 100644 --- a/strands-py/tests/strands/_context_manager/test_retrieval_tool.py +++ b/strands-py/tests/strands/_context_manager/test_retrieval_tool.py @@ -82,6 +82,13 @@ async def test_refuses_to_serve_a_stashed_preview_as_the_original(self, stash): assert result["status"] == "error" assert "no longer available" in result["content"][0]["text"] + @pytest.mark.asyncio + async def test_serves_an_original_that_cites_another_reference(self, stash): + ref = await _store_text(stash, "see the earlier output [ref: tool-9_0]") + tool = _create_retrieval_tool(stash) + result = await tool._tool_func({"toolUseId": "t1", "input": {"reference": ref}}) + assert result["status"] == "success" + @pytest.mark.asyncio async def test_retrieves_json_content(self, stash): data = json.dumps({"json": {"key": "value"}}).encode("utf-8") diff --git a/strands-py/tests/strands/_context_manager/test_stash.py b/strands-py/tests/strands/_context_manager/test_stash.py index 268c61bd65..1b7d51139b 100644 --- a/strands-py/tests/strands/_context_manager/test_stash.py +++ b/strands-py/tests/strands/_context_manager/test_stash.py @@ -241,12 +241,12 @@ class TestIsOffloadMarker: [ "[Truncated: 1 block, ~7,500 tokens]\n\nhead\n\n[... 100 chars elided ...]\n\n[Stashed] [ref: tu-1_0]", "[Dropped] [refs: tu-1_0, tu-1_1]", - "[Offloaded: ~500 tokens] [ref: track-1_2]", - "[Summarized: ~900 tokens]\n\nsummary [ref: tooluse_abc_0]\n", + "[Offloaded: ~500 tokens] [ref: tu-1_0]", + "[Summarized: ~900 tokens]\n\nsummary [ref: tu-1_0]\n", ], ) - def test_matches_offload_replacements(self, text): - assert _is_offload_marker(text) + def test_matches_offload_replacements_for_their_own_ref(self, text): + assert _is_offload_marker(text, ["tu-1_0"]) @pytest.mark.parametrize( "text", @@ -254,11 +254,16 @@ def test_matches_offload_replacements(self, text): "plain result", "see [ref: RFC-2119] for details", "[ref: tu-1_0] appears at the start, not the end", + "an original that happens to cite another entry [ref: tu-9_0]", "", ], ) def test_ignores_originals(self, text): - assert not _is_offload_marker(text) + assert not _is_offload_marker(text, ["tu-1_0"]) + + def test_bounded_on_large_bracket_heavy_text(self): + text = "[" * 200_000 + " [ref: tu-1_0]" + assert _is_offload_marker(text, ["tu-1_0"]) class TestStoreMessageSkipsOffloadMarkers: @@ -283,6 +288,14 @@ async def test_keep_existing_skips_marker_blocks_on_empty_stash(self, stash): assert await stash.retrieve("track-1_1") is None assert await stash.retrieve("track-1_2") == {"text": "an ordinary block"} + @pytest.mark.asyncio + async def test_keep_existing_stores_originals_that_cite_other_refs(self, stash): + block = ContentBlock( + toolResult=ToolResult(toolUseId="tu-2", status="success", content=[{"text": "cites [ref: tu-1_0]"}]) + ) + await stash.store_message(Message(role="user", content=[block]), keep_existing=True) + assert await stash.retrieve("tu-2_0") == {"text": "cites [ref: tu-1_0]"} + @pytest.mark.asyncio async def test_live_path_still_stores_marker_lookalikes(self, stash): await stash.store_message(self._preview_message())