refactor(utils): unify AWS session functions into single get_session

- Merge get_boto3_session and get_assumed_role_session into get_session(region, profile_name, role_arn)
- Rename clear_sessions to clear_session_cache
- Remove top-level exports from utils/__init__.py for clearer imports
- Users now import from specific modules: strands_env.utils.aws, strands_env.utils.sglang
- Update all usages in CLI, environments, tools, and documentation

Co-Authored-By: Claude Opus 4.5 <[email protected]>

Author: Lawhy

CLAUDE.md

@@ -85,7 +85,7 @@ The package lives in `src/strands_env/` with these modules:
 
 **sglang.py** — SGLang client caching with `@lru_cache`. `get_cached_client(base_url, max_connections)` for connection pooling. `get_cached_client_from_slime_args(args)` for slime RL training integration. `check_server_health(base_url)` for early validation.
 
-**aws.py** — AWS boto3 session caching. `get_boto3_session(region, profile_name)` with `@lru_cache` (boto3 handles credential refresh). `get_assumed_role_session(role_arn, region)` uses `RefreshableCredentials` for programmatic role assumption with auto-refresh.
+**aws.py** — AWS boto3 session caching. `get_session(region, profile_name, role_arn)` with `@lru_cache`. If `role_arn` provided, uses `RefreshableCredentials` for programmatic role assumption with auto-refresh; otherwise returns basic session.
 
 ### `tools/`
 

docs/rl-training.md

@@ -17,12 +17,11 @@ Customize the `generate` and `reward_func` methods to replace single generation
 ```python
 from strands_env.core import Action, TaskContext
 from strands_env.core.models import sglang_model_factory
-from strands_env.utils import get_cached_client_from_slime_args
+from strands_env.utils.sglang import get_cached_client_from_slime_args
 
 async def generate(args, sample, sampling_params):
     # Build model factory with cached client
     factory = sglang_model_factory(
-        model_id=args.hf_checkpoint,
         tokenizer=tokenizer,
         client=get_cached_client_from_slime_args(args),
         sampling_params=sampling_params,

src/strands_env/cli/utils.py

@@ -273,14 +273,15 @@ def _build_sglang_model_factory(config: ModelConfig, max_concurrency: int, sampl
 
 def _build_bedrock_model_factory(config: ModelConfig, sampling: dict) -> ModelFactory:
     """Build Bedrock model factory."""
-    from strands_env.utils.aws import get_assumed_role_session, get_boto3_session
+    from strands_env.utils.aws import get_session
 
     if not config.model_id:
         raise click.ClickException("--model-id is required for Bedrock backend")
 
-    if config.role_arn:
-        boto_session = get_assumed_role_session(config.role_arn, config.region)
-    else:
-        boto_session = get_boto3_session(config.region, config.profile_name)
+    boto_session = get_session(
+        region=config.region,
+        profile_name=config.profile_name,
+        role_arn=config.role_arn,
+    )
 
     return bedrock_model_factory(model_id=config.model_id, boto_session=boto_session, sampling_params=sampling)

src/strands_env/environments/code_sandbox/env.py

@@ -24,7 +24,7 @@
 
 from strands_env.core.environment import Environment
 from strands_env.tools import CodeInterpreterToolkit
-from strands_env.utils.aws import get_boto3_session
+from strands_env.utils.aws import get_session
 
 if TYPE_CHECKING:
     import boto3
@@ -53,9 +53,9 @@ class CodeSandboxEnv(Environment):
 
     Example:
         from strands_env.environments.code_sandbox import CodeSandboxEnv, CodeMode
-        from strands_env.utils import get_boto3_session
+        from strands_env.utils.aws import get_session
 
-        session = get_boto3_session(region="us-east-1")
+        session = get_session(region="us-east-1")
         env = CodeSandboxEnv(
             boto3_session=session,
             model_factory=model_factory,
@@ -99,7 +99,7 @@ def __init__(
         )
         self.mode = mode
         self._toolkit = CodeInterpreterToolkit(
-            boto3_session=boto3_session or get_boto3_session(), session_name="strands-env-code-sandbox"
+            boto3_session=boto3_session or get_session(), session_name="strands-env-code-sandbox"
         )
 
     @override

src/strands_env/tools/code_interpreter.py

@@ -33,9 +33,9 @@ class CodeInterpreterToolkit:
     and shell commands in a sandboxed environment.
 
     Example:
-        from strands_env.utils import get_boto3_session
+        from strands_env.utils.aws import get_session
 
-        session = get_boto3_session(region="us-east-1")
+        session = get_session(region="us-east-1")
         toolkit = CodeInterpreterToolkit(boto3_session=session)
 
         # In environment:

src/strands_env/utils/__init__.py

@@ -13,17 +13,3 @@
 # limitations under the License.
 
 """Utilities for `strands_env`."""
-
-from .aws import clear_sessions, get_assumed_role_session, get_boto3_session
-from .sglang import clear_clients, get_cached_client, get_cached_client_from_slime_args
-
-__all__ = [
-    # AWS
-    "clear_sessions",
-    "get_assumed_role_session",
-    "get_boto3_session",
-    # SGLang
-    "clear_clients",
-    "get_cached_client",
-    "get_cached_client_from_slime_args",
-]

src/strands_env/utils/aws.py

@@ -25,48 +25,40 @@
 
 
 @lru_cache(maxsize=None)
-def get_boto3_session(region: str = "us-east-1", profile_name: str | None = None) -> boto3.Session:
+def get_session(
+    region: str = "us-east-1",
+    profile_name: str | None = None,
+    role_arn: str | None = None,
+    session_name: str = "StrandsEnvSession",
+) -> boto3.Session:
     """Get a cached boto3 session.
 
     Credentials are managed by boto3's provider chain (env vars, ~/.aws/credentials,
     IAM instance role, etc.) and auto-refresh automatically.
 
-    For role assumption, either:
-    1. Configure a profile in ~/.aws/config with `role_arn` and `source_profile`, or
-    2. Use `get_assumed_role_session()` for programmatic role assumption
+    If `role_arn` is provided, assumes the role using STS with auto-refreshing
+    credentials via botocore's `RefreshableCredentials`.
 
     Args:
         region: AWS region name.
         profile_name: Optional AWS profile name from ~/.aws/config.
+        role_arn: Optional ARN of the IAM role to assume.
+        session_name: Session name for assumed role (only used if role_arn provided).
 
     Returns:
         Cached boto3 Session instance.
     """
-    logger.info(f"Creating boto3 session: region={region}, profile={profile_name}")
-    return boto3.Session(region_name=region, profile_name=profile_name)
-
-
-@lru_cache(maxsize=None)
-def get_assumed_role_session(
-    role_arn: str,
-    region: str = "us-east-1",
-    session_name: str = "StrandsEnvSession",
-) -> boto3.Session:
-    """Get a cached boto3 session with assumed role credentials.
+    if role_arn:
+        return _create_assumed_role_session(role_arn, region, session_name)
+    else:
+        logger.info(f"Creating boto3 session: region={region}, profile={profile_name}")
+        return boto3.Session(region_name=region, profile_name=profile_name)
 
-    Uses botocore's `RefreshableCredentials` so credentials auto-refresh when expired.
-    The session is cached by (role_arn, region, session_name).
 
-    Args:
-        role_arn: ARN of the IAM role to assume.
-        region: AWS region name.
-        session_name: Session name for the assumed role.
-
-    Returns:
-        Cached boto3 Session with auto-refreshing credentials.
-    """
+def _create_assumed_role_session(role_arn: str, region: str, session_name: str) -> boto3.Session:
+    """Create a boto3 session with assumed role credentials."""
     from botocore.credentials import RefreshableCredentials
-    from botocore.session import get_session
+    from botocore.session import get_session as get_botocore_session
 
     logger.info(f"Creating boto3 session with assumed role: role={role_arn}, region={region}")
 
@@ -87,12 +79,11 @@ def refresh() -> dict:
         method="sts-assume-role",
     )
 
-    botocore_session = get_session()
+    botocore_session = get_botocore_session()
     botocore_session._credentials = session_credentials
     return boto3.Session(botocore_session=botocore_session, region_name=region)
 
 
-def clear_sessions() -> None:
+def clear_session_cache() -> None:
     """Clear all cached boto3 sessions."""
-    get_boto3_session.cache_clear()
-    get_assumed_role_session.cache_clear()
+    get_session.cache_clear()

tests/unit/test_aws.py

@@ -18,40 +18,40 @@
 
 import boto3
 
-from strands_env.utils.aws import clear_sessions, get_assumed_role_session, get_boto3_session
+from strands_env.utils.aws import clear_session_cache, get_session
 
 
-class TestGetBoto3Session:
-    """Tests for get_boto3_session."""
+class TestGetSession:
+    """Tests for get_session (basic mode without role assumption)."""
 
     def setup_method(self):
         """Clear cache before each test."""
-        clear_sessions()
+        clear_session_cache()
 
     def test_returns_session(self):
         """Should return a boto3 Session."""
-        session = get_boto3_session(region="us-west-2")
+        session = get_session(region="us-west-2")
         assert isinstance(session, boto3.Session)
         assert session.region_name == "us-west-2"
 
     def test_cached_by_region(self):
         """Same region should return cached session."""
-        session1 = get_boto3_session(region="us-east-1")
-        session2 = get_boto3_session(region="us-east-1")
+        session1 = get_session(region="us-east-1")
+        session2 = get_session(region="us-east-1")
         assert session1 is session2
 
     def test_different_regions_different_sessions(self):
         """Different regions should return different sessions."""
-        session1 = get_boto3_session(region="us-east-1")
-        session2 = get_boto3_session(region="us-west-2")
+        session1 = get_session(region="us-east-1")
+        session2 = get_session(region="us-west-2")
         assert session1 is not session2
 
     @patch("strands_env.utils.aws.boto3.Session")
     def test_cached_by_profile(self, mock_session_cls):
         """Same profile should return cached session."""
         mock_session_cls.return_value = MagicMock()
-        session1 = get_boto3_session(region="us-east-1", profile_name="test-profile")
-        session2 = get_boto3_session(region="us-east-1", profile_name="test-profile")
+        session1 = get_session(region="us-east-1", profile_name="test-profile")
+        session2 = get_session(region="us-east-1", profile_name="test-profile")
         assert session1 is session2
         # Session should only be created once due to caching
         assert mock_session_cls.call_count == 1
@@ -60,22 +60,22 @@ def test_cached_by_profile(self, mock_session_cls):
     def test_different_profiles_different_sessions(self, mock_session_cls):
         """Different profiles should return different sessions."""
         mock_session_cls.side_effect = [MagicMock(), MagicMock()]
-        session1 = get_boto3_session(region="us-east-1", profile_name="profile-a")
-        session2 = get_boto3_session(region="us-east-1", profile_name="profile-b")
+        session1 = get_session(region="us-east-1", profile_name="profile-a")
+        session2 = get_session(region="us-east-1", profile_name="profile-b")
         assert session1 is not session2
 
 
-class TestGetAssumedRoleSession:
-    """Tests for get_assumed_role_session."""
+class TestGetSessionWithRoleAssumption:
+    """Tests for get_session with role_arn (role assumption mode)."""
 
     def setup_method(self):
         """Clear cache before each test."""
-        clear_sessions()
+        clear_session_cache()
 
     @patch("strands_env.utils.aws.boto3.client")
     @patch("botocore.session.get_session")
     def test_assumes_role(self, mock_get_session, mock_boto3_client):
-        """Should call STS assume_role."""
+        """Should call STS assume_role when role_arn provided."""
         from datetime import datetime, timezone
 
         # Mock STS response
@@ -95,7 +95,7 @@ def test_assumes_role(self, mock_get_session, mock_boto3_client):
         mock_get_session.return_value = mock_botocore_session
 
         role_arn = "arn:aws:iam::123456789:role/TestRole"
-        session = get_assumed_role_session(role_arn=role_arn, region="us-east-1")
+        session = get_session(region="us-east-1", role_arn=role_arn)
 
         # Verify STS was called
         mock_boto3_client.assert_called_with("sts", region_name="us-east-1")
@@ -123,8 +123,8 @@ def test_cached_by_role_arn(self, mock_get_session, mock_boto3_client):
         mock_get_session.return_value = MagicMock()
 
         role_arn = "arn:aws:iam::123456789:role/TestRole"
-        session1 = get_assumed_role_session(role_arn=role_arn)
-        session2 = get_assumed_role_session(role_arn=role_arn)
+        session1 = get_session(role_arn=role_arn)
+        session2 = get_session(role_arn=role_arn)
 
         assert session1 is session2
         # assume_role should only be called once due to caching
@@ -149,7 +149,7 @@ def test_has_refreshable_credentials(self, mock_boto3_client):
         mock_boto3_client.return_value = mock_sts
 
         role_arn = "arn:aws:iam::123456789:role/TestRole"
-        session = get_assumed_role_session(role_arn=role_arn)
+        session = get_session(role_arn=role_arn)
 
         # Get the underlying botocore credentials directly
         botocore_creds = session._session._credentials
@@ -159,17 +159,17 @@ def test_has_refreshable_credentials(self, mock_boto3_client):
         assert botocore_creds._refresh_using is not None
 
 
-class TestClearSessions:
-    """Tests for clear_sessions."""
+class TestClearSessionCache:
+    """Tests for clear_session_cache."""
 
-    def test_clears_all_caches(self):
-        """Should clear both session caches."""
-        # Create some cached sessions
-        session1 = get_boto3_session(region="us-east-1")
+    def test_clears_cache(self):
+        """Should clear the session cache."""
+        # Create a cached session
+        session1 = get_session(region="us-east-1")
 
         # Clear
-        clear_sessions()
+        clear_session_cache()
 
         # New call should create new session
-        session2 = get_boto3_session(region="us-east-1")
+        session2 = get_session(region="us-east-1")
         assert session1 is not session2