fix: Fix bug for deleting agents belonging to the user org (#2033)

This commit is contained in:
Matthew Zhou 2024-11-13 12:37:31 -08:00 committed by GitHub
parent 151ce6f655
commit 6d6ee93458
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 29 additions and 3 deletions

View File

@ -1512,12 +1512,16 @@ class SyncServer(Server):
if self.ms.get_agent(agent_id=agent_id, user_id=user_id) is None: if self.ms.get_agent(agent_id=agent_id, user_id=user_id) is None:
raise ValueError(f"Agent agent_id={agent_id} does not exist") raise ValueError(f"Agent agent_id={agent_id} does not exist")
# Verify that the agent exists and is owned by the user # Verify that the agent exists and belongs to the org of the user
agent_state = self.ms.get_agent(agent_id=agent_id, user_id=user_id) agent_state = self.ms.get_agent(agent_id=agent_id, user_id=user_id)
if not agent_state: if not agent_state:
raise ValueError(f"Could not find agent_id={agent_id} under user_id={user_id}") raise ValueError(f"Could not find agent_id={agent_id} under user_id={user_id}")
if agent_state.user_id != user_id:
raise ValueError(f"Could not authorize agent_id={agent_id} with user_id={user_id}") agent_state_user = self.user_manager.get_user_by_id(user_id=agent_state.user_id)
if agent_state_user.organization_id != actor.organization_id:
raise ValueError(
f"Could not authorize agent_id={agent_id} with user_id={user_id} because of differing organizations; agent_id was created in {agent_state_user.organization_id} while user belongs to {actor.organization_id}. How did they get the agent id?"
)
# First, if the agent is in the in-memory cache we should remove it # First, if the agent is in the in-memory cache we should remove it
# List of {'user_id': user_id, 'agent_id': agent_id, 'agent': agent_obj} dicts # List of {'user_id': user_id, 'agent_id': agent_id, 'agent': agent_obj} dicts

View File

@ -7,6 +7,7 @@ import pytest
import letta.utils as utils import letta.utils as utils
from letta.constants import BASE_TOOLS, DEFAULT_MESSAGE_TOOL, DEFAULT_MESSAGE_TOOL_KWARG from letta.constants import BASE_TOOLS, DEFAULT_MESSAGE_TOOL, DEFAULT_MESSAGE_TOOL_KWARG
from letta.schemas.enums import MessageRole from letta.schemas.enums import MessageRole
from letta.schemas.user import User
from .test_managers import DEFAULT_EMBEDDING_CONFIG from .test_managers import DEFAULT_EMBEDDING_CONFIG
@ -575,3 +576,24 @@ def test_load_agent_with_nonexistent_tool_names_does_not_error(server: SyncServe
# cleanup # cleanup
server.delete_agent(user_id, agent_state.id) server.delete_agent(user_id, agent_state.id)
def test_delete_agent_same_org(server: SyncServer, org_id: str, user_id: str):
agent_state = server.create_agent(
request=CreateAgent(
name="nonexistent_tools_agent",
memory=ChatMemory(
human="Sarah",
persona="I am a helpful assistant",
),
llm_config=LLMConfig.default_config("gpt-4"),
embedding_config=EmbeddingConfig.default_config(provider="openai"),
),
actor=server.get_user_or_default(user_id),
)
# create another user in the same org
another_user = server.user_manager.create_user(User(organization_id=org_id, name="another"))
# test that another user in the same org can delete the agent
server.delete_agent(another_user.id, agent_state.id)