From e56c5c5b4955cfcdcd4dc19e2fe9bc01eb03338b Mon Sep 17 00:00:00 2001 From: jnjpng Date: Mon, 5 Jan 2026 17:57:39 -0800 Subject: [PATCH] fix: global sandbox environment variables not injected on tool execution (#8310) * base * test --- letta/services/sandbox_config_manager.py | 24 ++++-- letta/services/tool_sandbox/base.py | 28 ++++-- letta/services/tool_sandbox/local_sandbox.py | 20 +---- letta/services/tool_sandbox/modal_sandbox.py | 51 +++++------ tests/managers/test_sandbox_manager.py | 89 ++++++++++++++++++++ 5 files changed, 152 insertions(+), 60 deletions(-) diff --git a/letta/services/sandbox_config_manager.py b/letta/services/sandbox_config_manager.py index cf26608e..365ae612 100644 --- a/letta/services/sandbox_config_manager.py +++ b/letta/services/sandbox_config_manager.py @@ -12,7 +12,9 @@ from letta.schemas.environment_variables import ( SandboxEnvironmentVariableUpdate, ) from letta.schemas.sandbox_config import ( + E2BSandboxConfig, LocalSandboxConfig, + ModalSandboxConfig, SandboxConfig as PydanticSandboxConfig, SandboxConfigCreate, SandboxConfigUpdate, @@ -35,13 +37,16 @@ class SandboxConfigManager: if not sandbox_config: logger.debug(f"Creating new sandbox config of type {sandbox_type}, none found for organization {actor.organization_id}.") - # TODO: Add more sandbox types later + # Create the appropriate config type based on sandbox_type + # Using the actual model classes ensures Pydantic's Union type resolution works correctly if sandbox_type == SandboxType.E2B: - default_config = {} # Empty + default_config = E2BSandboxConfig() + elif sandbox_type == SandboxType.MODAL: + default_config = ModalSandboxConfig() else: - # TODO: May want to move this to environment variables v.s. persisting in database + # LOCAL sandbox type default_local_sandbox_path = LETTA_TOOL_EXECUTION_DIR - default_config = LocalSandboxConfig(sandbox_dir=default_local_sandbox_path).model_dump(exclude_none=True) + default_config = LocalSandboxConfig(sandbox_dir=default_local_sandbox_path) sandbox_config = self.create_or_update_sandbox_config(SandboxConfigCreate(config=default_config), actor=actor) return sandbox_config @@ -53,13 +58,16 @@ class SandboxConfigManager: if not sandbox_config: logger.debug(f"Creating new sandbox config of type {sandbox_type}, none found for organization {actor.organization_id}.") - # TODO: Add more sandbox types later + # Create the appropriate config type based on sandbox_type + # Using the actual model classes ensures Pydantic's Union type resolution works correctly if sandbox_type == SandboxType.E2B: - default_config = {} # Empty + default_config = E2BSandboxConfig() + elif sandbox_type == SandboxType.MODAL: + default_config = ModalSandboxConfig() else: - # TODO: May want to move this to environment variables v.s. persisting in database + # LOCAL sandbox type default_local_sandbox_path = LETTA_TOOL_EXECUTION_DIR - default_config = LocalSandboxConfig(sandbox_dir=default_local_sandbox_path).model_dump(exclude_none=True) + default_config = LocalSandboxConfig(sandbox_dir=default_local_sandbox_path) sandbox_config = await self.create_or_update_sandbox_config_async(SandboxConfigCreate(config=default_config), actor=actor) return sandbox_config diff --git a/letta/services/tool_sandbox/base.py b/letta/services/tool_sandbox/base.py index 12f14df0..31c02bda 100644 --- a/letta/services/tool_sandbox/base.py +++ b/letta/services/tool_sandbox/base.py @@ -388,19 +388,35 @@ class AsyncToolSandboxBase(ABC): """ return False # Default to False for local execution - async def _gather_env_vars(self, agent_state: AgentState | None, additional_env_vars: dict[str, str], sbx_id: str, is_local: bool): + async def _gather_env_vars( + self, agent_state: AgentState | None, additional_env_vars: dict[str, str] | None, sbx_id: str, is_local: bool + ): + """ + Gather environment variables with proper layering: + 1. OS environment (for local sandboxes only) + 2. Global sandbox env vars from DB (always included) + 3. Provided sandbox env vars (agent-scoped, override global on key collision) + 4. Agent state env vars + 5. Additional runtime env vars (highest priority) + """ env = os.environ.copy() if is_local else {} + + # Always fetch and include global sandbox env vars from DB + global_env_vars = await self.sandbox_config_manager.get_sandbox_env_vars_as_dict_async( + sandbox_config_id=sbx_id, actor=self.user, limit=None + ) + env.update(global_env_vars) + + # Override with provided sandbox env vars if self.provided_sandbox_env_vars: env.update(self.provided_sandbox_env_vars) - else: - env_vars = await self.sandbox_config_manager.get_sandbox_env_vars_as_dict_async( - sandbox_config_id=sbx_id, actor=self.user, limit=None - ) - env.update(env_vars) + # TOOD: may be duplicative with provided sandbox env vars above + # Override with agent state env vars if agent_state: env.update(agent_state.get_agent_env_vars_as_dict()) + # Override with additional runtime env vars (highest priority) if additional_env_vars: env.update(additional_env_vars) diff --git a/letta/services/tool_sandbox/local_sandbox.py b/letta/services/tool_sandbox/local_sandbox.py index 3e88f052..0613c3f5 100644 --- a/letta/services/tool_sandbox/local_sandbox.py +++ b/letta/services/tool_sandbox/local_sandbox.py @@ -77,24 +77,8 @@ class AsyncToolSandboxLocal(AsyncToolSandboxBase): local_configs = sbx_config.get_local_config() use_venv = local_configs.use_venv - # Prepare environment variables - env = os.environ.copy() - if self.provided_sandbox_env_vars: - env.update(self.provided_sandbox_env_vars) - else: - env_vars = await self.sandbox_config_manager.get_sandbox_env_vars_as_dict_async( - sandbox_config_id=sbx_config.id, actor=self.user, limit=100 - ) - env.update(env_vars) - - if agent_state: - env.update(agent_state.get_agent_env_vars_as_dict()) - - if additional_env_vars: - env.update(additional_env_vars) - - # Filter out None values to prevent subprocess errors - env = {k: v for k, v in env.items() if v is not None} + # Prepare environment variables using standardized gathering logic + env = await self._gather_env_vars(agent_state, additional_env_vars, sbx_config.id, is_local=True) # Make sure sandbox directory exists sandbox_dir = os.path.expanduser(local_configs.sandbox_dir) diff --git a/letta/services/tool_sandbox/modal_sandbox.py b/letta/services/tool_sandbox/modal_sandbox.py index 34a403ae..4fd1cd6e 100644 --- a/letta/services/tool_sandbox/modal_sandbox.py +++ b/letta/services/tool_sandbox/modal_sandbox.py @@ -115,40 +115,35 @@ class AsyncToolSandboxModal(AsyncToolSandboxBase): else: letta_api_key = additional_env_vars.get("LETTA_SECRET_API_KEY", None) - # Construct dynamic env vars - # Priority order (later overrides earlier): - # 1. Sandbox-level env vars (from database) - # 2. Agent-specific env vars - # 3. Additional runtime env vars + # Construct dynamic env vars with proper layering: + # 1. Global sandbox env vars from DB (always included) + # 2. Provided sandbox env vars (agent-scoped, override global on key collision) + # 3. Agent-specific env vars from secrets + # 4. Additional runtime env vars (highest priority) env_vars = {} - # Load sandbox-level environment variables from the database - # These can be updated after deployment and will be available at runtime + # Always load global sandbox-level environment variables from the database + try: + sandbox_config = await self.sandbox_config_manager.get_or_create_default_sandbox_config_async( + sandbox_type=SandboxType.MODAL, actor=self.user + ) + if sandbox_config: + global_env_vars = await self.sandbox_config_manager.get_sandbox_env_vars_as_dict_async( + sandbox_config_id=sandbox_config.id, actor=self.user, limit=None + ) + env_vars.update(global_env_vars) + except Exception as e: + logger.warning(f"Could not load global sandbox env vars for tool {self.tool_name}: {e}") + + # Override with provided sandbox env vars (agent-scoped) if self.provided_sandbox_env_vars: env_vars.update(self.provided_sandbox_env_vars) - else: - try: - from letta.services.sandbox_config_manager import SandboxConfigManager - sandbox_config_manager = SandboxConfigManager() - sandbox_config = await sandbox_config_manager.get_or_create_default_sandbox_config_async( - sandbox_type=SandboxType.MODAL, actor=self.user - ) - if sandbox_config: - sandbox_env_vars = await sandbox_config_manager.get_sandbox_env_vars_as_dict_async( - sandbox_config_id=sandbox_config.id, actor=self.user, limit=None - ) - env_vars.update(sandbox_env_vars) - except Exception as e: - logger.warning(f"Could not load sandbox env vars for tool {self.tool_name}: {e}") + # Override with agent-specific environment variables from secrets + if agent_state: + env_vars.update(agent_state.get_agent_env_vars_as_dict()) - # Add agent-specific environment variables (these override sandbox-level) - # Use the pre-decrypted value field which was populated in from_orm_async() - if agent_state and agent_state.secrets: - for secret in agent_state.secrets: - env_vars[secret.key] = secret.value or "" - - # Add any additional env vars passed at runtime (highest priority) + # Override with additional env vars passed at runtime (highest priority) if additional_env_vars: env_vars.update(additional_env_vars) diff --git a/tests/managers/test_sandbox_manager.py b/tests/managers/test_sandbox_manager.py index be1dc56b..e52a7d78 100644 --- a/tests/managers/test_sandbox_manager.py +++ b/tests/managers/test_sandbox_manager.py @@ -299,3 +299,92 @@ async def test_get_sandbox_env_var_by_key(server: SyncServer, sandbox_env_var_fi # Assertions to verify correct retrieval assert retrieved_env_var.id == sandbox_env_var_fixture.id + + +@pytest.mark.asyncio +async def test_gather_env_vars_layering(server: SyncServer, sandbox_config_fixture, default_user): + """Test that _gather_env_vars properly layers env vars with correct priority. + + Priority order (later overrides earlier): + 1. Global sandbox env vars from DB (always included) + 2. Provided sandbox env vars (agent-scoped, override global on key collision) + 3. Agent state env vars + 4. Additional runtime env vars (highest priority) + """ + from unittest.mock import MagicMock + + from letta.services.tool_sandbox.local_sandbox import AsyncToolSandboxLocal + + # Create global sandbox env vars in the database + global_var1 = SandboxEnvironmentVariableCreate(key="GLOBAL_ONLY", value="global_value") + global_var2 = SandboxEnvironmentVariableCreate(key="OVERRIDE_BY_PROVIDED", value="global_will_be_overridden") + global_var3 = SandboxEnvironmentVariableCreate(key="OVERRIDE_BY_AGENT", value="global_will_be_overridden_by_agent") + global_var4 = SandboxEnvironmentVariableCreate(key="OVERRIDE_BY_ADDITIONAL", value="global_will_be_overridden_by_additional") + + await server.sandbox_config_manager.create_sandbox_env_var_async( + global_var1, sandbox_config_id=sandbox_config_fixture.id, actor=default_user + ) + await server.sandbox_config_manager.create_sandbox_env_var_async( + global_var2, sandbox_config_id=sandbox_config_fixture.id, actor=default_user + ) + await server.sandbox_config_manager.create_sandbox_env_var_async( + global_var3, sandbox_config_id=sandbox_config_fixture.id, actor=default_user + ) + await server.sandbox_config_manager.create_sandbox_env_var_async( + global_var4, sandbox_config_id=sandbox_config_fixture.id, actor=default_user + ) + + # Define provided sandbox env vars (agent-scoped) + provided_env_vars = { + "OVERRIDE_BY_PROVIDED": "provided_value", + "PROVIDED_ONLY": "provided_only_value", + } + + # Create a mock agent state with secrets + mock_agent_state = MagicMock() + mock_agent_state.get_agent_env_vars_as_dict.return_value = { + "OVERRIDE_BY_AGENT": "agent_value", + "AGENT_ONLY": "agent_only_value", + } + + # Define additional runtime env vars + additional_env_vars = { + "OVERRIDE_BY_ADDITIONAL": "additional_value", + "ADDITIONAL_ONLY": "additional_only_value", + } + + # Create a minimal sandbox instance to test _gather_env_vars + sandbox = AsyncToolSandboxLocal( + tool_name="test_tool", + args={}, + user=default_user, + tool_id="test-tool-id", + sandbox_env_vars=provided_env_vars, + ) + + # Call _gather_env_vars + result = await sandbox._gather_env_vars( + agent_state=mock_agent_state, + additional_env_vars=additional_env_vars, + sbx_id=sandbox_config_fixture.id, + is_local=False, # Use False to avoid copying os.environ + ) + + # Verify layering: + # 1. Global vars included + assert result["GLOBAL_ONLY"] == "global_value" + + # 2. Provided vars override global + assert result["OVERRIDE_BY_PROVIDED"] == "provided_value" + assert result["PROVIDED_ONLY"] == "provided_only_value" + + # 3. Agent vars override provided/global + assert result["OVERRIDE_BY_AGENT"] == "agent_value" + assert result["AGENT_ONLY"] == "agent_only_value" + + # 4. Additional vars have highest priority + assert result["OVERRIDE_BY_ADDITIONAL"] == "additional_value" + assert result["ADDITIONAL_ONLY"] == "additional_only_value" + + # Verify LETTA IDs are injected + assert result["LETTA_TOOL_ID"] == "test-tool-id"