feat(offloader): return relevant path type from storage

Author: lizradway

src/strands/vended_plugins/context_offloader/storage.py

@@ -131,15 +131,19 @@ def _extension_for(content_type: str) -> str:
         return f".{content_type.split('/')[-1]}"
 
     def store(self, key: str, content: bytes, content_type: str = "text/plain") -> str:
-        """Store content as a file and return the filename as reference.
+        """Store content as a file and return the path as reference.
+
+        The returned path preserves the form of ``artifact_dir`` passed to
+        the constructor: a relative ``artifact_dir`` yields a relative
+        reference, an absolute one yields an absolute reference.
 
         Args:
             key: A unique key for this content block.
             content: The raw content bytes to store.
             content_type: MIME type of the content.
 
         Returns:
-            The filename (not full path) used as the reference.
+            The file path (e.g., ``./artifacts/1234_1_key.txt``).
         """
         self._artifact_dir.mkdir(parents=True, exist_ok=True)
 
@@ -156,26 +160,33 @@ def store(self, key: str, content: bytes, content_type: str = "text/plain") -> s
         file_path = self._artifact_dir / filename
         file_path.write_bytes(content)
 
-        return filename
+        return str(file_path)
 
     def retrieve(self, reference: str) -> tuple[bytes, str]:
         """Retrieve content from a stored file.
 
+        Accepts both full paths (as returned by ``store()``) and bare
+        filenames for backward compatibility.
+
         Args:
-            reference: The filename reference returned by store().
+            reference: The file path or filename returned by store().
 
         Returns:
             A tuple of (content bytes, content type).
 
         Raises:
             KeyError: If the file does not exist.
         """
-        file_path = (self._artifact_dir / reference).resolve()
+        ref_path = Path(reference)
+        file_path = ref_path.resolve() if len(ref_path.parts) > 1 else (self._artifact_dir / reference).resolve()
+        if not file_path.is_relative_to(self._artifact_dir.resolve()):
+            file_path = (self._artifact_dir / reference).resolve()
         if not file_path.is_relative_to(self._artifact_dir.resolve()):
             raise KeyError(f"Reference not found: {reference}")
         if not file_path.is_file():
             raise KeyError(f"Reference not found: {reference}")
-        content_type = self._content_types.get(reference, "application/octet-stream")
+        filename = file_path.name
+        content_type = self._content_types.get(filename, "application/octet-stream")
         return file_path.read_bytes(), content_type
 
     def _load_metadata(self) -> dict[str, str]:
@@ -320,15 +331,15 @@ def __init__(
         self._lock = threading.Lock()
 
     def store(self, key: str, content: bytes, content_type: str = "text/plain") -> str:
-        """Store content as an S3 object and return the object key as reference.
+        """Store content as an S3 object and return an ``s3://`` URI as reference.
 
         Args:
             key: A unique key for this content block.
             content: The raw content bytes to store.
             content_type: MIME type of the content.
 
         Returns:
-            The S3 object key used as the reference.
+            An S3 URI (e.g., ``s3://bucket/prefix/1234_1_key``).
 
         Raises:
             botocore.exceptions.ClientError: If the S3 operation fails (e.g., bucket
@@ -348,22 +359,31 @@ def store(self, key: str, content: bytes, content_type: str = "text/plain") -> s
             ContentType=content_type,
         )
 
-        return s3_key
+        return f"s3://{self._bucket}/{s3_key}"
 
     def retrieve(self, reference: str) -> tuple[bytes, str]:
         """Retrieve content from an S3 object.
 
+        Accepts both ``s3://`` URIs (as returned by ``store()``) and raw
+        S3 keys for backward compatibility.
+
         Args:
-            reference: The S3 object key returned by store().
+            reference: The S3 URI or object key returned by store().
 
         Returns:
             A tuple of (content bytes, content type).
 
         Raises:
             KeyError: If the object does not exist.
         """
+        s3_key = reference
+        if reference.startswith("s3://"):
+            expected_prefix = f"s3://{self._bucket}/"
+            if not reference.startswith(expected_prefix):
+                raise KeyError(f"Reference not found: {reference}")
+            s3_key = reference[len(expected_prefix):]
         try:
-            response = self._client.get_object(Bucket=self._bucket, Key=reference)
+            response = self._client.get_object(Bucket=self._bucket, Key=s3_key)
             content: bytes = response["Body"].read()
             content_type: str = response.get("ContentType", "application/octet-stream")
             return content, content_type

tests/strands/vended_plugins/context_offloader/test_plugin.py

@@ -11,6 +11,7 @@
 from strands.types.tools import ToolContext, ToolUse
 from strands.vended_plugins.context_offloader import (
     ContextOffloader,
+    FileStorage,
     InMemoryStorage,
 )
 
@@ -537,3 +538,45 @@ async def test_guidance_does_not_mention_retrieval_tool_when_disabled(self, stor
         result_text = event.result["content"][0]["text"]
         assert "retrieve_offloaded_content" not in result_text
         assert "available tools" in result_text
+
+
+class TestActionableReferences:
+    """Tests that storage-specific references appear in the offloaded preview."""
+
+    @pytest.mark.asyncio
+    async def test_file_storage_path_in_preview(self, tmp_path, mock_agent):
+        storage = FileStorage(artifact_dir=str(tmp_path / "artifacts"))
+        plugin = ContextOffloader(storage=storage, max_result_tokens=25, preview_tokens=10)
+        event = _make_event(mock_agent, "a" * 200)
+
+        await plugin._handle_tool_result(event)
+
+        result_text = event.result["content"][0]["text"]
+        assert str(tmp_path / "artifacts") in result_text
+
+    @pytest.mark.asyncio
+    async def test_file_storage_image_placeholder_has_path(self, tmp_path, mock_agent):
+        storage = FileStorage(artifact_dir=str(tmp_path / "artifacts"))
+        plugin = ContextOffloader(storage=storage, max_result_tokens=25, preview_tokens=10)
+        img_bytes = b"\x89PNG" + b"\x00" * 100
+        content = [
+            {"text": "x" * 200},
+            {"image": {"format": "png", "source": {"bytes": img_bytes}}},
+        ]
+        event = _make_event(mock_agent, content)
+
+        await plugin._handle_tool_result(event)
+
+        placeholder = event.result["content"][1]["text"]
+        assert str(tmp_path / "artifacts") in placeholder
+
+    @pytest.mark.asyncio
+    async def test_inmemory_storage_opaque_reference_in_preview(self, mock_agent):
+        storage = InMemoryStorage()
+        plugin = ContextOffloader(storage=storage, max_result_tokens=25, preview_tokens=10)
+        event = _make_event(mock_agent, "a" * 200)
+
+        await plugin._handle_tool_result(event)
+
+        result_text = event.result["content"][0]["text"]
+        assert "mem_" in result_text

tests/strands/vended_plugins/context_offloader/test_storage.py

@@ -1,6 +1,7 @@
 """Tests for offload storage backends."""
 
 import threading
+from pathlib import Path
 from unittest.mock import MagicMock, patch
 
 import pytest
@@ -147,7 +148,30 @@ def test_sanitizes_path_traversal(self, tmp_path):
         storage = FileStorage(artifact_dir=str(tmp_path))
         ref = storage.store("../../etc/passwd", b"content")
         assert ".." not in ref
-        assert "/" not in ref
+        assert "/" not in Path(ref).name
+
+    def test_reference_includes_artifact_dir(self, tmp_path):
+        artifact_dir = str(tmp_path / "artifacts")
+        storage = FileStorage(artifact_dir=artifact_dir)
+        ref = storage.store("key_1", b"content")
+        assert ref.startswith(artifact_dir + "/")
+
+    def test_relative_artifact_dir_gives_relative_reference(self, tmp_path, monkeypatch):
+        monkeypatch.chdir(tmp_path)
+        storage = FileStorage(artifact_dir="./artifacts")
+        ref = storage.store("key_1", b"content")
+        assert ref.startswith("artifacts/")
+        content, content_type = storage.retrieve(ref)
+        assert content == b"content"
+        assert content_type == "text/plain"
+
+    def test_retrieve_accepts_bare_filename(self, tmp_path):
+        storage = FileStorage(artifact_dir=str(tmp_path))
+        ref = storage.store("key_1", b"hello world")
+        filename = Path(ref).name
+        content, content_type = storage.retrieve(filename)
+        assert content == b"hello world"
+        assert content_type == "text/plain"
 
     def test_metadata_survives_across_instances(self, tmp_path):
         artifact_dir = str(tmp_path / "artifacts")
@@ -233,9 +257,9 @@ def test_unique_references(self, storage):
         assert storage.retrieve(ref1)[0] == b"content a"
         assert storage.retrieve(ref2)[0] == b"content b"
 
-    def test_reference_includes_prefix(self, storage):
+    def test_reference_is_s3_uri(self, storage):
         ref = storage.store("tool_abc", b"content")
-        assert ref.startswith("artifacts/")
+        assert ref.startswith("s3://test-bucket/artifacts/")
 
     def test_empty_prefix(self, mock_s3_client):
         with patch("boto3.Session") as mock_session_cls:
@@ -245,9 +269,20 @@ def test_empty_prefix(self, mock_s3_client):
             storage = S3Storage(bucket="test-bucket", prefix="")
 
         ref = storage.store("tool_abc", b"content")
-        assert not ref.startswith("/")
+        assert ref.startswith("s3://test-bucket/")
         assert storage.retrieve(ref)[0] == b"content"
 
+    def test_retrieve_accepts_raw_key(self, storage, mock_s3_client):
+        ref = storage.store("key_1", b"hello world")
+        raw_key = ref.removeprefix("s3://test-bucket/")
+        content, content_type = storage.retrieve(raw_key)
+        assert content == b"hello world"
+        assert content_type == "text/plain"
+
+    def test_retrieve_rejects_wrong_bucket_uri(self, storage):
+        with pytest.raises(KeyError, match="Reference not found"):
+            storage.retrieve("s3://wrong-bucket/artifacts/some_key")
+
     def test_put_object_called_with_correct_params(self, storage, mock_s3_client):
         storage.store("key_1", b"test content", "application/json")