yingqi-z20-Agent-libOS

repository

loading code, commits, and activity

repositories

loading repo index

yingqi-z20-Agent-libOS — gitlawb

#1	from __future__ import annotations
#2
#3	import json
#4	import asyncio
#5	import tempfile
#6	import unittest
#7	from typing import Any
#8
#9	from agent_libos import Runtime
#10	from agent_libos.llm.client import LLMCompletion
#11	from agent_libos.models import ProcessMessageKind, ProcessStatus
#12	from agent_libos.runtime.syscalls import LibOSSyscallSession
#13
#14
#15	class ProcessMessageTests(unittest.TestCase):
#16	def test_process_message_tools_send_read_and_ack_related_processes(self) -> None:
#17	runtime = Runtime.open("local")
#18	try:
#19	parent = runtime.process.spawn(image="base-agent:v0", goal="parent")
#20	child = runtime.spawn_child_process(parent, "child")
#21
#22	sent = runtime.tools.call(
#23	parent,
#24	"send_process_message",
#25	{
#26	"recipient_pid": child,
#27	"kind": "normal",
#28	"subject": "status",
#29	"body": "send a status update",
#30	"payload": {"priority": 1},
#31	},
#32	)
#33	self.assertTrue(sent.ok, sent.error)
#34	self.assertEqual(len(runtime.messages.unread(child)), 1)
#35
#36	read = runtime.tools.call(child, "read_process_messages", {})
#37
#38	self.assertTrue(read.ok, read.error)
#39	self.assertEqual(read.payload["messages"][0]["subject"], "status")
#40	self.assertEqual(read.payload["messages"][0]["payload"], {"priority": 1})
#41	self.assertEqual(read.payload["messages"][0]["status"], "acked")
#42	self.assertEqual(read.payload["acked_message_ids"], [sent.payload["message_id"]])
#43	self.assertEqual(runtime.messages.unread(child), [])
#44	finally:
#45	runtime.close()
#46
#47	def test_unrelated_process_cannot_send_process_message(self) -> None:
#48	runtime = Runtime.open("local")
#49	try:
#50	first = runtime.process.spawn(image="base-agent:v0", goal="first")
#51	second = runtime.process.spawn(image="base-agent:v0", goal="second")
#52
#53	denied = runtime.tools.call(first, "send_process_message", {"recipient_pid": second, "body": "no"})
#54
#55	self.assertFalse(denied.ok)
#56	self.assertIn("can only message", denied.error or "")
#57	self.assertEqual(runtime.messages.unread(second), [])
#58	finally:
#59	runtime.close()
#60
#61	def test_human_can_send_normal_and_interrupt_process_messages(self) -> None:
#62	runtime = Runtime.open("local")
#63	try:
#64	pid = runtime.process.spawn(image="base-agent:v0", goal="listen to human")
#65
#66	normal = runtime.human.send_process_message(pid, "please check progress", subject="status")
#67	interrupt = runtime.human.send_process_message(
#68	pid,
#69	"stop current work and inspect this",
#70	kind=ProcessMessageKind.INTERRUPT,
#71	)
#72
#73	unread = runtime.messages.unread(pid)
#74
#75	self.assertEqual([message.message_id for message in unread], [normal.message_id, interrupt.message_id])
#76	self.assertEqual(unread[0].sender, "human:owner")
#77	self.assertEqual(unread[0].channel, "human")
#78	self.assertEqual(unread[0].payload["source"], "human_input")
#79	self.assertEqual(unread[1].kind, ProcessMessageKind.INTERRUPT)
#80	self.assertIn("human.message", _audit_actions(runtime))
#81	finally:
#82	runtime.close()
#83
#84	def test_process_message_syscalls_send_read_and_ack(self) -> None:
#85	runtime = Runtime.open("local")
#86	try:
#87	parent = runtime.process.spawn(image="base-agent:v0", goal="parent")
#88	child = runtime.spawn_child_process(parent, "child")
#89	parent_session = LibOSSyscallSession(runtime, parent)
#90	child_session = LibOSSyscallSession(runtime, child)
#91
#92	sent = asyncio.run(
#93	parent_session.handle(
#94	"process.send_message",
#95	{"recipient_pid": child, "kind": "normal", "subject": "via syscall", "body": "hello"},
#96	)
#97	)
#98	read = asyncio.run(child_session.handle("process.read_messages", {}))
#99
#100	self.assertEqual(sent["subject"], "via syscall")
#101	self.assertEqual(read["messages"][0]["message_id"], sent["message_id"])
#102	self.assertEqual(read["messages"][0]["status"], "acked")
#103	self.assertEqual(runtime.messages.unread(child), [])
#104	finally:
#105	runtime.close()
#106
#107	def test_process_message_filters_channel_correlation_reply_and_ids(self) -> None:
#108	runtime = Runtime.open("local")
#109	try:
#110	parent = runtime.process.spawn(image="base-agent:v0", goal="parent")
#111	child = runtime.spawn_child_process(parent, "child")
#112	first = runtime.messages.send_from_process(
#113	parent,
#114	child,
#115	channel="control",
#116	correlation_id="job-1",
#117	subject="request",
#118	body="start",
#119	)
#120	runtime.messages.send_from_process(
#121	parent,
#122	child,
#123	channel="noise",
#124	correlation_id="job-1",
#125	subject="ignore",
#126	)
#127	reply = runtime.messages.send_from_process(
#128	child,
#129	parent,
#130	channel="control",
#131	correlation_id="job-1",
#132	reply_to=first.message_id,
#133	subject="reply",
#134	)
#135
#136	selected = runtime.tools.call(
#137	child,
#138	"read_process_messages",
#139	{"channel": "control", "correlation_id": "job-1", "ack": False},
#140	)
#141	reply_selected = runtime.tools.call(
#142	parent,
#143	"read_process_messages",
#144	{"reply_to": first.message_id, "message_ids": [reply.message_id]},
#145	)
#146
#147	self.assertTrue(selected.ok, selected.error)
#148	self.assertEqual([message["message_id"] for message in selected.payload["messages"]], [first.message_id])
#149	self.assertEqual(selected.payload["messages"][0]["channel"], "control")
#150	self.assertEqual(selected.payload["messages"][0]["correlation_id"], "job-1")
#151	self.assertEqual(selected.payload["acked_message_ids"], [])
#152	self.assertEqual(len(runtime.messages.unread(child)), 2)
#153
#154	self.assertTrue(reply_selected.ok, reply_selected.error)
#155	self.assertEqual(reply_selected.payload["messages"][0]["reply_to"], first.message_id)
#156	self.assertEqual(reply_selected.payload["acked_message_ids"], [reply.message_id])
#157	self.assertEqual(runtime.messages.unread(parent), [])
#158	finally:
#159	runtime.close()
#160
#161	def test_receive_process_messages_blocks_until_matching_message_then_resumes(self) -> None:
#162	client = PlannedActionClient(
#163	[
#164	{
#165	"action": "receive_process_messages",
#166	"channel": "control",
#167	"correlation_id": "job-1",
#168	}
#169	]
#170	)
#171	runtime = Runtime.open("local")
#172	runtime.llm.client = client
#173	try:
#174	parent = runtime.process.spawn(image="base-agent:v0", goal="parent")
#175	child = runtime.spawn_child_process(parent, "wait for control message")
#176
#177	waiting = runtime.run_process_once(child)
#178
#179	self.assertTrue(waiting["waiting_message"])
#180	self.assertEqual(waiting["filters"]["channel"], "control")
#181	self.assertEqual(runtime.process.get(child).status, ProcessStatus.WAITING_EVENT)
#182	self.assertEqual(len(client.user_prompts), 1)
#183
#184	runtime.messages.send_from_process(
#185	parent,
#186	child,
#187	channel="noise",
#188	correlation_id="job-1",
#189	subject="not yet",
#190	)
#191
#192	self.assertEqual(runtime.process.get(child).status, ProcessStatus.WAITING_EVENT)
#193	skipped = runtime.run_process_once(child)
#194	self.assertTrue(skipped["skipped"])
#195	self.assertEqual(len(client.user_prompts), 1)
#196
#197	matching = runtime.messages.send_from_process(
#198	parent,
#199	child,
#200	channel="control",
#201	correlation_id="job-1",
#202	subject="resume",
#203	payload={"ready": True},
#204	)
#205
#206	self.assertEqual(runtime.process.get(child).status, ProcessStatus.RUNNABLE)
#207	resumed = runtime.run_process_once(child)
#208
#209	self.assertTrue(resumed["ok"])
#210	self.assertTrue(resumed["resumed_after_message"])
#211	self.assertEqual(resumed["action"]["action"], "receive_process_messages")
#212	self.assertEqual(resumed["result"]["payload"]["messages"][0]["message_id"], matching.message_id)
#213	self.assertEqual(resumed["result"]["payload"]["messages"][0]["payload"], {"ready": True})
#214	self.assertEqual(resumed["result"]["payload"]["acked_message_ids"], [matching.message_id])
#215	self.assertEqual(len(client.user_prompts), 1)
#216	self.assertEqual(
#217	[message.subject for message in runtime.messages.unread(child)],
#218	["not yet"],
#219	)
#220	self.assertIn("process.message.wait_wake", _audit_actions(runtime))
#221	finally:
#222	runtime.close()
#223
#224	def test_receive_message_syscall_waits_inside_single_syscall_until_matching_message(self) -> None:
#225	runtime = Runtime.open("local")
#226	try:
#227	runtime.scheduler.poll_interval_s = 0.001
#228	parent = runtime.process.spawn(image="base-agent:v0", goal="parent")
#229	child = runtime.spawn_child_process(parent, "wait via syscall")
#230	child_session = LibOSSyscallSession(runtime, child)
#231
#232	async def scenario() -> dict[str, Any]:
#233	task = asyncio.create_task(
#234	child_session.handle(
#235	"process.receive_messages",
#236	{"block": True, "channel": "control", "correlation_id": "job-1"},
#237	)
#238	)
#239	await asyncio.sleep(0.01)
#240	self.assertFalse(task.done())
#241	self.assertEqual(runtime.process.get(child).status, ProcessStatus.WAITING_EVENT)
#242
#243	runtime.messages.send_from_process(
#244	parent,
#245	child,
#246	channel="noise",
#247	correlation_id="job-1",
#248	subject="not matching",
#249	)
#250	await asyncio.sleep(0.01)
#251	self.assertFalse(task.done())
#252	self.assertEqual(runtime.process.get(child).status, ProcessStatus.WAITING_EVENT)
#253
#254	matching = runtime.messages.send_from_process(
#255	parent,
#256	child,
#257	channel="control",
#258	correlation_id="job-1",
#259	subject="matching",
#260	)
#261	result = await asyncio.wait_for(task, timeout=1.0)
#262	result["expected_message_id"] = matching.message_id
#263	return result
#264
#265	result = asyncio.run(scenario())
#266
#267	self.assertTrue(result["ready"])
#268	self.assertEqual(result["messages"][0]["message_id"], result["expected_message_id"])
#269	self.assertEqual(result["messages"][0]["status"], "acked")
#270	self.assertEqual(result["acked_message_ids"], [result["expected_message_id"]])
#271	self.assertEqual(runtime.process.get(child).status, ProcessStatus.RUNNABLE)
#272	finally:
#273	runtime.close()
#274
#275	def test_receive_message_syscall_blocks_by_default(self) -> None:
#276	runtime = Runtime.open("local")
#277	try:
#278	runtime.scheduler.poll_interval_s = 0.001
#279	parent = runtime.process.spawn(image="base-agent:v0", goal="parent")
#280	child = runtime.spawn_child_process(parent, "default receive")
#281	child_session = LibOSSyscallSession(runtime, child)
#282
#283	async def scenario() -> dict[str, Any]:
#284	task = asyncio.create_task(child_session.handle("process.receive_messages", {"channel": "control"}))
#285	await asyncio.sleep(0.01)
#286	self.assertFalse(task.done())
#287	self.assertEqual(runtime.process.get(child).status, ProcessStatus.WAITING_EVENT)
#288
#289	matching = runtime.messages.send_from_process(parent, child, channel="control")
#290	result = await asyncio.wait_for(task, timeout=1.0)
#291	result["expected_message_id"] = matching.message_id
#292	return result
#293
#294	result = asyncio.run(scenario())
#295
#296	self.assertTrue(result["ready"])
#297	self.assertEqual(result["messages"][0]["message_id"], result["expected_message_id"])
#298	self.assertEqual(runtime.process.get(child).status, ProcessStatus.RUNNABLE)
#299	finally:
#300	runtime.close()
#301
#302	def test_process_messages_are_durable_in_sqlite(self) -> None:
#303	with tempfile.TemporaryDirectory() as temp_dir:
#304	db = f"{temp_dir}/runtime.sqlite"
#305	runtime = Runtime.open(db)
#306	pid = runtime.process.spawn(image="base-agent:v0", goal="persist queue")
#307	message = runtime.messages.post(sender="test", recipient_pid=pid, subject="persisted")
#308	runtime.close()
#309
#310	reopened = Runtime.open(db)
#311	try:
#312	unread = reopened.messages.unread(pid)
#313
#314	self.assertEqual([item.message_id for item in unread], [message.message_id])
#315	self.assertEqual(unread[0].subject, "persisted")
#316	finally:
#317	reopened.close()
#318
#319	def test_interrupt_message_preempts_tool_call_until_read(self) -> None:
#320	client = PlannedActionClient(
#321	[
#322	{"action": "get_current_time", "timezone": "UTC"},
#323	{"action": "read_process_messages"},
#324	]
#325	)
#326	runtime = Runtime.open("local")
#327	runtime.llm.client = client
#328	try:
#329	pid = runtime.process.spawn(image="base-agent:v0", goal="handle interrupts")
#330	runtime.messages.post(
#331	sender="test",
#332	recipient_pid=pid,
#333	kind=ProcessMessageKind.INTERRUPT,
#334	subject="urgent",
#335	body="inspect this before other work",
#336	)
#337
#338	interrupted = runtime.run_process_once(pid)
#339
#340	self.assertTrue(interrupted["result"]["interrupted_by_message"])
#341	self.assertEqual(interrupted["result"]["message_notice"]["phase"], "before_tool_call")
#342	self.assertNotIn("primitive.clock.now", _audit_actions(runtime))
#343	self.assertIn("process_message_notice", client.user_prompts[0])
#344
#345	read = runtime.run_process_once(pid)
#346
#347	self.assertEqual(read["action"]["action"], "read_process_messages")
#348	self.assertEqual(read["result"]["payload"]["messages"][0]["kind"], "interrupt")
#349	self.assertEqual(runtime.messages.unread(pid, kind=ProcessMessageKind.INTERRUPT), [])
#350	finally:
#351	runtime.close()
#352
#353	def test_normal_message_notifies_after_tool_call_without_preempting(self) -> None:
#354	client = PlannedActionClient([{"action": "get_current_time", "timezone": "UTC"}])
#355	runtime = Runtime.open("local")
#356	runtime.llm.client = client
#357	try:
#358	pid = runtime.process.spawn(image="base-agent:v0", goal="handle normal messages")
#359	runtime.messages.post(
#360	sender="test",
#361	recipient_pid=pid,
#362	kind=ProcessMessageKind.NORMAL,
#363	subject="later",
#364	body="read after current tool",
#365	)
#366
#367	result = runtime.run_process_once(pid)
#368
#369	self.assertEqual(result["action"]["action"], "get_current_time")
#370	self.assertIn("primitive.clock.now", _audit_actions(runtime))
#371	self.assertEqual(result["result"]["message_notice"]["phase"], "after_tool_call")
#372	self.assertEqual(result["result"]["message_notice"]["kind"], "normal")
#373	self.assertEqual(len(runtime.messages.unread(pid, kind=ProcessMessageKind.NORMAL)), 1)
#374	finally:
#375	runtime.close()
#376
#377
#378	class PlannedActionClient:
#379	def __init__(self, actions: list[dict[str, Any]]):
#380	self.actions = list(actions)
#381	self.user_prompts: list[str] = []
#382
#383	def complete_action(self, messages: list[dict[str, Any]], tools: list[dict[str, Any]]) -> LLMCompletion:
#384	if not self.actions:
#385	raise AssertionError("no planned action remains")
#386	self.user_prompts.append(str(messages[-1]["content"]))
#387	action = self.actions.pop(0)
#388	name = str(action["action"])
#389	args = {key: value for key, value in action.items() if key != "action"}
#390	return LLMCompletion(
#391	content="",
#392	tool_calls=[{"id": f"message_{len(self.user_prompts)}", "name": name, "arguments": json.dumps(args)}],
#393	)
#394
#395
#396	def _audit_actions(runtime: Runtime) -> set[str]:
#397	return {record.action for record in runtime.audit.trace()}
#398
#399
#400	if __name__ == "__main__":
#401	unittest.main()
#402

z6MkqRzA/yingqi-z20-Agent-libOS