| 20 | |
| 21 | |
| 22 | class McpServer: |
| 23 | def __init__(self, binary, cache_dir=None, extra_env=None, cwd=None): |
| 24 | self.binary = binary |
| 25 | self._id = 0 |
| 26 | self.proc = None |
| 27 | self._stderr = [] |
| 28 | env = dict(os.environ) |
| 29 | if cache_dir: |
| 30 | env["CBM_CACHE_DIR"] = cache_dir # isolate the graph DB location |
| 31 | if extra_env: |
| 32 | env.update(extra_env) |
| 33 | self.env = env |
| 34 | self.cwd = cwd |
| 35 | |
| 36 | def __enter__(self): |
| 37 | self.start() |
| 38 | return self |
| 39 | |
| 40 | def __exit__(self, *a): |
| 41 | self.close() |
| 42 | |
| 43 | def start(self): |
| 44 | self.proc = subprocess.Popen( |
| 45 | [self.binary], |
| 46 | stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE, |
| 47 | env=self.env, cwd=self.cwd, bufsize=0) |
| 48 | self._stderr_thread = threading.Thread(target=self._drain_stderr, daemon=True) |
| 49 | self._stderr_thread.start() |
| 50 | |
| 51 | def _drain_stderr(self): |
| 52 | try: |
| 53 | for line in self.proc.stderr: |
| 54 | self._stderr.append(line.decode("utf-8", "replace")) |
| 55 | except Exception: |
| 56 | pass |
| 57 | |
| 58 | def stderr_text(self): |
| 59 | return "".join(self._stderr) |
| 60 | |
| 61 | def _send(self, obj): |
| 62 | data = json.dumps(obj, ensure_ascii=False).encode("utf-8") |
| 63 | self.proc.stdin.write(data + b"\n") |
| 64 | self.proc.stdin.flush() |
| 65 | |
| 66 | def _read_message(self, timeout=60): |
| 67 | result = {} |
| 68 | |
| 69 | def reader(): |
| 70 | try: |
| 71 | result["line"] = self.proc.stdout.readline() |
| 72 | except Exception as ex: |
| 73 | result["exc"] = ex |
| 74 | |
| 75 | th = threading.Thread(target=reader, daemon=True) |
| 76 | th.start() |
| 77 | th.join(timeout) |
| 78 | if th.is_alive(): |
| 79 | raise McpError("timeout after %ss (hang)" % timeout) |
no outgoing calls