From 59d1f9cc2560d0a30c42883ebd7ab3635cfbc508 Mon Sep 17 00:00:00 2001
From: Matthew Zhou <mattzh1314@gmail.com>
Date: Wed, 13 Nov 2024 12:37:31 -0800
Subject: [PATCH] fix: Fix bug for deleting agents belonging to the user org
 (#2033)

---
 letta/server/server.py | 10 +++++++---
 tests/test_server.py   | 22 ++++++++++++++++++++++
 2 files changed, 29 insertions(+), 3 deletions(-)

diff --git a/letta/server/server.py b/letta/server/server.py
index 25305f62..a23d0795 100644
--- a/letta/server/server.py
+++ b/letta/server/server.py
@@ -1512,12 +1512,16 @@ class SyncServer(Server):
         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")
 
-        # 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)
         if not agent_state:
             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
         # List of {'user_id': user_id, 'agent_id': agent_id, 'agent': agent_obj} dicts
diff --git a/tests/test_server.py b/tests/test_server.py
index bd68e45a..c4887606 100644
--- a/tests/test_server.py
+++ b/tests/test_server.py
@@ -7,6 +7,7 @@ import pytest
 import letta.utils as utils
 from letta.constants import BASE_TOOLS, DEFAULT_MESSAGE_TOOL, DEFAULT_MESSAGE_TOOL_KWARG
 from letta.schemas.enums import MessageRole
+from letta.schemas.user import User
 
 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
     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)