diff --git a/litellm/integrations/code_interpreter_interception/__init__.py b/litellm/integrations/code_interpreter_interception/__init__.py new file mode 100644 index 00000000000..2256356b6f5 --- /dev/null +++ b/litellm/integrations/code_interpreter_interception/__init__.py @@ -0,0 +1,15 @@ +""" +Code Interpreter Interception Module + +Converts the native OpenAI Responses ``code_interpreter`` tool into a function +tool, runs the model-emitted code in a sandbox, and feeds the result back into +the agentic loop. +""" + +from litellm.integrations.code_interpreter_interception.handler import ( + CodeInterpreterInterceptionLogger, +) + +__all__ = [ + "CodeInterpreterInterceptionLogger", +] diff --git a/litellm/integrations/code_interpreter_interception/handler.py b/litellm/integrations/code_interpreter_interception/handler.py new file mode 100644 index 00000000000..da8149eab9b --- /dev/null +++ b/litellm/integrations/code_interpreter_interception/handler.py @@ -0,0 +1,473 @@ +""" +Code Interpreter Interception Handler + +CustomLogger that swaps the native OpenAI Responses ``code_interpreter`` tool for +a function tool, executes the code the model emits inside a sandbox, and feeds the +captured stdout back through the typed agentic loop plan. +""" + +import json +import time +import uuid +from typing import Any, cast + +import litellm +from litellm._logging import verbose_logger +from litellm.integrations.custom_logger import CustomLogger +from litellm.types.integrations.code_interpreter_interception import ( + CodeInterpreterInterceptionConfig, +) +from litellm.types.integrations.custom_logger import ( + AgenticLoopPlan, + AgenticLoopRequestPatch, +) +from litellm.types.utils import CallTypes + +LITELLM_CODE_EXECUTION_TOOL_NAME = "litellm_code_execution" +_INTERCEPTION_ACTIVE_KEY = "_code_interpreter_interception_active" +_SANDBOX_KEY = "_code_interpreter_interception_sandbox_key" +_CACHE_TTL_SECONDS = 15 * 60 + + +def _resolve_sandbox_tool(sandbox_tool_name: str | None) -> dict[str, Any] | None: + try: + from litellm.sandbox.sandbox_tools import resolve_sandbox_tool + except ImportError: + return None + return resolve_sandbox_tool(sandbox_tool_name) + + +class CodeInterpreterInterceptionLogger(CustomLogger): + """ + CustomLogger that implements transparent code-interpreter execution loops. + + Flow: + 1. Replace the native ``code_interpreter`` tool with a function tool in the + pre-call hook so the model emits code as function-call arguments. + 2. Detect ``litellm_code_execution`` function calls in the model response. + 3. Run the emitted code in a sandbox (reused per request via a server-minted + sandbox key) and build a typed rerun plan that appends the + function_call_output. + """ + + def __init__( + self, + enabled: bool = True, + enabled_providers: list[str] | None = None, + sandbox_tool_name: str | None = None, + sandbox_config: Any | None = None, + ): + super().__init__() + self.enabled = enabled + 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]] = {} + + @classmethod + def from_config_yaml( + cls, config: CodeInterpreterInterceptionConfig + ) -> "CodeInterpreterInterceptionLogger": + return cls( + enabled=bool(config.get("enabled", True)), + enabled_providers=config.get("enabled_providers"), + sandbox_tool_name=config.get("sandbox_tool_name"), + ) + + @staticmethod + def initialize_from_proxy_config( + litellm_settings: dict[str, Any], + callback_specific_params: dict[str, Any], + ) -> "CodeInterpreterInterceptionLogger": + params: CodeInterpreterInterceptionConfig = {} + if "code_interpreter_interception_params" in litellm_settings: + params = litellm_settings["code_interpreter_interception_params"] + elif "code_interpreter_interception" in callback_specific_params and isinstance( + callback_specific_params["code_interpreter_interception"], dict + ): + params = cast( + CodeInterpreterInterceptionConfig, + callback_specific_params["code_interpreter_interception"], + ) + return CodeInterpreterInterceptionLogger.from_config_yaml(params) + + async def async_pre_call_deployment_hook( + self, kwargs: dict[str, Any], call_type: CallTypes | None + ) -> dict | None: + if not kwargs.get("_agentic_loop_depth"): + kwargs.pop(_INTERCEPTION_ACTIVE_KEY, None) + kwargs.pop(_SANDBOX_KEY, None) + if not self.enabled: + return None + if call_type not in (CallTypes.responses, CallTypes.aresponses): + return None + if ( + self.enabled_providers is not None + and self._resolve_provider(kwargs) not in self.enabled_providers + ): + return None + + tools = kwargs.get("tools") + if not isinstance(tools, list): + return None + if not any( + isinstance(tool, dict) and tool.get("type") == "code_interpreter" + for tool in tools + ): + return None + + kwargs[_INTERCEPTION_ACTIVE_KEY] = True + kwargs[_SANDBOX_KEY] = uuid.uuid4().hex + if kwargs.get("stream"): + kwargs["stream"] = False + kwargs["_code_interpreter_interception_converted_stream"] = True + + function_tool = { + "type": "function", + "name": LITELLM_CODE_EXECUTION_TOOL_NAME, + "description": "Execute python code in a sandbox and return stdout.", + "parameters": { + "type": "object", + "properties": {"code": {"type": "string"}}, + "required": ["code"], + }, + } + kwargs["tools"] = [ + ( + function_tool + if isinstance(tool, dict) and tool.get("type") == "code_interpreter" + else tool + ) + for tool in tools + ] + if self._tool_choice_targets_code_interpreter(kwargs.get("tool_choice")): + kwargs["tool_choice"] = { + "type": "function", + "name": LITELLM_CODE_EXECUTION_TOOL_NAME, + } + return kwargs + + @staticmethod + def _tool_choice_targets_code_interpreter(tool_choice: Any) -> bool: + if not isinstance(tool_choice, dict): + return False + return ( + tool_choice.get("type") == "code_interpreter" + or tool_choice.get("name") == "code_interpreter" + ) + + def _resolve_provider(self, kwargs: dict[str, Any]) -> str | None: + provider = kwargs.get("custom_llm_provider") + if provider: + return provider + model = kwargs.get("model") + if not isinstance(model, str): + return None + try: + return litellm.get_llm_provider(model=model)[1] + except Exception: + return None + + async def async_should_run_agentic_loop( + self, + response: Any, + model: str, + messages: list[dict], + tools: list[dict] | None, + stream: bool, + custom_llm_provider: str, + kwargs: dict, + ) -> tuple[bool, dict]: + if not self.enabled: + return False, {} + if not kwargs.get(_INTERCEPTION_ACTIVE_KEY): + return False, {} + if ( + self.enabled_providers is not None + and custom_llm_provider not in self.enabled_providers + ): + return False, {} + + tool_calls = self._extract_code_execution_tool_calls(response=response) + if not tool_calls: + return False, {} + + return True, {"tool_calls": tool_calls} + + async def async_build_agentic_loop_plan( + self, + tools: dict, + model: str, + messages: list[dict], + response: Any, + anthropic_messages_provider_config: Any, + anthropic_messages_optional_request_params: dict, + logging_obj: Any, + stream: bool, + kwargs: dict, + ) -> AgenticLoopPlan: + await self._prune_expired_cache() + tool_calls = cast(list[dict[str, Any]], tools.get("tool_calls", [])) + sandbox_key = kwargs.get(_SANDBOX_KEY) + container, params = await self._get_or_create_container(cache_key=sandbox_key) + + try: + container_id = getattr(container, "id", None) + input_list = self._normalize_messages(messages) + code_interpreter_calls = [] + for tool_call in tool_calls: + arguments = tool_call.get("arguments", "") + code = self._parse_code(arguments) + stdout = await self._run_tool_call( + container=container, params=params, arguments=arguments + ) + input_list.append( + { + "type": "function_call", + "call_id": tool_call.get("call_id"), + "name": LITELLM_CODE_EXECUTION_TOOL_NAME, + "arguments": arguments, + } + ) + input_list.append( + { + "type": "function_call_output", + "call_id": tool_call.get("call_id"), + "output": stdout, + } + ) + code_interpreter_calls.append( + { + "id": f"ci_{uuid.uuid4().hex}", + "type": "code_interpreter_call", + "status": "completed", + "code": code, + "container_id": container_id, + "outputs": ( + [{"type": "logs", "logs": stdout}] if stdout else [] + ), + } + ) + except Exception: + await self._delete_container_for_cache_key(sandbox_key) + raise + + optional_params = anthropic_messages_optional_request_params + request_patch = AgenticLoopRequestPatch( + model=model, + messages=input_list, + tools=optional_params.get("tools"), + optional_params={k: v for k, v in optional_params.items() if k != "tools"}, + kwargs={k: v for k, v in kwargs.items() if k != "litellm_logging_obj"}, + ) + + return AgenticLoopPlan( + run_agentic_loop=True, + request_patch=request_patch, + metadata={ + "tool_type": "code_interpreter", + "sandbox_key": sandbox_key or "", + "code_interpreter_calls": code_interpreter_calls, + }, + ) + + async def async_agentic_loop_cleanup_hook( + self, plan: AgenticLoopPlan, kwargs: dict + ) -> None: + metadata = plan.metadata or {} if plan else {} + await self._delete_container_for_cache_key(metadata.get("sandbox_key")) + + 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")) + + calls = metadata.get("code_interpreter_calls") + if not calls: + return response + + is_dict = isinstance(response, dict) + output = ( + response.get("output") if is_dict else getattr(response, "output", None) + ) + if not isinstance(output, list): + return response + + def _item_type(item: Any) -> Any: + return ( + item.get("type") + if isinstance(item, dict) + else getattr(item, "type", None) + ) + + insert_at = next( + (i for i, item in enumerate(output) if _item_type(item) == "message"), + len(output), + ) + new_output = output[:insert_at] + list(calls) + output[insert_at:] + if is_dict: + response["output"] = new_output + else: + response.output = new_output + return response + + @staticmethod + def _parse_code(arguments: str) -> str: + try: + return json.loads(arguments).get("code", "") if arguments else "" + except (json.JSONDecodeError, TypeError, AttributeError): + return "" + + async def _run_tool_call( + self, container: Any, params: dict[str, Any] | None, arguments: str + ) -> str: + try: + code = json.loads(arguments).get("code", "") if arguments else "" + except (json.JSONDecodeError, TypeError): + return "[invalid tool arguments: could not parse code]" + + result = await self._run_code(container=container, params=params, code=code) + if getattr(result, "error", None): + error = result.error + message = ( + error.get("value") or error.get("name") + if isinstance(error, dict) + else str(error) + ) + 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]: + if cache_key: + cached = self._container_cache.get(cache_key) + if cached is not None: + return cached[0], cached[1] + + container, params = await self._create_container() + if cache_key: + self._container_cache[cache_key] = (container, params, time.time()) + return container, params + + 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 + + params = _resolve_sandbox_tool(self.sandbox_tool_name) + if params is None: + raise ValueError( + "CodeInterpreterInterception: no sandbox available. Provide a " + "sandbox_config or configure a sandbox tool resolvable via " + "sandbox_tool_name." + ) + container = await litellm.acreate_sandbox( + provider=params["sandbox_provider"], + api_key=params.get("api_key"), + api_base=params.get("api_base"), + ) + return container, params + + async def _run_code( + self, container: Any, params: dict[str, Any] | None, code: str + ) -> Any: + if self.sandbox_config is not None: + return await self.sandbox_config.arun_code(container=container, code=code) + if params is None: + raise ValueError( + "CodeInterpreterInterception: no sandbox available to run code." + ) + return await litellm.arun_code( + provider=params["sandbox_provider"], + container=container, + code=code, + api_key=params.get("api_key"), + ) + + async def _delete_container( + self, container: Any, params: dict[str, Any] | None + ) -> None: + try: + if self.sandbox_config is not None: + await self.sandbox_config.adelete_sandbox(container=container) + return + if params is None: + return + await litellm.adelete_sandbox( + provider=params["sandbox_provider"], + container=container, + api_key=params.get("api_key"), + api_base=params.get("api_base"), + ) + except Exception: + verbose_logger.exception( + "CodeInterpreterInterception: failed to delete sandbox container" + ) + + async def _delete_container_for_cache_key(self, cache_key: str | None) -> None: + if not cache_key: + return + cached = self._container_cache.pop(cache_key, None) + if cached is None: + return + await self._delete_container(container=cached[0], params=cached[1]) + + def _normalize_messages(self, messages: Any) -> list[dict[str, Any]]: + if isinstance(messages, str): + return [{"role": "user", "content": messages}] + if isinstance(messages, list): + return list(messages) + return [] + + def _extract_code_execution_tool_calls(self, response: Any) -> list[dict[str, Any]]: + if isinstance(response, dict): + output = response.get("output", []) + else: + output = getattr(response, "output", []) or [] + if not isinstance(output, list): + return [] + + return [ + { + "call_id": ( + item.get("call_id") + if isinstance(item, dict) + else getattr(item, "call_id", None) + ), + "name": LITELLM_CODE_EXECUTION_TOOL_NAME, + "arguments": ( + item.get("arguments") + if isinstance(item, dict) + else getattr(item, "arguments", "") + ), + } + for item in output + if self._is_code_execution_call(item) + ] + + def _is_code_execution_call(self, item: Any) -> bool: + if isinstance(item, dict): + return ( + item.get("type") == "function_call" + and item.get("name") == LITELLM_CODE_EXECUTION_TOOL_NAME + ) + return ( + getattr(item, "type", None) == "function_call" + and getattr(item, "name", None) == LITELLM_CODE_EXECUTION_TOOL_NAME + ) + + async def _prune_expired_cache(self) -> None: + 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 in expired: + self._container_cache.pop(cache_key, None) + await self._delete_container(container=container, params=params) diff --git a/litellm/integrations/custom_logger.py b/litellm/integrations/custom_logger.py index 481cf7fce8e..94fb97dff53 100644 --- a/litellm/integrations/custom_logger.py +++ b/litellm/integrations/custom_logger.py @@ -718,6 +718,24 @@ async def async_post_agentic_loop_response_hook( """ return response + async def async_agentic_loop_cleanup_hook( + self, + plan: AgenticLoopPlan, + kwargs: dict, + ) -> None: + """ + Release resources held for an agentic-loop iteration. + + Runs in a ``finally`` around the follow-up provider call, so it fires + whether the rerun returns normally, hits a loop safety abort, or raises + an upstream error. Implementations must be idempotent because the + post-response hook may already have released the same resource on the + success path. Use ``plan.metadata`` to locate what to clean up. + + Default does nothing. + """ + return None + async def async_should_run_chat_completion_agentic_loop( self, response: Any, diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 8ac5b47c6e7..790bd0519d7 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -2407,11 +2407,17 @@ def response_api_handler( provider_config=responses_api_provider_config, ) - return responses_api_provider_config.transform_response_api_response( - model=model, - raw_response=response, - logging_obj=logging_obj, + initial_response = ( + responses_api_provider_config.transform_response_api_response( + model=model, + raw_response=response, + logging_obj=logging_obj, + ) ) + # Responses agentic interception (e.g. code interpreter) runs the follow-up + # loop via the async hook, so it is async-only for now; the sync path returns + # the initial response unchanged. + return initial_response async def async_response_api_handler( self, @@ -2570,12 +2576,44 @@ async def async_response_api_handler( provider_config=responses_api_provider_config, ) - return responses_api_provider_config.transform_response_api_response( + initial_response = ( + responses_api_provider_config.transform_response_api_response( + model=model, + raw_response=response, + logging_obj=logging_obj, + ) + ) + + final_response = await self._call_agentic_completion_hooks( + response=initial_response, model=model, - raw_response=response, + messages=( + input + if isinstance(input, list) + else [{"role": "user", "content": input}] + ), + anthropic_messages_provider_config=responses_api_provider_config, + anthropic_messages_optional_request_params=response_api_optional_request_params, logging_obj=logging_obj, + stream=False, + custom_llm_provider=custom_llm_provider, + kwargs=dict(litellm_params), + api_surface="responses", ) + result = final_response if final_response is not None else initial_response + if litellm_params.get( + "_code_interpreter_interception_converted_stream" + ) and not litellm_params.get("_agentic_loop_depth"): + return self._wrap_responses_response_as_fake_stream( + result=result, + model=model, + responses_api_provider_config=responses_api_provider_config, + logging_obj=logging_obj, + custom_llm_provider=custom_llm_provider, + ) + return result + async def async_delete_response_api_handler( self, response_id: str, @@ -4875,6 +4913,132 @@ async def _execute_anthropic_agentic_plan( return response + async def _execute_responses_agentic_plan( + self, + plan: AgenticLoopPlan, + model: str, + response_api_optional_request_params: dict, + logging_obj: "LiteLLMLoggingObj", + kwargs: dict, + depth: int, + max_loops: int, + fingerprints: list[str], + fingerprint: str, + callback: Any | None = None, + ) -> Any: + patch = plan.request_patch or AgenticLoopRequestPatch() + if patch.messages is None: + raise ValueError("Agentic loop plan missing patched responses input") + + optional_params = dict(response_api_optional_request_params) + optional_params.update(patch.optional_params) + if patch.tools is not None: + optional_params["tools"] = patch.tools + optional_params = { + k: v + for k, v in optional_params.items() + if k != "stream" and k != "_code_interpreter_interception_converted_stream" + } + + internal_keys = {"litellm_logging_obj"} + kwargs_for_followup = { + k: v + for k, v in kwargs.items() + if not k.startswith("_websearch_interception") + and not k.startswith("_compression_interception") + and k != "_code_interpreter_interception_converted_stream" + and k not in internal_keys + and k not in optional_params + } + kwargs_for_followup.update(patch.kwargs) + kwargs_for_followup["_agentic_loop_depth"] = depth + 1 + kwargs_for_followup["max_agentic_loops"] = max_loops + kwargs_for_followup["_agentic_loop_fingerprints"] = fingerprints + [fingerprint] + + try: + response = await litellm.aresponses( + model=patch.model or model, + input=patch.messages, + **optional_params, + **kwargs_for_followup, + ) + + if callback is not None: + try: + response = await callback.async_post_agentic_loop_response_hook( + response=response, plan=plan, kwargs=kwargs + ) + except Exception as e: + _call_id = getattr(logging_obj, "litellm_call_id", "unknown") + verbose_logger.exception( + "LiteLLM.AgenticHookError: Exception in " + "async_post_agentic_loop_response_hook [call_id=%s model=%s]: %s", + _call_id, + model, + str(e), + ) + + return response + finally: + if callback is not None: + await self._run_agentic_loop_cleanup( + callback=callback, + plan=plan, + kwargs=kwargs, + logging_obj=logging_obj, + model=model, + ) + + @staticmethod + async def _run_agentic_loop_cleanup( + callback: Any, + plan: AgenticLoopPlan, + kwargs: dict, + logging_obj: "LiteLLMLoggingObj", + model: str, + ) -> None: + try: + await callback.async_agentic_loop_cleanup_hook(plan=plan, kwargs=kwargs) + except Exception as e: + _call_id = getattr(logging_obj, "litellm_call_id", "unknown") + verbose_logger.exception( + "LiteLLM.AgenticHookError: Exception in " + "async_agentic_loop_cleanup_hook [call_id=%s model=%s]: %s", + _call_id, + model, + str(e), + ) + + def _wrap_responses_response_as_fake_stream( + self, + result: Any, + model: str, + responses_api_provider_config: Any, + logging_obj: "LiteLLMLoggingObj", + custom_llm_provider: str, + ) -> Any: + """ + Wrap a completed responses result as a synthetic stream. + + Used when an interceptor forced stream=False to run the agentic loop on + the non-streaming path, but the caller originally asked for streaming. + """ + import httpx + + from litellm.responses.streaming_iterator import ( + MockResponsesAPIStreamingIterator, + ) + + payload = result.model_dump() if hasattr(result, "model_dump") else result + raw_response = httpx.Response(status_code=200, json=payload) + return MockResponsesAPIStreamingIterator( + response=raw_response, + model=model, + responses_api_provider_config=responses_api_provider_config, + logging_obj=logging_obj, + custom_llm_provider=custom_llm_provider, + ) + async def _execute_chat_completion_agentic_plan( self, plan: AgenticLoopPlan, @@ -4940,6 +5104,7 @@ async def _call_agentic_completion_hooks( stream: bool, custom_llm_provider: str, kwargs: Dict, + api_surface: str = "anthropic_messages", ) -> Optional[Any]: """ Call agentic completion hooks for all custom loggers (Anthropic Messages API). @@ -5046,6 +5211,20 @@ async def _call_agentic_completion_hooks( if not plan.run_agentic_loop: continue + if api_surface == "responses": + return await self._execute_responses_agentic_plan( + plan=plan, + model=model, + response_api_optional_request_params=anthropic_messages_optional_request_params, + logging_obj=logging_obj, + kwargs=kwargs_with_provider, + depth=depth, + max_loops=max_loops, + fingerprints=fingerprints, + fingerprint=fingerprint, + callback=callback, + ) + return await self._execute_anthropic_agentic_plan( plan=plan, model=model, @@ -5083,7 +5262,7 @@ async def _call_agentic_completion_hooks( else False ) - if websearch_converted_stream: + if api_surface == "anthropic_messages" and websearch_converted_stream: from typing import cast from litellm._logging import verbose_logger diff --git a/litellm/llms/e2b/sandbox/transformation.py b/litellm/llms/e2b/sandbox/transformation.py index 1ce28bc55fb..c279fab22ab 100644 --- a/litellm/llms/e2b/sandbox/transformation.py +++ b/litellm/llms/e2b/sandbox/transformation.py @@ -51,11 +51,13 @@ async def acreate_sandbox( timeout: int | None = None, allow_internet_access: bool = True, api_key: str | None = None, + api_base: str | None = None, metadata: dict | None = None, client: AsyncHTTPHandler | None = None, **kwargs, ) -> ContainerHandle: key = self.validate_environment(api_key=api_key) + base = api_base or E2B_API_BASE body = { "templateID": template or E2B_DEFAULT_TEMPLATE, "timeout": timeout if timeout is not None else DEFAULT_SANDBOX_TIMEOUT, @@ -68,7 +70,7 @@ async def acreate_sandbox( response = cast( httpx.Response, await self._http(client).post( - url=f"{E2B_API_BASE}/sandboxes", + url=f"{base}/sandboxes", headers={"X-API-Key": key, "Content-Type": "application/json"}, json=body, ), @@ -84,6 +86,7 @@ async def acreate_sandbox( "envd_access_token": data.get("envdAccessToken"), "traffic_access_token": data.get("trafficAccessToken"), "api_key": key, + "api_base": base, } return handle @@ -130,6 +133,7 @@ async def adelete_sandbox( *, container: Union[ContainerHandle, str], api_key: str | None = None, + api_base: str | None = None, client: AsyncHTTPHandler | None = None, **kwargs, ) -> bool: @@ -139,11 +143,12 @@ async def adelete_sandbox( or handle._hidden_params.get("api_key") or self.validate_environment() ) + base = api_base or handle._hidden_params.get("api_base") or E2B_API_BASE try: response = cast( httpx.Response, await self._http(client).delete( - url=f"{E2B_API_BASE}/sandboxes/{handle.id}", + url=f"{base}/sandboxes/{handle.id}", headers={"X-API-Key": key}, ), ) diff --git a/litellm/proxy/common_utils/callback_utils.py b/litellm/proxy/common_utils/callback_utils.py index 71dce163b78..d48499af6f0 100644 --- a/litellm/proxy/common_utils/callback_utils.py +++ b/litellm/proxy/common_utils/callback_utils.py @@ -71,6 +71,23 @@ def initialize_callbacks_on_proxy( imported_list.append(compression_interception_obj) continue + if ( + isinstance(callback, str) + and callback == "code_interpreter_interception" + ): + from litellm.integrations.code_interpreter_interception.handler import ( + CodeInterpreterInterceptionLogger, + ) + + code_interpreter_interception_obj = ( + CodeInterpreterInterceptionLogger.initialize_from_proxy_config( + litellm_settings=litellm_settings, + callback_specific_params=callback_specific_params, + ) + ) + imported_list.append(code_interpreter_interception_obj) + continue + # check if callback is a custom logger compatible callback if isinstance(callback, str): callback = LoggingCallbackManager._add_custom_callback_generic_api_str( diff --git a/litellm/proxy/example_config_yaml/code_interpreter_interception_config.yaml b/litellm/proxy/example_config_yaml/code_interpreter_interception_config.yaml new file mode 100644 index 00000000000..9c85f5c5140 --- /dev/null +++ b/litellm/proxy/example_config_yaml/code_interpreter_interception_config.yaml @@ -0,0 +1,17 @@ +model_list: + - model_name: gpt-5 + litellm_params: + model: openai/gpt-5 + +# Sandbox tools configuration +sandbox_tools: + - sandbox_tool_name: "my-e2b" + litellm_params: + sandbox_provider: "e2b" + api_key: os.environ/E2B_API_KEY + +litellm_settings: + callbacks: ["code_interpreter_interception"] + code_interpreter_interception_params: + enabled_providers: ["openai"] + sandbox_tool_name: "my-e2b" diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 2c5937d9506..177fced5cd4 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -141,6 +141,20 @@ def parse_cache_control(cache_control): "service_callback", "logger_fn", "litellm_disabled_callbacks", + # Agentic-loop control fields. These bound or drive an interceptor's agentic + # loop (web search, compression, code interpreter) and are server-controlled. + # A client-supplied value would forge loop depth/cycle state, mark an + # interception as active (triggering sandbox code execution without the + # native tool ever being present), force the completed response to be + # re-wrapped as a synthetic stream the caller never asked for, or raise the + # loop ceiling to drive many upstream model calls and sandbox executions + # from a single request. + "_agentic_loop_depth", + "_agentic_loop_fingerprints", + "_code_interpreter_interception_active", + "_code_interpreter_interception_converted_stream", + "_code_interpreter_interception_sandbox_key", + "max_agentic_loops", ) _UNTRUSTED_METADATA_CONTROL_FIELDS = ( diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index c138626a272..9d391199ff0 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -4773,6 +4773,11 @@ async def load_config( config ) + ## SANDBOX TOOLS SETTINGS + from litellm.sandbox.sandbox_tools import register_sandbox_tools + + register_sandbox_tools(config.get("sandbox_tools") or []) + ## /fine_tuning/jobs endpoints config finetuning_config = config.get("finetune_settings", None) set_fine_tuning_config(config=finetuning_config) diff --git a/litellm/sandbox/main.py b/litellm/sandbox/main.py index a3f5f3b4665..45d3bffb4f9 100644 --- a/litellm/sandbox/main.py +++ b/litellm/sandbox/main.py @@ -70,6 +70,7 @@ async def acreate_sandbox( timeout: int | None = None, allow_internet_access: bool = True, api_key: str | None = None, + api_base: str | None = None, **kwargs, ) -> ContainerHandle: _update_logging(kwargs, provider, "create_sandbox") @@ -78,6 +79,7 @@ async def acreate_sandbox( timeout=timeout, allow_internet_access=allow_internet_access, api_key=api_key, + api_base=api_base, **_forward_kwargs(kwargs), ) @@ -104,12 +106,14 @@ async def adelete_sandbox( provider: str, container: Union[ContainerHandle, str], api_key: str | None = None, + api_base: str | None = None, **kwargs, ) -> bool: _update_logging(kwargs, provider, "delete_sandbox") return await _get_config(provider).adelete_sandbox( container=container, api_key=api_key, + api_base=api_base, **_forward_kwargs(kwargs), ) @@ -121,6 +125,7 @@ async def acode_interpreter_tool( template: str | None = None, timeout: int | None = None, api_key: str | None = None, + api_base: str | None = None, **kwargs, ) -> CodeExecutionResult: _update_logging(kwargs, provider, "code_interpreter_tool") @@ -128,7 +133,11 @@ async def acode_interpreter_tool( forwarded = _forward_kwargs(kwargs) container = await config.acreate_sandbox( - template=template, timeout=timeout, api_key=api_key, **forwarded + template=template, + timeout=timeout, + api_key=api_key, + api_base=api_base, + **forwarded, ) try: return await config.arun_code( @@ -137,7 +146,7 @@ async def acode_interpreter_tool( finally: try: await config.adelete_sandbox( - container=container, api_key=api_key, **forwarded + container=container, api_key=api_key, api_base=api_base, **forwarded ) except Exception as e: litellm._logging.verbose_logger.debug( diff --git a/litellm/sandbox/sandbox_tools.py b/litellm/sandbox/sandbox_tools.py new file mode 100644 index 00000000000..f4a6678f629 --- /dev/null +++ b/litellm/sandbox/sandbox_tools.py @@ -0,0 +1,60 @@ +""" +Registry for sandbox tools configured via the proxy's top-level `sandbox_tools`. + +A sandbox tool maps a name to a sandbox provider plus its credentials, so the +code interpreter interceptor can resolve a tool by name to provider/key/base. +""" + +from collections.abc import Iterator + +from litellm._logging import verbose_logger + +_SANDBOX_TOOL_REGISTRY: dict[str, dict] = {} + + +def _resolve_secret_value(value: str | None) -> str | None: + if not isinstance(value, str): + return None + if value.startswith("os.environ/"): + from litellm.secret_managers.main import get_secret_str + + return get_secret_str(value) + return value + + +def _iter_valid_tools(tools: list[dict]) -> Iterator[tuple[str, dict]]: + for tool in tools: + if not isinstance(tool, dict): + verbose_logger.warning("sandbox_tools: skipping non-dict entry %r", tool) + continue + name = tool.get("sandbox_tool_name") + if not name: + verbose_logger.warning( + "sandbox_tools: skipping entry missing 'sandbox_tool_name': %r", tool + ) + continue + params = tool.get("litellm_params") or {} + provider = params.get("sandbox_provider") + if not provider: + verbose_logger.warning( + "sandbox_tools: skipping entry missing 'sandbox_provider': %r", tool + ) + continue + yield name, { + "sandbox_provider": provider, + "api_key": _resolve_secret_value(params.get("api_key")), + "api_base": _resolve_secret_value(params.get("api_base")), + } + + +def register_sandbox_tools(tools: list[dict]) -> None: + global _SANDBOX_TOOL_REGISTRY + _SANDBOX_TOOL_REGISTRY = dict(_iter_valid_tools(tools)) + + +def resolve_sandbox_tool(name: str) -> dict | None: + return _SANDBOX_TOOL_REGISTRY.get(name) + + +def clear_sandbox_tools() -> None: + register_sandbox_tools([]) diff --git a/litellm/types/integrations/code_interpreter_interception.py b/litellm/types/integrations/code_interpreter_interception.py new file mode 100644 index 00000000000..2669c59db37 --- /dev/null +++ b/litellm/types/integrations/code_interpreter_interception.py @@ -0,0 +1,22 @@ +""" +Type definitions for Code Interpreter Interception integration. +""" + +from typing import List, TypedDict + + +class CodeInterpreterInterceptionConfig(TypedDict, total=False): + """ + Configuration parameters for CodeInterpreterInterceptionLogger. + + Used in proxy_config.yaml under litellm_settings: + litellm_settings: + code_interpreter_interception_params: + enabled: true + enabled_providers: ["openai"] + sandbox_tool_name: "my_e2b_sandbox" + """ + + enabled: bool + enabled_providers: List[str] + sandbox_tool_name: str diff --git a/tests/test_litellm/integrations/code_interpreter_interception/__init__.py b/tests/test_litellm/integrations/code_interpreter_interception/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/integrations/code_interpreter_interception/test_handler.py b/tests/test_litellm/integrations/code_interpreter_interception/test_handler.py new file mode 100644 index 00000000000..f33814b86df --- /dev/null +++ b/tests/test_litellm/integrations/code_interpreter_interception/test_handler.py @@ -0,0 +1,982 @@ +""" +Unit tests for CodeInterpreterInterceptionLogger. + +All sandbox dependencies are injected (dependency injection, no monkeypatch): +a FakeSandbox stands in for the real e2b config and records how it is called. +""" + +import time + +import pytest + +from litellm.integrations.code_interpreter_interception.handler import ( + CodeInterpreterInterceptionLogger, + LITELLM_CODE_EXECUTION_TOOL_NAME, +) +from litellm.llms.base_llm.sandbox.transformation import CodeExecutionResult +from litellm.types.utils import CallTypes + +_ACTIVE_KEY = "_code_interpreter_interception_active" +_SANDBOX_KEY = "_code_interpreter_interception_sandbox_key" + + +class FakeHandle: + def __init__(self, sandbox_id="sbx_fake"): + self.id = sandbox_id + + +class FakeSandbox: + """Records acreate_sandbox / arun_code / adelete_sandbox calls.""" + + def __init__(self, stdout="42"): + self.stdout = stdout + self.create_calls = [] + self.run_calls = [] + self.delete_calls = [] + + async def acreate_sandbox(self, **kwargs): + self.create_calls.append(kwargs) + return FakeHandle() + + async def arun_code(self, *, container, code, **kwargs): + self.run_calls.append({"container": container, "code": code}) + return CodeExecutionResult(stdout=self.stdout) + + async def adelete_sandbox(self, *, container, **kwargs): + self.delete_calls.append({"container": container}) + return True + + +class FakeLogging: + def __init__(self, litellm_call_id="k1"): + self.litellm_call_id = litellm_call_id + self.model_call_details = {} + + +def _function_call_item(call_id="c1", name=LITELLM_CODE_EXECUTION_TOOL_NAME): + return { + "type": "function_call", + "call_id": call_id, + "name": name, + "arguments": '{"code":"print(40 + 2)"}', + } + + +class FakeResponse: + def __init__(self, output): + self.output = output + + +def _iter_messages(plan): + patch = plan.request_patch + assert patch is not None, "plan.request_patch must be set" + assert patch.messages is not None, "plan.request_patch.messages must be set" + return patch.messages + + +@pytest.mark.asyncio +async def test_build_plan_runs_code_and_feeds_output_back(): + sandbox = FakeSandbox(stdout="42") + logger = CodeInterpreterInterceptionLogger(sandbox_config=sandbox) + response = FakeResponse(output=[_function_call_item()]) + + plan = await logger.async_build_agentic_loop_plan( + tools={ + "tool_calls": [ + { + "call_id": "c1", + "name": LITELLM_CODE_EXECUTION_TOOL_NAME, + "arguments": '{"code":"print(40 + 2)"}', + } + ] + }, + model="gpt-5", + messages=[{"role": "user", "content": "x"}], + response=response, + 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: "sbxkey1"}, + ) + + assert sandbox.run_calls, "sandbox.arun_code must be invoked" + 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" + ] + 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"]) + + +@pytest.mark.asyncio +async def test_pre_call_converts_code_interpreter_tool(): + 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.aresponses) + + 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" + names = [t.get("name") or (t.get("function") or {}).get("name") for t in tools] + assert LITELLM_CODE_EXECUTION_TOOL_NAME in names + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "tool_choice", + [ + {"type": "code_interpreter"}, + {"type": "hosted_tool", "name": "code_interpreter"}, + ], +) +async def test_pre_call_rewrites_forced_code_interpreter_tool_choice(tool_choice): + """A forced tool_choice targeting the native code_interpreter tool must be + rewritten to the generated function tool; otherwise the outbound request + references a tool that no longer exists and the provider rejects it.""" + logger = CodeInterpreterInterceptionLogger(sandbox_config=FakeSandbox()) + kwargs = { + "tools": [{"type": "code_interpreter", "container": {"type": "auto"}}], + "tool_choice": tool_choice, + "custom_llm_provider": "openai", + } + + result = await logger.async_pre_call_deployment_hook(kwargs, CallTypes.aresponses) + + assert result is not None + assert result["tool_choice"] == { + "type": "function", + "name": LITELLM_CODE_EXECUTION_TOOL_NAME, + } + + +@pytest.mark.asyncio +async def test_pre_call_leaves_unrelated_tool_choice_untouched(): + logger = CodeInterpreterInterceptionLogger(sandbox_config=FakeSandbox()) + kwargs = { + "tools": [{"type": "code_interpreter", "container": {"type": "auto"}}], + "tool_choice": "auto", + "custom_llm_provider": "openai", + } + + result = await logger.async_pre_call_deployment_hook(kwargs, CallTypes.aresponses) + + assert result is not None + assert result["tool_choice"] == "auto" + + +@pytest.mark.asyncio +async def test_pre_call_noop_on_non_responses(): + 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 None + + +@pytest.mark.asyncio +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)] + ) + should_run, payload = await logger.async_should_run_agentic_loop( + response=match, + model="gpt-5", + messages=[{"role": "user", "content": "x"}], + tools=[], + stream=False, + custom_llm_provider="openai", + kwargs=active_kwargs, + ) + assert should_run is True + assert payload.get("tool_calls") + + no_match = FakeResponse(output=[_function_call_item(name="something_else")]) + should_run2, payload2 = await logger.async_should_run_agentic_loop( + response=no_match, + model="gpt-5", + messages=[{"role": "user", "content": "x"}], + tools=[], + stream=False, + custom_llm_provider="openai", + kwargs=active_kwargs, + ) + assert should_run2 is False + assert payload2 == {} + + +@pytest.mark.asyncio +async def test_container_reused_within_request_via_server_sandbox_key(): + sandbox = FakeSandbox(stdout="42") + logger = CodeInterpreterInterceptionLogger(sandbox_config=sandbox) + response = FakeResponse(output=[_function_call_item()]) + + common = dict( + tools={ + "tool_calls": [ + { + "call_id": "c1", + "name": LITELLM_CODE_EXECUTION_TOOL_NAME, + "arguments": '{"code":"print(40 + 2)"}', + } + ] + }, + model="gpt-5", + messages=[{"role": "user", "content": "x"}], + response=response, + anthropic_messages_provider_config=None, + anthropic_messages_optional_request_params={"tools": []}, + stream=False, + ) + + await logger.async_build_agentic_loop_plan( + logging_obj=FakeLogging(litellm_call_id="k1"), + kwargs={"litellm_call_id": "k1", _SANDBOX_KEY: "server-nonce-1"}, + **common, + ) + await logger.async_build_agentic_loop_plan( + logging_obj=FakeLogging(litellm_call_id="k1"), + kwargs={"litellm_call_id": "k1", _SANDBOX_KEY: "server-nonce-1"}, + **common, + ) + + assert ( + len(sandbox.create_calls) == 1 + ), "the sandbox is reused across loop iterations sharing one server sandbox key" + + +@pytest.mark.asyncio +async def test_colliding_caller_call_id_does_not_share_sandbox(): + """Two requests with the same caller-controlled litellm_call_id but distinct + server-minted sandbox keys must NOT share a container; otherwise one user's + code could read another in-flight request's sandbox state.""" + sandbox = FakeSandbox(stdout="42") + logger = CodeInterpreterInterceptionLogger(sandbox_config=sandbox) + common = dict( + tools={ + "tool_calls": [ + { + "call_id": "c1", + "name": LITELLM_CODE_EXECUTION_TOOL_NAME, + "arguments": '{"code":"print(1)"}', + } + ] + }, + model="gpt-5", + messages=[{"role": "user", "content": "x"}], + response=FakeResponse(output=[_function_call_item()]), + anthropic_messages_provider_config=None, + anthropic_messages_optional_request_params={"tools": []}, + stream=False, + ) + + await logger.async_build_agentic_loop_plan( + logging_obj=FakeLogging(litellm_call_id="shared"), + kwargs={"litellm_call_id": "shared", _SANDBOX_KEY: "nonce-A"}, + **common, + ) + await logger.async_build_agentic_loop_plan( + logging_obj=FakeLogging(litellm_call_id="shared"), + kwargs={"litellm_call_id": "shared", _SANDBOX_KEY: "nonce-B"}, + **common, + ) + + assert ( + len(sandbox.create_calls) == 2 + ), "distinct server sandbox keys must isolate sandboxes despite a colliding call id" + + +@pytest.mark.asyncio +async def test_pre_call_mints_server_sandbox_key(): + """The interceptor mints a server-side sandbox key (not derived from the + caller-controlled call id) when it activates.""" + logger = CodeInterpreterInterceptionLogger(sandbox_config=FakeSandbox()) + kwargs = { + "tools": [{"type": "code_interpreter", "container": {"type": "auto"}}], + "custom_llm_provider": "openai", + "litellm_call_id": "caller-supplied", + _SANDBOX_KEY: "caller-forged", + } + + result = await logger.async_pre_call_deployment_hook(kwargs, CallTypes.aresponses) + + assert result is not None + assert result[_SANDBOX_KEY] not in ("caller-forged", "caller-supplied") + assert len(result[_SANDBOX_KEY]) >= 16 + + +@pytest.mark.asyncio +async def test_build_plan_records_code_interpreter_call_metadata(): + sandbox = FakeSandbox(stdout="42") + logger = CodeInterpreterInterceptionLogger(sandbox_config=sandbox) + plan = await logger.async_build_agentic_loop_plan( + tools={ + "tool_calls": [ + { + "call_id": "c1", + "name": LITELLM_CODE_EXECUTION_TOOL_NAME, + "arguments": '{"code":"print(40 + 2)"}', + } + ] + }, + model="gpt-5", + messages=[{"role": "user", "content": "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: "sbxkey1"}, + ) + + calls = plan.metadata["code_interpreter_calls"] + assert calls, "build_plan must record a code_interpreter_call for re-injection" + assert calls[0]["code"] == "print(40 + 2)" + assert calls[0]["container_id"] == "sbx_fake" + assert calls[0]["type"] == "code_interpreter_call" + assert calls[0]["status"] == "completed" + assert calls[0]["outputs"] == [{"type": "logs", "logs": "42"}], ( + "outputs must be an OpenAI-shaped logs array (not None) so clients that " + "iterate over code_interpreter_call.outputs do not break" + ) + + +@pytest.mark.asyncio +async def test_build_plan_outputs_empty_array_when_no_stdout(): + """No stdout must still yield an iteration-safe empty array, never None.""" + sandbox = FakeSandbox(stdout="") + logger = CodeInterpreterInterceptionLogger(sandbox_config=sandbox) + plan = await logger.async_build_agentic_loop_plan( + tools={ + "tool_calls": [ + { + "call_id": "c1", + "name": LITELLM_CODE_EXECUTION_TOOL_NAME, + "arguments": '{"code":"pass"}', + } + ] + }, + model="gpt-5", + messages=[{"role": "user", "content": "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: "sbxkey1"}, + ) + + assert plan.metadata["code_interpreter_calls"][0]["outputs"] == [] + + +@pytest.mark.asyncio +async def test_post_hook_injects_code_interpreter_call_matching_openai_shape(): + from litellm.types.integrations.custom_logger import AgenticLoopPlan + + logger = CodeInterpreterInterceptionLogger(sandbox_config=FakeSandbox()) + ci_item = { + "id": "ci_x", + "type": "code_interpreter_call", + "status": "completed", + "code": "print(1)", + "container_id": "sbx_fake", + "outputs": [{"type": "logs", "logs": "1"}], + } + plan = AgenticLoopPlan( + run_agentic_loop=True, + metadata={"code_interpreter_calls": [ci_item]}, + ) + response = FakeResponse(output=[{"type": "message", "content": []}]) + + 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" + ) + assert set(out.output[0].keys()) == { + "id", + "type", + "status", + "code", + "container_id", + "outputs", + }, "injected item must match OpenAI's code_interpreter_call keys exactly" + + +@pytest.mark.asyncio +async def test_post_hook_noop_without_recorded_calls(): + from litellm.types.integrations.custom_logger import AgenticLoopPlan + + logger = CodeInterpreterInterceptionLogger(sandbox_config=FakeSandbox()) + response = FakeResponse(output=[{"type": "message", "content": []}]) + out = await logger.async_post_agentic_loop_response_hook( + response=response, plan=AgenticLoopPlan(run_agentic_loop=True), kwargs={} + ) + assert [item.get("type") for item in out.output] == ["message"] + + +@pytest.mark.asyncio +async def test_pre_call_forces_non_stream_for_loop(): + logger = CodeInterpreterInterceptionLogger(sandbox_config=FakeSandbox()) + kwargs = { + "tools": [{"type": "code_interpreter", "container": {"type": "auto"}}], + "custom_llm_provider": "openai", + "stream": True, + } + + out = await logger.async_pre_call_deployment_hook(kwargs, CallTypes.aresponses) + + 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" + ) + + +async def _build_plan(logger, sandbox, call_id="k1", provider="openai"): + return await logger.async_build_agentic_loop_plan( + tools={ + "tool_calls": [ + { + "call_id": "c1", + "name": LITELLM_CODE_EXECUTION_TOOL_NAME, + "arguments": '{"code":"print(40 + 2)"}', + } + ] + }, + model="gpt-5", + messages=[{"role": "user", "content": "x"}], + response=FakeResponse(output=[_function_call_item()]), + anthropic_messages_provider_config=None, + anthropic_messages_optional_request_params={"tools": []}, + logging_obj=FakeLogging(litellm_call_id=call_id), + stream=False, + kwargs={"litellm_call_id": call_id, _SANDBOX_KEY: "sbxkey1"}, + ) + + +@pytest.mark.asyncio +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)] + ) + + should_run, payload = await logger.async_should_run_agentic_loop( + response=forged, + model="gpt-5", + messages=[{"role": "user", "content": "x"}], + tools=[], + stream=False, + custom_llm_provider="openai", + kwargs={}, + ) + + assert should_run is False + assert payload == {} + + +@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)] + ) + + should_run, _ = await logger.async_should_run_agentic_loop( + response=response, + model="claude-x", + messages=[{"role": "user", "content": "x"}], + tools=[], + stream=False, + custom_llm_provider="anthropic", + kwargs={_ACTIVE_KEY: True}, + ) + + assert should_run is False + + +@pytest.mark.asyncio +async def test_pre_call_strips_client_forged_marker_on_initial_request(): + """A client cannot pre-set the active marker on the original request.""" + logger = CodeInterpreterInterceptionLogger(sandbox_config=FakeSandbox()) + kwargs = { + "tools": [{"type": "web_search"}], + "custom_llm_provider": "openai", + _ACTIVE_KEY: True, + } + + 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" + ) + + +@pytest.mark.asyncio +async def test_pre_call_preserves_marker_on_server_followup(): + """On a server-driven followup (depth>0) the marker is trusted and kept.""" + logger = CodeInterpreterInterceptionLogger(sandbox_config=FakeSandbox()) + kwargs = { + "tools": [{"type": "function", "name": LITELLM_CODE_EXECUTION_TOOL_NAME}], + "custom_llm_provider": "openai", + "_agentic_loop_depth": 1, + _ACTIVE_KEY: True, + } + + 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" + ) + + +@pytest.mark.asyncio +async def test_sandbox_deleted_after_loop_completes(): + sandbox = FakeSandbox(stdout="42") + logger = CodeInterpreterInterceptionLogger(sandbox_config=sandbox) + + 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" + + await logger.async_post_agentic_loop_response_hook( + response=FakeResponse(output=[{"type": "message", "content": []}]), + plan=plan, + kwargs={}, + ) + + assert len(sandbox.delete_calls) == 1, ( + "the sandbox must be deleted once the final response is assembled, " + "otherwise it keeps running and billing" + ) + assert "sbxkey1" not in logger._container_cache + + +@pytest.mark.asyncio +async def test_post_hook_delete_is_idempotent_across_loop_levels(): + sandbox = FakeSandbox(stdout="42") + logger = CodeInterpreterInterceptionLogger(sandbox_config=sandbox) + 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={} + ) + + assert len(sandbox.delete_calls) == 1, ( + "deleting an already-removed container must be a no-op so unwinding " + "loop levels do not double-delete" + ) + + +@pytest.mark.asyncio +async def test_build_plan_deletes_sandbox_when_execution_raises(): + """If sandbox execution raises before a plan is built (e.g. E2B aborts + output over its cap), the cached sandbox must be deleted before re-raising, + otherwise a caller can leak paid containers until the prune TTL.""" + + class RaisingSandbox(FakeSandbox): + async def arun_code(self, *, container, code, **kwargs): + raise ValueError("output exceeded cap") + + sandbox = RaisingSandbox() + logger = CodeInterpreterInterceptionLogger(sandbox_config=sandbox) + + with pytest.raises(ValueError, match="exceeded cap"): + await _build_plan(logger, sandbox, call_id="k1") + + 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" + ) + assert "sbxkey1" not in logger._container_cache + + +@pytest.mark.asyncio +async def test_cleanup_hook_deletes_sandbox(): + sandbox = FakeSandbox(stdout="42") + logger = CodeInterpreterInterceptionLogger(sandbox_config=sandbox) + plan = await _build_plan(logger, sandbox, call_id="k1") + + 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" + ) + assert "sbxkey1" not in logger._container_cache + + +@pytest.mark.asyncio +async def test_cleanup_hook_is_idempotent_with_post_hook(): + 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={}, + ) + 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" + ) + + +@pytest.mark.asyncio +async def test_responses_plan_cleans_up_sandbox_when_followup_raises(): + """If the agentic rerun fails, _execute_responses_agentic_plan must still + invoke the cleanup hook so the sandbox is not left running.""" + import litellm + from litellm.integrations.custom_logger import CustomLogger + from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler + from litellm.types.integrations.custom_logger import ( + AgenticLoopPlan, + AgenticLoopRequestPatch, + ) + + cleanup_calls = [] + + class CleanupCallback(CustomLogger): + async def async_post_agentic_loop_response_hook(self, response, plan, kwargs): + return response + + async def async_agentic_loop_cleanup_hook(self, plan, kwargs): + cleanup_calls.append(plan) + + plan = AgenticLoopPlan( + run_agentic_loop=True, + request_patch=AgenticLoopRequestPatch( + model="gpt-5", messages=[{"role": "user", "content": "x"}] + ), + metadata={"sandbox_key": "sbxkey1"}, + ) + + original = litellm.aresponses + + async def _boom(*args, **kwargs): + raise RuntimeError("upstream blew up") + + litellm.aresponses = _boom + try: + with pytest.raises(RuntimeError, match="upstream blew up"): + await BaseLLMHTTPHandler()._execute_responses_agentic_plan( + plan=plan, + model="gpt-5", + response_api_optional_request_params={}, + logging_obj=FakeLogging(litellm_call_id="k1"), + kwargs={}, + depth=0, + max_loops=3, + fingerprints=[], + fingerprint="fp", + callback=CleanupCallback(), + ) + finally: + litellm.aresponses = original + + assert cleanup_calls == [plan], ( + "cleanup hook must run in finally even when the rerun raises, otherwise " + "the sandbox keeps running until the prune TTL" + ) + + +@pytest.mark.asyncio +async def test_run_code_does_not_re_resolve_registry(monkeypatch): + """Params resolved once at create time must be reused for running code, so a + registry clear between create and run cannot turn into a create-then-fail.""" + import litellm + from litellm.sandbox import sandbox_tools + + sandbox_tools.register_sandbox_tools( + [ + { + "sandbox_tool_name": "e2b_default", + "litellm_params": {"sandbox_provider": "e2b", "api_key": "sk-x"}, + } + ] + ) + + create_kwargs = {} + run_kwargs = {} + + async def fake_acreate_sandbox(**kwargs): + create_kwargs.update(kwargs) + return FakeHandle() + + async def fake_arun_code(**kwargs): + run_kwargs.update(kwargs) + return CodeExecutionResult(stdout="ok") + + monkeypatch.setattr(litellm, "acreate_sandbox", fake_acreate_sandbox) + monkeypatch.setattr(litellm, "arun_code", fake_arun_code) + + logger = CodeInterpreterInterceptionLogger(sandbox_tool_name="e2b_default") + try: + container, params = await logger._get_or_create_container(cache_key="k1") + assert params is not None and params["sandbox_provider"] == "e2b" + + sandbox_tools.clear_sandbox_tools() + + stdout = await logger._run_tool_call( + container=container, params=params, arguments='{"code":"print(1)"}' + ) + finally: + sandbox_tools.clear_sandbox_tools() + + assert stdout == "ok", "run must succeed using the params captured at create time" + assert run_kwargs["provider"] == "e2b" + + +@pytest.mark.asyncio +async def test_run_tool_call_surfaces_execution_error(): + """A sandbox execution error must be fed back to the model as a labelled + string, not raised, so the agentic loop can react to it.""" + + 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"} + ) + + sandbox = ErroringSandbox() + logger = CodeInterpreterInterceptionLogger(sandbox_config=sandbox) + container = await logger._create_container() + + stdout = await logger._run_tool_call( + container=container[0], params=None, arguments='{"code":"raise ValueError(1)"}' + ) + + assert stdout == "[execution error] boom" + + +@pytest.mark.asyncio +async def test_run_tool_call_reports_unparseable_arguments(): + """Malformed tool arguments must produce a parse error string the model can + see rather than crashing the interceptor.""" + sandbox = FakeSandbox() + logger = CodeInterpreterInterceptionLogger(sandbox_config=sandbox) + container = await logger._create_container() + + 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" + + +@pytest.mark.asyncio +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"] + ) + kwargs = { + "tools": [{"type": "code_interpreter", "container": {"type": "auto"}}], + "custom_llm_provider": "anthropic", + } + + result = await logger.async_pre_call_deployment_hook(kwargs, CallTypes.aresponses) + + assert result is None + assert kwargs["tools"][0]["type"] == "code_interpreter", "tool must be untouched" + assert _ACTIVE_KEY not in kwargs + + +@pytest.mark.asyncio +async def test_resolve_provider_falls_back_to_model_lookup(): + logger = CodeInterpreterInterceptionLogger(sandbox_config=FakeSandbox()) + + assert logger._resolve_provider({"custom_llm_provider": "openai"}) == "openai" + assert logger._resolve_provider({"model": "gpt-5"}) == "openai" + assert logger._resolve_provider({"model": 123}) is None + assert logger._resolve_provider({"model": "no-such-provider-xyz"}) is None + + +@pytest.mark.asyncio +async def test_create_container_without_sandbox_raises(): + """The registry path must raise a clear error when no sandbox is resolvable + instead of silently creating nothing.""" + logger = CodeInterpreterInterceptionLogger(sandbox_tool_name="missing") + + with pytest.raises(ValueError, match="no sandbox available"): + await logger._create_container() + + +@pytest.mark.asyncio +async def test_run_code_without_params_raises(): + logger = CodeInterpreterInterceptionLogger(sandbox_tool_name="missing") + + with pytest.raises(ValueError, match="no sandbox available to run code"): + await logger._run_code(container=FakeHandle(), params=None, code="print(1)") + + +@pytest.mark.asyncio +async def test_delete_container_swallows_errors(): + """A delete failure must not propagate; the request already succeeded.""" + + class FailingDeleteSandbox(FakeSandbox): + async def adelete_sandbox(self, *, container, **kwargs): + raise RuntimeError("e2b unreachable") + + sandbox = FailingDeleteSandbox() + logger = CodeInterpreterInterceptionLogger(sandbox_config=sandbox) + container, params = await logger._create_container() + + await logger._delete_container(container=container, params=params) + + +@pytest.mark.asyncio +async def test_prune_expired_cache_deletes_underlying_container(): + """Expired cache entries must have their sandbox deleted, not just dropped, + otherwise an orphaned sandbox keeps running.""" + import litellm.integrations.code_interpreter_interception.handler as handler_mod + + sandbox = FakeSandbox() + logger = CodeInterpreterInterceptionLogger(sandbox_config=sandbox) + container, params = await logger._create_container() + logger._container_cache["old"] = ( + container, + params, + time.time() - handler_mod._CACHE_TTL_SECONDS - 1, + ) + + await logger._prune_expired_cache() + + assert "old" not in logger._container_cache + assert len(sandbox.delete_calls) == 1, "expired sandbox must be deleted" + + +@pytest.mark.asyncio +async def test_normalize_messages_handles_str_and_unknown(): + logger = CodeInterpreterInterceptionLogger(sandbox_config=FakeSandbox()) + + assert logger._normalize_messages("hi") == [{"role": "user", "content": "hi"}] + assert logger._normalize_messages([{"role": "user"}]) == [{"role": "user"}] + assert logger._normalize_messages(42) == [] + + +def test_from_config_yaml_reads_fields(): + cfg = { + "enabled": False, + "enabled_providers": ["openai"], + "sandbox_tool_name": "e2b_default", + } + logger = CodeInterpreterInterceptionLogger.from_config_yaml(cfg) + + assert logger.enabled is False + assert logger.enabled_providers == ["openai"] + assert logger.sandbox_tool_name == "e2b_default" + + +def test_initialize_from_proxy_config_prefers_litellm_settings(): + logger = CodeInterpreterInterceptionLogger.initialize_from_proxy_config( + litellm_settings={ + "code_interpreter_interception_params": { + "enabled_providers": ["openai"], + "sandbox_tool_name": "e2b_default", + } + }, + callback_specific_params={}, + ) + + assert logger.enabled_providers == ["openai"] + assert logger.sandbox_tool_name == "e2b_default" + + +@pytest.mark.asyncio +async def test_build_plan_handles_dict_shaped_response(): + """A responses payload delivered as a plain dict (not an object) must flow + through detection, execution, and re-injection the same as the typed form.""" + sandbox = FakeSandbox(stdout="42") + logger = CodeInterpreterInterceptionLogger(sandbox_config=sandbox) + dict_response = {"output": [_function_call_item()]} + + plan = await logger.async_build_agentic_loop_plan( + tools={"tool_calls": logger._extract_code_execution_tool_calls(dict_response)}, + model="gpt-5", + messages=[{"role": "user", "content": "x"}], + response=dict_response, + 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: "sbxkey1"}, + ) + + assert sandbox.run_calls, "code must run for a dict-shaped response" + assert plan.metadata["code_interpreter_calls"][0]["code"] == "print(40 + 2)" + + out = await logger.async_post_agentic_loop_response_hook( + response={"output": [{"type": "message", "content": []}]}, + plan=plan, + kwargs={}, + ) + + assert [item.get("type") for item in out["output"]] == [ + "code_interpreter_call", + "message", + ], "the dict-shaped response must get the code_interpreter_call re-injected" + + +@pytest.mark.asyncio +async def test_extract_tool_calls_reads_object_attributes(): + """Detection must work when output items are objects with attributes, not + only dicts.""" + + class Item: + def __init__(self): + self.type = "function_call" + self.name = LITELLM_CODE_EXECUTION_TOOL_NAME + self.call_id = "c9" + self.arguments = '{"code":"print(1)"}' + + logger = CodeInterpreterInterceptionLogger(sandbox_config=FakeSandbox()) + calls = logger._extract_code_execution_tool_calls(FakeResponse(output=[Item()])) + + assert len(calls) == 1 + assert calls[0]["call_id"] == "c9" + assert calls[0]["arguments"] == '{"code":"print(1)"}' diff --git a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py index 6b692180559..ec262d75ab8 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -718,7 +718,18 @@ async def test_add_litellm_data_to_request_strips_user_control_fields(): @pytest.mark.asyncio @pytest.mark.parametrize( "control_field", - ["callbacks", "service_callback", "logger_fn", "litellm_disabled_callbacks"], + [ + "callbacks", + "service_callback", + "logger_fn", + "litellm_disabled_callbacks", + "_agentic_loop_depth", + "_agentic_loop_fingerprints", + "_code_interpreter_interception_active", + "_code_interpreter_interception_converted_stream", + "_code_interpreter_interception_sandbox_key", + "max_agentic_loops", + ], ) async def test_add_litellm_data_to_request_strips_callback_control_fields( control_field, @@ -741,12 +752,19 @@ async def test_add_litellm_data_to_request_strips_callback_control_fields( request_mock.client = MagicMock() request_mock.client.host = "127.0.0.1" - sample_value = ( - ["langfuse"] - if control_field - in ("callbacks", "service_callback", "litellm_disabled_callbacks") - else "module.func" - ) + sample_values = { + "callbacks": ["langfuse"], + "service_callback": ["langfuse"], + "litellm_disabled_callbacks": ["langfuse"], + "logger_fn": "module.func", + "_agentic_loop_depth": 5, + "_agentic_loop_fingerprints": ["forged"], + "_code_interpreter_interception_active": True, + "_code_interpreter_interception_converted_stream": True, + "_code_interpreter_interception_sandbox_key": "forged-key", + "max_agentic_loops": 9999, + } + sample_value = sample_values[control_field] updated = await add_litellm_data_to_request( data={ @@ -4150,7 +4168,9 @@ async def test_string_metadata_does_not_bypass_tag_max_budget_check(self): litellm_budget_table=LiteLLM_BudgetTable(max_budget=0.10), ) - async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): + async def mock_get_current_spend( + counter_key, fallback_spend, max_budget=None, **kwargs + ): if counter_key == "spend:tag:paid": return 0.50 return fallback_spend @@ -4207,7 +4227,9 @@ async def test_header_tags_visible_to_tag_max_budget_check(self): litellm_budget_table=LiteLLM_BudgetTable(max_budget=0.10), ) - async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): + async def mock_get_current_spend( + counter_key, fallback_spend, max_budget=None, **kwargs + ): if counter_key == "spend:tag:tenant:acme": return 0.50 return fallback_spend @@ -4362,7 +4384,9 @@ async def test_key_tags_visible_to_tag_max_budget_check(self): litellm_budget_table=LiteLLM_BudgetTable(max_budget=0.10), ) - async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): + async def mock_get_current_spend( + counter_key, fallback_spend, max_budget=None, **kwargs + ): if counter_key == "spend:tag:engineering": return 0.50 return fallback_spend @@ -4413,7 +4437,9 @@ async def test_key_tags_within_budget_passes_check(self): litellm_budget_table=LiteLLM_BudgetTable(max_budget=0.10), ) - async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): + async def mock_get_current_spend( + counter_key, fallback_spend, max_budget=None, **kwargs + ): if counter_key == "spend:tag:engineering": return 0.05 return fallback_spend diff --git a/tests/test_litellm/sandbox/test_e2b_sandbox.py b/tests/test_litellm/sandbox/test_e2b_sandbox.py index 948d39fac8f..cc5b12156a1 100644 --- a/tests/test_litellm/sandbox/test_e2b_sandbox.py +++ b/tests/test_litellm/sandbox/test_e2b_sandbox.py @@ -295,3 +295,24 @@ async def test_public_lifecycle_create_run_delete(): async def test_unsupported_provider_raises(): with pytest.raises(ValueError): await litellm.acreate_sandbox(provider="not-a-provider") + + +# ---------- api_base override ---------- + + +@pytest.mark.asyncio +async def test_create_uses_api_base_override(): + client = FakeHTTPClient() + await E2BSandboxConfig().acreate_sandbox( + api_base="http://my-sandbox:8080", api_key="k", client=client + ) + _, url, _, _ = client.calls[0] + assert url == "http://my-sandbox:8080/sandboxes" + + +@pytest.mark.asyncio +async def test_create_defaults_to_e2b_api_base(): + client = FakeHTTPClient() + await E2BSandboxConfig().acreate_sandbox(api_key="k", client=client) + _, url, _, _ = client.calls[0] + assert url == "https://api.e2b.app/sandboxes" diff --git a/tests/test_litellm/sandbox/test_sandbox_tools.py b/tests/test_litellm/sandbox/test_sandbox_tools.py new file mode 100644 index 00000000000..06136534b13 --- /dev/null +++ b/tests/test_litellm/sandbox/test_sandbox_tools.py @@ -0,0 +1,181 @@ +"""Unit tests for the sandbox-tool registry.""" + +from litellm.sandbox import sandbox_tools + + +def _reset(): + sandbox_tools.clear_sandbox_tools() + + +def test_register_resolves_provider_key_and_base(): + _reset() + sandbox_tools.register_sandbox_tools( + [ + { + "sandbox_tool_name": "e2b_default", + "litellm_params": { + "sandbox_provider": "e2b", + "api_key": "sk-literal", + "api_base": "https://sandbox.internal", + }, + } + ] + ) + + resolved = sandbox_tools.resolve_sandbox_tool("e2b_default") + assert resolved == { + "sandbox_provider": "e2b", + "api_key": "sk-literal", + "api_base": "https://sandbox.internal", + } + _reset() + + +def test_register_clears_stale_entries_on_reload(): + """A tool removed from the config must not survive a re-registration.""" + _reset() + sandbox_tools.register_sandbox_tools( + [ + { + "sandbox_tool_name": "old", + "litellm_params": {"sandbox_provider": "e2b"}, + } + ] + ) + assert sandbox_tools.resolve_sandbox_tool("old") is not None + + sandbox_tools.register_sandbox_tools( + [ + { + "sandbox_tool_name": "new", + "litellm_params": {"sandbox_provider": "e2b"}, + } + ] + ) + + assert sandbox_tools.resolve_sandbox_tool("new") is not None + assert ( + sandbox_tools.resolve_sandbox_tool("old") is None + ), "stale tool must be gone after the config is reloaded" + _reset() + + +def test_register_empty_list_clears_removed_tools(): + """Reloading a config with sandbox_tools removed (the proxy passes an empty + list) must drop previously registered credentials from the process.""" + _reset() + sandbox_tools.register_sandbox_tools( + [ + { + "sandbox_tool_name": "e2b_default", + "litellm_params": {"sandbox_provider": "e2b", "api_key": "sk-x"}, + } + ] + ) + assert sandbox_tools.resolve_sandbox_tool("e2b_default") is not None + + sandbox_tools.register_sandbox_tools([]) + + assert ( + sandbox_tools.resolve_sandbox_tool("e2b_default") is None + ), "removing sandbox_tools from config must clear stale credentials" + _reset() + + +def test_register_resolves_secret_from_env(monkeypatch): + _reset() + monkeypatch.setenv("MY_SANDBOX_KEY", "sk-from-env") + sandbox_tools.register_sandbox_tools( + [ + { + "sandbox_tool_name": "e2b_default", + "litellm_params": { + "sandbox_provider": "e2b", + "api_key": "os.environ/MY_SANDBOX_KEY", + }, + } + ] + ) + + resolved = sandbox_tools.resolve_sandbox_tool("e2b_default") + assert resolved is not None + assert resolved["api_key"] == "sk-from-env" + assert resolved["api_base"] is None + _reset() + + +def test_resolve_unknown_returns_none(): + _reset() + assert sandbox_tools.resolve_sandbox_tool("nope") is None + + +def test_register_skips_malformed_entries_without_crashing(): + """A single malformed entry (missing sandbox_tool_name, or not a dict) must + not crash registration during proxy startup/hot-reload; valid entries in the + same list must still register.""" + _reset() + sandbox_tools.register_sandbox_tools( + [ + {"litellm_params": {"sandbox_provider": "e2b"}}, # missing name + "not-a-dict", # wrong type + {"sandbox_tool_name": "", "litellm_params": {}}, # empty name + { + "sandbox_tool_name": "good", + "litellm_params": {"sandbox_provider": "e2b"}, + }, + ] + ) + + assert sandbox_tools.resolve_sandbox_tool("good") is not None + assert sandbox_tools.resolve_sandbox_tool("") is None + assert set(sandbox_tools._SANDBOX_TOOL_REGISTRY) == {"good"} + _reset() + + +def test_register_skips_entry_missing_sandbox_provider(): + """An entry with a name but no sandbox_provider must be skipped at + registration so it cannot later resolve and call acreate_sandbox(provider=None), + which fails with a cryptic runtime error instead of a clear startup warning.""" + _reset() + sandbox_tools.register_sandbox_tools( + [ + {"sandbox_tool_name": "no_provider", "litellm_params": {"api_key": "sk-x"}}, + { + "sandbox_tool_name": "null_provider", + "litellm_params": {"sandbox_provider": None}, + }, + { + "sandbox_tool_name": "good", + "litellm_params": {"sandbox_provider": "e2b"}, + }, + ] + ) + + assert sandbox_tools.resolve_sandbox_tool("no_provider") is None + assert sandbox_tools.resolve_sandbox_tool("null_provider") is None + assert set(sandbox_tools._SANDBOX_TOOL_REGISTRY) == {"good"} + _reset() + + +def test_register_swaps_registry_atomically(): + """register_sandbox_tools must replace the registry in one rebind so a + concurrent resolve never observes a half-populated or transiently empty + registry between clearing and repopulating.""" + _reset() + sandbox_tools.register_sandbox_tools( + [{"sandbox_tool_name": "a", "litellm_params": {"sandbox_provider": "e2b"}}] + ) + before = sandbox_tools._SANDBOX_TOOL_REGISTRY + + sandbox_tools.register_sandbox_tools( + [ + {"sandbox_tool_name": "b", "litellm_params": {"sandbox_provider": "e2b"}}, + {"sandbox_tool_name": "c", "litellm_params": {"sandbox_provider": "e2b"}}, + ] + ) + after = sandbox_tools._SANDBOX_TOOL_REGISTRY + + assert after is not before, "the registry must be replaced, not mutated in place" + assert set(after) == {"b", "c"} + assert "a" not in after + _reset()