fix: global sandbox environment variables not injected on tool execution (#8310)
* base * test
This commit is contained in:
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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"
|
||||
|
||||
Reference in New Issue
Block a user