From 23bc1882709c1ac7b1be73afc68e352296810ea2 Mon Sep 17 00:00:00 2001
From: Matthew Zhou <mattzh1314@gmail.com>
Date: Thu, 9 Jan 2025 10:48:21 -1000
Subject: [PATCH] feat: Add sandbox_type filter when listing sandbox
 configurations (#567)

---
 .../rest_api/routers/v1/sandbox_configs.py      |  3 ++-
 letta/services/sandbox_config_manager.py        | 13 ++++++-------
 tests/test_managers.py                          | 17 +++++++++++++----
 3 files changed, 21 insertions(+), 12 deletions(-)

diff --git a/letta/server/rest_api/routers/v1/sandbox_configs.py b/letta/server/rest_api/routers/v1/sandbox_configs.py
index d5c16c04a..edd4383ca 100644
--- a/letta/server/rest_api/routers/v1/sandbox_configs.py
+++ b/letta/server/rest_api/routers/v1/sandbox_configs.py
@@ -69,11 +69,12 @@ def delete_sandbox_config(
 def list_sandbox_configs(
     limit: int = Query(1000, description="Number of results to return"),
     cursor: Optional[str] = Query(None, description="Pagination cursor to fetch the next set of results"),
+    sandbox_type: Optional[SandboxType] = Query(None, description="Filter for this specific sandbox type"),
     server: SyncServer = Depends(get_letta_server),
     user_id: str = Depends(get_user_id),
 ):
     actor = server.user_manager.get_user_or_default(user_id=user_id)
-    return server.sandbox_config_manager.list_sandbox_configs(actor, limit=limit, cursor=cursor)
+    return server.sandbox_config_manager.list_sandbox_configs(actor, limit=limit, cursor=cursor, sandbox_type=sandbox_type)
 
 
 ### Sandbox Environment Variable Routes
diff --git a/letta/services/sandbox_config_manager.py b/letta/services/sandbox_config_manager.py
index 0511d3ec6..391d4aad7 100644
--- a/letta/services/sandbox_config_manager.py
+++ b/letta/services/sandbox_config_manager.py
@@ -111,16 +111,15 @@ class SandboxConfigManager:
 
     @enforce_types
     def list_sandbox_configs(
-        self, actor: PydanticUser, cursor: Optional[str] = None, limit: Optional[int] = 50
+        self, actor: PydanticUser, cursor: Optional[str] = None, limit: Optional[int] = 50, sandbox_type: Optional[SandboxType] = None
     ) -> List[PydanticSandboxConfig]:
         """List all sandbox configurations with optional pagination."""
+        kwargs = {"organization_id": actor.organization_id}
+        if sandbox_type:
+            kwargs.update({"type": sandbox_type})
+
         with self.session_maker() as session:
-            sandboxes = SandboxConfigModel.list(
-                db_session=session,
-                cursor=cursor,
-                limit=limit,
-                organization_id=actor.organization_id,
-            )
+            sandboxes = SandboxConfigModel.list(db_session=session, cursor=cursor, limit=limit, **kwargs)
             return [sandbox.to_pydantic() for sandbox in sandboxes]
 
     @enforce_types
diff --git a/tests/test_managers.py b/tests/test_managers.py
index 44afd991b..76001f3a1 100644
--- a/tests/test_managers.py
+++ b/tests/test_managers.py
@@ -2057,16 +2057,16 @@ def test_get_sandbox_config_by_type(server: SyncServer, sandbox_config_fixture,
 
 def test_list_sandbox_configs(server: SyncServer, default_user):
     # Creating multiple sandbox configs
-    config_a = SandboxConfigCreate(
+    config_e2b_create = SandboxConfigCreate(
         config=E2BSandboxConfig(),
     )
-    config_b = SandboxConfigCreate(
+    config_local_create = SandboxConfigCreate(
         config=LocalSandboxConfig(sandbox_dir=""),
     )
-    server.sandbox_config_manager.create_or_update_sandbox_config(config_a, actor=default_user)
+    config_e2b = server.sandbox_config_manager.create_or_update_sandbox_config(config_e2b_create, actor=default_user)
     if USING_SQLITE:
         time.sleep(CREATE_DELAY_SQLITE)
-    server.sandbox_config_manager.create_or_update_sandbox_config(config_b, actor=default_user)
+    config_local = server.sandbox_config_manager.create_or_update_sandbox_config(config_local_create, actor=default_user)
 
     # List configs without pagination
     configs = server.sandbox_config_manager.list_sandbox_configs(actor=default_user)
@@ -2080,6 +2080,15 @@ def test_list_sandbox_configs(server: SyncServer, default_user):
     assert len(next_page) == 1
     assert next_page[0].id != paginated_configs[0].id
 
+    # List configs using sandbox_type filter
+    configs = server.sandbox_config_manager.list_sandbox_configs(actor=default_user, sandbox_type=SandboxType.E2B)
+    assert len(configs) == 1
+    assert configs[0].id == config_e2b.id
+
+    configs = server.sandbox_config_manager.list_sandbox_configs(actor=default_user, sandbox_type=SandboxType.LOCAL)
+    assert len(configs) == 1
+    assert configs[0].id == config_local.id
+
 
 # ======================================================================================================================
 # SandboxConfigManager Tests - Environment Variables