from __future__ import annotations import asyncio import tempfile import unittest from pathlib import Path from peer import ROOT, Peer, behavior_copy, output class WorkbenchTests(unittest.IsolatedAsyncioTestCase): async def test_persistent_namespace_and_last_expression(self): async with Peer() as peer: first = await peer.execute("answer = 6 * 7\nanswer") self.assertEqual(first["status"], "ok") self.assertEqual(output(first), "42") self.assertEqual(output(await peer.execute("answer + _")), "84") async def test_top_level_await_and_background_computation(self): async with Peer() as peer: first = await peer.execute("import asyncio\ntask = asyncio.create_task(asyncio.sleep(0.03, result=7))") self.assertEqual(first["status"], "ok") second = await peer.execute("await task") self.assertEqual(output(second), "7") async def test_future_flags_survive_cells(self): async with Peer() as peer: self.assertEqual((await peer.execute("from __future__ import annotations\ndef f(x: Unknown): return x"))["status"],"ok") result = await peer.execute("def g(x: AlsoUnknown): return x\ng.__annotations__") self.assertEqual(result["status"],"ok") self.assertIn("AlsoUnknown",output(result)) async def test_syntax_and_runtime_errors_do_not_destroy_namespace(self): async with Peer() as peer: await peer.execute("saved = 9") for code in ["if :", "1/0", "raise SystemExit(5)", "raise KeyboardInterrupt()"]: with self.subTest(code=code): result = await peer.execute(code) self.assertEqual(result["status"],"error") self.assertTrue(result["error"]) self.assertEqual(output(await peer.execute("saved")),"9") async def test_print_and_native_stdout_cannot_send_or_forge_control(self): async with Peer() as peer: result = await peer.execute("import os,sys\nprint('scratch-like output')\nprint('error lane',file=sys.stderr)\nos.write(1, b'{\"kind\":\"host_request\"}\\n')\nNone") self.assertEqual(result["status"],"ok") self.assertIn("scratch-like output",output(result,"stdout")) self.assertIn("error lane",output(result,"stderr")) self.assertEqual(peer.requests,[]) await asyncio.sleep(0.01) self.assertIn(b'host_request',peer.raw["stdout"]) async def test_output_is_bounded_and_truncation_explicit(self): async with Peer() as peer: result = await peer.execute("print('🌙' * 100000)") self.assertLessEqual(sum(len(x["text"].encode()) for x in result["output"]),65536) self.assertGreater(result["truncated_bytes"],300000) self.assertEqual(output(await peer.execute("21*2")),"42") async def test_many_alternating_writes_bound_record_count(self): async with Peer() as peer: result = await peer.execute("import sys\nfor i in range(10000):\n sys.stdout.write('a'); sys.stderr.write('b')") self.assertLessEqual(len(result["output"]),256) self.assertGreater(result["truncated_bytes"],0) async def test_explicit_send_roundtrip(self): async def service(message): self.assertEqual(message["request"],{"method":"local.send","text":"hello"}) return {"kind":"ok","value":{"effect_id":"effect-1","status":"confirmed","event_id":7,"reason":None}} async with Peer() as peer: result = await peer.execute("from klbr import local\nreceipt = await local.send('hello')\nreceipt.status",service) self.assertEqual(output(result),"'confirmed'") self.assertEqual(len(peer.requests),1) async def test_host_error_is_inspectable_and_kernel_recovers(self): async def service(_): return {"kind":"error","code":"scope_closed","message":"execution expired"} async with Peer() as peer: result = await peer.execute("from klbr import local\nawait local.send('hello')",service) self.assertEqual(result["status"],"error") self.assertIn("scope_closed",result["error"]) self.assertEqual(output(await peer.execute("2")),"2") async def test_wait_is_terminal_not_a_sleep(self): async def service(message): self.assertEqual(message["request"]["method"],"runtime.wait") return {"kind":"ok","value":{"wake_at_ms":None}} async with Peer() as peer: result = await peer.execute("from klbr import runtime\nawait runtime.wait(reason='quiet')\nprint('must not execute')",service) self.assertEqual(result["status"],"yielded") self.assertNotIn("must not execute",output(result)) self.assertEqual(len(peer.requests),1) async def test_catching_yield_does_not_restore_effect_scope(self): async def service(_): return {"kind":"ok","value":{}} code = "from klbr import runtime,local\ntry:\n await runtime.wait(reason='quiet')\nexcept BaseException:\n pass\ntry:\n await local.send('not allowed')\nexcept RuntimeError as e:\n print(type(e).__name__)" async with Peer() as peer: result = await peer.execute(code,service) self.assertEqual(result["status"],"yielded") self.assertIn("ScopeClosed",output(result)) self.assertEqual(len(peer.requests),1) async def test_invalid_wait_values_never_reach_host(self): async with Peer() as peer: for value in ["0", "-1", "601", "True", "float('nan')", "float('inf')"]: with self.subTest(value=value): result = await peer.execute(f"from klbr import runtime\nawait runtime.wait({value},reason='bad')") self.assertEqual(result["status"],"error") self.assertFalse(peer.requests) async def test_cooperative_interrupt_preserves_globals(self): async with Peer() as peer: operation = await peer.begin("import asyncio\nsaved=42\nawait asyncio.sleep(60)") await asyncio.sleep(0.03) await peer.wire.send({"kind":"interrupt","id":operation}) result = await asyncio.wait_for(peer.finish(operation),1) self.assertEqual(result["status"],"interrupted") self.assertEqual(output(await peer.execute("saved")),"42") async def test_late_reply_after_interrupt_is_discarded(self): async with Peer() as peer: operation = await peer.begin("from klbr import local\nawait local.send('pending')") request = await asyncio.wait_for(peer.wire.receive(),1) self.assertEqual(request["kind"],"host_request") await peer.wire.send({"kind":"interrupt","id":operation}) self.assertEqual((await peer.finish(operation))["status"],"interrupted") await peer.wire.send({"kind":"host_reply","id":request["id"],"execution_id":operation, "reply":{"kind":"ok","value":{"ignored":True}}}) self.assertEqual(output(await peer.execute("42")),"42") async def test_background_effects_expire_with_originating_cell(self): async with Peer() as peer: await peer.execute("import asyncio\nfrom klbr import local\nresult='pending'\nasync def later():\n global result\n await asyncio.sleep(.03)\n try:\n await local.send('late')\n except RuntimeError as e:\n result=type(e).__name__\ntask=asyncio.create_task(later())") result=await peer.execute("await task\nresult") self.assertEqual(output(result),"'ScopeClosed'") self.assertFalse(peer.requests) async def test_background_output_is_not_reattributed_to_next_cell(self): async with Peer() as peer: await peer.execute("import asyncio\nasync def later():\n await asyncio.sleep(.02)\n print('from old cell')\ntask=asyncio.create_task(later())") result=await peer.execute("await task\nprint('new cell')") self.assertEqual(output(result),"new cell\n") await asyncio.sleep(.02) self.assertIn(b'background origin=',peer.raw["stdout"]) self.assertIn(b'from old cell',peer.raw["stdout"]) async def test_stale_generation_is_connection_fatal(self): from klbr_runtime.wire import encode async with Peer() as peer: peer.wire.writer.write(encode({"version":1,"session_id":peer.session,"generation":"old-generation", "message":{"kind":"execute","id":"old","code":"42"}})) await peer.wire.writer.drain() status=await asyncio.wait_for(peer.process.wait(),2) self.assertNotEqual(status,0) async def test_native_exit_does_not_claim_completion(self): async with Peer() as peer: await peer.begin("import os; os._exit(19)") with self.assertRaises(asyncio.IncompleteReadError): await asyncio.wait_for(peer.wire.receive(),1) self.assertEqual(await peer.process.wait(),19) class PolicyProcessTests(unittest.IsolatedAsyncioTestCase): async def test_editable_defaults_match_attention_contract(self): async with Peer("hooks",ROOT/"behavior") as peer: cases=[("operator",{"action":"wake","priority":"operator"}), ("discord.dm",{"action":"wake","priority":"directed"}), ("discord.mention",{"action":"wake","priority":"directed"}), ("discord.ambient",{"action":"defer","seconds":8})] for source,expected in cases: result=await peer.invoke("attention.plan",{"event_id":1,"source":source,"received_ms":0}) self.assertEqual(result,{"kind":"hook","value":expected}) async def test_send_waits_for_independent_hook_without_circular_wait(self): published=[] async with Peer("hooks",ROOT/"behavior") as hooks, Peer() as workbench: async def service(message): request=message["request"] decision=await hooks.invoke("delivery.review",{"effect_id":"effect-1","destination":"operator","text":request["text"]}) self.assertEqual(decision["value"],{"action":"allow"}) published.append(request["text"]) return {"kind":"ok","value":{"effect_id":"effect-1","status":"confirmed","event_id":1,"reason":None}} result=await workbench.execute("from klbr import local\n(await local.send('hello')).status",service) self.assertEqual(output(result),"'confirmed'") self.assertEqual(published,["hello"]) async def test_hook_worker_available_while_workbench_cpu_blocked(self): async with Peer("hooks",ROOT/"behavior") as hooks, Peer() as workbench: await workbench.begin("while True: pass") result=await hooks.invoke("delivery.review",{"effect_id":"e","destination":"operator","text":"hello"},timeout=1) self.assertEqual(result["value"],{"action":"allow"}) await workbench.close() async def test_wrong_hook_return_is_failure_not_approval(self): with tempfile.TemporaryDirectory() as directory: behavior=behavior_copy(Path(directory)/"candidate","def delivery(request, context):\n return True") async with Peer("hooks",behavior) as hooks: result=await hooks.invoke("delivery.review",{"effect_id":"e","destination":"operator","text":"hello"}) self.assertEqual(result["kind"],"failed") self.assertIn("DeliveryDecision",result["error"]) async def test_hook_cannot_use_workbench_effect_sdk(self): source="async def delivery(request, context):\n from klbr import local\n return await local.send('not a proposal')" with tempfile.TemporaryDirectory() as directory: behavior=behavior_copy(Path(directory)/"candidate",source) async with Peer("hooks",behavior) as hooks: result=await hooks.invoke("delivery.review",{"effect_id":"e","destination":"operator","text":"hello"}) self.assertEqual(result["kind"],"failed") self.assertIn("ScopeClosed",result["error"]) self.assertFalse(hooks.requests) async def test_replacing_hook_worker_preserves_workbench_heap(self): with tempfile.TemporaryDirectory() as directory: candidate=behavior_copy(Path(directory)/"candidate","def delivery(request, context):\n return DeliveryDecision.hold('new rule')") async with Peer() as workbench: await workbench.execute("kept = object()\nidentity = id(kept)") async with Peer("hooks",ROOT/"behavior") as old: self.assertEqual((await old.invoke("delivery.review",{"effect_id":"e","destination":"operator","text":"hello"}))["value"]["action"],"allow") async with Peer("hooks",candidate) as new: self.assertEqual((await new.invoke("delivery.review",{"effect_id":"e","destination":"operator","text":"hello"}))["value"]["action"],"hold") self.assertEqual(output(await workbench.execute("id(kept) == identity")),"True") async def test_bad_candidate_does_not_affect_existing_worker(self): with tempfile.TemporaryDirectory() as directory: candidate=behavior_copy(Path(directory)/"candidate","this is a syntax error !!!") async with Peer("hooks",ROOT/"behavior") as good: with self.assertRaises(RuntimeError): await Peer("hooks",candidate).start() self.assertEqual((await good.invoke("delivery.review",{"effect_id":"e","destination":"operator","text":"hello"}))["value"]["action"],"allow") async def test_blocked_policy_does_not_block_workbench_and_can_be_killed(self): with tempfile.TemporaryDirectory() as directory: candidate=behavior_copy(Path(directory)/"candidate","def delivery(request, context):\n while True: pass") async with Peer("hooks",candidate) as hooks, Peer() as workbench: with self.assertRaises(asyncio.TimeoutError): await hooks.invoke("delivery.review",{"effect_id":"e","destination":"operator","text":"hello"},timeout=.05) self.assertEqual(output(await workbench.execute("2+2")),"4") await hooks.close() self.assertIsNotNone(hooks.process.returncode)