From 6c21029cb7e8fe827ea9f8d108f1ff18ae3b9e4b Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Tue, 30 Jun 2026 18:58:09 -0700 Subject: [PATCH] feat(sandbox): reuse e2b container across requests when metadata.session_id is set (#31688) * feat(sandbox): reuse e2b container across requests when metadata.session_id is set When a client passes `metadata.session_id` in a /chat/completions request alongside a code_interpreter tool, the proxy now routes all requests sharing that session_id to the same sandbox container. State (variables, imports, installed packages) persists across requests within the session. Without a session_id the existing ephemeral behavior is unchanged: one container per agentic loop, deleted immediately after. The sandbox key is derived from session_id rather than a per-request UUID. The cleanup and post-loop hooks skip deletion for session-scoped containers. TTL-based pruning (15 min idle) still applies and refreshes on every use, so an active session never expires mid-use. The session_id-scoped key is registered in all_litellm_params and the proxy strip-list so it never leaks to the upstream LLM provider. * fix(sandbox): scope session sandbox key to API key identity; add per-identity LRU cap Two security issues addressed: 1. Cross-user sandbox isolation: the session_id supplied by the client is now combined with the server-minted user_api_key_hash to form the cache key (format: "{hash}:{session_id}" when authenticated, bare session_id for non-proxy use). Two tenants sharing the same session_id no longer share a sandbox. 2. Bounded session allocation: each API key identity is capped at _SESSION_SCOPED_PER_IDENTITY_CAP (10) live session-scoped containers. When a new session is opened beyond the cap, the least-recently-used entry for that identity is evicted and its sandbox deleted, preventing unbounded accumulation via rotating session IDs. The container cache tuple gains a fourth element (identity: str | None) so eviction can filter by identity without parsing key formats. Tests added for both properties. --- .../code_interpreter_interception/handler.py | 71 +++- litellm/proxy/dev_config.yaml | 15 + litellm/proxy/litellm_pre_call_utils.py | 1 + litellm/types/utils.py | 1 + qa_sticky_session.sh | 59 +++ .../test_handler.py | 369 ++++++++++++++---- 6 files changed, 421 insertions(+), 95 deletions(-) create mode 100755 qa_sticky_session.sh diff --git a/litellm/integrations/code_interpreter_interception/handler.py b/litellm/integrations/code_interpreter_interception/handler.py index cd7b211f1a5..759b2be3a84 100644 --- a/litellm/integrations/code_interpreter_interception/handler.py +++ b/litellm/integrations/code_interpreter_interception/handler.py @@ -40,9 +40,11 @@ from litellm.types.utils import ( LITELLM_CODE_EXECUTION_TOOL_NAME = "litellm_code_execution" _INTERCEPTION_ACTIVE_KEY = "_code_interpreter_interception_active" _SANDBOX_KEY = "_code_interpreter_interception_sandbox_key" +_SESSION_SCOPED_KEY = "_code_interpreter_interception_session_scoped" _CONVERTED_STREAM_KEY = "_code_interpreter_interception_converted_stream" _LITELLM_METADATA_KEY = "litellm_metadata" _CACHE_TTL_SECONDS = 15 * 60 +_SESSION_SCOPED_PER_IDENTITY_CAP = 10 class CodeExecutionToolCall(TypedDict, total=False): @@ -107,6 +109,20 @@ class ChatCompletionFunctionToolChoice(TypedDict): CodeExecutionFunctionToolChoice = ResponsesFunctionToolChoice | ChatCompletionFunctionToolChoice +def _extract_session_id(kwargs: dict[str, Any]) -> str | None: + for meta_key in ("metadata", "litellm_metadata"): + meta = kwargs.get(meta_key) + if isinstance(meta, dict): + sid = meta.get("session_id") + if sid and isinstance(sid, str): + return sid + return None + + +def _extract_identity(kwargs: dict[str, Any]) -> str: + return kwargs.get("user_api_key_hash") or "" + + def _resolve_sandbox_tool(sandbox_tool_name: str | None) -> dict[str, Any] | None: try: from litellm.sandbox.sandbox_tools import resolve_sandbox_tool @@ -140,7 +156,7 @@ class CodeInterpreterInterceptionLogger(CustomLogger): self.enabled_providers = enabled_providers self.sandbox_tool_name = sandbox_tool_name self.sandbox_config = sandbox_config - self._container_cache: dict[str, tuple[Any, dict[str, Any] | None, float]] = {} + self._container_cache: dict[str, tuple[Any, dict[str, Any] | None, float, str | None]] = {} @classmethod def from_config_yaml(cls, config: CodeInterpreterInterceptionConfig) -> "CodeInterpreterInterceptionLogger": @@ -191,7 +207,13 @@ class CodeInterpreterInterceptionLogger(CustomLogger): return None kwargs[_INTERCEPTION_ACTIVE_KEY] = True - kwargs[_SANDBOX_KEY] = uuid.uuid4().hex + session_id = _extract_session_id(kwargs) + if session_id: + identity = _extract_identity(kwargs) + kwargs[_SANDBOX_KEY] = f"{identity}:{session_id}" if identity else session_id + kwargs[_SESSION_SCOPED_KEY] = True + else: + kwargs[_SANDBOX_KEY] = uuid.uuid4().hex if kwargs.get("stream"): kwargs["stream"] = False kwargs[_CONVERTED_STREAM_KEY] = True @@ -217,6 +239,7 @@ class CodeInterpreterInterceptionLogger(CustomLogger): if not is_interception_internal_key(key) and not key.startswith("_agentic_loop") and key != "max_agentic_loops" + and key != _SESSION_SCOPED_KEY } if filtered_metadata: kwargs[_LITELLM_METADATA_KEY] = filtered_metadata @@ -227,7 +250,7 @@ class CodeInterpreterInterceptionLogger(CustomLogger): def _write_interception_metadata(kwargs: dict[str, Any]) -> None: metadata = kwargs.get(_LITELLM_METADATA_KEY) metadata = dict(metadata) if isinstance(metadata, dict) else {} - for key in (_INTERCEPTION_ACTIVE_KEY, _SANDBOX_KEY, _CONVERTED_STREAM_KEY): + for key in (_INTERCEPTION_ACTIVE_KEY, _SANDBOX_KEY, _SESSION_SCOPED_KEY, _CONVERTED_STREAM_KEY): if key in kwargs: metadata[key] = kwargs[key] kwargs[_LITELLM_METADATA_KEY] = metadata @@ -347,7 +370,9 @@ class CodeInterpreterInterceptionLogger(CustomLogger): await self._prune_expired_cache() tool_calls = cast(list[CodeExecutionToolCall], tools.get("tool_calls", [])) sandbox_key = kwargs.get(_SANDBOX_KEY) - container, params = await self._get_or_create_container(cache_key=sandbox_key) + is_session = bool(kwargs.get(_SESSION_SCOPED_KEY)) + identity = _extract_identity(kwargs) if is_session else None + container, params = await self._get_or_create_container(cache_key=sandbox_key, identity=identity) try: container_id = cast(str | None, getattr(container, "id", None)) @@ -404,6 +429,7 @@ class CodeInterpreterInterceptionLogger(CustomLogger): metadata={ "tool_type": "code_interpreter", "sandbox_key": sandbox_key or "", + "is_session_scoped": bool(kwargs.get(_SESSION_SCOPED_KEY)), "code_interpreter_calls": code_interpreter_calls, }, ) @@ -419,7 +445,9 @@ class CodeInterpreterInterceptionLogger(CustomLogger): await self._prune_expired_cache() tool_calls = cast(list[CodeExecutionToolCall], tools.get("tool_calls", [])) sandbox_key = cast(str | None, kwargs.get(_SANDBOX_KEY)) - container, params = await self._get_or_create_container(cache_key=sandbox_key) + is_session = bool(kwargs.get(_SESSION_SCOPED_KEY)) + identity = _extract_identity(cast(dict[str, Any], kwargs)) if is_session else None + container, params = await self._get_or_create_container(cache_key=sandbox_key, identity=identity) try: container_id = cast(str | None, getattr(container, "id", None)) @@ -455,6 +483,7 @@ class CodeInterpreterInterceptionLogger(CustomLogger): metadata={ "tool_type": "code_interpreter", "sandbox_key": sandbox_key or "", + "is_session_scoped": bool(kwargs.get(_SESSION_SCOPED_KEY)), "code_interpreter_calls": code_interpreter_calls, "response_format": "openai", }, @@ -489,6 +518,8 @@ class CodeInterpreterInterceptionLogger(CustomLogger): async def async_agentic_loop_cleanup_hook(self, plan: AgenticLoopPlan, kwargs: dict) -> None: metadata = plan.metadata or {} if plan else {} + if metadata.get("is_session_scoped"): + return await self._delete_container_for_cache_key(metadata.get("sandbox_key")) @staticmethod @@ -520,7 +551,8 @@ class CodeInterpreterInterceptionLogger(CustomLogger): async def async_post_agentic_loop_response_hook(self, response: Any, plan: AgenticLoopPlan, kwargs: dict) -> Any: metadata = plan.metadata or {} if plan else {} - await self._delete_container_for_cache_key(metadata.get("sandbox_key")) + if not metadata.get("is_session_scoped"): + await self._delete_container_for_cache_key(metadata.get("sandbox_key")) calls = metadata.get("code_interpreter_calls") if not calls: @@ -565,17 +597,32 @@ class CodeInterpreterInterceptionLogger(CustomLogger): return f"[execution error] {message}" return getattr(result, "stdout", "") or "" - async def _get_or_create_container(self, cache_key: str | None) -> tuple[Any, dict[str, Any] | None]: + async def _get_or_create_container( + self, + cache_key: str | None, + identity: str | None = None, + ) -> tuple[Any, dict[str, Any] | None]: if cache_key: cached = self._container_cache.get(cache_key) if cached is not None: + self._container_cache[cache_key] = (cached[0], cached[1], time.time(), cached[3]) return cached[0], cached[1] container, params = await self._create_container() if cache_key: - self._container_cache[cache_key] = (container, params, time.time()) + if identity is not None: + await self._evict_lru_session_if_over_cap(identity) + self._container_cache[cache_key] = (container, params, time.time(), identity) return container, params + async def _evict_lru_session_if_over_cap(self, identity: str) -> None: + identity_entries = [(k, v) for k, v in self._container_cache.items() if v[3] == identity] + if len(identity_entries) < _SESSION_SCOPED_PER_IDENTITY_CAP: + return + lru_key, lru_entry = min(identity_entries, key=lambda item: item[1][2]) + self._container_cache.pop(lru_key, None) + await self._delete_container(container=lru_entry[0], params=lru_entry[1]) + async def _create_container(self) -> tuple[Any, dict[str, Any] | None]: if self.sandbox_config is not None: return await self.sandbox_config.acreate_sandbox(), None @@ -739,12 +786,8 @@ class CodeInterpreterInterceptionLogger(CustomLogger): now = time.time() expired = [ (cache_key, container, params) - for cache_key, ( - container, - params, - created_at, - ) in self._container_cache.items() - if now - created_at > _CACHE_TTL_SECONDS + for cache_key, (container, params, last_accessed, *_) in self._container_cache.items() + if now - last_accessed > _CACHE_TTL_SECONDS ] for cache_key, container, params in expired: self._container_cache.pop(cache_key, None) diff --git a/litellm/proxy/dev_config.yaml b/litellm/proxy/dev_config.yaml index e437ed7a118..6078161c780 100644 --- a/litellm/proxy/dev_config.yaml +++ b/litellm/proxy/dev_config.yaml @@ -182,10 +182,25 @@ model_list: litellm_params: model: openai/gpt-5.5 api_key: os.environ/OPENAI_API_KEY + - model_name: gpt-4o-mini + litellm_params: + model: openai/gpt-4o-mini + api_key: os.environ/OPENAI_API_KEY general_settings: master_key: sk-1234 +sandbox_tools: + - sandbox_tool_name: e2b_sandbox + litellm_params: + sandbox_provider: e2b + api_key: os.environ/E2B_API_KEY + litellm_settings: drop_params: True telemetry: False + code_interpreter_interception_params: + enabled: true + sandbox_tool_name: e2b_sandbox + callbacks: + - code_interpreter_interception diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 0ffb0337545..6f5c82530e3 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -153,6 +153,7 @@ _UNTRUSTED_ROOT_CONTROL_FIELDS = ( "_code_interpreter_interception_active", "_code_interpreter_interception_converted_stream", "_code_interpreter_interception_sandbox_key", + "_code_interpreter_interception_session_scoped", "max_agentic_loops", ) diff --git a/litellm/types/utils.py b/litellm/types/utils.py index dff4e4af89e..4f0c0c21bce 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3062,6 +3062,7 @@ agentic_loop_internal_litellm_params = [ "max_agentic_loops", "_code_interpreter_interception_active", "_code_interpreter_interception_sandbox_key", + "_code_interpreter_interception_session_scoped", "_code_interpreter_interception_converted_stream", ] diff --git a/qa_sticky_session.sh b/qa_sticky_session.sh new file mode 100755 index 00000000000..326bb8117c7 --- /dev/null +++ b/qa_sticky_session.sh @@ -0,0 +1,59 @@ +#!/usr/bin/env bash +# QA: code interpreter sandbox stickiness via metadata.session_id +# bash qa_sticky_session.sh +# LITELLM_BASE_URL=http://localhost:4000 LITELLM_KEY=sk-1234 bash qa_sticky_session.sh + +set -euo pipefail + +BASE="${LITELLM_BASE_URL:-http://localhost:4000}" +KEY="${LITELLM_KEY:-sk-1234}" +MODEL="${LITELLM_MODEL:-gpt-4o-mini}" +# proxy running at http://localhost:4000 (master key: sk-1234) +SESSION_A="qa-session-$(date +%s)-A" +SESSION_B="qa-session-$(date +%s)-B" + +content() { + echo "$1" | python3 -c "import sys,json; d=json.load(sys.stdin); print(d.get('choices',[{}])[0].get('message',{}).get('content',''))" +} + +call() { + local session="${1:-}" code="$2" meta="" + [[ -n "$session" ]] && meta=", \"metadata\": {\"session_id\": \"$session\"}" + curl -s -X POST "$BASE/chat/completions" \ + -H "Content-Type: application/json" \ + -H "Authorization: Bearer $KEY" \ + -d "{\"model\":\"$MODEL\"$meta,\"tools\":[{\"type\":\"code_interpreter\"}],\"messages\":[{\"role\":\"user\",\"content\":\"Run this Python code and tell me the result: $code\"}]}" +} + +assert_match() { + local label="$1" body="$2" pattern="$3" + if echo "$body" | grep -qiE "$pattern"; then + echo "PASS $label" + else + echo "FAIL $label (expected /$pattern/)" + echo " $(content "$body")" + exit 1 + fi +} + +echo "=== Sticky Session Sandbox QA ===" +echo "base: $BASE session A: $SESSION_A session B: $SESSION_B" +echo + +R=$(call "$SESSION_A" "x = 42; print(x)") +assert_match "same session_id reuses sandbox (set x=42)" "$R" "42" + +R=$(call "$SESSION_A" "print(x)") +assert_match "same session_id keeps state (x still 42)" "$R" "42" + +R=$(call "$SESSION_B" "print(x)") +assert_match "different session_id is isolated" "$R" "not defined|NameError|undefined|error" + +R=$(call "" "y = 99; print(y)") +assert_match "no session_id runs code" "$R" "99" + +R=$(call "" "print(y)") +assert_match "no session_id gets fresh sandbox each request" "$R" "not defined|NameError|undefined|error" + +echo +echo "All checks passed." diff --git a/tests/test_litellm/integrations/code_interpreter_interception/test_handler.py b/tests/test_litellm/integrations/code_interpreter_interception/test_handler.py index 7ff58ba6324..9d308ac1989 100644 --- a/tests/test_litellm/integrations/code_interpreter_interception/test_handler.py +++ b/tests/test_litellm/integrations/code_interpreter_interception/test_handler.py @@ -14,6 +14,7 @@ from litellm.integrations.code_interpreter_interception.handler import ( LITELLM_CODE_EXECUTION_TOOL_NAME, _INTERCEPTION_ACTIVE_KEY as _ACTIVE_KEY, _SANDBOX_KEY, + _SESSION_SCOPED_KEY, ) from litellm.types.integrations.custom_logger import ( CHAT_COMPLETION_AGENTIC_SURFACE, @@ -138,11 +139,7 @@ async def test_build_plan_runs_code_and_feeds_output_back(): assert sandbox.run_calls[0]["code"] == "print(40 + 2)" messages = _iter_messages(plan) - outputs = [ - m - for m in messages - if isinstance(m, dict) and m.get("type") == "function_call_output" - ] + outputs = [m for m in messages if isinstance(m, dict) and m.get("type") == "function_call_output"] assert outputs, "expected a function_call_output item appended" output_item = next(m for m in outputs if m.get("call_id") == "c1") assert "42" in str(output_item["output"]) @@ -160,9 +157,7 @@ async def test_pre_call_converts_code_interpreter_tool(): assert result is not None tools = result["tools"] - assert not any( - t.get("type") == "code_interpreter" for t in tools - ), "code_interpreter tool must be removed" + assert not any(t.get("type") == "code_interpreter" for t in tools), "code_interpreter tool must be removed" names = [t.get("name") or (t.get("function") or {}).get("name") for t in tools] assert LITELLM_CODE_EXECUTION_TOOL_NAME in names @@ -267,9 +262,7 @@ async def test_should_run_detects_only_matching_function_call(): logger = CodeInterpreterInterceptionLogger(sandbox_config=FakeSandbox()) active_kwargs = {"_code_interpreter_interception_active": True} - match = FakeResponse( - output=[_function_call_item(name=LITELLM_CODE_EXECUTION_TOOL_NAME)] - ) + match = FakeResponse(output=[_function_call_item(name=LITELLM_CODE_EXECUTION_TOOL_NAME)]) should_run, payload = await logger.async_should_run_agentic_loop( response=match, model="gpt-5", @@ -331,9 +324,7 @@ async def test_container_reused_within_request_via_server_sandbox_key(): **common, ) - assert ( - len(sandbox.create_calls) == 1 - ), "the sandbox is reused across loop iterations sharing one server sandbox key" + assert len(sandbox.create_calls) == 1, "the sandbox is reused across loop iterations sharing one server sandbox key" @pytest.mark.asyncio @@ -372,9 +363,9 @@ async def test_colliding_caller_call_id_does_not_share_sandbox(): **common, ) - assert ( - len(sandbox.create_calls) == 2 - ), "distinct server sandbox keys must isolate sandboxes despite a colliding call id" + assert len(sandbox.create_calls) == 2, ( + "distinct server sandbox keys must isolate sandboxes despite a colliding call id" + ) @pytest.mark.asyncio @@ -479,14 +470,11 @@ async def test_post_hook_injects_code_interpreter_call_matching_openai_shape(): ) response = FakeResponse(output=[{"type": "message", "content": []}]) - out = await logger.async_post_agentic_loop_response_hook( - response=response, plan=plan, kwargs={} - ) + out = await logger.async_post_agentic_loop_response_hook(response=response, plan=plan, kwargs={}) types = [item.get("type") for item in out.output] assert types == ["code_interpreter_call", "message"], ( - "code_interpreter_call must be re-injected before the message, matching " - "OpenAI's native output ordering" + "code_interpreter_call must be re-injected before the message, matching OpenAI's native output ordering" ) assert set(out.output[0].keys()) == { "id", @@ -524,8 +512,7 @@ async def test_pre_call_forces_non_stream_for_loop(): assert out is not None assert out["stream"] is False, "loop requires a non-streaming upstream call" assert out["_code_interpreter_interception_converted_stream"] is True, ( - "the converted-stream flag must be set so the final response is wrapped " - "back into a stream for the caller" + "the converted-stream flag must be set so the final response is wrapped back into a stream for the caller" ) @@ -556,9 +543,7 @@ async def test_gate_refuses_without_server_active_marker(): """A forged litellm_code_execution call must not trigger the loop unless the pre-call hook actually converted a native code_interpreter tool.""" logger = CodeInterpreterInterceptionLogger(sandbox_config=FakeSandbox()) - forged = FakeResponse( - output=[_function_call_item(name=LITELLM_CODE_EXECUTION_TOOL_NAME)] - ) + forged = FakeResponse(output=[_function_call_item(name=LITELLM_CODE_EXECUTION_TOOL_NAME)]) should_run, payload = await logger.async_should_run_agentic_loop( response=forged, @@ -577,12 +562,8 @@ async def test_gate_refuses_without_server_active_marker(): @pytest.mark.asyncio async def test_gate_rechecks_provider_scope(): """enabled_providers must be re-enforced at the gate, not only in pre-call.""" - logger = CodeInterpreterInterceptionLogger( - sandbox_config=FakeSandbox(), enabled_providers=["openai"] - ) - response = FakeResponse( - output=[_function_call_item(name=LITELLM_CODE_EXECUTION_TOOL_NAME)] - ) + logger = CodeInterpreterInterceptionLogger(sandbox_config=FakeSandbox(), enabled_providers=["openai"]) + response = FakeResponse(output=[_function_call_item(name=LITELLM_CODE_EXECUTION_TOOL_NAME)]) should_run, _ = await logger.async_should_run_agentic_loop( response=response, @@ -600,11 +581,7 @@ async def test_gate_rechecks_provider_scope(): @pytest.mark.asyncio async def test_chat_completion_gate_detects_code_execution_tool_call(): logger = CodeInterpreterInterceptionLogger(sandbox_config=FakeSandbox()) - response = { - "choices": [ - {"message": {"tool_calls": [_chat_function_call_item(call_id="call_123")]}} - ] - } + response = {"choices": [{"message": {"tool_calls": [_chat_function_call_item(call_id="call_123")]}}]} should_run, payload = await logger.async_should_run_agentic_loop( response=response, @@ -661,9 +638,7 @@ async def test_chat_completion_build_plan_runs_code_and_appends_tool_message(): }, model="gpt-5", messages=[{"role": "user", "content": "x"}], - response={ - "choices": [{"message": {"tool_calls": [_chat_function_call_item()]}}] - }, + response={"choices": [{"message": {"tool_calls": [_chat_function_call_item()]}}]}, anthropic_messages_provider_config=None, anthropic_messages_optional_request_params={ "tools": [native_chat_tool], @@ -738,8 +713,7 @@ async def test_pre_call_strips_client_forged_marker_on_initial_request(): await logger.async_pre_call_deployment_hook(kwargs, CallTypes.aresponses) assert _ACTIVE_KEY not in kwargs, ( - "no native code_interpreter tool was present, so a client-supplied " - "active marker must be cleared" + "no native code_interpreter tool was present, so a client-supplied active marker must be cleared" ) assert kwargs["litellm_metadata"] == {"safe_user_value": "kept"} @@ -774,8 +748,7 @@ async def test_pre_call_strips_forged_loop_controls_then_mints_own_markers(): assert metadata[_ACTIVE_KEY] is True assert metadata[_SANDBOX_KEY] == result[_SANDBOX_KEY] assert metadata[_SANDBOX_KEY] != "client-forged", ( - "the surviving sandbox key must be the server-minted one, not the forged " - "value the client supplied" + "the surviving sandbox key must be the server-minted one, not the forged value the client supplied" ) @@ -793,8 +766,7 @@ async def test_pre_call_preserves_marker_on_server_followup(): await logger.async_pre_call_deployment_hook(kwargs, CallTypes.aresponses) assert kwargs.get(_ACTIVE_KEY) is True, ( - "the server-set marker must survive followup requests so multi-round " - "code execution keeps working" + "the server-set marker must survive followup requests so multi-round code execution keeps working" ) @@ -805,9 +777,7 @@ async def test_sandbox_deleted_after_loop_completes(): plan = await _build_plan(logger, sandbox, call_id="k1") assert sandbox.create_calls, "sandbox must be created during the loop" - assert ( - not sandbox.delete_calls - ), "sandbox must outlive the loop until the final hook" + assert not sandbox.delete_calls, "sandbox must outlive the loop until the final hook" await logger.async_post_agentic_loop_response_hook( response=FakeResponse(output=[{"type": "message", "content": []}]), @@ -816,8 +786,7 @@ async def test_sandbox_deleted_after_loop_completes(): ) assert len(sandbox.delete_calls) == 1, ( - "the sandbox must be deleted once the final response is assembled, " - "otherwise it keeps running and billing" + "the sandbox must be deleted once the final response is assembled, otherwise it keeps running and billing" ) assert "sbxkey1" not in logger._container_cache @@ -829,16 +798,11 @@ async def test_post_hook_delete_is_idempotent_across_loop_levels(): plan = await _build_plan(logger, sandbox, call_id="k1") response = FakeResponse(output=[{"type": "message", "content": []}]) - await logger.async_post_agentic_loop_response_hook( - response=response, plan=plan, kwargs={} - ) - await logger.async_post_agentic_loop_response_hook( - response=response, plan=plan, kwargs={} - ) + await logger.async_post_agentic_loop_response_hook(response=response, plan=plan, kwargs={}) + await logger.async_post_agentic_loop_response_hook(response=response, plan=plan, kwargs={}) assert len(sandbox.delete_calls) == 1, ( - "deleting an already-removed container must be a no-op so unwinding " - "loop levels do not double-delete" + "deleting an already-removed container must be a no-op so unwinding loop levels do not double-delete" ) @@ -860,8 +824,7 @@ async def test_build_plan_deletes_sandbox_when_execution_raises(): assert len(sandbox.create_calls) == 1, "the sandbox must have been created" assert len(sandbox.delete_calls) == 1, ( - "a build failure must delete the cached sandbox so it does not keep " - "running and billing" + "a build failure must delete the cached sandbox so it does not keep running and billing" ) assert "sbxkey1" not in logger._container_cache @@ -875,8 +838,7 @@ async def test_cleanup_hook_deletes_sandbox(): await logger.async_agentic_loop_cleanup_hook(plan=plan, kwargs={}) assert len(sandbox.delete_calls) == 1, ( - "the cleanup hook must delete the sandbox so a rerun failure cannot " - "leak a running container" + "the cleanup hook must delete the sandbox so a rerun failure cannot leak a running container" ) assert "sbxkey1" not in logger._container_cache @@ -895,8 +857,7 @@ async def test_cleanup_hook_is_idempotent_with_post_hook(): await logger.async_agentic_loop_cleanup_hook(plan=plan, kwargs={}) assert len(sandbox.delete_calls) == 1, ( - "cleanup running in finally after the success-path post hook already " - "deleted the sandbox must not double-delete" + "cleanup running in finally after the success-path post hook already deleted the sandbox must not double-delete" ) @@ -923,9 +884,7 @@ async def test_responses_plan_cleans_up_sandbox_when_followup_raises(): plan = AgenticLoopPlan( run_agentic_loop=True, - request_patch=AgenticLoopRequestPatch( - model="gpt-5", messages=[{"role": "user", "content": "x"}] - ), + request_patch=AgenticLoopRequestPatch(model="gpt-5", messages=[{"role": "user", "content": "x"}]), metadata={"sandbox_key": "sbxkey1"}, ) @@ -995,9 +954,7 @@ async def test_run_code_does_not_re_resolve_registry(monkeypatch): sandbox_tools.clear_sandbox_tools() - stdout = await logger._run_tool_call( - container=container, params=params, arguments='{"code":"print(1)"}' - ) + stdout = await logger._run_tool_call(container=container, params=params, arguments='{"code":"print(1)"}') finally: sandbox_tools.clear_sandbox_tools() @@ -1013,9 +970,7 @@ async def test_run_tool_call_surfaces_execution_error(): class ErroringSandbox(FakeSandbox): async def arun_code(self, *, container, code, **kwargs): self.run_calls.append({"container": container, "code": code}) - return CodeExecutionResult( - stdout="", error={"name": "ValueError", "value": "boom"} - ) + return CodeExecutionResult(stdout="", error={"name": "ValueError", "value": "boom"}) sandbox = ErroringSandbox() logger = CodeInterpreterInterceptionLogger(sandbox_config=sandbox) @@ -1036,9 +991,7 @@ async def test_run_tool_call_reports_unparseable_arguments(): logger = CodeInterpreterInterceptionLogger(sandbox_config=sandbox) container = await logger._create_container() - stdout = await logger._run_tool_call( - container=container[0], params=None, arguments="not-json" - ) + stdout = await logger._run_tool_call(container=container[0], params=None, arguments="not-json") assert stdout == "[invalid tool arguments: could not parse code]" assert not sandbox.run_calls, "code must not run when arguments cannot be parsed" @@ -1048,9 +1001,7 @@ async def test_run_tool_call_reports_unparseable_arguments(): async def test_pre_call_skips_provider_outside_scope(): """enabled_providers must filter the pre-call conversion so a request to an out-of-scope provider is left untouched.""" - logger = CodeInterpreterInterceptionLogger( - sandbox_config=FakeSandbox(), enabled_providers=["openai"] - ) + logger = CodeInterpreterInterceptionLogger(sandbox_config=FakeSandbox(), enabled_providers=["openai"]) kwargs = { "tools": [{"type": "code_interpreter", "container": {"type": "auto"}}], "custom_llm_provider": "anthropic", @@ -1119,6 +1070,7 @@ async def test_prune_expired_cache_deletes_underlying_container(): container, params, time.time() - handler_mod._CACHE_TTL_SECONDS - 1, + None, ) await logger._prune_expired_cache() @@ -1217,3 +1169,258 @@ async def test_extract_tool_calls_reads_object_attributes(): assert len(calls) == 1 assert calls[0]["call_id"] == "c9" assert calls[0]["arguments"] == '{"code":"print(1)"}' + + +# --------------------------------------------------------------------------- +# Sticky session tests +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_pre_call_uses_session_id_from_metadata_as_sandbox_key(): + """When session_id is in request metadata, it becomes the sandbox key so the + container is shared across requests in the same session.""" + logger = CodeInterpreterInterceptionLogger(sandbox_config=FakeSandbox()) + session_id = "conv-abc-123" + kwargs = { + "tools": [{"type": "code_interpreter", "container": {"type": "auto"}}], + "custom_llm_provider": "openai", + "metadata": {"session_id": session_id}, + } + + result = await logger.async_pre_call_deployment_hook(kwargs, CallTypes.acompletion) + + assert result is not None + assert result[_SANDBOX_KEY] == session_id + assert result[_SESSION_SCOPED_KEY] is True + assert result["litellm_metadata"][_SANDBOX_KEY] == session_id + assert result["litellm_metadata"][_SESSION_SCOPED_KEY] is True + + +@pytest.mark.asyncio +async def test_pre_call_uses_session_id_from_litellm_metadata(): + """session_id in litellm_metadata also works as the sticky key.""" + logger = CodeInterpreterInterceptionLogger(sandbox_config=FakeSandbox()) + session_id = "sess-xyz-789" + kwargs = { + "tools": [{"type": "code_interpreter", "container": {"type": "auto"}}], + "custom_llm_provider": "openai", + "litellm_metadata": {"session_id": session_id}, + } + + result = await logger.async_pre_call_deployment_hook(kwargs, CallTypes.acompletion) + + assert result is not None + assert result[_SANDBOX_KEY] == session_id + assert result[_SESSION_SCOPED_KEY] is True + + +@pytest.mark.asyncio +async def test_pre_call_without_session_id_still_mints_random_key(): + """Requests without a session_id still get a server-minted random sandbox key.""" + logger = CodeInterpreterInterceptionLogger(sandbox_config=FakeSandbox()) + kwargs = { + "tools": [{"type": "code_interpreter", "container": {"type": "auto"}}], + "custom_llm_provider": "openai", + } + + result = await logger.async_pre_call_deployment_hook(kwargs, CallTypes.acompletion) + + assert result is not None + assert _SESSION_SCOPED_KEY not in result or result[_SESSION_SCOPED_KEY] is False + assert len(result[_SANDBOX_KEY]) >= 16 + + +@pytest.mark.asyncio +async def test_session_scoped_sandbox_survives_agentic_loop_cleanup(): + """A session-scoped sandbox must NOT be deleted by the cleanup or post hooks; + it needs to persist across requests within the same session.""" + sandbox = FakeSandbox(stdout="42") + logger = CodeInterpreterInterceptionLogger(sandbox_config=sandbox) + session_id = "conv-persist-me" + + plan = await logger.async_build_agentic_loop_plan( + tools={ + "tool_calls": [ + { + "call_id": "c1", + "name": LITELLM_CODE_EXECUTION_TOOL_NAME, + "arguments": '{"code":"x = 10"}', + } + ] + }, + model="gpt-4o-mini", + messages=[{"role": "user", "content": "set x"}], + response=FakeResponse(output=[_function_call_item()]), + anthropic_messages_provider_config=None, + anthropic_messages_optional_request_params={"tools": []}, + logging_obj=FakeLogging(litellm_call_id="k1"), + stream=False, + kwargs={ + "litellm_call_id": "k1", + _SANDBOX_KEY: session_id, + _SESSION_SCOPED_KEY: True, + }, + ) + + assert plan.metadata["is_session_scoped"] is True + + await logger.async_post_agentic_loop_response_hook( + response=FakeResponse(output=[{"type": "message", "content": []}]), + plan=plan, + kwargs={}, + ) + await logger.async_agentic_loop_cleanup_hook(plan=plan, kwargs={}) + + assert not sandbox.delete_calls, ( + "session-scoped sandbox must not be deleted after a single agentic loop; " + "it must persist for the next request in the session" + ) + assert session_id in logger._container_cache, "session-scoped container must remain in cache after loop ends" + + +@pytest.mark.asyncio +async def test_session_scoped_sandbox_reused_across_sequential_requests(): + """Two sequential requests with the same session_id must share one container, + confirming state (e.g. assigned variables) can persist across HTTP requests.""" + sandbox = FakeSandbox(stdout="42") + logger = CodeInterpreterInterceptionLogger(sandbox_config=sandbox) + session_id = "conv-reuse-me" + + common_plan_args = dict( + tools={ + "tool_calls": [ + { + "call_id": "c1", + "name": LITELLM_CODE_EXECUTION_TOOL_NAME, + "arguments": '{"code":"print(1)"}', + } + ] + }, + model="gpt-4o-mini", + messages=[{"role": "user", "content": "x"}], + response=FakeResponse(output=[_function_call_item()]), + anthropic_messages_provider_config=None, + anthropic_messages_optional_request_params={"tools": []}, + stream=False, + ) + session_kwargs = {_SANDBOX_KEY: session_id, _SESSION_SCOPED_KEY: True} + + plan1 = await logger.async_build_agentic_loop_plan( + logging_obj=FakeLogging(litellm_call_id="req1"), + kwargs={"litellm_call_id": "req1", **session_kwargs}, + **common_plan_args, + ) + await logger.async_post_agentic_loop_response_hook( + response=FakeResponse(output=[{"type": "message", "content": []}]), + plan=plan1, + kwargs={}, + ) + + plan2 = await logger.async_build_agentic_loop_plan( + logging_obj=FakeLogging(litellm_call_id="req2"), + kwargs={"litellm_call_id": "req2", **session_kwargs}, + **common_plan_args, + ) + await logger.async_post_agentic_loop_response_hook( + response=FakeResponse(output=[{"type": "message", "content": []}]), + plan=plan2, + kwargs={}, + ) + + assert len(sandbox.create_calls) == 1, ( + "a single container must serve both requests in the same session; " + "two creates means state cannot persist between requests" + ) + assert len(sandbox.delete_calls) == 0, "the session container must still be alive after both requests complete" + + +@pytest.mark.asyncio +async def test_non_session_sandbox_still_deleted_after_loop(): + """Without a session_id, the existing per-request ephemeral behavior is unchanged.""" + sandbox = FakeSandbox(stdout="42") + logger = CodeInterpreterInterceptionLogger(sandbox_config=sandbox) + + plan = await _build_plan(logger, sandbox, call_id="k1") + await logger.async_post_agentic_loop_response_hook( + response=FakeResponse(output=[{"type": "message", "content": []}]), + plan=plan, + kwargs={}, + ) + + assert len(sandbox.delete_calls) == 1, "non-session sandbox must still be cleaned up after each request" + + +@pytest.mark.asyncio +async def test_sandbox_key_scoped_to_api_key_hash_isolates_users(): + """Two callers supplying the same session_id but different API key hashes must + each get their own sandbox; sharing across tenants would let one read or mutate + the other's interpreter state.""" + logger = CodeInterpreterInterceptionLogger(sandbox_config=FakeSandbox()) + session_id = "same-session-id" + + result_a = await logger.async_pre_call_deployment_hook( + { + "tools": [{"type": "code_interpreter"}], + "custom_llm_provider": "openai", + "metadata": {"session_id": session_id}, + "user_api_key_hash": "hash-for-tenant-a", + }, + CallTypes.acompletion, + ) + result_b = await logger.async_pre_call_deployment_hook( + { + "tools": [{"type": "code_interpreter"}], + "custom_llm_provider": "openai", + "metadata": {"session_id": session_id}, + "user_api_key_hash": "hash-for-tenant-b", + }, + CallTypes.acompletion, + ) + + assert result_a is not None and result_b is not None + assert result_a[_SANDBOX_KEY] != result_b[_SANDBOX_KEY], ( + "same session_id from different API keys must yield different sandbox keys; " + "otherwise tenant A can read tenant B's sandbox state" + ) + assert "hash-for-tenant-a" in result_a[_SANDBOX_KEY] + assert "hash-for-tenant-b" in result_b[_SANDBOX_KEY] + + +@pytest.mark.asyncio +async def test_per_identity_cap_evicts_lru_session(): + """When a single identity holds the cap limit of session sandboxes and opens a + new one, the least-recently-used session is evicted so the allocation stays + bounded. Without this, rotating session IDs is an unbounded sandbox leak.""" + from litellm.integrations.code_interpreter_interception.handler import _SESSION_SCOPED_PER_IDENTITY_CAP + + sandbox = FakeSandbox(stdout="ok") + logger = CodeInterpreterInterceptionLogger(sandbox_config=sandbox) + identity = "hash-for-identity-x" + + for i in range(_SESSION_SCOPED_PER_IDENTITY_CAP): + await logger._get_or_create_container( + cache_key=f"{identity}:session-{i}", + identity=identity, + ) + logger._container_cache[f"{identity}:session-{i}"] = ( + logger._container_cache[f"{identity}:session-{i}"][0], + logger._container_cache[f"{identity}:session-{i}"][1], + float(i), + identity, + ) + + assert len(logger._container_cache) == _SESSION_SCOPED_PER_IDENTITY_CAP + + await logger._get_or_create_container( + cache_key=f"{identity}:session-new", + identity=identity, + ) + + assert len(logger._container_cache) == _SESSION_SCOPED_PER_IDENTITY_CAP, ( + "adding a new session beyond the cap must evict one entry so total stays bounded" + ) + assert f"{identity}:session-0" not in logger._container_cache, ( + "the entry with the oldest last_accessed timestamp must be evicted first (LRU)" + ) + assert len(sandbox.delete_calls) == 1, "evicted sandbox must be deleted, not just removed from cache"