fix: bound hosted binary tool responses (#112)

* fix: bound hosted binary tool responses

* fix: guard compressed upstream error bodies

Author: ChiragAgg5k

src/mcp_server_appwrite/constants.py

@@ -49,6 +49,9 @@ def _resolve_server_version() -> str:
 EXCLUDED_SERVICES: frozenset[str] = frozenset()
 
 MAX_FETCH_BYTES = 25 * 1024 * 1024  # 25 MB cap on server-fetched files
+# MCP embeds binary tool results as base64 in one JSON-RPC response. Bound the
+# source bytes before base64 and JSON encoding create additional in-memory copies.
+MAX_HOSTED_BINARY_RESPONSE_BYTES = 25 * 1024 * 1024
 # Match Cloud/agent chat attachment max (10 MB). Hosted uploads resolve
 # turn attachments to inline base64; keep this at least that large.
 MAX_INLINE_BYTES = 10 * 1024 * 1024  # 10 MB cap on decoded inline content

src/mcp_server_appwrite/error_classification.py

@@ -7,8 +7,9 @@
 
 from __future__ import annotations
 
+import json
 from collections.abc import Iterator
-from typing import Literal
+from typing import Any, Literal
 
 from appwrite_console.exception import AppwriteException
 
@@ -17,6 +18,7 @@
     "appwrite_4xx",
     "appwrite_5xx",
     "sdk_validation",
+    "response_too_large",
     "internal",
 ]
 
@@ -26,6 +28,7 @@
         "appwrite_4xx",
         "appwrite_5xx",
         "sdk_validation",
+        "response_too_large",
         "internal",
     }
 )
@@ -35,13 +38,43 @@ class WriteConfirmationRequired(RuntimeError):
     """A mutating hidden tool was called without explicit confirmation."""
 
 
+class HostedBinaryResponseTooLarge(ValueError):
+    """A binary Appwrite response exceeded the hosted MCP memory-safe limit."""
+
+    def __init__(
+        self,
+        tool_name: str,
+        limit_bytes: int,
+        *,
+        content_length: int | None = None,
+        observed_bytes: int | None = None,
+    ) -> None:
+        error: dict[str, Any] = {
+            "code": "hosted_response_too_large",
+            "tool": tool_name,
+            "limitBytes": limit_bytes,
+            "message": (
+                "The binary response is too large to return through hosted MCP. "
+                "Use an Appwrite SDK or REST API for larger content."
+            ),
+        }
+        if content_length is not None:
+            error["contentLength"] = content_length
+        if observed_bytes is not None:
+            error["observedBytes"] = observed_bytes
+        super().__init__(json.dumps({"error": error}, separators=(",", ":")))
+
+
 def classify_tool_error(exc: BaseException) -> ErrorCategory:
     """Return the bounded operational category for an exception chain."""
     chain = tuple(_exception_chain(exc))
 
     if any(isinstance(item, WriteConfirmationRequired) for item in chain):
         return "write_confirmation"
 
+    if any(isinstance(item, HostedBinaryResponseTooLarge) for item in chain):
+        return "response_too_large"
+
     if any(_is_sdk_validation_error(item) for item in chain):
         return "sdk_validation"
 

src/mcp_server_appwrite/server.py

@@ -60,6 +60,7 @@
     FETCH_TIMEOUT_SECONDS,
     HOSTED_PATH_GUIDANCE,
     MAX_FETCH_BYTES,
+    MAX_HOSTED_BINARY_RESPONSE_BYTES,
     MAX_INLINE_BYTES,
     SERVER_ICON_URL,
     SERVER_VERSION,
@@ -73,7 +74,7 @@
     get_appwrite_context,
 )
 from .docs_search import DocsSearch
-from .error_classification import is_response_parse_error
+from .error_classification import HostedBinaryResponseTooLarge, is_response_parse_error
 from .operator import Operator, _parse_tool_name
 from .service import Service
 from .tool_manager import ToolManager
@@ -381,6 +382,11 @@ def register_services(
                 name,
                 allowed_methods=allowed_methods,
                 context_scope=context_scope(name),
+                binary_response_limit=(
+                    MAX_HOSTED_BINARY_RESPONSE_BYTES
+                    if profile == OAUTH_PROFILE
+                    else None
+                ),
             )
         )
     return tools_manager
@@ -820,6 +826,126 @@ def _prepare_arguments(tool_info: dict, arguments: dict[str, Any]) -> dict[str,
     return prepared_arguments
 
 
+def _raise_bounded_response_error(response: httpx.Response) -> None:
+    """Translate an upstream streaming error into the SDK's public exception."""
+    body = bytearray()
+    for chunk in response.iter_bytes():
+        remaining = MAX_INLINE_BYTES - len(body)
+        if remaining <= 0:
+            break
+        body.extend(chunk[:remaining])
+    text = bytes(body).decode("utf-8", errors="replace")
+    message = text or response.reason_phrase
+    error_type = None
+    try:
+        payload = json.loads(text)
+        if isinstance(payload, dict):
+            message = str(payload.get("message") or message)
+            raw_type = payload.get("type")
+            error_type = str(raw_type) if raw_type is not None else None
+    except (TypeError, ValueError):
+        pass
+    raise AppwriteException(message, response.status_code, error_type, text)
+
+
+def _perform_bounded_binary_client_call(
+    client: Client,
+    tool_name: str,
+    method: str,
+    path: str = "",
+    headers: dict[str, Any] | None = None,
+    params: dict[str, Any] | None = None,
+    response_type: str = "json",
+) -> bytes:
+    """Stream one SDK binary call into a bounded buffer for hosted HTTP."""
+    if method.lower() != "get" or response_type != "json":
+        raise RuntimeError(f"Unsupported bounded binary request for {tool_name}.")
+
+    request_headers = {
+        key: value
+        for key, value in {**client._global_headers, **(headers or {})}.items()
+        if value
+    }
+    # Prevent HTTPX from transparently inflating a compressed response into one
+    # oversized chunk before the decoded-byte limit can run.
+    request_headers["accept-encoding"] = "identity"
+    request_params = client.flatten(params or {})
+    endpoint = client._endpoint.rstrip("/")
+
+    with httpx.Client(
+        verify=not client._self_signed,
+        timeout=FETCH_TIMEOUT_SECONDS,
+        follow_redirects=True,
+    ) as http_client:
+        with http_client.stream(
+            method, endpoint + path, headers=request_headers, params=request_params
+        ) as response:
+            # Check before reading success or error bodies: HTTPX decodes
+            # ``iter_bytes()`` chunks, so either path could otherwise inflate a
+            # compressed response beyond the limit before we can count it.
+            content_encoding = response.headers.get("content-encoding", "identity")
+            if content_encoding.lower().strip() not in {"", "identity"}:
+                raise ValueError(
+                    "Hosted MCP cannot safely return a compressed binary response. "
+                    "Use an Appwrite SDK or REST API for this content."
+                )
+
+            if response.status_code >= 400:
+                _raise_bounded_response_error(response)
+
+            warning = response.headers.get("x-appwrite-warning")
+            if warning:
+                for item in warning.split(";"):
+                    print(f"Warning: {item}", file=sys.stderr)
+
+            declared = response.headers.get("content-length")
+            if declared:
+                try:
+                    content_length = int(declared)
+                except ValueError:
+                    content_length = None
+                if (
+                    content_length is not None
+                    and content_length > MAX_HOSTED_BINARY_RESPONSE_BYTES
+                ):
+                    raise HostedBinaryResponseTooLarge(
+                        tool_name,
+                        MAX_HOSTED_BINARY_RESPONSE_BYTES,
+                        content_length=content_length,
+                    )
+
+            body = bytearray()
+            for chunk in response.iter_bytes():
+                observed_bytes = len(body) + len(chunk)
+                if observed_bytes > MAX_HOSTED_BINARY_RESPONSE_BYTES:
+                    raise HostedBinaryResponseTooLarge(
+                        tool_name,
+                        MAX_HOSTED_BINARY_RESPONSE_BYTES,
+                        observed_bytes=observed_bytes,
+                    )
+                body.extend(chunk)
+            return bytes(body)
+
+
+def _bounded_binary_client_call(
+    client: Client,
+    tool_name: str,
+    method: str,
+    path: str = "",
+    headers: dict[str, Any] | None = None,
+    params: dict[str, Any] | None = None,
+    response_type: str = "json",
+) -> bytes:
+    try:
+        return _perform_bounded_binary_client_call(
+            client, tool_name, method, path, headers, params, response_type
+        )
+    except httpx.HTTPError as exc:
+        # Match the generated SDK contract so callers receive the existing
+        # Appwrite-formatted tool error instead of an internal HTTPX exception.
+        raise AppwriteException(str(exc)) from exc
+
+
 def execute_registered_tool(
     tools_manager: ToolManager,
     name: str,
@@ -843,13 +969,39 @@ def execute_registered_tool(
     # Re-bind the SDK method to a client authenticated for the current request.
     # An explicit client takes precedence (used by tests); otherwise it is resolved
     # from the request's OAuth access token.
+    hosted = client is None
     if client is None:
         client = resolve_client(target_project, organization_id)
     bound_method = getattr(service_cls(client), method_name)
+    bounded_binary = (
+        hosted and inspect.signature(bound_method).return_annotation is bytes
+    )
 
     parsed = _parse_tool_name(name)
     try:
-        result = bound_method(**prepared_arguments)
+        if bounded_binary:
+            original_call = client.call
+            setattr(
+                client,
+                "call",
+                lambda method, path="", headers=None, params=None, response_type="json": _bounded_binary_client_call(
+                    client,
+                    name,
+                    method,
+                    path,
+                    headers,
+                    params,
+                    response_type,
+                ),
+            )
+            try:
+                result = bound_method(**prepared_arguments)
+            finally:
+                setattr(client, "call", original_call)
+        else:
+            result = bound_method(**prepared_arguments)
+    except HostedBinaryResponseTooLarge:
+        raise
     except AppwriteException as exc:
         error_monitoring.capture_appwrite_exception(
             exc,

src/mcp_server_appwrite/service.py

@@ -23,11 +23,13 @@ def __init__(
         *,
         allowed_methods: frozenset[str] | None = None,
         context_scope: str = "console",
+        binary_response_limit: int | None = None,
     ):
         self.service = service_instance
         self.service_name = service_name
         self.allowed_methods = allowed_methods
         self.context_scope = context_scope
+        self.binary_response_limit = binary_response_limit
         self._method_name_overrides = self.get_method_name_overrides()
 
     def get_method_name_overrides(self) -> Dict[str, str]:
@@ -204,9 +206,20 @@ def list_tools(self) -> Dict[str, Dict]:
                 if param.default is param.empty:
                     required.append(param_name)
 
+            description = docstring.short_description or "No description available"
+            if (
+                self.binary_response_limit is not None
+                and type_hints.get("return") is bytes
+            ):
+                limit_mib = self.binary_response_limit // (1024 * 1024)
+                description = (
+                    f"{description} Hosted MCP returns binary responses up to "
+                    f"{limit_mib} MiB; use an Appwrite SDK or REST API for larger content."
+                )
+
             tool_definition = Tool(
                 name=tool_name,
-                description=docstring.short_description or "No description available",
+                description=description,
                 input_schema={
                     "type": "object",
                     "properties": properties,

tests/unit/test_error_classification.py

@@ -4,6 +4,7 @@
 from pydantic import BaseModel, ValidationError
 
 from mcp_server_appwrite.error_classification import (
+    HostedBinaryResponseTooLarge,
     WriteConfirmationRequired,
     classify_tool_error,
     is_response_parse_error,
@@ -17,6 +18,10 @@ def test_write_confirmation(self):
             "write_confirmation",
         )
 
+    def test_hosted_binary_response_too_large(self):
+        error = HostedBinaryResponseTooLarge("storage_get_file_download", 1024)
+        self.assertEqual(classify_tool_error(error), "response_too_large")
+
     def test_wrapped_appwrite_4xx(self):
         for code in (400, 401, 404, 409, 429, 499):
             with self.subTest(code=code):