@aws/agentcore 0.9.1 → 1.0.0-preview.2

package/dist/assets/python/http/strands/base/main.py

@@ -9,6 +9,9 @@ from mcp_client.client import get_streamable_http_mcp_client
 {{#if hasMemory}}
 from memory.session import get_memory_session_manager
 {{/if}}
+{{#if sessionStorageMountPath}}
+import os
+{{/if}}
 
 app = BedrockAgentCoreApp()
 log = app.logger
@@ -30,11 +33,70 @@ def add_numbers(a: int, b: int) -> int:
     return a+b
 tools.append(add_numbers)
 
+{{#if sessionStorageMountPath}}
+SESSION_STORAGE_PATH = "{{sessionStorageMountPath}}"
+
+def _safe_resolve(path: str) -> str:
+    """Resolve path safely within the storage boundary."""
+    resolved = os.path.realpath(os.path.join(SESSION_STORAGE_PATH, path.lstrip("/")))
+    if not resolved.startswith(os.path.realpath(SESSION_STORAGE_PATH)):
+        raise ValueError(f"Path '{path}' is outside the storage boundary")
+    return resolved
+
+@tool
+def file_read(path: str) -> str:
+    """Read a file from persistent storage. The path is relative to the storage root."""
+    try:
+        full_path = _safe_resolve(path)
+        with open(full_path) as f:
+            return f.read()
+    except ValueError as e:
+        return str(e)
+    except OSError as e:
+        return f"Error reading '{path}': {e.strerror}"
+
+@tool
+def file_write(path: str, content: str) -> str:
+    """Write content to a file in persistent storage. The path is relative to the storage root."""
+    try:
+        full_path = _safe_resolve(path)
+        parent = os.path.dirname(full_path)
+        if parent:
+            os.makedirs(parent, exist_ok=True)
+        with open(full_path, "w") as f:
+            f.write(content)
+        return f"Written to {path}"
+    except ValueError as e:
+        return str(e)
+    except OSError as e:
+        return f"Error writing '{path}': {e.strerror}"
+
+@tool
+def list_files(directory: str = "") -> str:
+    """List files in persistent storage. The directory is relative to the storage root."""
+    try:
+        target = _safe_resolve(directory)
+        entries = os.listdir(target)
+        return "\n".join(entries) if entries else "(empty directory)"
+    except ValueError as e:
+        return str(e)
+    except OSError as e:
+        return f"Error listing '{directory}': {e.strerror}"
+
+tools.extend([file_read, file_write, list_files])
+{{/if}}
+
 # Add MCP client to tools if available
 for mcp_client in mcp_clients:
     if mcp_client:
         tools.append(mcp_client)
 
+SYSTEM_PROMPT = """
+You are a helpful assistant. Use tools when appropriate.
+{{#if sessionStorageMountPath}}
+You have persistent storage at {{sessionStorageMountPath}}. Use file tools to read and write files. Data persists across sessions.
+{{/if}}
+"""
 
 {{#if hasMemory}}
 def agent_factory():
@@ -46,9 +108,7 @@ def agent_factory():
             cache[key] = Agent(
                 model=load_model(),
                 session_manager=get_memory_session_manager(session_id, user_id),
-                system_prompt="""
-                    You are a helpful assistant. Use tools when appropriate.
-                """,
+                system_prompt=SYSTEM_PROMPT,
                 tools=tools
             )
         return cache[key]
@@ -62,9 +122,7 @@ def get_or_create_agent():
     if _agent is None:
         _agent = Agent(
             model=load_model(),
-            system_prompt="""
-                You are a helpful assistant. Use tools when appropriate.
-            """,
+            system_prompt=SYSTEM_PROMPT,
             tools=tools
         )
     return _agent