diff --git a/letta/schemas/agent.py b/letta/schemas/agent.py index 57268645..cbd9e204 100644 --- a/letta/schemas/agent.py +++ b/letta/schemas/agent.py @@ -83,6 +83,13 @@ class AgentState(OrmMetadataBase, validate_assignment=True): ..., description="The environment variables for tool execution specific to this agent." ) + def get_agent_env_vars_as_dict(self) -> Dict[str, str]: + # Get environment variables for this agent specifically + per_agent_env_vars = {} + for agent_env_var_obj in self.tool_exec_environment_variables: + per_agent_env_vars[agent_env_var_obj.key] = agent_env_var_obj.value + return per_agent_env_vars + class CreateAgent(BaseModel, validate_assignment=True): # # all optional as server can generate defaults diff --git a/letta/services/tool_execution_sandbox.py b/letta/services/tool_execution_sandbox.py index d0a6f406..9eca9fcb 100644 --- a/letta/services/tool_execution_sandbox.py +++ b/letta/services/tool_execution_sandbox.py @@ -98,7 +98,7 @@ class ToolExecutionSandbox: os.environ.clear() os.environ.update(original_env) # Restore original environment variables - def run_local_dir_sandbox(self, agent_state: AgentState) -> SandboxRunResult: + def run_local_dir_sandbox(self, agent_state: Optional[AgentState] = None) -> SandboxRunResult: sbx_config = self.sandbox_config_manager.get_or_create_default_sandbox_config(sandbox_type=SandboxType.LOCAL, actor=self.user) local_configs = sbx_config.get_local_config() @@ -107,6 +107,10 @@ class ToolExecutionSandbox: env = os.environ.copy() env.update(env_vars) + # Get environment variables for this agent specifically + if agent_state: + env.update(agent_state.get_agent_env_vars_as_dict()) + # Safety checks if not os.path.isdir(local_configs.sandbox_dir): raise FileNotFoundError(f"Sandbox directory does not exist: {local_configs.sandbox_dir}") @@ -273,7 +277,7 @@ class ToolExecutionSandbox: # e2b sandbox specific functions - def run_e2b_sandbox(self, agent_state: AgentState) -> SandboxRunResult: + def run_e2b_sandbox(self, agent_state: Optional[AgentState] = None) -> SandboxRunResult: sbx_config = self.sandbox_config_manager.get_or_create_default_sandbox_config(sandbox_type=SandboxType.E2B, actor=self.user) sbx = self.get_running_e2b_sandbox_with_same_state(sbx_config) if not sbx or self.force_recreate: @@ -292,6 +296,10 @@ class ToolExecutionSandbox: # Get environment variables for the sandbox # TODO: We set limit to 100 here, but maybe we want it uncapped? Realistically this should be fine. env_vars = self.sandbox_config_manager.get_sandbox_env_vars_as_dict(sandbox_config_id=sbx_config.id, actor=self.user, limit=100) + # Get environment variables for this agent specifically + if agent_state: + env_vars.update(agent_state.get_agent_env_vars_as_dict()) + code = self.generate_execution_script(agent_state=agent_state) execution = sbx.run_code(code, envs=env_vars) diff --git a/tests/integration_test_tool_execution_sandbox.py b/tests/integration_test_tool_execution_sandbox.py index 4a497d8e..04785c64 100644 --- a/tests/integration_test_tool_execution_sandbox.py +++ b/tests/integration_test_tool_execution_sandbox.py @@ -12,7 +12,7 @@ from letta.functions.function_sets.base import core_memory_append, core_memory_r from letta.orm.sandbox_config import SandboxConfig, SandboxEnvironmentVariable from letta.schemas.agent import AgentState from letta.schemas.embedding_config import EmbeddingConfig -from letta.schemas.environment_variables import SandboxEnvironmentVariableCreate +from letta.schemas.environment_variables import AgentEnvironmentVariable, SandboxEnvironmentVariableCreate from letta.schemas.llm_config import LLMConfig from letta.schemas.memory import ChatMemory from letta.schemas.organization import Organization @@ -323,6 +323,41 @@ def test_local_sandbox_env(mock_e2b_api_key_none, get_env_tool, test_user): assert long_random_string in result.func_return +@pytest.mark.local_sandbox +def test_local_sandbox_per_agent_env(mock_e2b_api_key_none, get_env_tool, agent_state, test_user): + manager = SandboxConfigManager(tool_settings) + key = "secret_word" + + # Make a custom local sandbox config + sandbox_dir = str(Path(__file__).parent / "test_tool_sandbox") + config_create = SandboxConfigCreate(config=LocalSandboxConfig(sandbox_dir=sandbox_dir).model_dump()) + config = manager.create_or_update_sandbox_config(config_create, test_user) + + # Make a environment variable with a long random string + # Note: This has an overlapping key with agent state's environment variables + # We expect that the agent's env var supersedes this + wrong_long_random_string = "".join(secrets.choice(string.ascii_letters + string.digits) for _ in range(20)) + manager.create_sandbox_env_var( + SandboxEnvironmentVariableCreate(key=key, value=wrong_long_random_string), sandbox_config_id=config.id, actor=test_user + ) + + # Make a environment variable with a long random string and put into agent state + correct_long_random_string = "".join(secrets.choice(string.ascii_letters + string.digits) for _ in range(20)) + agent_state.tool_exec_environment_variables = [ + AgentEnvironmentVariable(key=key, value=correct_long_random_string, agent_id=agent_state.id) + ] + + # Create tool and args + args = {} + + # Run the custom sandbox + sandbox = ToolExecutionSandbox(get_env_tool.name, args, user=test_user) + result = sandbox.run(agent_state=agent_state) + + assert wrong_long_random_string not in result.func_return + assert correct_long_random_string in result.func_return + + @pytest.mark.local_sandbox def test_local_sandbox_e2e_composio_star_github(mock_e2b_api_key_none, check_composio_key_set, composio_github_star_tool, test_user): # Add the composio key @@ -470,6 +505,42 @@ def test_e2b_sandbox_inject_env_var_existing_sandbox(check_e2b_key_is_set, get_e assert long_random_string in result.func_return +# TODO: There is a near dupe of this test above for local sandbox - we should try to make it parameterized tests to minimize code bloat +@pytest.mark.e2b_sandbox +def test_e2b_sandbox_per_agent_env(check_e2b_key_is_set, get_env_tool, agent_state, test_user): + manager = SandboxConfigManager(tool_settings) + key = "secret_word" + + # Make a custom local sandbox config + sandbox_dir = str(Path(__file__).parent / "test_tool_sandbox") + config_create = SandboxConfigCreate(config=LocalSandboxConfig(sandbox_dir=sandbox_dir).model_dump()) + config = manager.create_or_update_sandbox_config(config_create, test_user) + + # Make a environment variable with a long random string + # Note: This has an overlapping key with agent state's environment variables + # We expect that the agent's env var supersedes this + wrong_long_random_string = "".join(secrets.choice(string.ascii_letters + string.digits) for _ in range(20)) + manager.create_sandbox_env_var( + SandboxEnvironmentVariableCreate(key=key, value=wrong_long_random_string), sandbox_config_id=config.id, actor=test_user + ) + + # Make a environment variable with a long random string and put into agent state + correct_long_random_string = "".join(secrets.choice(string.ascii_letters + string.digits) for _ in range(20)) + agent_state.tool_exec_environment_variables = [ + AgentEnvironmentVariable(key=key, value=correct_long_random_string, agent_id=agent_state.id) + ] + + # Create tool and args + args = {} + + # Run the custom sandbox + sandbox = ToolExecutionSandbox(get_env_tool.name, args, user=test_user) + result = sandbox.run(agent_state=agent_state) + + assert wrong_long_random_string not in result.func_return + assert correct_long_random_string in result.func_return + + @pytest.mark.e2b_sandbox def test_e2b_sandbox_config_change_force_recreates_sandbox(check_e2b_key_is_set, list_tool, test_user): manager = SandboxConfigManager(tool_settings) @@ -506,7 +577,7 @@ def test_e2b_sandbox_with_list_rv(check_e2b_key_is_set, list_tool, test_user): assert len(result.func_return) == 5 -@pytest.mark.e2b_sandboxfunc +@pytest.mark.e2b_sandbox def test_e2b_e2e_composio_star_github(check_e2b_key_is_set, check_composio_key_set, composio_github_star_tool, test_user): # Add the composio key manager = SandboxConfigManager(tool_settings)