okami.mom/albedo

git clone https://git.nkp.pet/okami.mom/albedo

nekomimi.petcodemode: mcp callable from python, 256 MiB kernel memory cap, cell output spilled to a file4359cba

main
branchesmainall refs →
16.8 KiB409 linesraw
1"""MCP stdio discovery, credentials, namespaced calls, and teardown through the daemon."""
2
3from typing import Any
4from concurrent.futures import ThreadPoolExecutor
5import errno
6import json
7import os
8import shutil
9import sys
10import time
11import unittest
12from pathlib import Path
13
14from harness import Albedo, Provider, exclusive, operation_id, Reply, text
15
16SERVER = Path(__file__).with_name("fake_mcp_server.py")
17
18
19# exclusive: fixture configures global MCP servers, credentials, and daemon environment
20@exclusive
21class McpTests(unittest.TestCase):
22    def setUp(self):
23        self.message = "ping from albedo"
24        self.python_mode = False
25
26        def script(request):
27            tools = sorted(tool["name"] for tool in request.get("tools", []))
28            answered = any(
29                item.get("type") == "function_call_output" for item in request["input"]
30            )
31            mcp_tool = next((name for name in tools if name.startswith("mcp_")), None)
32            if mcp_tool and not answered:
33                if self.python_mode:
34                    code = (
35                        "import asyncio\n"
36                        'r = await mcp.call("fake/echo", message="ping from python")\n'
37                        "print(r.content[0].text)\n"
38                        "print(await mcp.tools())\n"
39                        "results = await asyncio.gather(\n"
40                        '    mcp.call("fake/echo", message="batch one"),\n'
41                        '    mcp.call("fake/echo", message="batch two"),\n'
42                        "    return_exceptions=True,\n"
43                        ")\n"
44                        "print(results)"
45                    )
46                    return Reply("python", tool_arguments={"code": code})
47                return Reply(
48                    "python",
49                    tool_name=mcp_tool,
50                    tool_arguments={"message": self.message},
51                )
52            return text("done")
53
54        self.provider = Provider(script)
55        self.addCleanup(self.provider.close)
56
57        def prepare(app):
58            app.env.update(
59                ALBEDO_MCP_SECRET="configured-secret",
60                ALBEDO_MCP_CLOSED=str(app.root / "closed"),
61                ALBEDO_MCP_AMBIENT="must-not-reach-the-server",
62                ALBEDO_MCP_CONTROL=str(app.root / "control"),
63            )
64            (app.root / "control").mkdir()
65            self.configure(app, [str(SERVER)])
66            app.store_secrets(
67                "mcp", {"fake": {"env": {"FAKE_SECRET": "stored-secret"}}}
68            )
69
70        self.app = Albedo(self.provider, protocol="responses", prepare=prepare)
71        self.app.__enter__()
72        self.addCleanup(self.app.__exit__, None, None, None)
73        self.session = self.app.session()
74
75    def configure(self, app, args, startup=20000, retry=None):
76        path = app.home / "extensions.json"
77        settings: dict[str, Any] = (
78            json.loads(path.read_text())
79            if path.exists()
80            else {"models": {"refreshHours": 0}}
81        )
82        settings["mcp"] = {
83            "servers": {
84                "fake": {
85                    "type": "stdio",
86                    "command": sys.executable,
87                    "args": args,
88                    "env": {
89                        "FAKE_SECRET": {"env": "ALBEDO_MCP_SECRET"},
90                        "FAKE_MCP_CLOSED": {"env": "ALBEDO_MCP_CLOSED"},
91                        "FAKE_MCP_CONTROL": {"env": "ALBEDO_MCP_CONTROL"},
92                    },
93                    "startupTimeoutMs": startup,
94                }
95            }
96        }
97        if retry is not None:
98            settings["mcp"]["retryMs"] = retry
99        path.write_text(json.dumps(settings))
100
101    def tools(self, requests):
102        return [
103            tool["name"]
104            for tool in requests[0]["tools"]
105            if tool["name"].startswith("mcp_")
106        ]
107
108    def mcp(self):
109        with self.app.api(f"/sessions/{self.session}/catalog") as response:
110            catalog = json.load(response)
111        return next(
112            row
113            for row in catalog["discovery"]["candidates"]
114            if row["kind"] == "extension" and row["preference_key"] == "mcp"
115        )
116
117    def select_mcp(self, session, enabled):
118        route = f"/sessions/{session}?view=configuration"
119        with self.app.api(route) as response:
120            json.load(response)
121            etag = response.headers["ETag"]
122        self.app.api(
123            route,
124            {"selection": {"extensions": {"mcp": enabled}}},
125            method="PATCH",
126            headers={"If-Match": etag},
127        ).close()
128
129    def set_mcp(self, enabled):
130        self.select_mcp(self.session, enabled)
131        with self.app.api(
132            f"/sessions/{self.session}/reload", {"target": "session"}
133        ) as response:
134            return json.load(response)
135
136    def turn(self, prompt, session=None):
137        session = session or self.session
138        before = len(self.provider.requests)
139        self.app.prompt(session, prompt).close()
140        self.app.idle(session, timeout=60)
141        return [entry["request"] for entry in self.provider.requests[before:]]
142
143    def test_cold_command_and_recorded_recovery_do_not_block_another_session(self):
144        other = self.app.session()
145        for session, enabled in ((self.session, True), (other, False)):
146            route = f"/sessions/{session}?view=configuration"
147            with self.app.api(route) as response:
148                json.load(response)
149                etag = response.headers["ETag"]
150            self.app.api(
151                route,
152                {"selection": {"extensions": {"mcp": enabled}}},
153                method="PATCH",
154                headers={"If-Match": etag},
155            ).close()
156        with self.app.api(f"/sessions/{self.session}?tail=0") as response:
157            self.assertEqual(json.load(response)["kernel"]["state"], "none")
158        marker = self.app.workspace / "cold-mcp-started"
159        gate = self.app.workspace / "cold-mcp-release"
160        os.mkfifo(gate)
161        script = (
162            "import pathlib; "
163            + f"pathlib.Path({str(marker)!r}).touch(); "
164            + f"open({str(gate)!r}, 'rb').read(1); "
165            + f"exec(pathlib.Path({str(SERVER)!r}).read_text())"
166        )
167        self.configure(self.app, ["-c", script], startup=30000)
168        input_id = operation_id()
169
170        def submit_command():
171            with self.app.api(
172                f"/sessions/{self.session}/inputs/{input_id}",
173                {
174                    "kind": "command",
175                    "command_id": "/research",
176                    "arguments": {"arguments": "cold composition fixture"},
177                },
178                method="PUT",
179            ) as response:
180                return json.load(response)
181
182        def release_preparation():
183            deadline = time.monotonic() + 10
184            while True:
185                try:
186                    descriptor = os.open(gate, os.O_WRONLY | os.O_NONBLOCK)
187                    break
188                except OSError as error:
189                    if error.errno != errno.ENXIO or time.monotonic() >= deadline:
190                        raise
191                    time.sleep(0.01)
192            try:
193                os.write(descriptor, b"x")
194            finally:
195                os.close(descriptor)
196
197        def complete_independent_turn():
198            with self.app.api(f"/sessions/{other}?tail=0") as response:
199                self.assertEqual(json.load(response)["id"], other)
200            with self.app.prompt(other, "independent turn") as response:
201                other_input = json.load(response)["id"]
202            self.app.idle(other)
203            with self.app.api(f"/sessions/{other}/inputs/{other_input}") as response:
204                self.assertEqual(json.load(response)["turn"]["state"], "completed")
205
206        with ThreadPoolExecutor(max_workers=2) as pool:
207            command = pool.submit(submit_command)
208            deadline = time.monotonic() + 10
209            while not marker.exists() and time.monotonic() < deadline:
210                time.sleep(0.01)
211            self.assertTrue(marker.exists(), "cold MCP preparation did not start")
212            try:
213                complete_independent_turn()
214                self.assertFalse(command.done(), "command passed the held preparation")
215            finally:
216                release_preparation()
217            self.assertEqual(command.result(timeout=40)["id"], input_id)
218        self.app.idle(self.session)
219        with self.app.api(f"/sessions/{self.session}/inputs/{input_id}") as response:
220            self.assertEqual(json.load(response)["turn"]["state"], "completed")
221        with self.app.api(f"/sessions/{self.session}?tail=0") as response:
222            original_kernel = json.load(response)["kernel"]["instance_id"]
223        self.assertIsNotNone(original_kernel)
224        original_history = self.app.history(self.session)["items"]
225        marker.unlink()
226        self.app.restart()
227        deadline = time.monotonic() + 10
228        while not marker.exists() and time.monotonic() < deadline:
229            time.sleep(0.01)
230        self.assertTrue(marker.exists(), "recorded kernel preparation did not start")
231        try:
232            complete_independent_turn()
233        finally:
234            release_preparation()
235        deadline = time.monotonic() + 20
236        while True:
237            with self.app.api(f"/sessions/{self.session}?tail=0") as response:
238                recovered = json.load(response)["kernel"]
239            if recovered["state"] == "attached" or time.monotonic() >= deadline:
240                break
241            time.sleep(0.01)
242        self.assertEqual(recovered["state"], "attached")
243        self.assertEqual(recovered["instance_id"], original_kernel)
244        self.assertEqual(self.app.history(self.session)["items"], original_history)
245
246    def test_discovery_namespaced_call_and_scrubbed_credentials(self):
247        self.assertEqual(self.set_mcp(True)["session"]["state"], "applied")
248        self.assertTrue(self.mcp()["effective_enabled"])
249        self.assertEqual(self.mcp()["extension"]["plugins"], ["managed"])
250        requests = self.turn("use the mcp server")
251        self.assertEqual(len(requests), 2)
252        context = requests[0]["instructions"]
253        self.assertNotIn("<extension-context", json.dumps(requests[0]["input"]))
254        self.assertIn('name="mcp"', context)
255        self.assertIn("echo", context)
256        advertised = [
257            tool["name"]
258            for tool in requests[0]["tools"]
259            if tool["name"].startswith("mcp_")
260        ]
261        self.assertEqual(len(advertised), 1)
262        self.assertTrue(advertised[0].startswith("mcp_fake_echo_"))
263        output = next(
264            item["output"]
265            for item in requests[1]["input"]
266            if item.get("type") == "function_call_output"
267        )
268        self.assertEqual(
269            json.loads(json.loads(output)["content"][0]["text"]),
270            {"echoed": "ping from albedo", "secret": "stored-secret", "ambient": None},
271        )
272
273    def test_python_binding_calls_and_batches_mcp_operations(self):
274        self.python_mode = True
275        self.set_mcp(True)
276        requests = self.turn("compose mcp calls in python")
277        self.assertEqual(len(requests), 2)
278        output = next(
279            item["output"]
280            for item in requests[1]["input"]
281            if item.get("type") == "function_call_output"
282        )
283        rendered = json.dumps(output)
284        self.assertIn("ping from python", rendered)
285        self.assertIn("fake", rendered)
286        self.assertIn("echo", rendered)
287        self.assertIn("batch one", rendered)
288        self.assertIn("batch two", rendered)
289
290    def test_a_failed_tool_call_is_refused_without_ending_the_turn(self):
291        self.set_mcp(True)
292        self.message = "fail"
293        requests = self.turn("the server will refuse")
294        self.assertEqual(len(requests), 2)
295        output = next(
296            item["output"]
297            for item in requests[1]["input"]
298            if item.get("type") == "function_call_output"
299        )
300        self.assertIn("MCP request failed", json.loads(output)["error"])
301        self.assertIn("done", json.dumps(self.app.history(self.session)))
302
303    def test_disable_closes_server_and_removes_tools(self):
304        self.set_mcp(True)
305        self.assertEqual(self.set_mcp(False)["session"]["state"], "applied")
306        self.assertFalse(self.mcp()["effective_enabled"])
307        closed = self.app.root / "closed"
308        deadline = time.monotonic() + 15
309        while time.monotonic() < deadline and not closed.exists():
310            time.sleep(0.05)
311        self.assertTrue(closed.exists(), "disabling must close the server")
312        requests = self.turn("the server is gone")
313        self.assertFalse(
314            any(tool["name"].startswith("mcp_") for tool in requests[0]["tools"])
315        )
316
317    def test_insecure_credentials_fail_closed(self):
318        credentials = self.app.home / "creds.json"
319        credentials.chmod(0o644)
320        self.addCleanup(credentials.chmod, 0o600)
321        outcome = self.set_mcp(True)
322        self.assertEqual(outcome["session"]["state"], "failed")
323        self.assertIn("0600", outcome["session"]["failure"]["detail"])
324        self.assertFalse(self.provider.requests)
325        credentials.chmod(0o600)
326        self.assertEqual(self.set_mcp(True)["session"]["state"], "applied")
327
328    def test_an_unavailable_server_is_skipped_and_the_rest_load(self):
329        path = self.app.home / "extensions.json"
330        settings = json.loads(path.read_text())
331        settings["mcp"]["servers"]["dead"] = {
332            "type": "stdio",
333            "command": str(self.app.root / "missing-server"),
334            "startupTimeoutMs": 3000,
335        }
336        path.write_text(json.dumps(settings))
337        self.assertEqual(self.set_mcp(True)["session"]["state"], "applied")
338        advertised = self.tools(self.turn("one server is down"))
339        self.assertEqual(len(advertised), 1)
340        self.assertTrue(advertised[0].startswith("mcp_fake_echo_"))
341
342    def test_a_server_that_comes_online_joins_the_session_after_a_turn(self):
343        late = self.app.root / "late_server.py"
344        self.configure(self.app, [str(late)], startup=3000, retry=200)
345        self.assertEqual(self.set_mcp(True)["session"]["state"], "applied")
346        self.assertEqual(self.tools(self.turn("nothing is up yet")), [])
347        shutil.copy(SERVER, late)
348        deadline = time.monotonic() + 40
349        advertised = []
350        while time.monotonic() < deadline and not advertised:
351            time.sleep(0.5)
352            advertised = self.tools(self.turn("is it up now?"))
353        self.assertEqual(len(advertised), 1, "the server never joined the session")
354        self.assertTrue(advertised[0].startswith("mcp_fake_echo_"))
355        history = json.dumps(self.app.history(self.session))
356        self.assertIn("capabilities changed", history)
357        self.assertIn("fake", history)
358
359    def test_a_known_server_does_not_hold_up_opening_a_session(self):
360        self.set_mcp(True)
361        self.assertEqual(len(self.tools(self.turn("learn the catalogue"))), 1)
362        (self.app.root / "control" / "delay").write_text("3")
363
364        session = self.app.session()
365        self.select_mcp(session, True)
366        started = time.monotonic()
367        with self.app.api(
368            f"/sessions/{session}/reload", {"target": "session"}
369        ) as response:
370            self.assertEqual(json.load(response)["session"]["state"], "applied")
371        self.assertLess(time.monotonic() - started, 1.5)
372
373        # The saved tools are advertised at once; the call waits for the dial.
374        requests = self.turn("use the mcp server", session)
375        self.assertEqual(len(self.tools(requests)), 1)
376        output = next(
377            item["output"]
378            for item in requests[-1]["input"]
379            if item.get("type") == "function_call_output"
380        )
381        self.assertIn(self.message, output)
382        self.assertNotIn("capabilities changed", json.dumps(self.app.history(session)))
383
384    def test_a_changed_catalogue_refreshes_the_session_after_its_turn(self):
385        self.set_mcp(True)
386        self.turn("learn the catalogue")
387        schema = {
388            "type": "object",
389            "required": ["message"],
390            "properties": {"message": {"type": "string"}},
391        }
392        tools = [
393            {"name": name, "description": name, "inputSchema": schema}
394            for name in ("echo", "shout")
395        ]
396        (self.app.root / "control" / "tools").write_text(json.dumps(tools))
397
398        session = self.app.session()
399        self.select_mcp(session, True)
400        self.assertEqual(
401            len(self.tools(self.turn("still the saved tools", session))), 1
402        )
403        deadline = time.monotonic() + 20
404        while "capabilities changed" not in json.dumps(self.app.history(session)):
405            self.assertLess(time.monotonic(), deadline, "the session never refreshed")
406            time.sleep(0.1)
407        advertised = self.tools(self.turn("now the new ones", session))
408        self.assertEqual(len(advertised), 2)
409        self.assertTrue(any(name.startswith("mcp_fake_shout_") for name in advertised))