From 1027e24f7b2d860ad976153f3bddb5c92076bb53 Mon Sep 17 00:00:00 2001 From: ciregenz Date: Mon, 6 Jul 2026 04:27:46 -0700 Subject: [PATCH] [eric] google-workspace: cap shim tool-result text under the CLI's MCP-spill threshold (gmail-dump thrash) --- .../cap_tool_result.py | 44 ++++++++++++ backend/apps/google_workspace_mcp_shim/run.py | 23 +++++++ backend/tests/test_gws_cap_tool_result.py | 67 +++++++++++++++++++ 3 files changed, 134 insertions(+) create mode 100644 backend/apps/google_workspace_mcp_shim/cap_tool_result.py create mode 100644 backend/tests/test_gws_cap_tool_result.py diff --git a/backend/apps/google_workspace_mcp_shim/cap_tool_result.py b/backend/apps/google_workspace_mcp_shim/cap_tool_result.py new file mode 100644 index 00000000..b8e1e8a8 --- /dev/null +++ b/backend/apps/google_workspace_mcp_shim/cap_tool_result.py @@ -0,0 +1,44 @@ +"""Cap the cumulative text of a FastMCP call_tool return so one Gmail/Drive dump can't +blow the model's context. Pure + stdlib-only (no upstream imports) so it's importable +and unit-testable outside the shim's ephemeral uv env. + +The bundled Claude CLI hard-rejects any MCP result over ~25K tokens and spills it to a +file, which the model then re-reads back in, refilling the context and tripping the CLI's +autocompact-thrash. Capping under that spill threshold keeps the result inline and the +model out of the re-read loop.""" + +from typing import Any + +MAX_RESULT_CHARS = 48_000 +P_TRUNCATION_NOTE = ( + "\n\n[Truncated: this tool returned more than {cap} characters, too much to fit " + "in context at once. Narrow the request (add a search filter, a date range, or a " + "smaller max_results / page size) or fetch the next page.]" +) + + +def cap_tool_result(result: Any, max_chars: int = MAX_RESULT_CHARS) -> Any: + """Cap the text content blocks of a call_tool return in place. Duck-typed and + fail-open: any shape we don't recognize passes through unchanged, so an upstream + contract change degrades to no-cap, never a crash.""" + try: + blocks = result[0] if isinstance(result, tuple) else result + if not isinstance(blocks, list): + return result + used = 0 + truncated = False + for b in blocks: + if getattr(b, "type", None) != "text" or getattr(b, "text", None) is None: + continue + if truncated: + b.text = "" + continue + text = b.text + if used + len(text) <= max_chars: + used += len(text) + continue + b.text = text[: max(0, max_chars - used)] + P_TRUNCATION_NOTE.format(cap=max_chars) + truncated = True + return result + except Exception: + return result diff --git a/backend/apps/google_workspace_mcp_shim/run.py b/backend/apps/google_workspace_mcp_shim/run.py index 87bd44fb..a2fa27ba 100644 --- a/backend/apps/google_workspace_mcp_shim/run.py +++ b/backend/apps/google_workspace_mcp_shim/run.py @@ -17,11 +17,21 @@ CLIENT_ID/SECRET become unused placeholders. """ import functools +import importlib.util import os +import sys import google_workspace_mcp.auth.gauth as gauth from google.oauth2.credentials import Credentials +# Load the cap helper as a loose sibling file (not `from backend...`): the shim runs in uv's ephemeral env where the project isn't a package, and a path-load can't drag in backend's transitive deps. Kept next to run.py so the bundle always ships them together. +def p_load_cap(): + path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "cap_tool_result.py") + spec = importlib.util.spec_from_file_location("gws_cap_tool_result", path) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod.cap_tool_result + @functools.lru_cache(maxsize=1) def p_patched_get_credentials(): @@ -47,6 +57,19 @@ from google_workspace_mcp import __main__ as p_gw_main # noqa: E402,F401 from google_workspace_mcp.app import mcp # noqa: E402 +# Fail-open: if the cap helper can't load (unexpected bundle layout), run uncapped rather than break the whole Google Workspace tool. +try: + p_cap = p_load_cap() + p_orig_call_tool = mcp.call_tool + + async def p_capped_call_tool(name, arguments): + return p_cap(await p_orig_call_tool(name, arguments)) + + mcp.call_tool = p_capped_call_tool +except Exception as p_e: + print(f"[gws-shim] result cap disabled ({p_e}); running uncapped", file=sys.stderr) + + if __name__ == "__main__": # Upstream google_workspace_mcp.__main__.main() wraps a synchronous mcp.run() in asyncio.run() which throws "a coroutine was expected, got None" against current FastMCP. Skip it and invoke FastMCP's stdio loop directly. The `_gw_main` import above is what actually registers every tool/prompt/resource module against the shared `mcp` instance via its top-level imports. mcp.run("stdio") diff --git a/backend/tests/test_gws_cap_tool_result.py b/backend/tests/test_gws_cap_tool_result.py new file mode 100644 index 00000000..6b185ef7 --- /dev/null +++ b/backend/tests/test_gws_cap_tool_result.py @@ -0,0 +1,67 @@ +"""Google-workspace shim result cap invariant. + +The bug class (1.5.4 field report, Alex's query_gmail_emails thrash): a single +oversized Gmail/Drive dump exceeds the CLI's ~25K-token MCP cap, gets spilled to +a file, the model re-reads it back, and the context refills into the CLI's +autocompact-thrash. The seal: the shim caps its own tool-result text under that +spill threshold, with a clear paginate marker, and fails open on any shape it +doesn't recognize so an upstream contract change never crashes the shim. +""" + +from types import SimpleNamespace + +from backend.apps.google_workspace_mcp_shim.cap_tool_result import ( + MAX_RESULT_CHARS, + cap_tool_result, +) + + +def block(text: str) -> SimpleNamespace: + return SimpleNamespace(type="text", text=text) + + +def test_small_result_untouched() -> None: + b = block("one short email") + cap_tool_result(([b], {"result": "one short email"})) + assert b.text == "one short email" + + +def test_oversized_single_block_capped_with_marker() -> None: + b = block("E" * 300_000) + cap_tool_result(([b], {"result": "E" * 300_000})) + assert len(b.text) < MAX_RESULT_CHARS + 400 + assert b.text.startswith("E") + assert "Truncated" in b.text + assert len(b.text) // 4 < 25_000 + + +def test_budget_spans_multiple_blocks() -> None: + a, b, c = block("A" * 40_000), block("B" * 40_000), block("C" * 40_000) + cap_tool_result([a, b, c]) + assert a.text == "A" * 40_000 + assert "Truncated" in b.text and b.text.startswith("B") + assert c.text == "" + + +def test_non_text_blocks_pass_through() -> None: + img = SimpleNamespace(type="image", data="zzz") + txt = block("hello") + cap_tool_result([img, txt]) + assert img.data == "zzz" + assert txt.text == "hello" + + +def test_bare_list_return_shape() -> None: + b = block("Z" * 100_000) + out = cap_tool_result([b]) + assert out is not None + assert "Truncated" in b.text + + +def test_fail_open_on_unexpected_shapes() -> None: + assert cap_tool_result(None) is None + assert cap_tool_result({"structured": "only"}) == {"structured": "only"} + assert cap_tool_result("raw string") == "raw string" + junk = [SimpleNamespace(nope=1)] + cap_tool_result(junk) # no .type/.text -> untouched, no raise + assert junk[0].nope == 1