feat(sagemaker/sessions): add validation layer so if session_manager is None but request is a session request, raise 400 error

- change session api transform's transform_request to never return
  request field in output (previously used to pass session_manager
- change handler code to only take raw_request as a parameter and use
  new utility function get_session_manager
- add new integration tests for expected errors when sessions is
  disabled

Author: zhaozuy

python/model_hosting_container_standards/sagemaker/sessions/handlers.py

@@ -5,8 +5,13 @@
 from fastapi import Request, Response
 from fastapi.exceptions import HTTPException
 
-from .manager import SessionManager
-from .models import SageMakerSessionHeader, SessionRequestType
+from .manager import get_session_manager
+from .models import (
+    SESSION_DISABLED_ERROR_DETAIL,
+    SESSION_DISABLED_LOG_MESSAGE,
+    SageMakerSessionHeader,
+    SessionRequestType,
+)
 from .utils import get_session_id_from_request
 
 logger = logging.getLogger(__name__)
@@ -29,11 +34,10 @@ def get_handler_for_request_type(request_type: SessionRequestType):
         return None
 
 
-async def close_session(session_manager: SessionManager, raw_request: Request):
+async def close_session(raw_request: Request):
     """Close an existing session and clean up its resources.
 
     Args:
-        session_manager: SessionManager instance to manage the session lifecycle
         raw_request: FastAPI Request object containing session ID in headers
 
     Returns:
@@ -43,6 +47,13 @@ async def close_session(session_manager: SessionManager, raw_request: Request):
         HTTPException: If session closure fails with 424 FAILED_DEPENDENCY status
     """
     session_id = get_session_id_from_request(raw_request)
+    session_manager = get_session_manager()
+    if session_manager is None:
+        logger.error(SESSION_DISABLED_LOG_MESSAGE)
+        raise HTTPException(
+            status_code=HTTPStatus.BAD_REQUEST.value,
+            detail=SESSION_DISABLED_ERROR_DETAIL,
+        )
     try:
         session_manager.close_session(session_id)
         logger.info(f"Session {session_id} closed")
@@ -59,11 +70,10 @@ async def close_session(session_manager: SessionManager, raw_request: Request):
         )
 
 
-async def create_session(session_manager: SessionManager, raw_request: Request):
+async def create_session(raw_request: Request):
     """Create a new stateful session with expiration tracking.
 
     Args:
-        session_manager: SessionManager instance to manage the session lifecycle
         raw_request: FastAPI Request object (unused but part of handler signature)
 
     Returns:
@@ -72,6 +82,13 @@ async def create_session(session_manager: SessionManager, raw_request: Request):
     Raises:
         HTTPException: If session creation fails with 424 FAILED_DEPENDENCY status
     """
+    session_manager = get_session_manager()
+    if session_manager is None:
+        logger.error(SESSION_DISABLED_LOG_MESSAGE)
+        raise HTTPException(
+            status_code=HTTPStatus.BAD_REQUEST.value,
+            detail=SESSION_DISABLED_ERROR_DETAIL,
+        )
     try:
         session = session_manager.create_session()
         # expiration_ts is guaranteed to be set for newly created sessions

python/model_hosting_container_standards/sagemaker/sessions/models.py

@@ -37,3 +37,11 @@ class SageMakerSessionHeader:
     SESSION_ID = "X-Amzn-SageMaker-Session-Id"
     NEW_SESSION_ID = "X-Amzn-SageMaker-New-Session-Id"
     CLOSED_SESSION_ID = "X-Amzn-SageMaker-Closed-Session-Id"
+
+
+# Error messages for session management
+SESSION_DISABLED_ERROR_DETAIL = "Invalid payload. stateful sessions not enabled"
+SESSION_DISABLED_LOG_MESSAGE = (
+    f"Invalid payload. stateful sessions not enabled, "
+    f"{SageMakerSessionHeader.SESSION_ID} header not supported"
+)

python/model_hosting_container_standards/sagemaker/sessions/transform.py

@@ -10,7 +10,11 @@
 from ...common import BaseApiTransform, BaseTransformRequestOutput
 from .handlers import get_handler_for_request_type
 from .manager import SessionManager, get_session_manager
-from .models import SessionRequest
+from .models import (
+    SESSION_DISABLED_ERROR_DETAIL,
+    SESSION_DISABLED_LOG_MESSAGE,
+    SessionRequest,
+)
 from .utils import get_session, get_session_id_from_request
 
 logger = logging.getLogger(__name__)
@@ -63,7 +67,7 @@ def _validate_session_if_present(raw_request: Request, session_manager: SessionM
 
 
 def process_session_request(
-    request_data: dict, raw_request: Request, session_manager: SessionManager
+    request_data: dict, raw_request: Request, session_manager: Optional[SessionManager]
 ):
     """Process a potential session management request.
 
@@ -92,16 +96,22 @@ def process_session_request(
     # Not a session request - pass through for normal processing
     if session_request is None:
         return BaseTransformRequestOutput(
-            request=None,
             raw_request=raw_request,
             intercept_func=None,
         )
 
+    if session_manager is None:
+        logger.error(SESSION_DISABLED_LOG_MESSAGE)
+        raise HTTPException(
+            status_code=HTTPStatus.BAD_REQUEST.value,
+            detail=SESSION_DISABLED_ERROR_DETAIL,
+        )
+
     # Route to appropriate session management handler
     intercept_func = get_handler_for_request_type(session_request.requestType)
 
     return BaseTransformRequestOutput(
-        request=session_manager, raw_request=raw_request, intercept_func=intercept_func
+        raw_request=raw_request, intercept_func=intercept_func
     )
 
 

python/tests/integration/test_sagemaker_sessions_integration.py

@@ -33,6 +33,7 @@
     init_session_manager_from_env,
 )
 from model_hosting_container_standards.sagemaker.sessions.models import (
+    SESSION_DISABLED_ERROR_DETAIL,
     SageMakerSessionHeader,
 )
 
@@ -542,5 +543,97 @@ def test_interleaved_session_operations(self):
         assert response.status_code == 200
 
 
+class TestSessionsDisabled:
+    """Test behavior when stateful sessions are disabled.
+
+    These tests verify that session management requests fail gracefully
+    when the SAGEMAKER_ENABLE_STATEFUL_SESSIONS flag is not set.
+    """
+
+    @pytest.fixture
+    def app_with_sessions_disabled(self, monkeypatch):
+        """Create app with sessions disabled."""
+        # Explicitly disable sessions
+        monkeypatch.delenv("SAGEMAKER_ENABLE_STATEFUL_SESSIONS", raising=False)
+        monkeypatch.delenv("SAGEMAKER_SESSIONS_PATH", raising=False)
+        monkeypatch.delenv("SAGEMAKER_SESSIONS_EXPIRATION", raising=False)
+
+        # Reinitialize the global session manager (should be None)
+        init_session_manager_from_env()
+
+        # Now create the app with sessions disabled
+        app = FastAPI()
+        router = APIRouter()
+
+        @router.post("/invocations")
+        @sagemaker_standards.stateful_session_manager()
+        async def invocations(request: Request):
+            """Stateful invocation handler."""
+            body_bytes = await request.body()
+            body = json.loads(body_bytes.decode())
+
+            return Response(
+                status_code=200,
+                content=json.dumps({"message": "success", "echo": body}),
+            )
+
+        app.include_router(router)
+        sagemaker_standards.bootstrap(app)
+
+        return TestClient(app)
+
+    def test_new_session_request_fails_when_disabled(self, app_with_sessions_disabled):
+        """Test that NEW_SESSION request fails when sessions are disabled."""
+        response = app_with_sessions_disabled.post(
+            "/invocations", json={"requestType": "NEW_SESSION"}
+        )
+
+        # Should fail with 400 BAD_REQUEST since sessions are not enabled
+        assert response.status_code == 400
+        assert SESSION_DISABLED_ERROR_DETAIL in response.text
+
+    def test_close_session_request_fails_when_disabled(
+        self, app_with_sessions_disabled
+    ):
+        """Test that CLOSE request fails when sessions are disabled."""
+        response = app_with_sessions_disabled.post(
+            "/invocations",
+            json={"requestType": "CLOSE"},
+            headers={SageMakerSessionHeader.SESSION_ID: "some-session-id"},
+        )
+
+        # Should fail with 400 BAD_REQUEST due to session header when sessions disabled
+        assert response.status_code == 400
+        assert SESSION_DISABLED_ERROR_DETAIL in response.text
+
+    def test_regular_requests_work_when_sessions_disabled(
+        self, app_with_sessions_disabled
+    ):
+        """Test that regular requests still work when sessions are disabled."""
+        response = app_with_sessions_disabled.post(
+            "/invocations", json={"prompt": "test request"}
+        )
+
+        # Regular requests should still work
+        assert response.status_code == 200
+        data = json.loads(response.text)
+        assert data["message"] == "success"
+        assert data["echo"]["prompt"] == "test request"
+
+    def test_regular_requests_with_session_header_when_disabled(
+        self, app_with_sessions_disabled
+    ):
+        """Test that requests with session headers fail validation when sessions disabled."""
+        response = app_with_sessions_disabled.post(
+            "/invocations",
+            json={"prompt": "test"},
+            headers={SageMakerSessionHeader.SESSION_ID: "invalid-session"},
+        )
+
+        # Should fail with 400 BAD_REQUEST since sessions are not enabled
+        assert response.status_code == 400
+        assert SESSION_DISABLED_ERROR_DETAIL in response.text
+
+
 if __name__ == "__main__":
     pytest.main([__file__, "-v"])

python/tests/sagemaker/sessions/test_handlers.py

@@ -2,7 +2,7 @@
 
 import time
 from http import HTTPStatus
-from unittest.mock import Mock
+from unittest.mock import Mock, patch
 
 import pytest
 from fastapi import Response
@@ -61,7 +61,11 @@ async def test_creates_session_successfully(
         """Test successfully creates a session and returns response."""
         mock_session_manager.create_session.return_value = mock_session_with_expiration
 
-        response = await create_session(mock_session_manager, mock_request)
+        with patch(
+            "model_hosting_container_standards.sagemaker.sessions.handlers.get_session_manager",
+            return_value=mock_session_manager,
+        ):
+            response = await create_session(mock_request)
 
         assert isinstance(response, Response)
         assert response.status_code == HTTPStatus.OK.value
@@ -80,7 +84,11 @@ async def test_calls_session_manager_create_session(
         """Test calls session_manager.create_session method."""
         mock_session_manager.create_session.return_value = mock_session_with_expiration
 
-        await create_session(mock_session_manager, mock_request)
+        with patch(
+            "model_hosting_container_standards.sagemaker.sessions.handlers.get_session_manager",
+            return_value=mock_session_manager,
+        ):
+            await create_session(mock_request)
 
         mock_session_manager.create_session.assert_called_once()
 
@@ -91,8 +99,12 @@ async def test_raises_http_exception_on_session_creation_failure(
         """Test raises HTTPException when session creation fails."""
         mock_session_manager.create_session.side_effect = Exception("Creation failed")
 
-        with pytest.raises(HTTPException) as exc_info:
-            await create_session(mock_session_manager, mock_request)
+        with patch(
+            "model_hosting_container_standards.sagemaker.sessions.handlers.get_session_manager",
+            return_value=mock_session_manager,
+        ):
+            with pytest.raises(HTTPException) as exc_info:
+                await create_session(mock_request)
 
         assert exc_info.value.status_code == HTTPStatus.FAILED_DEPENDENCY.value
         assert "Failed to create session" in exc_info.value.detail
@@ -109,7 +121,11 @@ async def test_closes_session_successfully(
         session_id = "test-session-123"
         mock_session_manager.close_session.return_value = None
 
-        response = await close_session(mock_session_manager, mock_request_with_session)
+        with patch(
+            "model_hosting_container_standards.sagemaker.sessions.handlers.get_session_manager",
+            return_value=mock_session_manager,
+        ):
+            response = await close_session(mock_request_with_session)
 
         assert isinstance(response, Response)
         assert response.status_code == HTTPStatus.OK.value
@@ -124,8 +140,12 @@ async def test_raises_http_exception_on_close_failure(
         """Test raises HTTPException when session close fails."""
         mock_session_manager.close_session.side_effect = ValueError("Session not found")
 
-        with pytest.raises(HTTPException) as exc_info:
-            await close_session(mock_session_manager, mock_request_with_session)
+        with patch(
+            "model_hosting_container_standards.sagemaker.sessions.handlers.get_session_manager",
+            return_value=mock_session_manager,
+        ):
+            with pytest.raises(HTTPException) as exc_info:
+                await close_session(mock_request_with_session)
 
         assert exc_info.value.status_code == HTTPStatus.FAILED_DEPENDENCY.value
         assert "Failed to close session" in exc_info.value.detail

python/tests/sagemaker/sessions/test_transform.py

@@ -183,7 +183,7 @@ def test_returns_create_handler_for_new_session_request(
         )
 
         assert isinstance(result, BaseTransformRequestOutput)
-        assert result.request == mock_session_manager
+        assert result.request is None
         assert result.raw_request == mock_request
         assert result.intercept_func == create_session
 
@@ -198,7 +198,7 @@ def test_returns_close_handler_for_close_request(
         )
 
         assert isinstance(result, BaseTransformRequestOutput)
-        assert result.request == mock_session_manager
+        assert result.request is None
         assert result.raw_request == mock_request
         assert result.intercept_func == close_session
 
@@ -335,10 +335,10 @@ async def test_end_to_end_new_session_flow(self, transform):
 
         # Verify we get an intercept function
         assert result.intercept_func == create_session
-        assert result.request == transform._session_manager
+        assert result.request is None
 
         # Verify we can call the handler
-        response = await result.intercept_func(result.request, mock_request)
+        response = await result.intercept_func(mock_request)
         assert response.status_code == HTTPStatus.OK.value
         assert SageMakerSessionHeader.NEW_SESSION_ID in response.headers
 

python/tests/sagemaker/sessions/test_utils.py

@@ -8,6 +8,7 @@
 from fastapi.exceptions import HTTPException
 
 from model_hosting_container_standards.sagemaker.sessions.models import (
+    SESSION_DISABLED_ERROR_DETAIL,
     SageMakerSessionHeader,
 )
 from model_hosting_container_standards.sagemaker.sessions.utils import (
@@ -92,7 +93,7 @@ def test_raises_http_exception_when_sessions_not_enabled_but_header_present(self
             get_session(None, raw_request)
 
         assert exc_info.value.status_code == HTTPStatus.BAD_REQUEST.value
-        assert "stateful sessions not enabled" in exc_info.value.detail
+        assert SESSION_DISABLED_ERROR_DETAIL in exc_info.value.detail
         assert SageMakerSessionHeader.SESSION_ID in exc_info.value.detail
 
     def test_returns_none_when_sessions_not_enabled_and_no_header(self):