fix: paginate large collection reads and surface errors in MCP tools (#339, #338)

mempalace/mcp_server.py

@@ -96,16 +96,22 @@ def tool_status():
     count = col.count()
     wings = {}
     rooms = {}
-    try:
-        all_meta = col.get(include=["metadatas"], limit=10000)["metadatas"]
-        for m in all_meta:
-            w = m.get("wing", "unknown")
-            r = m.get("room", "unknown")
-            wings[w] = wings.get(w, 0) + 1
-            rooms[r] = rooms.get(r, 0) + 1
-    except Exception:
-        pass
-    return {
+    batch_size = 5000
+    offset = 0
+    error_info = None
+    while offset < count:
+        try:
+            batch = col.get(include=["metadatas"], limit=batch_size, offset=offset)
+            for m in batch["metadatas"]:
+                w = m.get("wing", "unknown")
+                r = m.get("room", "unknown")
+                wings[w] = wings.get(w, 0) + 1
+                rooms[r] = rooms.get(r, 0) + 1
+            offset += batch_size
+        except Exception as e:
+            error_info = f"Partial result, failed at offset {offset}: {str(e)}"
+            break
+    result = {
         "total_drawers": count,
         "wings": wings,
         "rooms": rooms,
@@ -113,6 +119,10 @@ def tool_status():
         "protocol": PALACE_PROTOCOL,
         "aaak_dialect": AAAK_SPEC,
     }
+    if error_info:
+        result["error"] = error_info
+        result["partial"] = True
+    return result
 
 
 # ── AAAK Dialect Spec ─────────────────────────────────────────────────────────
@@ -153,13 +163,21 @@ def tool_list_wings():
     if not col:
         return _no_palace()
     wings = {}
+    batch_size = 5000
+    offset = 0
     try:
-        all_meta = col.get(include=["metadatas"], limit=10000)["metadatas"]
-        for m in all_meta:
-            w = m.get("wing", "unknown")
-            wings[w] = wings.get(w, 0) + 1
-    except Exception:
-        pass
+        total = col.count()
+    except Exception as e:
+        return {"wings": {}, "error": str(e)}
+    while offset < total:
+        try:
+            batch = col.get(include=["metadatas"], limit=batch_size, offset=offset)
+            for m in batch["metadatas"]:
+                w = m.get("wing", "unknown")
+                wings[w] = wings.get(w, 0) + 1
+            offset += batch_size
+        except Exception as e:
+            return {"wings": wings, "error": f"Partial result, failed at offset {offset}: {str(e)}", "partial": True}
     return {"wings": wings}
 
 
@@ -168,16 +186,25 @@ def tool_list_rooms(wing: str = None):
     if not col:
         return _no_palace()
     rooms = {}
+    batch_size = 5000
+    offset = 0
+    where = {"wing": wing} if wing else None
     try:
-        kwargs = {"include": ["metadatas"], "limit": 10000}
-        if wing:
-            kwargs["where"] = {"wing": wing}
-        all_meta = col.get(**kwargs)["metadatas"]
-        for m in all_meta:
-            r = m.get("room", "unknown")
-            rooms[r] = rooms.get(r, 0) + 1
-    except Exception:
-        pass
+        total = col.count()
+    except Exception as e:
+        return {"wing": wing or "all", "rooms": {}, "error": str(e)}
+    while offset < total:
+        try:
+            kwargs = {"include": ["metadatas"], "limit": batch_size, "offset": offset}
+            if where:
+                kwargs["where"] = where
+            batch = col.get(**kwargs)
+            for m in batch["metadatas"]:
+                r = m.get("room", "unknown")
+                rooms[r] = rooms.get(r, 0) + 1
+            offset += batch_size
+        except Exception as e:
+            return {"wing": wing or "all", "rooms": rooms, "error": f"Partial result, failed at offset {offset}: {str(e)}", "partial": True}
     return {"wing": wing or "all", "rooms": rooms}
 
 
@@ -186,16 +213,24 @@ def tool_get_taxonomy():
     if not col:
         return _no_palace()
     taxonomy = {}
+    batch_size = 5000
+    offset = 0
     try:
-        all_meta = col.get(include=["metadatas"], limit=10000)["metadatas"]
-        for m in all_meta:
-            w = m.get("wing", "unknown")
-            r = m.get("room", "unknown")
-            if w not in taxonomy:
-                taxonomy[w] = {}
-            taxonomy[w][r] = taxonomy[w].get(r, 0) + 1
-    except Exception:
-        pass
+        total = col.count()
+    except Exception as e:
+        return {"taxonomy": {}, "error": str(e)}
+    while offset < total:
+        try:
+            batch = col.get(include=["metadatas"], limit=batch_size, offset=offset)
+            for m in batch["metadatas"]:
+                w = m.get("wing", "unknown")
+                r = m.get("room", "unknown")
+                if w not in taxonomy:
+                    taxonomy[w] = {}
+                taxonomy[w][r] = taxonomy[w].get(r, 0) + 1
+            offset += batch_size
+        except Exception as e:
+            return {"taxonomy": taxonomy, "error": f"Partial result, failed at offset {offset}: {str(e)}", "partial": True}
     return {"taxonomy": taxonomy}