test/harness/mcpdrv.py

1#!/usr/bin/env python3

2"""Drive apiary (an MCP-stdio server) as a scriptable client.

4Apiary has no CLI; running the binary speaks MCP JSON-RPC on

5stdin/stdout. This module spawns one apiary process per bot, performs

6the MCP handshake, and lets a test script call tools and observe

7inbound `notifications/claude/channel` events -- which is exactly the

8signal that says "this bot's trusted-set filter delivered that

9message".

11Everything is line-delimited JSON on stdout; stderr carries apiary's

12own structured log.

14Configuration comes from the environment so a single experiment script

15can be pointed at two different binaries for a matched pair:

17 APIARY_BIN path to the apiary release bundle to drive (required)

18 HARNESS_RUN state/log directory (default /tmp/apiary-harness)

19"""

21import json

22import os

23import queue

24import subprocess

25import sys

26import threading

27import time

29RUN_DIR = os.environ.get("HARNESS_RUN", "/tmp/apiary-harness")

31# Identifiers used by env-up.sh's seeded database. Generic on purpose:

32# this is a public repo and these must not encode any particular

33# deployment's nicks, hosts or channels.

34OWNER = os.environ.get("HARNESS_OWNER", "owner-1")

35LEADER_NICK = os.environ.get("HARNESS_LEADER", "leader-1")

36CHANNEL = os.environ.get("HARNESS_CHANNEL", "#coord")

37WORKER_GROUP = os.environ.get("HARNESS_GROUP", "test-workers")

40def apiary_bin():

41 b = os.environ.get("APIARY_BIN")

42 if not b:

43 raise SystemExit(

44 "APIARY_BIN is not set. Point it at the apiary release "

45 "bundle to drive, e.g.\n"

46 " APIARY_BIN=$PWD/build/release/bin/apiary python3 "

47 "test/harness/exp_silent_loss.py\n"

48 "Refusing to guess: a matched pair is worthless if both "

49 "arms silently run the same binary.")

50 if not os.access(b, os.X_OK):

51 raise SystemExit(f"APIARY_BIN is not executable: {b}")

52 return b

55def load_state(run_dir=None):

56 """Read env-up.sh's state.env."""

57 run_dir = run_dir or RUN_DIR

58 path = os.path.join(run_dir, "state.env")

59 if not os.path.exists(path):

60 raise SystemExit(

61 f"no {path} -- run test/harness/env-up.sh first")

62 state = {}

63 with open(path) as f:

64 for line in f:

65 k, _, v = line.strip().partition("=")

66 if k:

67 state[k] = v

68 return state

71class Bot:

72 def __init__(self, label, env, bin=None, log_dir=None):

73 self.label = label

74 self.env = dict(os.environ)

75 # Scrub any inherited APIARY_* so a stray value from the shell

76 # cannot silently change what we think we are testing.

77 for k in list(self.env):

78 if k.startswith("APIARY_"):

79 del self.env[k]

80 self.env.update({k: str(v) for k, v in env.items()})

81 log_dir = log_dir or RUN_DIR

82 os.makedirs(log_dir, exist_ok=True)

83 self.stderr_path = os.path.join(log_dir, f"{label}.stderr")

84 self._stderr_f = open(self.stderr_path, "w")

85 self.bin = bin or apiary_bin()

86 self.proc = subprocess.Popen(

87 [self.bin],

88 stdin=subprocess.PIPE,

89 stdout=subprocess.PIPE,

90 stderr=self._stderr_f,

91 env=self.env,

92 text=True,

93 bufsize=1,

94 )

95 self._id = 0

96 self._responses = {}

97 self._resp_lock = threading.Condition()

98 self.notifications = []

99 self._notif_lock = threading.Condition()

100 self.raw = []

101 self._reader = threading.Thread(target=self._read_loop, daemon=True)

102 self._reader.start()

103

104 # ---- plumbing -------------------------------------------------

105

106 def _read_loop(self):

107 for line in self.proc.stdout:

108 line = line.strip()

109 if not line:

110 continue

111 self.raw.append(line)

112 try:

113 msg = json.loads(line)

114 except json.JSONDecodeError:

115 continue

116 if "id" in msg and ("result" in msg or "error" in msg):

117 with self._resp_lock:

118 self._responses[msg["id"]] = msg

119 self._resp_lock.notify_all()

120 elif msg.get("method") == "notifications/claude/channel":

121 with self._notif_lock:

122 self.notifications.append(msg["params"])

123 self._notif_lock.notify_all()

124

125 def _send(self, obj):

126 self.proc.stdin.write(json.dumps(obj) + "\n")

127 self.proc.stdin.flush()

128

129 def _request(self, method, params=None, timeout=60):

130 self._id += 1

131 rid = self._id

132 self._send({"jsonrpc": "2.0", "id": rid, "method": method,

133 "params": params or {}})

134 deadline = time.time() + timeout

135 with self._resp_lock:

136 while rid not in self._responses:

137 remaining = deadline - time.time()

138 if remaining <= 0:

139 raise TimeoutError(f"{self.label}: {method} timed out")

140 self._resp_lock.wait(remaining)

141 return self._responses.pop(rid)

142

143 # ---- MCP ------------------------------------------------------

144

145 def initialize(self, timeout=60):

146 r = self._request("initialize", {

147 "protocolVersion": "2024-11-05",

148 "capabilities": {},

149 "clientInfo": {"name": "apiary-harness", "version": "0.0"},

150 }, timeout=timeout)

151 self._send({"jsonrpc": "2.0", "method": "notifications/initialized",

152 "params": {}})

153 return r

154

155 def tools(self):

156 r = self._request("tools/list")

157 return sorted(t["name"] for t in r["result"]["tools"])

158

159 def call(self, tool, timeout=60, **args):

160 r = self._request("tools/call", {"name": tool, "arguments": args},

161 timeout=timeout)

162 if "error" in r:

163 return f"<JSONRPC-ERROR {json.dumps(r['error'])}>"

164 parts = [c.get("text", "") for c in r["result"].get("content", [])]

165 return "\n".join(parts)

166

167 def wait_connected(self, timeout=25.0):

168 """Block until the bridge logs a successful connect."""

169 deadline = time.time() + timeout

170 while time.time() < deadline:

171 if self.grep_stderr("Apiary bridge connected"):

172 return True

173 time.sleep(0.1)

174 return False

175

176 # ---- observation ---------------------------------------------

177

178 def wait_notification(self, match, timeout=20):

179 """Wait for an inbound channel notification whose content

180 contains `match`. Returns the params dict, or None on timeout."""

181 deadline = time.time() + timeout

182 seen = 0

183 with self._notif_lock:

184 while True:

185 while seen < len(self.notifications):

186 n = self.notifications[seen]

187 seen += 1

188 if match in n.get("content", ""):

189 return n

190 remaining = deadline - time.time()

191 if remaining <= 0:

192 return None

193 self._notif_lock.wait(remaining)

194

195 def contents(self):

196 return [n.get("content", "") for n in self.notifications]

197

198 def alive(self):

199 return self.proc.poll() is None

200

201 def stderr_text(self):

202 self._stderr_f.flush()

203 with open(self.stderr_path) as f:

204 return f.read()

205

206 def grep_stderr(self, pat):

207 return [l for l in self.stderr_text().splitlines() if pat in l]

208

209 def close(self):

210 try:

211 self.proc.terminate()

212 self.proc.wait(timeout=10)

213 except Exception:

214 try:

215 self.proc.kill()

216 except Exception:

217 pass

218 try:

219 self._stderr_f.close()

220 except Exception:

221 pass

222

223

224def leader_env(port, token, nick=None, host="127.0.0.1",

225 channel=None, owner=None, **extra):

226 e = {

227 "APIARY_ENCLAVE_HOST": host,

228 "APIARY_ENCLAVE_PORT": port,

229 "APIARY_ENCLAVE_TLS": "no",

230 "APIARY_USER": nick or LEADER_NICK,

231 "APIARY_TOKEN": token,

232 "APIARY_CHANNEL": channel or CHANNEL,

233 "APIARY_MODE": "leader",

234 "APIARY_OWNER_NICK": owner or OWNER,

235 "APIARY_REPORTS_TO": owner or OWNER,

236 "APIARY_SERVICES_TIMEOUT": "8.0",

237 "APIARY_WORKER_GROUP": WORKER_GROUP,

238 }

239 e.update(extra)

240 return e

241

242

243def worker_env(port, token, nick, reports_to=None,

244 host="127.0.0.1", channel=None, owner=None, **extra):

245 e = {

246 "APIARY_ENCLAVE_HOST": host,

247 "APIARY_ENCLAVE_PORT": port,

248 "APIARY_ENCLAVE_TLS": "no",

249 "APIARY_USER": nick,

250 "APIARY_TOKEN": token,

251 "APIARY_CHANNEL": channel or CHANNEL,

252 "APIARY_MODE": "worker",

253 "APIARY_OWNER_NICK": owner or OWNER,

254 "APIARY_REPORTS_TO": reports_to or LEADER_NICK,

255 "APIARY_SERVICES_TIMEOUT": "8.0",

256 }

257 e.update(extra)

258 return e

259

260

261def unique_nick(prefix):

262 """A nick that has not been used on this server before.

263

264 Revoking a bot DISABLES the row rather than deleting it, so the

265 nick stays burned: re-running an experiment with a fixed nick fails

266 with `name-taken`. That failure is loud, but in a matched pair it

267 lands on the SECOND arm, which is the arm you are testing -- so a

268 careless run reports the pre-change result and nothing else.

269 """

270 return f"{prefix}{int(time.time()) % 100000}"

271

272

273def say(*a):

274 print(*a, flush=True)

275

276

277def check(label, ok, detail=""):

278 say(f" [{'PASS' if ok else 'FAIL'}] {label}"

279 + (f" -- {detail}" if detail else ""))

280 return bool(ok)