perf(status): aggregate wing/room counts in the backend instead of client-side fetch - #1824
perf(status): aggregate wing/room counts in the backend instead of client-side fetch#1824gavinmcfall wants to merge 1 commit into
Conversation
… fetch mempalace_status tallied wings/rooms by pulling EVERY drawer's metadata to the client and counting in Python — O(n) memory and time. On a 418k-drawer pgvector palace this took ~344s (and OOM'd at 2Gi). Add an optional BaseCollection-style hook count_by_metadata_field(field) implemented by the pgvector backend as a single SQL GROUP BY, and use it in tool_status when the collection provides it; other backends keep the existing client-side path. Result on the same 418k palace: status drops from ~344s to ~0.4s. Assisted-by: Claude Code (claude-opus-4-8) Agentically-Engineered: https://nerdz.cloud/agentic-engineering
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a server-side aggregation method count_by_metadata_field for the pgvector backend, optimizing the tool_status endpoint to fetch wing and room counts in a single query instead of streaming all metadata to the client. A review comment identifies a correctness bug in pgvector.py where explicit "unknown" values and null values would overwrite each other in the dictionary comprehension, and suggests resolving this by using COALESCE directly in the SQL query.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| sql = ( | ||
| f"SELECT metadata->>%s AS k, count(*) AS c " | ||
| f"FROM {_quote_identifier(self._table)} GROUP BY 1" | ||
| ) | ||
| rows = self._client._execute(sql, [field], fetch=True) | ||
| return {(r[0] if r[0] is not None else "unknown"): int(r[1]) for r in rows} |
There was a problem hiding this comment.
There is a correctness bug here when a collection contains both explicit "unknown" values and missing/null values for the metadata field. Because both map to the key "unknown" in the dict comprehension, one will overwrite the other, resulting in an incorrect count. Pushing the COALESCE aggregation to SQL solves this cleanly and correctly.
| sql = ( | |
| f"SELECT metadata->>%s AS k, count(*) AS c " | |
| f"FROM {_quote_identifier(self._table)} GROUP BY 1" | |
| ) | |
| rows = self._client._execute(sql, [field], fetch=True) | |
| return {(r[0] if r[0] is not None else "unknown"): int(r[1]) for r in rows} | |
| sql = ( | |
| f"SELECT COALESCE(metadata->>%s, 'unknown') AS k, count(*) AS c " | |
| f"FROM {_quote_identifier(self._table)} GROUP BY 1" | |
| ) | |
| rows = self._client._execute(sql, [field], fetch=True) | |
| return {r[0]: int(r[1]) for r in rows} |
|
Hey @gavinmcfall — thanks for filing this, the SQL is essentially the right shape. A couple of things have landed in adjacent threads since you opened this PR that you might not have seen:
I'd suggest aligning this PR's surface area with #1835's spec (rename Two concrete additions from building the same thing for our pgvector deployment that might be worth folding in:
Happy to send a follow-up PR with these on top of yours once the design is settled, or you can pick them up directly — whatever's easier. |
|
Cross-link follow-up to my earlier comment: the wire-byte half I mentioned (the |
igorls
left a comment
igorls
left a comment
There was a problem hiding this comment.
Thanks for this, and especially for the concrete incident data — a 419k-drawer pgvector palace taking ~344s and OOM-killing a 2Gi pod on status is exactly the kind of real-world pressure that justifies server-side aggregation. The SQL here is sound and injection-safe (bound %s field, _quote_identifier on the table, NULL → "unknown" matching the client-side ground truth).
The blocker isn't correctness — it's interface convergence. This PR predates the agreed design in #1835, which a parallel PR (#1868) now implements as the canonical contract:
- a
BaseCollection.facet_counts(field, where=None) -> dict[str, int]method (default raisesUnsupportedCapabilityError, mirroringlexical_search), - a
supports_metadata_facetscapability token on the backend, - all four call sites (
tool_status,tool_list_wings,tool_list_rooms,tool_get_taxonomy) checking the token and falling back to client-side counting.
This PR instead adds count_by_metadata_field(field) (no where, duck-typed via getattr/callable, only tool_status), which would entrench a second, non-conforming aggregation interface — and it textually conflicts with #1868 (both rewrite the same wing/room loop in tool_status).
Proposed path so your work lands cleanly:
- We merge #1868 first (it establishes the base method, the capability token, and all four call sites).
- Rebase this PR on top as the pgvector implementation of
facet_counts: renamecount_by_metadata_field→facet_counts(field, where=None), add thewhere → SQL WHEREmapping, registersupports_metadata_facetsonPgVectorBackend, drop thetool_statusedit entirely (those call sites already live in #1868), and add unit tests for the new method.
After that rebase this shrinks to roughly one method on pgvector.py + one token line, with no mcp_server.py changes and no conflict. The SQL, the NULL bucketing, and the incident motivation all carry straight over. Would you be up for reshaping it that way? Happy to help line it up once #1868 is in.
One thing worth confirming for your scenario: the wing/room tools already short-circuit through a _sqlite_taxonomy() cache when present, so it's worth verifying the OOM path actually reached the client-side counting branch (i.e. the SQLite cache wasn't populated) so the fix lands where the incident occurred.
|
@gavinmcfall — thank you for the original SQL and the concrete incident data; both carry forward. Filed #2038 as a take-over of this PR's territory, implementing the rebase path @igorls proposed above: renamed to Filing separately rather than pushing to your branch since the changes are substantial and you may still want to iterate on your own version. If maintainers prefer to close this in favor of the new PR that seems reasonable to us; if you'd rather reshape this one instead, let us know and we can withdraw the new PR. |
|
Agreed — #2038 is the right shape and supersedes this. The |
3 participants
Problem
mempalace_statusbuilds its wing/room breakdown by pulling every drawer's metadata to the client (_get_cached_metadata) and counting in Python. That's O(n) memory and time.On a 418,997-drawer pgvector palace this took ~344s and OOM-killed a 2Gi pod. Since the MCP wake-up protocol calls
statusfirst, large palaces effectively can't wake up.Change
count_by_metadata_field(field) -> {value: count}.SELECT metadata->>field, count(*) ... GROUP BY 1.tool_statususes it when the collection provides it; backends without it keep the existing client-side path unchanged (duck-typed, fully backward compatible).Result
Same 418k palace:
statusdrops from ~344s → ~0.4s, flat memory. Counts verified identical to the client-side tally (7 wings / 50 rooms; e.g.my_other_repos=169,587).Notes
list_wings/list_roomshave the same client-side pattern and could adopt this hook in a follow-up; kept this PR focused on the wake-up path.