Snowflake-managed MCP or the Cortex Agents REST API: code¶

#!/usr/bin/env python3
"""
Stream diagnostic for a Snowflake-managed MCP server.

Sends one tools/call the way a spec-compliant client does and times every frame
of the Server-Sent Events stream as it arrives. It reports:

  transport          Content-Type, plus the no-buffering headers the server sets
  keep-alives        count of `: keep-alive` comment frames and their interval
  longest gap        the longest silence on the wire, from the request onward,
                     against your proxy's idle timeout
  progress           notifications/progress events received for the progressToken it
                     sends (--no-progress-token sends none, for comparison)
  result             result events, and whether the stream ended with `data: [DONE]`
  request id         x-snowflake-request-id, which joins CORTEX_AGENT_USAGE_HISTORY

Use it to find out why a client times out. If the longest gap is under the idle
timeout of every proxy between the client and Snowflake, keep-alives are doing
their job and a timeout comes from the client's own per-request limit instead.
If progress events arrive, a client that resets its timeout on progress (the MCP
specification allows it) can wait longer than that limit.

Standard library only.

    export SNOWFLAKE_ACCOUNT_URL=https://<account_identifier>.snowflakecomputing.com
    export SNOWFLAKE_TOKEN=<OAuth access token or programmatic access token>
    python client/mcp_stream_check.py
    python client/mcp_stream_check.py --idle-timeout 30 --compare-json
    python client/mcp_stream_check.py --tool revenue_agent --text "Total revenue by region?"

Use hyphens, not underscores, in the account hostname.
"""

import argparse
import json
import os
import sys
import time
import urllib.error
import urllib.request

DATABASE = os.environ.get("MCP_DATABASE", "COOKBOOK_MCP_OR_AGENT_REST_API")
SCHEMA = os.environ.get("MCP_SCHEMA", "APP")
SERVER = os.environ.get("MCP_SERVER", "REVENUE_MCP")

SSE_ACCEPT = "application/json, text/event-stream"
PROGRESS_TOKEN = "stream-check-1"


def _url(account_url):
    return f"{account_url.rstrip('/')}/api/v2/databases/{DATABASE}/schemas/{SCHEMA}/mcp-servers/{SERVER}"


def _request(account_url, token, payload, accept):
    return urllib.request.Request(
        _url(account_url),
        data=json.dumps(payload).encode(),
        method="POST",
        headers={
            "Authorization": f"Bearer {token}",
            "Content-Type": "application/json",
            "Accept": accept,
        },
    )


def _open(request, timeout):
    try:
        return urllib.request.urlopen(request, timeout=timeout)
    except urllib.error.HTTPError as err:
        detail = err.read().decode(errors="replace")[:500]
        sys.exit(f"HTTP {err.code}: {detail}")


def list_tools(account_url, token, timeout):
    """tools/list answers with one JSON body whatever the Accept header says."""
    payload = {"jsonrpc": "2.0", "id": 1, "method": "tools/list"}
    with _open(_request(account_url, token, payload, SSE_ACCEPT), timeout) as response:
        content_type = response.headers.get("Content-Type", "")
        raw = response.read().decode()
    if "text/event-stream" in content_type:
        raw = next((ln[5:].strip() for ln in raw.splitlines() if ln.startswith("data:")), "{}")
    tools = (json.loads(raw).get("result") or {}).get("tools") or []
    return tools, content_type


def _text_argument(tool):
    """Name of the argument that carries the question: `text` for agent tools,
    `message` for Cortex Analyst, `query` for Cortex Search."""
    schema = tool.get("inputSchema") or {}
    for name in schema.get("required") or []:
        if name in ("text", "message", "query"):
            return name
    for name in ("text", "message", "query"):
        if name in (schema.get("properties") or {}):
            return name
    return "text"


def stream_call(account_url, token, tool, arguments, accept, timeout, progress_token=True):
    """One tools/call. Reads the body line by line and timestamps every frame."""
    params = {"name": tool, "arguments": arguments}
    if progress_token:
        params["_meta"] = {"progressToken": PROGRESS_TOKEN}
    payload = {"jsonrpc": "2.0", "id": 2, "method": "tools/call", "params": params}
    report = {
        "accept": accept,
        "progress_token": progress_token,
        "keepalives": [],
        "progress": [],
        "results": 0,
        "done": False,
        "other": [],
        "error": None,
        "result_bytes": 0,
    }
    started = time.monotonic()
    frame_times = []
    with _open(_request(account_url, token, payload, accept), timeout) as response:
        report["headers_s"] = time.monotonic() - started
        report["status"] = response.status
        for name in ("Content-Type", "Cache-Control", "X-Accel-Buffering", "x-snowflake-request-id"):
            report[name] = response.headers.get(name, "")
        streaming = "text/event-stream" in report["Content-Type"]
        if not streaming:
            body = response.read()
            report["first_byte_s"] = time.monotonic() - started
            frame_times.append(report["first_byte_s"])
            report["result_bytes"] = len(body)
            report["results"] = 1
            rpc = json.loads(body or b"{}")
            report["error"] = rpc.get("error")
        else:
            while True:
                line = response.readline()
                if not line:
                    break  # stream closed
                at = time.monotonic() - started
                text = line.decode(errors="replace").rstrip("\r\n")
                if not text:
                    continue  # event separator
                report.setdefault("first_byte_s", at)
                frame_times.append(at)
                if text.startswith(":"):
                    report["keepalives"].append(at)
                elif text.startswith("data:"):
                    _classify(text[5:].strip(), at, report)
                elif not text.startswith(("event:", "id:", "retry:")):
                    report["other"].append(text[:80])
    report["total_s"] = time.monotonic() - started
    report["streaming"] = streaming
    # Measure from the request, not from the headers: a proxy's idle clock starts
    # when the request is sent, and the server sends headers with the first frame.
    edges = [0.0, *frame_times, report["total_s"]]
    report["longest_gap_s"] = max(b - a for a, b in zip(edges, edges[1:]))
    return report


def _classify(data, at, report):
    if data == "[DONE]":
        report["done"] = True
        return
    try:
        message = json.loads(data)
    except json.JSONDecodeError:
        report["other"].append(data[:80])
        return
    if message.get("method") == "notifications/progress":
        report["progress"].append((at, (message.get("params") or {}).get("progress")))
    elif "result" in message or "error" in message:
        report["results"] += 1
        report["result_bytes"] += len(data)
        report["error"] = report["error"] or message.get("error")
    else:
        report["other"].append(data[:80])


def _interval(times):
    if len(times) < 2:
        return "n/a"
    gaps = [b - a for a, b in zip(times, times[1:])]
    return f"every {sum(gaps) / len(gaps):.1f}s (min {min(gaps):.1f}s, max {max(gaps):.1f}s)"


def print_report(report, idle_timeout):
    print(f"\n  Accept             {report['accept']}")
    print(f"  progressToken      {'sent' if report['progress_token'] else 'not sent'}")
    print(f"  HTTP status        {report['status']}")
    print(f"  Content-Type       {report['Content-Type'] or '(none)'}")
    print(f"  Cache-Control      {report['Cache-Control'] or '(none)'}")
    print(f"  X-Accel-Buffering  {report['X-Accel-Buffering'] or '(none)'}")
    print(f"  request id         {report['x-snowflake-request-id'] or '(none)'}")
    print(f"  headers after      {report['headers_s']:6.1f}s")
    print(f"  first byte after   {report.get('first_byte_s', report['total_s']):6.1f}s")
    if report["streaming"]:
        print(f"  keep-alives        {len(report['keepalives'])}, {_interval(report['keepalives'])}")
        print(f"  progress events    {len(report['progress'])}, {_interval([t for t, _ in report['progress']])}")
        print(f"  result events      {report['results']} ({report['result_bytes'] / 1024:.0f} KB)")
        print(f"  ended with [DONE]  {'yes' if report['done'] else 'no, the stream closed'}")
        if report["other"]:
            print(f"  other frames       {len(report['other'])}, first: {report['other'][0]}")
    else:
        print(f"  body               one JSON document ({report['result_bytes'] / 1024:.0f} KB)")
    print(f"  total              {report['total_s']:6.1f}s")
    flag = "  OVER the idle timeout" if report["longest_gap_s"] >= idle_timeout else ""
    print(f"  longest silence    {report['longest_gap_s']:6.1f}s  (idle timeout {idle_timeout:.0f}s){flag}")
    if report["error"]:
        print(f"  JSON-RPC error     {report['error']}")


def verdict(report, idle_timeout, client_deadline):
    """Plain-language reading of one streamed call."""
    lines = []
    if not report["streaming"]:
        lines.append(
            "The server sent one JSON body. Managed MCP streams tools/call by default, so the "
            "Accept header that reached Snowflake did not list text/event-stream. Check the "
            "client and any proxy that rewrites headers."
        )
    elif report["longest_gap_s"] < idle_timeout:
        lines.append(
            f"Keep-alives kept every silence under {idle_timeout:.0f}s, so a proxy with that idle "
            "timeout will not drop this call."
        )
    else:
        lines.append(
            f"A {report['longest_gap_s']:.0f}s silence exceeds the {idle_timeout:.0f}s idle timeout. "
            "Look for a proxy that buffers the response (it would hide keep-alives)."
        )
    if report["total_s"] >= client_deadline:
        lines.append(
            f"The call took {report['total_s']:.0f}s, past the {client_deadline:.0f}s client limit. "
            "Streaming does not deliver the answer earlier. Reset the timeout on progress if your "
            "client supports it, make the agent faster, or split the work into submit and fetch "
            "calls (see the long-running agents over MCP recipe)."
        )
    if report["streaming"] and report["progress_token"]:
        if report["progress"]:
            lines.append(
                f"{len(report['progress'])} progress notifications arrived. A client that resets its "
                "per-request timeout on progress keeps waiting while the tool runs; one that does not "
                "still times out at its own limit."
            )
        else:
            lines.append(
                "No progress notifications arrived for the progressToken. Rely on keep-alives "
                "for idle timeouts and on your client's per-request limit for everything else."
            )
    if report["streaming"] and not report["done"]:
        lines.append("The stream ended on close, without data: [DONE]. End reading on close, not on the sentinel.")
    if report["results"] != 1:
        lines.append(f"Expected one result event, got {report['results']}.")
    return lines


def main():
    parser = argparse.ArgumentParser(description="Time the SSE stream of one managed MCP tools/call.")
    parser.add_argument("--tool", help="tool name; defaults to the first tool the server lists")
    parser.add_argument("--text", default="What was total revenue by region in August 2026?")
    parser.add_argument(
        "--idle-timeout",
        type=float,
        default=float(os.environ.get("MCP_IDLE_TIMEOUT_S", "60")),
        help="shortest idle timeout of any proxy or load balancer on the path (seconds)",
    )
    parser.add_argument(
        "--client-deadline",
        type=float,
        default=float(os.environ.get("MCP_CLIENT_CEILING_S", "120")),
        help="your MCP client's absolute per-call limit (seconds)",
    )
    parser.add_argument(
        "--no-progress-token",
        action="store_true",
        help="omit _meta.progressToken, as a client that does not track progress would",
    )
    parser.add_argument(
        "--compare-json",
        action="store_true",
        help="repeat the call with Accept: application/json to show what a non-streaming client sees",
    )
    parser.add_argument("--timeout", type=float, default=600, help="socket read timeout (seconds)")
    args = parser.parse_args()

    account_url = os.environ.get("SNOWFLAKE_ACCOUNT_URL")
    token = os.environ.get("SNOWFLAKE_TOKEN")
    if not account_url or not token:
        sys.exit("Set SNOWFLAKE_ACCOUNT_URL and SNOWFLAKE_TOKEN.")

    tools, list_type = list_tools(account_url, token, args.timeout)
    if not tools:
        sys.exit("tools/list returned no tools: check USAGE on the server for the session's role.")
    tool = next((t for t in tools if t.get("name") == args.tool), None) if args.tool else tools[0]
    if tool is None:
        sys.exit(f"no tool named {args.tool}; the server lists {[t.get('name') for t in tools]}")
    print(f"\nserver   {DATABASE}.{SCHEMA}.{SERVER}")
    print(f"tools    {[t.get('name') for t in tools]}  (tools/list Content-Type: {list_type})")
    print(f"calling  {tool['name']}")

    arguments = {_text_argument(tool): args.text}
    progress_token = not args.no_progress_token
    report = stream_call(account_url, token, tool["name"], arguments, SSE_ACCEPT, args.timeout, progress_token)
    print_report(report, args.idle_timeout)
    if args.compare_json:
        print("\nsame call, Accept: application/json only")
        json_report = stream_call(
            account_url, token, tool["name"], arguments, "application/json", args.timeout, progress_token
        )
        print_report(json_report, args.idle_timeout)

    print("\nverdict")
    for line in verdict(report, args.idle_timeout, args.client_deadline):
        print(f"  - {line}")
    healthy = (
        report["streaming"]
        and report["results"] == 1
        and report["longest_gap_s"] < args.idle_timeout
        and report["total_s"] < args.client_deadline
    )
    sys.exit(0 if healthy else 1)


if __name__ == "__main__":
    main()