fix: global sandbox environment variables not injected on tool execution (#8310)

* base

* test
This commit is contained in:
jnjpng
2026-01-05 17:57:39 -08:00
committed by Caren Thomas
parent e57adc0a6e
commit e56c5c5b49
5 changed files with 152 additions and 60 deletions

View File

@@ -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

View File

@@ -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)

View File

@@ -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)

View File

@@ -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)

View File

@@ -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"