From 5c1dfbc3680c81d27b1993fb6c2693fff9cd13d5 Mon Sep 17 00:00:00 2001 From: Mukunda Katta Date: Sat, 16 May 2026 12:00:14 -0700 Subject: [PATCH] feat(session): allow S3SessionManager to reuse a pre-built S3 client Adds a new optional `s3_client` keyword argument to `S3SessionManager`. When provided, the manager reuses the caller's boto3 S3 client directly instead of constructing a new `boto3.Session` and S3 client, avoiding the per-instance HTTP connection pool + endpoint discovery overhead. Callers also retain full control of the client (credentials, retry config, custom endpoints). When `s3_client` is supplied, the existing `boto_session`, `boto_client_config`, and `region_name` parameters are ignored (documented in the docstring). Backward-compatible: callers that do not pass `s3_client` get the existing behavior unchanged, including the `strands-agents` user-agent tag on the auto-built client. Tests added (moto-backed): - test_s3_client_kwarg_reuses_supplied_client - test_s3_client_kwarg_ignores_session_and_config - test_s3_client_kwarg_supports_session_round_trip - test_default_path_still_works Full session/s3 suite: 50 passed. Refs strands-agents/sdk-python#1163 --- src/strands/session/s3_session_manager.py | 46 ++++++++----- .../session/test_s3_session_manager.py | 66 +++++++++++++++++++ 2 files changed, 96 insertions(+), 16 deletions(-) diff --git a/src/strands/session/s3_session_manager.py b/src/strands/session/s3_session_manager.py index fad5e4fd0..99ec76e16 100644 --- a/src/strands/session/s3_session_manager.py +++ b/src/strands/session/s3_session_manager.py @@ -51,6 +51,7 @@ def __init__( boto_session: boto3.Session | None = None, boto_client_config: BotocoreConfig | None = None, region_name: str | None = None, + s3_client: Any = None, **kwargs: Any, ): """Initialize S3SessionManager with S3 storage. @@ -60,29 +61,42 @@ def __init__( ID is not allowed to contain path separators (e.g., a/b). bucket: S3 bucket name (required) prefix: S3 key prefix for storage organization - boto_session: Optional boto3 session - boto_client_config: Optional boto3 client configuration - region_name: AWS region for S3 storage + boto_session: Optional boto3 session. Ignored if ``s3_client`` is supplied. + boto_client_config: Optional boto3 client configuration. Ignored if ``s3_client`` is supplied. + region_name: AWS region for S3 storage. Ignored if ``s3_client`` is supplied. + s3_client: Optional pre-built boto3 S3 client. When provided, S3SessionManager + reuses it directly instead of constructing a new boto3.Session and S3 client. + This avoids per-instance boto initialization overhead (HTTP connection pool, + endpoint discovery) when many managers are created in the same process, and + gives callers full control over the underlying client (credentials, retry + config, custom endpoints). When ``s3_client`` is provided, ``boto_session``, + ``boto_client_config``, and ``region_name`` are ignored. **kwargs: Additional keyword arguments for future extensibility. """ self.bucket = bucket self.prefix = prefix - session = boto_session or boto3.Session(region_name=region_name) - - # Add strands-agents to the request user agent - if boto_client_config: - existing_user_agent = getattr(boto_client_config, "user_agent_extra", None) - # Append 'strands-agents' to existing user_agent_extra or set it if not present - if existing_user_agent: - new_user_agent = f"{existing_user_agent} strands-agents" - else: - new_user_agent = "strands-agents" - client_config = boto_client_config.merge(BotocoreConfig(user_agent_extra=new_user_agent)) + if s3_client is not None: + # Reuse the caller's pre-built client. We deliberately do not modify + # its user_agent_extra; the caller owns the client's configuration. + self.client = s3_client else: - client_config = BotocoreConfig(user_agent_extra="strands-agents") + session = boto_session or boto3.Session(region_name=region_name) + + # Add strands-agents to the request user agent + if boto_client_config: + existing_user_agent = getattr(boto_client_config, "user_agent_extra", None) + # Append 'strands-agents' to existing user_agent_extra or set it if not present + if existing_user_agent: + new_user_agent = f"{existing_user_agent} strands-agents" + else: + new_user_agent = "strands-agents" + client_config = boto_client_config.merge(BotocoreConfig(user_agent_extra=new_user_agent)) + else: + client_config = BotocoreConfig(user_agent_extra="strands-agents") + + self.client = session.client(service_name="s3", config=client_config) - self.client = session.client(service_name="s3", config=client_config) super().__init__(session_id=session_id, session_repository=self) def _get_session_path(self, session_id: str) -> str: diff --git a/tests/strands/session/test_s3_session_manager.py b/tests/strands/session/test_s3_session_manager.py index 29bc97ab5..aa59ddd54 100644 --- a/tests/strands/session/test_s3_session_manager.py +++ b/tests/strands/session/test_s3_session_manager.py @@ -510,3 +510,69 @@ def test_update_nonexistent_multi_agent(s3_manager, sample_session): nonexistent_mock.id = "nonexistent" with pytest.raises(SessionException): s3_manager.update_multi_agent(sample_session.session_id, nonexistent_mock) + + +# --- s3_client reuse (issue #1163) --- + + +def test_s3_client_kwarg_reuses_supplied_client(mocked_aws, s3_bucket): + """When ``s3_client`` is passed, S3SessionManager must reuse it directly + instead of building a new boto3.Session + client. + """ + supplied = boto3.client("s3", region_name="us-west-2") + manager = S3SessionManager( + session_id="test-reuse", + bucket=s3_bucket, + prefix="sessions/", + s3_client=supplied, + ) + assert manager.client is supplied + + +def test_s3_client_kwarg_ignores_session_and_config(mocked_aws, s3_bucket): + """When ``s3_client`` is supplied, boto_session / boto_client_config / + region_name are ignored. We assert by checking that no extra boto3.Session + is constructed when those args are also passed. + """ + supplied = boto3.client("s3", region_name="us-west-2") + + # boto_session is the sentinel that would otherwise become self.client; + # supplying it together with s3_client should NOT override s3_client. + bogus_session = Mock(spec=boto3.Session) + manager = S3SessionManager( + session_id="test-reuse-2", + bucket=s3_bucket, + prefix="sessions/", + boto_session=bogus_session, + boto_client_config=BotocoreConfig(retries={"max_attempts": 99}), + region_name="us-east-1", + s3_client=supplied, + ) + assert manager.client is supplied + bogus_session.client.assert_not_called() + + +def test_s3_client_kwarg_supports_session_round_trip(mocked_aws, s3_bucket, sample_session): + """End-to-end smoke: a manager built with s3_client= can write and read.""" + supplied = boto3.client("s3", region_name="us-west-2") + manager = S3SessionManager( + session_id="test-roundtrip", + bucket=s3_bucket, + prefix="sessions/", + s3_client=supplied, + ) + manager.create_session(sample_session) + fetched = manager.read_session(sample_session.session_id) + assert fetched.session_id == sample_session.session_id + + +def test_default_path_still_works(mocked_aws, s3_bucket): + """Sanity: omitting s3_client still goes through the boto3.Session path.""" + manager = S3SessionManager( + session_id="test-default", + bucket=s3_bucket, + prefix="sessions/", + region_name="us-west-2", + ) + # client is built by session.client("s3", ...); only check it's there + assert manager.client is not None