feat: Add automatic session management to Code Interpreter tool

src/strands_tools/code_interpreter/__init__.py

@@ -31,9 +31,4 @@
 from .agent_core_code_interpreter import AgentCoreCodeInterpreter
 from .code_interpreter import CodeInterpreter
 
-__all__ = [
-    # Base classes
-    "CodeInterpreter",
-    # Platform implementations
-    "AgentCoreCodeInterpreter",
-]
+__all__ = ["CodeInterpreter", "AgentCoreCodeInterpreter"]

src/strands_tools/code_interpreter/agent_core_code_interpreter.py

@@ -1,4 +1,5 @@
 import logging
+import uuid
 from dataclasses import dataclass
 from typing import Any, Dict, List, Optional
 
@@ -42,88 +43,59 @@ class SessionInfo:
 
 
 class AgentCoreCodeInterpreter(CodeInterpreter):
-    """
-    Bedrock AgentCore implementation of the CodeInterpreter.
-
-    This class provides a code interpreter interface using AWS Bedrock AgentCore services.
-    It supports executing Python, JavaScript, and TypeScript code in isolated sandbox
-    environments with custom code interpreter identifiers.
-
-    The class maintains session state and provides methods for code execution, file
-    operations, and session management. It supports both default AWS code interpreter
-    environments and custom environments specified by identifier.
-
-    Examples:
-        Basic usage with default identifier:
-
-        >>> interpreter = AgentCoreCodeInterpreter(region="us-west-2")
-        >>> # Uses default identifier: "aws.codeinterpreter.v1"
-
-        Using a custom code interpreter identifier:
-
-        >>> custom_id = "my-custom-interpreter-abc123"
-        >>> interpreter = AgentCoreCodeInterpreter(
-        ...     region="us-west-2",
-        ...     identifier=custom_id
-        ... )
-
-        Environment-specific usage:
-
-        >>> # For testing environments
-        >>> test_interpreter = AgentCoreCodeInterpreter(
-        ...     region="us-east-1",
-        ...     identifier="test-interpreter-xyz789"
-        ... )
-
-        >>> # For production environments
-        >>> prod_interpreter = AgentCoreCodeInterpreter(
-        ...     region="us-west-2",
-        ...     identifier="prod-interpreter-def456"
-        ... )
-
-    Attributes:
-        region (str): The AWS region where the code interpreter service is hosted.
-        identifier (str): The code interpreter identifier being used for sessions.
-    """
-
-    def __init__(self, region: Optional[str] = None, identifier: Optional[str] = None) -> None:
+    def __init__(
+        self,
+        region: Optional[str] = None,
+        identifier: Optional[str] = None,
+        session_name: Optional[str] = None,
+        auto_create: bool = True,
+    ) -> None:
         """
         Initialize the Bedrock AgentCore code interpreter.
 
         Args:
-            region (Optional[str]): AWS region for the sandbox service. If not provided,
-                the region will be resolved from AWS configuration (environment variables,
-                AWS config files, or instance metadata). Defaults to None.
-            identifier (Optional[str]): Custom code interpreter identifier to use
-                for code execution sessions. This allows you to specify custom code
-                interpreter environments instead of the default AWS-provided one.
-
-                Valid formats include:
-                - Default identifier: "aws.codeinterpreter.v1" (used when None)
-                - Custom identifier: "my-custom-interpreter-abc123"
-                - Environment-specific: "test-interpreter-xyz789"
-
-                Note: Use the code interpreter ID, not the full ARN. The AWS service
-                expects the identifier portion only (e.g., "my-interpreter-123" rather
-                than "arn:aws:bedrock-agentcore:region:account:code-interpreter-custom/my-interpreter-123").
-
-                If not provided, defaults to "aws.codeinterpreter.v1" for backward
-                compatibility. Defaults to None.
-
-        Note:
-            This constructor maintains full backward compatibility. Existing code
-            that doesn't specify the identifier parameter will continue to work
-            unchanged with the default AWS code interpreter environment.
-
-        Raises:
-            Exception: If there are issues with AWS region resolution or client
-                initialization during session creation.
+            region: AWS region for the sandbox service.
+            identifier: Custom code interpreter identifier.
+            session_name: Session name or strategy:
+                - None (default): Generate random session ID per instance
+                - "runtime": Use AgentCore runtime session_id (set at invocation time)
+                - Specific string: Use this exact session name
+            auto_create: Automatically create sessions if they don't exist.
+                - True (default): Create sessions on-demand
+                - False: Fail if session doesn't exist (strict mode)
+
+        Examples:
+            # Case 1: Random session per instance (default)
+            interpreter = AgentCoreCodeInterpreter()
+
+            # Case 2: Bind to runtime session (recommended for production)
+            session_id = getattr(context, 'session_id', None)
+            interpreter = AgentCoreCodeInterpreter(session_name=session_id)
+
+            # Case 3: Named session with auto-create
+            interpreter = AgentCoreCodeInterpreter(session_name="my-analysis")
+
+            # Case 4: Strict mode - must pre-initialize
+            interpreter = AgentCoreCodeInterpreter(
+                session_name="must-exist",
+                auto_create=False
+            )
         """
         super().__init__()
         self.region = resolve_region(region)
         self.identifier = identifier or "aws.codeinterpreter.v1"
+        self.auto_create = auto_create
+
+        # Generate session name strategy
+        if session_name is None:
+            self.default_session = f"session-{uuid.uuid4().hex[:12]}"
+        else:
+            self.default_session = session_name
+
         self._sessions: Dict[str, SessionInfo] = {}
 
+        logger.info(f"Initialized CodeInterpreter with session='{self.default_session}', " f"auto_create={auto_create}")
+
     def start_platform(self) -> None:
         """Initialize the Bedrock AgentCoreplatform connection."""
         pass
@@ -240,84 +212,127 @@ def list_local_sessions(self) -> Dict[str, Any]:
             "content": [{"json": {"sessions": sessions_info, "totalSessions": len(sessions_info)}}],
         }
 
+    def _ensure_session(self, session_name: Optional[str]) -> tuple[str, Optional[Dict[str, Any]]]:
+        """
+        Ensure a session exists based on configuration.
+
+        Behavior matrix:
+        | session_name | auto_create | Behavior |
+        |--------------|-------------|----------|
+        | None         | True        | Use default_session, create if needed |
+        | None         | False       | Use default_session, error if missing |
+        | "my-session" | True        | Use "my-session", create if needed |
+        | "my-session" | False       | Use "my-session", error if missing |
+
+        Args:
+            session_name: Explicit session name from action, or None to use default
+
+        Returns:
+            Tuple of (session_name, error_dict or None)
+        """
+        # Determine which session to use
+        target_session = session_name if session_name else self.default_session
+
+        # Session already exists - use it
+        if target_session in self._sessions:
+            logger.debug(f"Using existing session: {target_session}")
+            return target_session, None
+
+        # Session doesn't exist - check auto_create
+        if self.auto_create:
+            # Auto-create the session
+            logger.info(f"Auto-creating session: {target_session}")
+            init_action = InitSessionAction(
+                type="initSession", session_name=target_session, description="Auto-initialized session"
+            )
+            result = self.init_session(init_action)
+
+            if result.get("status") != "success":
+                return target_session, result
+
+            logger.info(f"Successfully auto-created session: {target_session}")
+            return target_session, None
+
+        # auto_create=False and session doesn't exist - raise exception
+        logger.debug(f"Session '{target_session}' not found (auto_create disabled)")
+        raise ValueError(f"Session '{target_session}' not found. Create it first using initSession.")
+
     def execute_code(self, action: ExecuteCodeAction) -> Dict[str, Any]:
-        """Execute code in a Bedrock AgentCoresession."""
-        if action.session_name not in self._sessions:
-            return {"status": "error", "content": [{"text": f"Session '{action.session_name}' not found"}]}
+        """Execute code in a Bedrock AgentCore session with automatic session management."""
+        session_name, error = self._ensure_session(action.session_name)
+        if error:
+            return error
 
-        logger.debug(f"Executing {action.language} code in session '{action.session_name}'")
+        logger.debug(f"Executing {action.language} code in session '{session_name}'")
 
-        # Use the invoke method with proper parameters as shown in the example
         params = {"code": action.code, "language": action.language.value, "clearContext": action.clear_context}
-        response = self._sessions[action.session_name].client.invoke("executeCode", params)
+        response = self._sessions[session_name].client.invoke("executeCode", params)
 
         return self._create_tool_result(response)
 
     def execute_command(self, action: ExecuteCommandAction) -> Dict[str, Any]:
-        """Execute a command in a Bedrock AgentCoresession."""
-        if action.session_name not in self._sessions:
-            return {"status": "error", "content": [{"text": f"Session '{action.session_name}' not found"}]}
+        """Execute a command in a Bedrock AgentCore session with automatic session management."""
+        session_name, error = self._ensure_session(action.session_name)
+        if error:
+            return error
 
-        logger.debug(f"Executing command in session '{action.session_name}': {action.command}")
+        logger.debug(f"Executing command in session '{session_name}': {action.command}")
 
-        # Use the invoke method with proper parameters as shown in the example
         params = {"command": action.command}
-        response = self._sessions[action.session_name].client.invoke("executeCommand", params)
+        response = self._sessions[session_name].client.invoke("executeCommand", params)
 
         return self._create_tool_result(response)
 
     def read_files(self, action: ReadFilesAction) -> Dict[str, Any]:
-        """Read files from a Bedrock AgentCoresession."""
-        if action.session_name not in self._sessions:
-            return {"status": "error", "content": [{"text": f"Session '{action.session_name}' not found"}]}
+        """Read files from a Bedrock AgentCore session with automatic session management."""
+        session_name, error = self._ensure_session(action.session_name)
+        if error:
+            return error
 
-        logger.debug(f"Reading files from session '{action.session_name}': {action.paths}")
+        logger.debug(f"Reading files from session '{session_name}': {action.paths}")
 
-        # Use the invoke method with proper parameters as shown in the example
         params = {"paths": action.paths}
-        response = self._sessions[action.session_name].client.invoke("readFiles", params)
+        response = self._sessions[session_name].client.invoke("readFiles", params)
 
         return self._create_tool_result(response)
 
     def list_files(self, action: ListFilesAction) -> Dict[str, Any]:
-        """List files in a Bedrock AgentCoresession directory."""
-        if action.session_name not in self._sessions:
-            return {"status": "error", "content": [{"text": f"Session '{action.session_name}' not found"}]}
+        """List files in a Bedrock AgentCore session directory with automatic session management."""
+        session_name, error = self._ensure_session(action.session_name)
+        if error:
+            return error
 
-        logger.debug(f"Listing files in session '{action.session_name}' at path: {action.path}")
+        logger.debug(f"Listing files in session '{session_name}' at path: {action.path}")
 
-        # Use the invoke method with proper parameters as shown in the example
         params = {"path": action.path}
-        response = self._sessions[action.session_name].client.invoke("listFiles", params)
+        response = self._sessions[session_name].client.invoke("listFiles", params)
 
         return self._create_tool_result(response)
 
     def remove_files(self, action: RemoveFilesAction) -> Dict[str, Any]:
-        """Remove files from a Bedrock AgentCoresession."""
-        if action.session_name not in self._sessions:
-            return {"status": "error", "content": [{"text": f"Session '{action.session_name}' not found"}]}
+        """Remove files from a Bedrock AgentCore session with automatic session management."""
+        session_name, error = self._ensure_session(action.session_name)
+        if error:
+            return error
 
-        logger.debug(f"Removing files from session '{action.session_name}': {action.paths}")
+        logger.debug(f"Removing files from session '{session_name}': {action.paths}")
 
-        # Use the invoke method with proper parameters as shown in the example
         params = {"paths": action.paths}
-        response = self._sessions[action.session_name].client.invoke("removeFiles", params)
+        response = self._sessions[session_name].client.invoke("removeFiles", params)
 
         return self._create_tool_result(response)
 
     def write_files(self, action: WriteFilesAction) -> Dict[str, Any]:
-        """Write files to a Bedrock AgentCoresession."""
-        if action.session_name not in self._sessions:
-            return {"status": "error", "content": [{"text": f"Session '{action.session_name}' not found"}]}
+        """Write files to a Bedrock AgentCore session with automatic session management."""
+        session_name, error = self._ensure_session(action.session_name)
+        if error:
+            return error
 
-        logger.debug(f"Writing {len(action.content)} files to session '{action.session_name}'")
+        logger.debug(f"Writing {len(action.content)} files to session '{session_name}'")
 
-        # Convert FileContent objects to dictionaries for the API
         content_dicts = [{"path": fc.path, "text": fc.text} for fc in action.content]
-
-        # Use the invoke method with proper parameters as shown in the example
         params = {"content": content_dicts}
-        response = self._sessions[action.session_name].client.invoke("writeFiles", params)
+        response = self._sessions[session_name].client.invoke("writeFiles", params)
 
         return self._create_tool_result(response)