97 lines
3.5 KiB
Python
97 lines
3.5 KiB
Python
"""Minimal MCP streamable-http client. One session per worker."""
|
|
import json
|
|
import os
|
|
import random
|
|
import time
|
|
import urllib.request
|
|
import urllib.error
|
|
|
|
|
|
class McpError(Exception):
|
|
pass
|
|
|
|
|
|
class McpClient:
|
|
def __init__(self, url=None, token=None, timeout=300):
|
|
self.url = url or os.environ.get("MCP_URL", "http://127.0.0.1:3000/mcp")
|
|
self.token = token or os.environ.get("MCP_TOKEN", "")
|
|
self.timeout = timeout
|
|
self.sid = None
|
|
self._id = 0
|
|
self._lock_retries = 0
|
|
self.lock_retry_total = 0
|
|
|
|
def _post(self, body, method="POST"):
|
|
h = {"Content-Type": "application/json",
|
|
"Accept": "application/json, text/event-stream"}
|
|
if self.token:
|
|
h["Authorization"] = f"Bearer {self.token}"
|
|
if self.sid:
|
|
h["Mcp-Session-Id"] = self.sid
|
|
data = json.dumps(body).encode() if body is not None else None
|
|
req = urllib.request.Request(self.url, data, h, method=method)
|
|
with urllib.request.urlopen(req, timeout=self.timeout) as r:
|
|
sid = r.headers.get("Mcp-Session-Id")
|
|
raw = r.read().decode()
|
|
if sid:
|
|
self.sid = sid
|
|
if "data:" in raw[:40]:
|
|
raw = "".join(l[5:].strip() for l in raw.splitlines() if l.startswith("data:"))
|
|
return json.loads(raw) if raw.strip() else None
|
|
|
|
def open(self):
|
|
self._id += 1
|
|
self._post({"jsonrpc": "2.0", "id": self._id, "method": "initialize",
|
|
"params": {"protocolVersion": "2025-03-26", "capabilities": {},
|
|
"clientInfo": {"name": "abap-llm-harness", "version": "0.1"}}})
|
|
self._post({"jsonrpc": "2.0", "method": "notifications/initialized"})
|
|
return self
|
|
|
|
def close(self):
|
|
if self.sid:
|
|
try:
|
|
self._post(None, method="DELETE")
|
|
except Exception:
|
|
pass
|
|
self.sid = None
|
|
|
|
def list_tools(self):
|
|
self._id += 1
|
|
res = self._post({"jsonrpc": "2.0", "id": self._id, "method": "tools/list"})
|
|
return res["result"]["tools"]
|
|
|
|
def call(self, name, args, retry=True):
|
|
"""Return (is_error, text)."""
|
|
self._id += 1
|
|
body = {"jsonrpc": "2.0", "id": self._id, "method": "tools/call",
|
|
"params": {"name": name, "arguments": args}}
|
|
try:
|
|
res = self._post(body)
|
|
except urllib.error.HTTPError as e:
|
|
if e.code == 404 and retry: # session dropped (capacity limit)
|
|
self.sid = None
|
|
self.open()
|
|
return self.call(name, args, retry=False)
|
|
raise McpError(f"{name}: HTTP {e.code}") from e
|
|
if "error" in res:
|
|
return True, json.dumps(res["error"])
|
|
r = res["result"]
|
|
text = "\n".join(c.get("text", "") for c in r.get("content", []))
|
|
# The server shares one RFC connection between sessions. A parallel call gets
|
|
# "Concurrent call detected" instead of waiting. Retry; the caller never sees it.
|
|
if "Concurrent call detected" in text and self._lock_retries < 8:
|
|
self._lock_retries += 1
|
|
time.sleep(0.5 * self._lock_retries + random.random())
|
|
try:
|
|
return self.call(name, args, retry)
|
|
finally:
|
|
self._lock_retries = 0
|
|
self.lock_retry_total += self._lock_retries
|
|
return bool(r.get("isError")), text
|
|
|
|
def __enter__(self):
|
|
return self.open()
|
|
|
|
def __exit__(self, *a):
|
|
self.close()
|