feat: replace tool_exec_environtment_variables with secrets (#2953)

This commit is contained in:
cthomas
2025-09-16 13:22:48 -07:00
committed by GitHub
parent 86f6b81e67
commit edb6c5e14e
15 changed files with 77 additions and 49 deletions

View File

@@ -1628,7 +1628,7 @@ class Agent(BaseAgent):
action_name = generate_composio_action_from_func_name(target_letta_tool.name)
# Get entity ID from the agent_state
entity_id = None
for env_var in self.agent_state.tool_exec_environment_variables:
for env_var in self.agent_state.secrets:
if env_var.key == COMPOSIO_ENTITY_ENV_VAR_KEY:
entity_id = env_var.value
# Get composio_api_key

View File

@@ -1874,7 +1874,7 @@ class LettaAgent(BaseAgent):
start_time = get_utc_timestamp_ns()
agent_step_span.add_event(name="tool_execution_started")
sandbox_env_vars = {var.key: var.value for var in agent_state.tool_exec_environment_variables}
sandbox_env_vars = {var.key: var.value for var in agent_state.secrets}
tool_execution_manager = ToolExecutionManager(
agent_state=agent_state,
message_manager=self.message_manager,

View File

@@ -1106,7 +1106,7 @@ class LettaAgentV2(BaseAgentV2):
start_time = get_utc_timestamp_ns()
agent_step_span.add_event(name="tool_execution_started")
sandbox_env_vars = {var.key: var.value for var in agent_state.tool_exec_environment_variables}
sandbox_env_vars = {var.key: var.value for var in agent_state.secrets}
tool_execution_manager = ToolExecutionManager(
agent_state=agent_state,
message_manager=self.message_manager,

View File

@@ -441,7 +441,7 @@ class VoiceAgent(BaseAgent):
)
# Use ToolExecutionManager for modern tool execution
sandbox_env_vars = {var.key: var.value for var in agent_state.tool_exec_environment_variables}
sandbox_env_vars = {var.key: var.value for var in agent_state.secrets}
tool_execution_manager = ToolExecutionManager(
agent_state=agent_state,
message_manager=self.message_manager,

View File

@@ -233,6 +233,7 @@ class Agent(SqlalchemyBase, OrganizationMixin, ProjectMixin, TemplateEntityMixin
"identity_ids": [],
"multi_agent_group": None,
"tool_exec_environment_variables": [],
"secrets": [],
}
# Optional fields: only included if requested
@@ -253,6 +254,7 @@ class Agent(SqlalchemyBase, OrganizationMixin, ProjectMixin, TemplateEntityMixin
"identity_ids": lambda: [i.id for i in self.identities],
"multi_agent_group": lambda: self.multi_agent_group,
"tool_exec_environment_variables": lambda: self.tool_exec_environment_variables,
"secrets": lambda: self.tool_exec_environment_variables,
}
include_relationships = set(optional_fields.keys() if include_relationships is None else include_relationships)
@@ -324,6 +326,7 @@ class Agent(SqlalchemyBase, OrganizationMixin, ProjectMixin, TemplateEntityMixin
"identity_ids": [],
"multi_agent_group": None,
"tool_exec_environment_variables": [],
"secrets": [],
}
# Initialize include_relationships to an empty set if it's None
@@ -344,7 +347,7 @@ class Agent(SqlalchemyBase, OrganizationMixin, ProjectMixin, TemplateEntityMixin
multi_agent_group = self.awaitable_attrs.multi_agent_group if "multi_agent_group" in include_relationships else none_async()
tool_exec_environment_variables = (
self.awaitable_attrs.tool_exec_environment_variables
if "tool_exec_environment_variables" in include_relationships
if "tool_exec_environment_variables" in include_relationships or "secrets" in include_relationships
else empty_list_async()
)
file_agents = self.awaitable_attrs.file_agents if "memory" in include_relationships else empty_list_async()
@@ -368,5 +371,6 @@ class Agent(SqlalchemyBase, OrganizationMixin, ProjectMixin, TemplateEntityMixin
state["identity_ids"] = [i.id for i in identities]
state["multi_agent_group"] = multi_agent_group
state["tool_exec_environment_variables"] = tool_exec_environment_variables
state["secrets"] = tool_exec_environment_variables
return self.__pydantic_model__(**state)

View File

@@ -86,6 +86,11 @@ class AgentState(OrmMetadataBase, validate_assignment=True):
sources: List[Source] = Field(..., description="The sources used by the agent.")
tags: List[str] = Field(..., description="The tags associated with the agent.")
tool_exec_environment_variables: List[AgentEnvironmentVariable] = Field(
default_factory=list,
description="Deprecated: use `secrets` field instead.",
deprecated=True,
)
secrets: List[AgentEnvironmentVariable] = Field(
default_factory=list, description="The environment variables for tool execution specific to this agent."
)
project_id: Optional[str] = Field(None, description="The id of the project the agent belongs to.")
@@ -133,7 +138,7 @@ class AgentState(OrmMetadataBase, validate_assignment=True):
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:
for agent_env_var_obj in self.secrets:
per_agent_env_vars[agent_env_var_obj.key] = agent_env_var_obj.value
return per_agent_env_vars
@@ -222,9 +227,8 @@ class CreateAgent(BaseModel, validate_assignment=True): #
deprecated=True,
description="Deprecated: Project should now be passed via the X-Project header instead of in the request body. If using the sdk, this can be done via the new x_project field below.",
)
tool_exec_environment_variables: Optional[Dict[str, str]] = Field(
None, description="The environment variables for tool execution specific to this agent."
)
tool_exec_environment_variables: Optional[Dict[str, str]] = Field(None, description="Deprecated: use `secrets` field instead.")
secrets: Optional[Dict[str, str]] = Field(None, description="The environment variables for tool execution specific to this agent.")
memory_variables: Optional[Dict[str, str]] = Field(None, description="The variables that should be set for the agent.")
project_id: Optional[str] = Field(None, description="The id of the project the agent belongs to.")
template_id: Optional[str] = Field(None, description="The id of the template the agent belongs to.")
@@ -328,9 +332,8 @@ class UpdateAgent(BaseModel):
message_ids: Optional[List[str]] = Field(None, description="The ids of the messages in the agent's in-context memory.")
description: Optional[str] = Field(None, description="The description of the agent.")
metadata: Optional[Dict] = Field(None, description="The metadata of the agent.")
tool_exec_environment_variables: Optional[Dict[str, str]] = Field(
None, description="The environment variables for tool execution specific to this agent."
)
tool_exec_environment_variables: Optional[Dict[str, str]] = Field(None, description="Deprecated: use `secrets` field instead")
secrets: Optional[Dict[str, str]] = Field(None, description="The environment variables for tool execution specific to this agent.")
project_id: Optional[str] = Field(None, description="The id of the project the agent belongs to.")
template_id: Optional[str] = Field(None, description="The id of the template the agent belongs to.")
base_template_id: Optional[str] = Field(None, description="The base template id of the agent.")

View File

@@ -40,6 +40,7 @@ class MarshmallowAgentSchema(BaseSchema):
core_memory = fields.List(fields.Nested(SerializedBlockSchema))
tools = fields.List(fields.Nested(SerializedToolSchema))
tool_exec_environment_variables = fields.List(fields.Nested(SerializedAgentEnvironmentVariableSchema))
secrets = fields.List(fields.Nested(SerializedAgentEnvironmentVariableSchema))
tags = fields.List(fields.Nested(SerializedAgentTagSchema))
def __init__(self, *args, session: sessionmaker, actor: User, max_steps: Optional[int] = None, **kwargs):
@@ -214,6 +215,9 @@ class MarshmallowAgentSchema(BaseSchema):
for env_var in data.get("tool_exec_environment_variables", []):
# need to be re-set at load time
env_var["value"] = ""
for env_var in data.get("secrets", []):
# need to be re-set at load time
env_var["value"] = ""
return data
@pre_load

View File

@@ -455,7 +455,8 @@ class AgentManager:
[{"agent_id": aid, "identity_id": iid} for iid in identity_ids],
)
if agent_create.tool_exec_environment_variables:
agent_secrets = agent_create.secrets or agent_create.tool_exec_environment_variables
if agent_secrets:
env_rows = [
{
"agent_id": aid,
@@ -463,7 +464,7 @@ class AgentManager:
"value": val,
"organization_id": actor.organization_id,
}
for key, val in agent_create.tool_exec_environment_variables.items()
for key, val in agent_secrets.items()
]
session.execute(insert(AgentEnvironmentVariable).values(env_rows))
@@ -674,7 +675,8 @@ class AgentManager:
)
env_rows = []
if agent_create.tool_exec_environment_variables:
agent_secrets = agent_create.secrets or agent_create.tool_exec_environment_variables
if agent_secrets:
env_rows = [
{
"agent_id": aid,
@@ -682,7 +684,7 @@ class AgentManager:
"value": val,
"organization_id": actor.organization_id,
}
for key, val in agent_create.tool_exec_environment_variables.items()
for key, val in agent_secrets.items()
]
result = await session.execute(insert(AgentEnvironmentVariable).values(env_rows).returning(AgentEnvironmentVariable.id))
env_rows = [{**row, "id": env_var_id} for row, env_var_id in zip(env_rows, result.scalars().all())]
@@ -701,8 +703,9 @@ class AgentManager:
result = await new_agent.to_pydantic_async(include_relationships=include_relationships)
if agent_create.tool_exec_environment_variables and env_rows:
if agent_secrets and env_rows:
result.tool_exec_environment_variables = [AgentEnvironmentVariable(**row) for row in env_rows]
result.secrets = [AgentEnvironmentVariable(**row) for row in env_rows]
# initial message sequence (skip if _init_with_no_messages is True)
if not _init_with_no_messages:
@@ -894,7 +897,8 @@ class AgentManager:
)
session.expire(agent, ["tags"])
if agent_update.tool_exec_environment_variables is not None:
agent_secrets = agent_update.secrets or agent_update.tool_exec_environment_variables
if agent_secrets is not None:
session.execute(delete(AgentEnvironmentVariable).where(AgentEnvironmentVariable.agent_id == aid))
env_rows = [
{
@@ -903,7 +907,7 @@ class AgentManager:
"value": v,
"organization_id": agent.organization_id,
}
for k, v in agent_update.tool_exec_environment_variables.items()
for k, v in agent_secrets.items()
]
if env_rows:
self._bulk_insert_pivot(session, AgentEnvironmentVariable.__table__, env_rows)
@@ -1019,7 +1023,8 @@ class AgentManager:
)
session.expire(agent, ["tags"])
if agent_update.tool_exec_environment_variables is not None:
agent_secrets = agent_update.secrets or agent_update.tool_exec_environment_variables
if agent_secrets is not None:
await session.execute(delete(AgentEnvironmentVariable).where(AgentEnvironmentVariable.agent_id == aid))
env_rows = [
{
@@ -1028,7 +1033,7 @@ class AgentManager:
"value": v,
"organization_id": agent.organization_id,
}
for k, v in agent_update.tool_exec_environment_variables.items()
for k, v in agent_secrets.items()
]
if env_rows:
await self._bulk_insert_pivot_async(session, AgentEnvironmentVariable.__table__, env_rows)
@@ -1544,6 +1549,8 @@ class AgentManager:
if env_vars:
for var in agent.tool_exec_environment_variables:
var.value = env_vars.get(var.key, "")
for var in agent.secrets:
var.value = env_vars.get(var.key, "")
agent = agent.create(session, actor=actor)
@@ -1627,6 +1634,7 @@ class AgentManager:
# Remove stale variables
stale_keys = set(existing_vars) - set(env_vars)
agent.tool_exec_environment_variables = [var for var in updated_vars if var.key not in stale_keys]
agent.secrets = [var for var in updated_vars if var.key not in stale_keys]
# Update the agent in the database
agent.update(session, actor=actor)

View File

@@ -209,8 +209,10 @@ class AgentSerializationManager:
agent_schema.id = agent_file_id
# wipe the values of tool_exec_environment_variables (they contain secrets)
if agent_schema.tool_exec_environment_variables:
agent_schema.tool_exec_environment_variables = {key: "" for key in agent_schema.tool_exec_environment_variables}
agent_secrets = agent_schema.secrets or agent_schema.tool_exec_environment_variables
if agent_secrets:
agent_schema.tool_exec_environment_variables = {key: "" for key in agent_secrets}
agent_schema.secrets = {key: "" for key in agent_secrets}
if agent_schema.messages:
for message in agent_schema.messages:
@@ -655,10 +657,16 @@ class AgentSerializationManager:
if agent_data.get("source_ids"):
agent_data["source_ids"] = [file_to_db_ids[file_id] for file_id in agent_data["source_ids"]]
if env_vars and agent_data.get("tool_exec_environment_variables"):
if env_vars and agent_data.get("secrets"):
# update environment variable values from the provided env_vars dict
for key in agent_data["secrets"]:
agent_data["secrets"][key] = env_vars.get(key, "")
agent_data["tool_exec_environment_variables"][key] = env_vars.get(key, "")
elif env_vars and agent_data.get("tool_exec_environment_variables"):
# also handle tool_exec_environment_variables for backwards compatibility
for key in agent_data["tool_exec_environment_variables"]:
agent_data["tool_exec_environment_variables"][key] = env_vars.get(key, "")
agent_data["secrets"][key] = env_vars.get(key, "")
# Override project_id if provided
if project_id:

View File

@@ -51,7 +51,7 @@ class ExternalComposioToolExecutor(ToolExecutor):
def _get_entity_id(self, agent_state: AgentState) -> Optional[str]:
"""Extract the entity ID from environment variables."""
for env_var in agent_state.tool_exec_environment_variables:
for env_var in agent_state.secrets:
if env_var.key == COMPOSIO_ENTITY_ENV_VAR_KEY:
return env_var.value
return None

View File

@@ -108,11 +108,16 @@ def comprehensive_agent_checks(agent: AgentState, request: Union[CreateAgent, Up
assert agent.metadata == request.metadata, f"Metadata mismatch: {agent.metadata} != {request.metadata}"
# Assert agent env vars
if hasattr(request, "tool_exec_environment_variables"):
if hasattr(request, "tool_exec_environment_variables") and request.tool_exec_environment_variables:
for agent_env_var in agent.tool_exec_environment_variables:
assert agent_env_var.key in request.tool_exec_environment_variables
assert request.tool_exec_environment_variables[agent_env_var.key] == agent_env_var.value
assert agent_env_var.organization_id == actor.organization_id
if hasattr(request, "secrets") and request.secrets:
for agent_env_var in agent.secrets:
assert agent_env_var.key in request.secrets
assert request.secrets[agent_env_var.key] == agent_env_var.value
assert agent_env_var.organization_id == actor.organization_id
# Assert agent type
if hasattr(request, "agent_type"):

View File

@@ -589,7 +589,7 @@ async def test_local_sandbox_per_agent_env(disable_e2b_api_key, get_env_tool, ag
manager.create_sandbox_env_var(SandboxEnvironmentVariableCreate(key=key, value=wrong_val), sandbox_config_id=config.id, actor=test_user)
correct_val = "".join(secrets.choice(string.ascii_letters + string.digits) for _ in range(20))
agent_state.tool_exec_environment_variables = [AgentEnvironmentVariable(key=key, value=correct_val, agent_id=agent_state.id)]
agent_state.secrets = [AgentEnvironmentVariable(key=key, value=correct_val, agent_id=agent_state.id)]
sandbox = AsyncToolSandboxLocal(get_env_tool.name, {}, user=test_user)
result = await sandbox.run(agent_state=agent_state)
@@ -812,7 +812,7 @@ async def test_e2b_sandbox_per_agent_env(check_e2b_key_is_set, get_env_tool, age
actor=test_user,
)
agent_state.tool_exec_environment_variables = [AgentEnvironmentVariable(key=key, value=correct_val, agent_id=agent_state.id)]
agent_state.secrets = [AgentEnvironmentVariable(key=key, value=correct_val, agent_id=agent_state.id)]
sandbox = AsyncToolSandboxE2B(get_env_tool.name, {}, user=test_user)
result = await sandbox.run(agent_state=agent_state)
@@ -1116,7 +1116,7 @@ async def test_local_sandbox_async_per_agent_env(disable_e2b_api_key, async_get_
manager.create_sandbox_env_var(SandboxEnvironmentVariableCreate(key=key, value=wrong_val), sandbox_config_id=config.id, actor=test_user)
correct_val = "correct_async_local_value"
agent_state.tool_exec_environment_variables = [AgentEnvironmentVariable(key=key, value=correct_val, agent_id=agent_state.id)]
agent_state.secrets = [AgentEnvironmentVariable(key=key, value=correct_val, agent_id=agent_state.id)]
sandbox = AsyncToolSandboxLocal(async_get_env_tool.name, {}, user=test_user)
result = await sandbox.run(agent_state=agent_state)
@@ -1141,7 +1141,7 @@ async def test_e2b_sandbox_async_per_agent_env(check_e2b_key_is_set, async_get_e
actor=test_user,
)
agent_state.tool_exec_environment_variables = [AgentEnvironmentVariable(key=key, value=correct_val, agent_id=agent_state.id)]
agent_state.secrets = [AgentEnvironmentVariable(key=key, value=correct_val, agent_id=agent_state.id)]
sandbox = AsyncToolSandboxE2B(async_get_env_tool.name, {}, user=test_user)
result = await sandbox.run(agent_state=agent_state)

View File

@@ -359,9 +359,7 @@ def test_local_sandbox_per_agent_env(disable_e2b_api_key, get_env_tool, agent_st
# 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)
]
agent_state.secrets = [AgentEnvironmentVariable(key=key, value=correct_long_random_string, agent_id=agent_state.id)]
# Create tool and args
args = {}
@@ -569,9 +567,7 @@ def test_e2b_sandbox_per_agent_env(check_e2b_key_is_set, get_env_tool, agent_sta
# 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)
]
agent_state.secrets = [AgentEnvironmentVariable(key=key, value=correct_long_random_string, agent_id=agent_state.id)]
# Create tool and args
args = {}

View File

@@ -298,7 +298,7 @@ def _compare_agent_state_model_dump(d1: Dict[str, Any], d2: Dict[str, Any], log:
if not isinstance(v1, list) or not isinstance(v2, list) or len(v1) != len(v2):
_log_mismatch(key, v1, v2, log)
return False
elif key == "tool_exec_environment_variables":
elif key == "tool_exec_environment_variables" or key == "secrets":
if not isinstance(v1, list) or not isinstance(v2, list) or len(v1) != len(v2):
_log_mismatch(key, v1, v2, log)
return False

View File

@@ -503,7 +503,7 @@ def _compare_agents(orig: AgentSchema, imp: AgentSchema, index: int) -> List[str
errors.append(f"Agent {index}: tool_rules mismatch")
# Environment variables
if orig.tool_exec_environment_variables != imp.tool_exec_environment_variables:
if orig.secrets != imp.secrets:
errors.append(f"Agent {index}: tool_exec_environment_variables mismatch")
# Messages
@@ -1084,16 +1084,16 @@ class TestAgentFileExport:
exported_agent = agent_file.agents[0]
# verify environment variables exist but values are scrubbed (empty strings)
assert exported_agent.tool_exec_environment_variables is not None
assert len(exported_agent.tool_exec_environment_variables) == 3
assert "API_KEY" in exported_agent.tool_exec_environment_variables
assert "DATABASE_PASSWORD" in exported_agent.tool_exec_environment_variables
assert "TOKEN" in exported_agent.tool_exec_environment_variables
assert exported_agent.secrets is not None
assert len(exported_agent.secrets) == 3
assert "API_KEY" in exported_agent.secrets
assert "DATABASE_PASSWORD" in exported_agent.secrets
assert "TOKEN" in exported_agent.secrets
# most importantly: verify all secret values have been wiped
assert exported_agent.tool_exec_environment_variables["API_KEY"] == ""
assert exported_agent.tool_exec_environment_variables["DATABASE_PASSWORD"] == ""
assert exported_agent.tool_exec_environment_variables["TOKEN"] == ""
assert exported_agent.secrets["API_KEY"] == ""
assert exported_agent.secrets["DATABASE_PASSWORD"] == ""
assert exported_agent.secrets["TOKEN"] == ""
# verify no secret values appear anywhere in the exported data
assert "super-secret-api-key-12345" not in str(agent_file)
@@ -1287,9 +1287,9 @@ class TestAgentFileImport:
# verify values are scrubbed in export
exported_agent = agent_file.agents[0]
assert exported_agent.tool_exec_environment_variables["API_KEY"] == ""
assert exported_agent.tool_exec_environment_variables["DATABASE_URL"] == ""
assert exported_agent.tool_exec_environment_variables["SECRET_TOKEN"] == ""
assert exported_agent.secrets["API_KEY"] == ""
assert exported_agent.secrets["DATABASE_URL"] == ""
assert exported_agent.secrets["SECRET_TOKEN"] == ""
# import with new environment variable values
new_env_vars = {