feat: Pass in entity_id explicitly to auto-gen composio tool (#535)

tests/integration_test_tool_execution_sandbox.py

@@ -8,6 +8,7 @@ import pytest
 from sqlalchemy import delete
 
 from letta import create_client
+from letta.constants import COMPOSIO_ENTITY_ENV_VAR_KEY
 from letta.functions.function_sets.base import core_memory_append, core_memory_replace
 from letta.orm.sandbox_config import SandboxConfig, SandboxEnvironmentVariable
 from letta.schemas.agent import AgentState
@@ -186,6 +187,14 @@ def composio_github_star_tool(test_user):
     yield tool
 
 
+@pytest.fixture
+def composio_gmail_get_profile_tool(test_user):
+    tool_manager = ToolManager()
+    tool_create = ToolCreate.from_composio(action_name="GMAIL_GET_PROFILE")
+    tool = tool_manager.create_or_update_tool(pydantic_tool=Tool(**tool_create.model_dump()), actor=test_user)
+    yield tool
+
+
 @pytest.fixture
 def clear_core_memory_tool(test_user):
     def clear_memory(agent_state: AgentState):
@@ -374,6 +383,29 @@ def test_local_sandbox_e2e_composio_star_github(mock_e2b_api_key_none, check_com
     assert result.func_return["details"] == "Action executed successfully"
 
 
+@pytest.mark.local_sandbox
+def test_local_sandbox_multiple_composio_entities(
+    mock_e2b_api_key_none, check_composio_key_set, composio_gmail_get_profile_tool, agent_state, test_user
+):
+    # Agent state with no composio entity ID
+    result = ToolExecutionSandbox(composio_gmail_get_profile_tool.name, {}, user=test_user).run(agent_state=agent_state)
+    assert result.func_return["response_data"]["emailAddress"] == "sarah@letta.com"
+
+    # Agent state with the composio entity set to 'matt'
+    agent_state.tool_exec_environment_variables = [
+        AgentEnvironmentVariable(key=COMPOSIO_ENTITY_ENV_VAR_KEY, value="matt", agent_id=agent_state.id)
+    ]
+    result = ToolExecutionSandbox(composio_gmail_get_profile_tool.name, {}, user=test_user).run(agent_state=agent_state)
+    assert result.func_return["response_data"]["emailAddress"] == "matt@letta.com"
+
+    # Agent state with composio entity ID set to default
+    agent_state.tool_exec_environment_variables = [
+        AgentEnvironmentVariable(key=COMPOSIO_ENTITY_ENV_VAR_KEY, value="default", agent_id=agent_state.id)
+    ]
+    result = ToolExecutionSandbox(composio_gmail_get_profile_tool.name, {}, user=test_user).run(agent_state=agent_state)
+    assert result.func_return["response_data"]["emailAddress"] == "sarah@letta.com"
+
+
 @pytest.mark.local_sandbox
 def test_local_sandbox_e2e_composio_star_github_without_setting_db_env_vars(
     mock_e2b_api_key_none, check_composio_key_set, composio_github_star_tool, test_user
@@ -593,6 +625,38 @@ def test_e2b_e2e_composio_star_github(check_e2b_key_is_set, check_composio_key_s
     assert result.func_return["details"] == "Action executed successfully"
 
 
+@pytest.mark.e2b_sandbox
+def test_e2b_multiple_composio_entities(
+    check_e2b_key_is_set, check_composio_key_set, composio_gmail_get_profile_tool, agent_state, test_user
+):
+    manager = SandboxConfigManager(tool_settings)
+    config = manager.get_or_create_default_sandbox_config(sandbox_type=SandboxType.E2B, actor=test_user)
+
+    manager.create_sandbox_env_var(
+        SandboxEnvironmentVariableCreate(key="COMPOSIO_API_KEY", value=tool_settings.composio_api_key),
+        sandbox_config_id=config.id,
+        actor=test_user,
+    )
+
+    # Agent state with no composio entity ID
+    result = ToolExecutionSandbox(composio_gmail_get_profile_tool.name, {}, user=test_user).run(agent_state=agent_state)
+    assert result.func_return["response_data"]["emailAddress"] == "sarah@letta.com"
+
+    # Agent state with the composio entity set to 'matt'
+    agent_state.tool_exec_environment_variables = [
+        AgentEnvironmentVariable(key=COMPOSIO_ENTITY_ENV_VAR_KEY, value="matt", agent_id=agent_state.id)
+    ]
+    result = ToolExecutionSandbox(composio_gmail_get_profile_tool.name, {}, user=test_user).run(agent_state=agent_state)
+    assert result.func_return["response_data"]["emailAddress"] == "matt@letta.com"
+
+    # Agent state with composio entity ID set to default
+    agent_state.tool_exec_environment_variables = [
+        AgentEnvironmentVariable(key=COMPOSIO_ENTITY_ENV_VAR_KEY, value="default", agent_id=agent_state.id)
+    ]
+    result = ToolExecutionSandbox(composio_gmail_get_profile_tool.name, {}, user=test_user).run(agent_state=agent_state)
+    assert result.func_return["response_data"]["emailAddress"] == "sarah@letta.com"
+
+
 # Core memory integration tests
 class TestCoreMemoryTools:
     """