Merge pull request #371 from RhettOP/fix/issue-339-338-silent-exceptions-pagination

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

mempalace/mcp_server.py

@@ -143,16 +143,25 @@ 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 True:
+        try:
+            batch = col.get(include=["metadatas"], limit=batch_size, offset=offset)
+            rows = batch["metadatas"]
+            for m in rows:
+                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 += len(rows)
+            if len(rows) < batch_size:
+                break
+        except Exception as e:
+            error_info = f"Partial result, failed at offset {offset}: {str(e)}"
+            break
+    result = {
         "total_drawers": count,
         "wings": wings,
         "rooms": rooms,
@@ -160,6 +169,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 ─────────────────────────────────────────────────────────
@@ -200,13 +213,28 @@ 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
+        col.count()  # verify collection is accessible
+    except Exception as e:
+        return {"wings": {}, "error": str(e)}
+    while True:
+        try:
+            batch = col.get(include=["metadatas"], limit=batch_size, offset=offset)
+            rows = batch["metadatas"]
+            for m in rows:
+                w = m.get("wing", "unknown")
+                wings[w] = wings.get(w, 0) + 1
+            offset += len(rows)
+            if len(rows) < batch_size:
+                break
+        except Exception as e:
+            return {
+                "wings": wings,
+                "error": f"Partial result, failed at offset {offset}: {str(e)}",
+                "partial": True,
+            }
     return {"wings": wings}
 
 
@@ -215,16 +243,33 @@ 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
+        col.count()  # verify collection is accessible
+    except Exception as e:
+        return {"wing": wing or "all", "rooms": {}, "error": str(e)}
+    while True:
+        try:
+            kwargs = {"include": ["metadatas"], "limit": batch_size, "offset": offset}
+            if where:
+                kwargs["where"] = where
+            batch = col.get(**kwargs)
+            rows = batch["metadatas"]
+            for m in rows:
+                r = m.get("room", "unknown")
+                rooms[r] = rooms.get(r, 0) + 1
+            offset += len(rows)
+            if len(rows) < batch_size:
+                break
+        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}
 
 
@@ -233,16 +278,31 @@ 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
+        col.count()  # verify collection is accessible
+    except Exception as e:
+        return {"taxonomy": {}, "error": str(e)}
+    while True:
+        try:
+            batch = col.get(include=["metadatas"], limit=batch_size, offset=offset)
+            rows = batch["metadatas"]
+            for m in rows:
+                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 += len(rows)
+            if len(rows) < batch_size:
+                break
+        except Exception as e:
+            return {
+                "taxonomy": taxonomy,
+                "error": f"Partial result, failed at offset {offset}: {str(e)}",
+                "partial": True,
+            }
     return {"taxonomy": taxonomy}