main
py 1,180 lines 38.3 KB
Raw
1 from __future__ import annotations
2
3 import json
4 import time
5 import sys
6 from types import SimpleNamespace
7 from pathlib import Path
8
9 import pytest
10
11 PROJECT_ROOT = Path(__file__).resolve().parents[1]
12 if str(PROJECT_ROOT) not in sys.path:
13 sys.path.insert(0, str(PROJECT_ROOT))
14
15 from helpers import parallel_tools
16 from helpers.tool import Response
17
18
19 class _FakeLogItem:
20 def __init__(self, type_, heading="", content="", kvps=None, id_=None, **kwargs) -> None:
21 self.type = type_
22 self.heading = heading
23 self.content = content
24 self.kvps = dict(kvps or {})
25 self.kvps.update(kwargs)
26 self.id = id_
27
28 def update(self, content=None, kvps=None, **kwargs):
29 if content is not None:
30 self.content = content
31 if kvps:
32 self.kvps.update(kvps)
33 self.kvps.update(kwargs)
34
35
36 class _FakeLog:
37 def __init__(self) -> None:
38 self.items = []
39
40 def log(self, type, heading="", content="", kvps=None, id=None, **kwargs):
41 item = _FakeLogItem(type, heading, content, kvps, id, **kwargs)
42 self.items.append(item)
43 return item
44
45
46 class _FakeContext:
47 def __init__(self) -> None:
48 self.id = "ctx"
49 self.data = {}
50 self.log = _FakeLog()
51 self.task = None
52
53 def get_data(self, key: str, recursive: bool = True):
54 return self.data.get(key)
55
56 def set_data(self, key: str, value, recursive: bool = True):
57 self.data[key] = value
58
59
60 class _FakeAgent:
61 def __init__(self) -> None:
62 self.context = _FakeContext()
63 self.agent_name = "A0"
64 self.number = 0
65
66
67 class _FakeDeferredTask:
68 def __init__(
69 self,
70 *,
71 ready: bool = False,
72 alive: bool = True,
73 result=None,
74 thread_name=None,
75 ) -> None:
76 self.ready = ready
77 self.alive = alive
78 self._result = result
79 self.killed = 0
80 self.thread_name = thread_name
81 self.started = None
82 self.children = []
83
84 def start_task(self, func, *args):
85 self.started = (func, args)
86 return self
87
88 def is_ready(self):
89 return self.ready
90
91 def is_alive(self):
92 return self.alive
93
94 async def result(self):
95 return self._result
96
97 def kill(self):
98 self.killed += 1
99 self.alive = False
100 for child in self.children:
101 child.kill()
102 self.children = []
103
104 def add_child_task(self, task, terminate_thread=False):
105 self.children.append(task)
106
107
108 def test_normalize_parallel_tool_calls_accepts_normal_tool_request_shapes() -> None:
109 calls = parallel_tools.normalize_parallel_tool_calls(
110 [
111 {
112 "tool_name": "text_editor:read",
113 "tool_args": {"path": "README.md"},
114 },
115 {
116 "tool": "scheduler",
117 "args": {"method": "list_tasks"},
118 },
119 ]
120 )
121
122 assert calls[0].tool_name == "text_editor"
123 assert calls[0].tool_args == {"path": "README.md", "action": "read"}
124 assert calls[1].tool_name == "scheduler"
125 assert calls[1].tool_args == {"method": "list_tasks", "action": "list_tasks"}
126
127
128 def test_normalize_parallel_tool_calls_accepts_json_string_array() -> None:
129 calls = parallel_tools.normalize_parallel_tool_calls(
130 json.dumps(
131 [
132 {
133 "tool_name": "call_subordinate",
134 "tool_args": {
135 "profile": "researcher",
136 "reset": True,
137 "message": "Research nuclear fusion news in French.",
138 },
139 "headline": "Researching nuclear fusion news in French",
140 },
141 {
142 "tool_name": "call_subordinate",
143 "tool_args": {
144 "profile": "researcher",
145 "reset": True,
146 "message": "Research nuclear fusion news in Italian.",
147 },
148 "headline": "Researching nuclear fusion news in Italian",
149 },
150 ]
151 )
152 )
153
154 assert [call.tool_name for call in calls] == [
155 "call_subordinate",
156 "call_subordinate",
157 ]
158 assert calls[0].tool_args["profile"] == "researcher"
159 assert calls[0].tool_args["reset"] is True
160 assert calls[1].tool_args["message"] == "Research nuclear fusion news in Italian."
161
162
163 def test_subordinate_prompts_share_reusable_tree_contract() -> None:
164 call_prompt = (PROJECT_ROOT / "prompts/agent.system.tool.call_sub.md").read_text(
165 encoding="utf-8"
166 )
167 parallel_prompt = (PROJECT_ROOT / "prompts/agent.system.tool.parallel.md").read_text(
168 encoding="utf-8"
169 )
170
171 assert "A0 creates A1 children, A1 creates A2 children" in call_prompt
172 assert "stable child ID" in call_prompt
173 assert "same child lifecycle here as it does top-level" in parallel_prompt
174 assert "each job's `context_id`" in parallel_prompt
175
176
177 def test_normalize_parallel_tool_calls_rejects_nested_parallel() -> None:
178 with pytest.raises(ValueError, match="cannot be nested"):
179 parallel_tools.normalize_parallel_tool_calls(
180 [{"tool_name": "parallel", "tool_args": {"tool_calls": []}}]
181 )
182
183
184 @pytest.mark.parametrize("tool_name", ["document_query", "response"])
185 def test_normalize_parallel_tool_calls_rejects_disallowed_tools(tool_name: str) -> None:
186 with pytest.raises(ValueError, match=rf"{tool_name}.*parallel"):
187 parallel_tools.normalize_parallel_tool_calls(
188 [
189 {
190 "tool_name": tool_name,
191 "tool_args": {},
192 }
193 ]
194 )
195
196
197 @pytest.mark.asyncio
198 async def test_parallel_jobs_extras_lists_running_and_ready_jobs() -> None:
199 agent = _FakeAgent()
200 running = parallel_tools.ParallelJob(
201 id="search-1234abcd",
202 parent_context_id="ctx",
203 index=0,
204 tool_name="search_engine",
205 tool_args={"query": "Agent Zero"},
206 kind="tool",
207 state="running",
208 started_at=time.time() - 2,
209 )
210 ready = parallel_tools.ParallelJob(
211 id="callsubordin-5678efgh",
212 parent_context_id="ctx",
213 index=1,
214 tool_name="call_subordinate",
215 tool_args={"message": "Summarize"},
216 kind="subordinate",
217 state="success",
218 started_at=time.time() - 4,
219 completed_at=time.time() - 1,
220 result="done",
221 )
222 agent.context.set_data(
223 parallel_tools.PARALLEL_JOBS_KEY,
224 {running.id: running, ready.id: ready},
225 )
226
227 extras = await parallel_tools.build_parallel_jobs_extras(agent) # type: ignore[arg-type]
228
229 assert "search-1234abcd" in extras
230 assert "callsubordin-5678efgh" in extras
231 assert "ready to collect" in extras
232
233
234 @pytest.mark.asyncio
235 async def test_parallel_await_timeout_keeps_running_jobs_awaitable(monkeypatch) -> None:
236 agent = _FakeAgent()
237 task = _FakeDeferredTask(alive=True)
238 job = parallel_tools.ParallelJob(
239 id="wait-1234abcd",
240 parent_context_id="ctx",
241 index=0,
242 tool_name="wait",
243 tool_args={"seconds": 60},
244 kind="tool",
245 state="running",
246 created_at=99.0,
247 started_at=99.0,
248 deferred_task=task, # type: ignore[arg-type]
249 )
250 agent.context.set_data(parallel_tools.PARALLEL_JOBS_KEY, {job.id: job})
251 times = iter([100.0, 102.0])
252 monkeypatch.setattr(
253 parallel_tools.time,
254 "time",
255 lambda: next(times, 102.0),
256 )
257
258 results = await parallel_tools.await_parallel_jobs( # type: ignore[arg-type]
259 agent,
260 [job.id],
261 timeout=1,
262 collect=True,
263 wait=True,
264 )
265 payload = json.loads(parallel_tools.format_parallel_results(results))
266
267 assert results[0]["state"] == "running"
268 assert results[0]["wait_timed_out"] is True
269 assert payload["status"] == "waiting"
270 assert payload["wait_timeout"] is True
271 assert "await" in payload["instruction"]
272 assert task.killed == 0
273 assert agent.context.get_data(parallel_tools.PARALLEL_JOBS_KEY)[job.id] is job
274
275
276 @pytest.mark.asyncio
277 async def test_parallel_collect_returns_running_jobs_without_waiting_or_canceling() -> None:
278 from tools.parallel import ParallelTool
279
280 agent = _FakeAgent()
281 task = _FakeDeferredTask(alive=True)
282 job = parallel_tools.ParallelJob(
283 id="wait-collect",
284 parent_context_id="ctx",
285 index=0,
286 tool_name="wait",
287 tool_args={"seconds": 60},
288 kind="tool",
289 state="running",
290 deferred_task=task, # type: ignore[arg-type]
291 )
292 agent.context.set_data(parallel_tools.PARALLEL_JOBS_KEY, {job.id: job})
293 tool = ParallelTool(
294 agent, # type: ignore[arg-type]
295 "parallel",
296 None,
297 {"action": "collect", "job_ids": [job.id]},
298 "",
299 None,
300 )
301
302 response = await tool.execute(**tool.args)
303 payload = json.loads(response.message)
304
305 assert payload["status"] == "running"
306 assert payload["jobs"][0]["job_id"] == job.id
307 assert "wait_timeout" not in payload
308 assert task.killed == 0
309 assert agent.context.get_data(parallel_tools.PARALLEL_JOBS_KEY)[job.id] is job
310
311
312 @pytest.mark.asyncio
313 async def test_parallel_cancel_still_stops_and_removes_running_jobs() -> None:
314 agent = _FakeAgent()
315 task = _FakeDeferredTask(alive=True)
316 job = parallel_tools.ParallelJob(
317 id="wait-cancel",
318 parent_context_id="ctx",
319 index=0,
320 tool_name="wait",
321 tool_args={"seconds": 60},
322 kind="tool",
323 state="running",
324 deferred_task=task, # type: ignore[arg-type]
325 )
326 agent.context.set_data(parallel_tools.PARALLEL_JOBS_KEY, {job.id: job})
327
328 results = await parallel_tools.cancel_parallel_jobs(agent, [job.id]) # type: ignore[arg-type]
329 payload = json.loads(parallel_tools.format_parallel_results(results))
330
331 assert results[0]["state"] == "cancelled"
332 assert payload["status"] == "cancelled"
333 assert task.killed == 1
334 assert job.id not in agent.context.get_data(parallel_tools.PARALLEL_JOBS_KEY)
335
336
337 @pytest.mark.asyncio
338 async def test_parallel_remove_context_deletes_persisted_worker_chat(monkeypatch) -> None:
339 removed = []
340
341 from helpers import persist_chat
342
343 monkeypatch.setattr(persist_chat, "remove_chat", removed.append)
344
345 await parallel_tools._remove_context("missing-worker")
346
347 assert removed == ["missing-worker"]
348
349
350 @pytest.mark.asyncio
351 async def test_direct_parallel_worker_inherits_chat_model_override(monkeypatch) -> None:
352 from agent import AgentConfig, AgentContext
353
354 parent_id = "ctx-parallel-model-override"
355 AgentContext.remove(parent_id)
356 parent = AgentContext(
357 AgentConfig(mcp_servers="", profile="agent0"),
358 id=parent_id,
359 set_current=False,
360 )
361 override = {"preset_name": "Text only"}
362 parent.set_data("chat_model_override", override)
363 current_user_message = object()
364 parent.agent0.last_user_message = current_user_message
365 observed = {}
366
367 async def fake_execute_tool_call(agent, *_args, **_kwargs):
368 observed["override"] = agent.context.get_data("chat_model_override")
369 observed["last_user_message"] = agent.last_user_message
370 return "done"
371
372 async def remove_context(context_id):
373 AgentContext.remove(context_id)
374
375 monkeypatch.setattr(parallel_tools, "execute_tool_call", fake_execute_tool_call)
376 monkeypatch.setattr(parallel_tools, "_remove_context", remove_context)
377 job = parallel_tools.ParallelJob(
378 id="vision-load-override",
379 parent_context_id=parent_id,
380 index=0,
381 tool_name="vision_load",
382 tool_args={"paths": ["/tmp/example.png"]},
383 kind="tool",
384 )
385
386 try:
387 assert await parallel_tools._run_direct_tool_job(parent_id, job) == "done"
388 assert observed["override"] == override
389 assert observed["last_user_message"] is current_user_message
390 finally:
391 AgentContext.remove(parent_id)
392
393
394 @pytest.mark.asyncio
395 async def test_parallel_recursion_guard_allows_subordinate_children_but_blocks_tool_workers() -> None:
396 from extensions.python.tool_execute_before._20_block_parallel_recursion import (
397 BlockParallelRecursion,
398 )
399 from helpers.errors import RepairableException
400
401 agent = _FakeAgent()
402 agent.context.set_data(parallel_tools.PARALLEL_WORKER_JOB_KEY, "legacy-job")
403 assert parallel_tools.is_parallel_worker(agent) is True # type: ignore[arg-type]
404
405 agent.context.set_data(parallel_tools.PARALLEL_WORKER_KIND_KEY, "subordinate")
406 assert parallel_tools.is_parallel_worker(agent) is False # type: ignore[arg-type]
407 await BlockParallelRecursion(agent=agent).execute(tool_name="parallel") # type: ignore[arg-type]
408
409 agent.context.set_data(parallel_tools.PARALLEL_WORKER_KIND_KEY, "tool")
410 assert parallel_tools.is_parallel_worker(agent) is True # type: ignore[arg-type]
411 with pytest.raises(RepairableException, match="cannot be used inside a parallel worker"):
412 await BlockParallelRecursion(agent=agent).execute(tool_name="parallel") # type: ignore[arg-type]
413
414
415 @pytest.mark.asyncio
416 async def test_parallel_subordinate_jobs_are_visible_child_logs_not_scheduler_tasks(monkeypatch) -> None:
417 class FakeDeferredTask:
418 def __init__(self, thread_name=None) -> None:
419 self.thread_name = thread_name
420 self.started = None
421
422 def start_task(self, func, *args):
423 self.started = (func, args)
424 return self
425
426 def is_ready(self):
427 return False
428
429 def is_alive(self):
430 return True
431
432 def kill(self):
433 pass
434
435 monkeypatch.setattr(parallel_tools, "DeferredTask", FakeDeferredTask)
436 agent = _FakeAgent()
437
438 jobs = await parallel_tools.start_parallel_jobs(
439 agent, # type: ignore[arg-type]
440 [
441 parallel_tools.NormalizedToolCall(
442 index=0,
443 tool_name="call_subordinate",
444 tool_args={
445 "profile": "developer",
446 "message": "Return ALPHA=1",
447 "reset": True,
448 },
449 )
450 ],
451 )
452
453 assert jobs[0].kind == "subordinate"
454 snapshot = parallel_tools._job_snapshot(jobs[0], include_result=False)
455 assert "scheduler_task_uuid" not in snapshot
456 assert agent.context.log.items[0].type == "subagent"
457 assert agent.context.log.items[0].kvps == {
458 "profile": "developer",
459 "message": "Return ALPHA=1",
460 "reset": True,
461 }
462 assert "id" not in agent.context.log.items[0].kvps
463 assert "tool_name" not in agent.context.log.items[0].kvps
464 assert "parallel_child" not in agent.context.log.items[0].kvps
465
466
467 @pytest.mark.asyncio
468 async def test_parallel_subordinate_enforces_parent_delegation_policy(
469 monkeypatch,
470 ) -> None:
471 from agent import AgentContext
472 from helpers import tool_policy
473 from helpers.errors import RepairableException
474
475 parent_agent = SimpleNamespace(
476 config=SimpleNamespace(profile="restricted"),
477 context=_FakeContext(),
478 )
479 parent_context = SimpleNamespace(agent0=parent_agent)
480 monkeypatch.setattr(
481 AgentContext,
482 "get",
483 staticmethod(lambda _context_id: parent_context),
484 )
485 monkeypatch.setattr(tool_policy.subagents, "get_paths", lambda *args, **kwargs: [])
486 monkeypatch.setattr(
487 tool_policy,
488 "get_policy",
489 lambda agent: {
490 "mode": "custom",
491 "default": "allow",
492 "allowed": [],
493 "blocked": ["local:call_subordinate"],
494 },
495 )
496 job = parallel_tools.ParallelJob(
497 id="callsubordin-blocked",
498 parent_context_id="ctx",
499 index=0,
500 tool_name="call_subordinate",
501 tool_args={"profile": "developer", "message": "Work"},
502 kind="subordinate",
503 )
504
505 with pytest.raises(
506 RepairableException,
507 match='Tool "call_subordinate" is blocked for agent profile "restricted"',
508 ):
509 await parallel_tools._run_subordinate_context_job("ctx", job)
510
511
512 @pytest.mark.asyncio
513 async def test_parallel_subordinate_reuses_profile_validation(monkeypatch) -> None:
514 from agent import AgentContext
515 from helpers import tool_policy
516 from helpers.errors import RepairableException
517 from tools import call_subordinate
518
519 parent_agent = SimpleNamespace(
520 config=SimpleNamespace(profile="agent0"),
521 context=_FakeContext(),
522 )
523 parent_context = SimpleNamespace(agent0=parent_agent)
524 monkeypatch.setattr(
525 AgentContext,
526 "get",
527 staticmethod(lambda _context_id: parent_context),
528 )
529 monkeypatch.setattr(tool_policy.subagents, "get_paths", lambda *args, **kwargs: [])
530 monkeypatch.setattr(
531 tool_policy,
532 "get_policy",
533 lambda agent: {
534 "mode": "inherit",
535 "default": "allow",
536 "allowed": [],
537 "blocked": [],
538 },
539 )
540 monkeypatch.setattr(
541 call_subordinate.subagents,
542 "get_available_agents_dict",
543 lambda project_name: {"developer": SimpleNamespace(title="Developer")},
544 )
545 job = parallel_tools.ParallelJob(
546 id="callsubordin-invalid",
547 parent_context_id="ctx",
548 index=0,
549 tool_name="call_subordinate",
550 tool_args={"profile": "ghost", "message": "Work"},
551 kind="subordinate",
552 )
553
554 with pytest.raises(RepairableException, match="Agent profile 'ghost' not found"):
555 await parallel_tools._run_subordinate_context_job("ctx", job)
556
557
558 @pytest.mark.asyncio
559 async def test_parallel_subordinates_are_distinct_reusable_a1_children(monkeypatch) -> None:
560 from agent import Agent, AgentConfig, AgentContext
561 from helpers import message_queue, persist_chat, tool_policy
562
563 parent_id = "ctx-parallel-a1-tree"
564 AgentContext.remove(parent_id)
565 parent = AgentContext(
566 AgentConfig(mcp_servers="", profile="agent0"),
567 id=parent_id,
568 set_current=False,
569 )
570
571 async def fake_monologue(agent):
572 return agent.agent_name
573
574 monkeypatch.setattr(Agent, "monologue", fake_monologue)
575 monkeypatch.setattr(tool_policy, "ensure_tool_allowed", lambda *_args, **_kwargs: None)
576 monkeypatch.setattr(message_queue, "log_user_message", lambda *_args, **_kwargs: None)
577 monkeypatch.setattr(persist_chat, "save_tmp_chat", lambda _context: None)
578
579 child_ids = []
580 try:
581 jobs = await parallel_tools.start_parallel_jobs(
582 parent.agent0,
583 [
584 parallel_tools.NormalizedToolCall(
585 index=0,
586 tool_name="call_subordinate",
587 tool_args={"message": "left branch", "reset": True},
588 ),
589 parallel_tools.NormalizedToolCall(
590 index=1,
591 tool_name="call_subordinate",
592 tool_args={"message": "right branch", "reset": True},
593 ),
594 ],
595 )
596 results = await parallel_tools.await_parallel_jobs(
597 parent.agent0,
598 [job.id for job in jobs],
599 timeout=10,
600 )
601 child_ids = [result["context_id"] for result in results]
602
603 assert [result["state"] for result in results] == ["success", "success"]
604 assert [result["result"] for result in results] == ["A1", "A1"]
605 assert len(set(child_ids)) == 2
606 assert set(parent.agent0.get_data("_subordinates")) == set(child_ids)
607 for child_id in child_ids:
608 child = AgentContext.get(child_id)
609 assert child is not None
610 assert child.agent0.number == 1
611 assert child.get_output_data("parent_context_id") == parent.id
612 assert child.get_output_data("parent_agent_number") == 0
613 assert child.get_output_data("parent_context_kind") == "subordinate"
614 finally:
615 for child_id in child_ids:
616 AgentContext.remove(child_id)
617 AgentContext.remove(parent_id)
618
619
620 @pytest.mark.asyncio
621 async def test_failed_parallel_subordinate_continues_directly_or_in_parallel(
622 monkeypatch,
623 ) -> None:
624 from agent import Agent, AgentConfig, AgentContext
625 from helpers import message_queue, persist_chat, tool_policy
626 from tools.call_subordinate import Delegation
627
628 parent_id = "ctx-parallel-resume-tree"
629 AgentContext.remove(parent_id)
630 parent = AgentContext(
631 AgentConfig(mcp_servers="", profile="agent0"),
632 id=parent_id,
633 set_current=False,
634 )
635 calls = {}
636
637 async def flaky_monologue(agent):
638 count = calls.get(agent.context.id, 0) + 1
639 calls[agent.context.id] = count
640 if count == 1:
641 raise RuntimeError("simulated API failure")
642 return f"{agent.agent_name} continuation {count}"
643
644 monkeypatch.setattr(Agent, "monologue", flaky_monologue)
645 monkeypatch.setattr(tool_policy, "ensure_tool_allowed", lambda *_args, **_kwargs: None)
646 monkeypatch.setattr(message_queue, "log_user_message", lambda *_args, **_kwargs: None)
647 monkeypatch.setattr(persist_chat, "save_tmp_chat", lambda _context: None)
648
649 child_id = ""
650 try:
651 failed = parallel_tools.ParallelJob(
652 id="callsubordin-failed",
653 parent_context_id=parent.id,
654 index=0,
655 tool_name="call_subordinate",
656 tool_args={"message": "remember ALPHA", "reset": True},
657 kind="subordinate",
658 parent_agent=parent.agent0,
659 )
660 parallel_tools._jobs_for_context(parent)[failed.id] = failed
661 await parallel_tools._run_parallel_job(parent.id, failed.id)
662 child_id = failed.worker_context_id or ""
663
664 assert failed.state == "error"
665 assert failed.error == "simulated API failure"
666 assert child_id
667 assert AgentContext.get(child_id).agent0.number == 1 # type: ignore[union-attr]
668
669 direct = Delegation(
670 parent.agent0,
671 "call_subordinate",
672 None,
673 {},
674 "",
675 None,
676 )
677 direct_result = await direct.execute(
678 message="continue after the API failure",
679 context_id=child_id,
680 reset=False,
681 )
682 assert direct_result.message == "A1 continuation 2"
683 assert direct_result.additional == {"context_id": child_id}
684
685 continued = parallel_tools.ParallelJob(
686 id="callsubordin-continued",
687 parent_context_id=parent.id,
688 index=0,
689 tool_name="call_subordinate",
690 tool_args={
691 "message": "continue once more",
692 "context_id": child_id,
693 "reset": False,
694 },
695 kind="subordinate",
696 parent_agent=parent.agent0,
697 )
698 parallel_tools._jobs_for_context(parent)[continued.id] = continued
699 await parallel_tools._run_parallel_job(parent.id, continued.id)
700
701 assert continued.state == "success"
702 assert continued.worker_context_id == child_id
703 assert continued.result == "A1 continuation 3"
704 assert calls == {child_id: 3}
705 finally:
706 if child_id:
707 AgentContext.remove(child_id)
708 AgentContext.remove(parent_id)
709
710
711 @pytest.mark.asyncio
712 async def test_parallel_a1_spawns_a2_with_same_lifecycle(monkeypatch) -> None:
713 from agent import Agent, AgentConfig, AgentContext
714 from helpers import message_queue, persist_chat, tool_policy
715
716 parent_id = "ctx-parallel-a2-tree"
717 AgentContext.remove(parent_id)
718 parent = AgentContext(
719 AgentConfig(mcp_servers="", profile="agent0"),
720 id=parent_id,
721 set_current=False,
722 )
723
724 async def fake_monologue(agent):
725 return agent.agent_name
726
727 monkeypatch.setattr(Agent, "monologue", fake_monologue)
728 monkeypatch.setattr(tool_policy, "ensure_tool_allowed", lambda *_args, **_kwargs: None)
729 monkeypatch.setattr(message_queue, "log_user_message", lambda *_args, **_kwargs: None)
730 monkeypatch.setattr(persist_chat, "save_tmp_chat", lambda _context: None)
731
732 child_ids = []
733 try:
734 a1_job = parallel_tools.ParallelJob(
735 id="callsubordin-a1",
736 parent_context_id=parent.id,
737 index=0,
738 tool_name="call_subordinate",
739 tool_args={"message": "be A1", "reset": True},
740 kind="subordinate",
741 parent_agent=parent.agent0,
742 )
743 a1_result = await parallel_tools._run_subordinate_context_job(parent.id, a1_job)
744 a1 = AgentContext.get(a1_job.worker_context_id or "").agent0 # type: ignore[union-attr]
745 child_ids.append(a1.context.id)
746
747 a2_job = parallel_tools.ParallelJob(
748 id="callsubordin-a2",
749 parent_context_id=a1.context.id,
750 index=0,
751 tool_name="call_subordinate",
752 tool_args={"message": "be A2", "reset": True},
753 kind="subordinate",
754 parent_agent=a1,
755 )
756 a2_result = await parallel_tools._run_subordinate_context_job(
757 a1.context.id, a2_job
758 )
759 a2_context = AgentContext.get(a2_job.worker_context_id or "")
760 child_ids.append(a2_context.id) # type: ignore[union-attr]
761
762 assert a1_result == "A1"
763 assert a1.number == 1
764 assert a2_result == "A2"
765 assert a2_context.agent0.number == 2 # type: ignore[union-attr]
766 assert a2_context.get_output_data("parent_context_id") == a1.context.id # type: ignore[union-attr]
767 assert a2_context.get_output_data("parent_agent_number") == 1 # type: ignore[union-attr]
768 finally:
769 for child_id in reversed(child_ids):
770 AgentContext.remove(child_id)
771 AgentContext.remove(parent_id)
772
773
774 @pytest.mark.asyncio
775 async def test_parallel_subordinate_owns_nested_parallel_tasks(monkeypatch) -> None:
776 monkeypatch.setattr(parallel_tools, "DeferredTask", _FakeDeferredTask)
777 agent = _FakeAgent()
778 parent_task = _FakeDeferredTask()
779 agent.context.task = parent_task
780 agent.context.set_data(parallel_tools.PARALLEL_WORKER_KIND_KEY, "subordinate")
781
782 jobs = await parallel_tools.start_parallel_jobs(
783 agent, # type: ignore[arg-type]
784 [
785 parallel_tools.NormalizedToolCall(
786 index=0,
787 tool_name="call_subordinate",
788 tool_args={"message": "nested", "reset": True},
789 )
790 ],
791 )
792
793 assert parent_task.children == [jobs[0].deferred_task]
794 parent_task.kill()
795 assert jobs[0].deferred_task.killed == 1 # type: ignore[union-attr]
796
797
798 @pytest.mark.asyncio
799 async def test_parallel_direct_tool_jobs_fallback_to_generic_tool_log_type(monkeypatch) -> None:
800 class FakeDeferredTask:
801 def __init__(self, thread_name=None) -> None:
802 self.thread_name = thread_name
803 self.started = None
804
805 def start_task(self, func, *args):
806 self.started = (func, args)
807 return self
808
809 def is_ready(self):
810 return False
811
812 def is_alive(self):
813 return True
814
815 def kill(self):
816 pass
817
818 monkeypatch.setattr(parallel_tools, "DeferredTask", FakeDeferredTask)
819 monkeypatch.setattr(parallel_tools, "_resolve_parallel_tool", lambda *_args, **_kwargs: None)
820 agent = _FakeAgent()
821
822 jobs = await parallel_tools.start_parallel_jobs(
823 agent, # type: ignore[arg-type]
824 [
825 parallel_tools.NormalizedToolCall(
826 index=0,
827 tool_name="wait",
828 tool_args={"seconds": 1},
829 )
830 ],
831 )
832
833 assert jobs[0].kind == "tool"
834 assert agent.context.log.items[0].type == "tool"
835 assert agent.context.log.items[0].kvps == {"seconds": 1, "_tool_name": "wait"}
836
837 parallel_tools._finish_job(jobs[0], "success", result="done")
838
839 assert agent.context.log.items[0].content == "done"
840 assert agent.context.log.items[0].kvps == {"seconds": 1, "_tool_name": "wait"}
841
842
843 @pytest.mark.asyncio
844 async def test_parallel_code_execution_child_uses_code_exe_log_type(monkeypatch) -> None:
845 class FakeDeferredTask:
846 def __init__(self, thread_name=None) -> None:
847 self.thread_name = thread_name
848
849 def start_task(self, func, *args):
850 return self
851
852 def is_ready(self):
853 return False
854
855 def is_alive(self):
856 return True
857
858 def kill(self):
859 pass
860
861 class FakeCodeExecutionTool:
862 def __init__(self, agent, args):
863 self.agent = agent
864 self.args = args
865
866 def get_log_object(self):
867 runtime = self.args.get("runtime", "unknown")
868 session = self.args.get("session", None)
869 session_text = f"[{session}] " if session or session == 0 else ""
870 return self.agent.context.log.log(
871 type="code_exe",
872 heading=f"icon://terminal {session_text}code_execution_tool - {runtime}",
873 content="",
874 kvps=self.args,
875 )
876
877 monkeypatch.setattr(parallel_tools, "DeferredTask", FakeDeferredTask)
878 monkeypatch.setattr(
879 parallel_tools,
880 "_resolve_parallel_tool",
881 lambda _agent, _tool_name, args: FakeCodeExecutionTool(_agent, args),
882 )
883 agent = _FakeAgent()
884
885 jobs = await parallel_tools.start_parallel_jobs(
886 agent, # type: ignore[arg-type]
887 [
888 parallel_tools.NormalizedToolCall(
889 index=0,
890 tool_name="code_execution_tool",
891 tool_args={
892 "runtime": "terminal",
893 "session": 0,
894 "code": "pwd",
895 },
896 )
897 ],
898 )
899
900 assert jobs[0].kind == "tool"
901 assert agent.context.log.items[0].type == "code_exe"
902 assert agent.context.log.items[0].heading == "icon://terminal [0] code_execution_tool - terminal"
903 assert agent.context.log.items[0].kvps == {
904 "runtime": "terminal",
905 "session": 0,
906 "code": "pwd",
907 }
908
909
910 @pytest.mark.asyncio
911 async def test_parallel_wait_child_uses_wait_log_type(monkeypatch) -> None:
912 class FakeDeferredTask:
913 def __init__(self, thread_name=None) -> None:
914 self.thread_name = thread_name
915
916 def start_task(self, func, *args):
917 return self
918
919 def is_ready(self):
920 return False
921
922 def is_alive(self):
923 return True
924
925 def kill(self):
926 pass
927
928 class FakeWaitTool:
929 def __init__(self, agent, args):
930 self.agent = agent
931 self.args = args
932
933 def get_log_object(self):
934 return self.agent.context.log.log(
935 type="progress",
936 heading="icon://timer Wait: Waiting...",
937 content="",
938 kvps=self.args,
939 )
940
941 monkeypatch.setattr(parallel_tools, "DeferredTask", FakeDeferredTask)
942 monkeypatch.setattr(
943 parallel_tools,
944 "_resolve_parallel_tool",
945 lambda _agent, _tool_name, args: FakeWaitTool(_agent, args),
946 )
947 agent = _FakeAgent()
948
949 jobs = await parallel_tools.start_parallel_jobs(
950 agent, # type: ignore[arg-type]
951 [
952 parallel_tools.NormalizedToolCall(
953 index=0,
954 tool_name="wait",
955 tool_args={"seconds": 1},
956 )
957 ],
958 )
959
960 assert jobs[0].kind == "tool"
961 assert agent.context.log.items[0].type == "progress"
962 assert agent.context.log.items[0].heading == "icon://timer Wait: Waiting..."
963 assert agent.context.log.items[0].kvps == {"seconds": 1}
964
965
966 @pytest.mark.asyncio
967 async def test_parallel_execute_reuses_child_log_object(monkeypatch) -> None:
968 class FakeTool:
969 def __init__(self, agent, args):
970 self.agent = agent
971 self.args = args
972
973 def get_log_object(self):
974 return self.agent.context.log.log(
975 type="tool",
976 heading="generic tool log",
977 content="",
978 kvps=self.args,
979 )
980
981 async def before_execution(self, **kwargs):
982 self.log = self.get_log_object()
983
984 async def execute(self, **kwargs):
985 return Response(message="done", break_loop=False)
986
987 async def after_execution(self, response):
988 self.log.update(content=response.message)
989
990 class FakeWorkerAgent(_FakeAgent):
991 def __init__(self) -> None:
992 super().__init__()
993 self.loop_data = SimpleNamespace(current_tool=None)
994
995 def get_tool(self, **kwargs):
996 return FakeTool(self, kwargs["args"])
997
998 async def handle_intervention(self):
999 pass
1000
1001 async def noop_extensions(*_args, **_kwargs):
1002 pass
1003
1004 monkeypatch.setattr(parallel_tools, "call_extensions_async", noop_extensions)
1005
1006 agent = FakeWorkerAgent()
1007 child_log = agent.context.log.log(
1008 type="progress",
1009 heading="icon://timer Wait: Waiting...",
1010 content="",
1011 kvps={"seconds": 1},
1012 )
1013
1014 result = await parallel_tools.execute_tool_call(
1015 agent, # type: ignore[arg-type]
1016 "wait",
1017 {"seconds": 1},
1018 log_item=child_log,
1019 )
1020
1021 assert result == "done"
1022 assert agent.context.log.items == [child_log]
1023 assert child_log.type == "progress"
1024 assert child_log.content == "done"
1025
1026
1027 @pytest.mark.asyncio
1028 async def test_parallel_tool_keeps_wrapper_out_of_visible_log() -> None:
1029 from tools.parallel import ParallelTool
1030
1031 class HistoryAgent(_FakeAgent):
1032 def __init__(self) -> None:
1033 super().__init__()
1034 self.tool_results = []
1035
1036 def hist_add_tool_result(self, tool_name, tool_result, **kwargs):
1037 self.tool_results.append((tool_name, tool_result, kwargs))
1038
1039 agent = HistoryAgent()
1040 tool = ParallelTool(agent, "parallel", None, {}, "", None) # type: ignore[arg-type]
1041
1042 await tool.before_execution()
1043 await tool.after_execution(Response(message="done", break_loop=False, additional={"extra": "value"}))
1044
1045 assert agent.context.log.items == []
1046 assert agent.tool_results == [("parallel", "done", {"extra": "value"})]
1047
1048
1049 @pytest.mark.asyncio
1050 async def test_parallel_collect_promotes_history_after_wrapper_result() -> None:
1051 from tools.parallel import ParallelTool
1052
1053 class HistoryAgent(_FakeAgent):
1054 def __init__(self) -> None:
1055 super().__init__()
1056 self.history_events = []
1057
1058 def hist_add_tool_result(self, tool_name, tool_result, **kwargs):
1059 self.history_events.append(("tool", tool_name, tool_result, kwargs))
1060
1061 def hist_add_message(self, ai, content, tokens=0, **kwargs):
1062 self.history_events.append(("message", ai, content, tokens, kwargs))
1063
1064 agent = HistoryAgent()
1065 job = parallel_tools.ParallelJob(
1066 id="vision-ready",
1067 parent_context_id=agent.context.id,
1068 index=0,
1069 tool_name="vision_load",
1070 tool_args={"paths": ["/image.png"]},
1071 kind="tool",
1072 state="success",
1073 result="Loaded images (1)",
1074 parent_history=[("raw-image-message", 1500)],
1075 )
1076 agent.context.set_data(parallel_tools.PARALLEL_JOBS_KEY, {job.id: job})
1077 tool = ParallelTool(
1078 agent, # type: ignore[arg-type]
1079 "parallel",
1080 None,
1081 {"action": "collect", "job_ids": [job.id]},
1082 "",
1083 None,
1084 )
1085
1086 response = await tool.execute(**tool.args)
1087
1088 assert job.id in agent.context.get_data(parallel_tools.PARALLEL_JOBS_KEY)
1089 assert agent.history_events == []
1090
1091 await tool.after_execution(response)
1092
1093 assert agent.history_events[0][0] == "tool"
1094 assert agent.history_events[1] == (
1095 "message",
1096 False,
1097 "raw-image-message",
1098 1500,
1099 {},
1100 )
1101 assert job.id not in agent.context.get_data(parallel_tools.PARALLEL_JOBS_KEY)
1102
1103
1104 @pytest.mark.asyncio
1105 async def test_parallel_child_contexts_are_chats_not_tasks(monkeypatch) -> None:
1106 from agent import AgentContext
1107 from initialize import initialize_agent
1108 from helpers import state_snapshot
1109
1110 class NoTaskScheduler:
1111 def get_task_by_uuid(self, _task_id):
1112 return None
1113
1114 monkeypatch.setattr(
1115 state_snapshot,
1116 "TaskScheduler",
1117 SimpleNamespace(get=lambda: NoTaskScheduler()),
1118 )
1119
1120 parent_id = "ctx-par-parent"
1121 child_id = "ctx-par-child"
1122 parent = AgentContext(config=initialize_agent(), id=parent_id, name="Parent", set_current=False)
1123 child = AgentContext(config=initialize_agent(), id=child_id, name="Child", set_current=False)
1124 try:
1125 child.set_output_data(parallel_tools.CHILD_PARENT_CONTEXT_ID_KEY, parent.id)
1126 child.set_output_data(parallel_tools.CHILD_PARENT_CONTEXT_KIND_KEY, "parallel")
1127 child.set_output_data(parallel_tools.CHILD_PARENT_CONTEXT_LABEL_KEY, "Child task")
1128 child.set_output_data(parallel_tools.CHILD_PARALLEL_JOB_ID_KEY, "job-123")
1129
1130 payload = await state_snapshot.build_snapshot(
1131 context=parent.id,
1132 log_from=0,
1133 notifications_from=0,
1134 timezone="UTC",
1135 )
1136
1137 contexts_by_id = {ctx["id"]: ctx for ctx in payload["contexts"]}
1138 task_ids = {task["id"] for task in payload["tasks"]}
1139 assert parent_id in contexts_by_id
1140 assert child_id in contexts_by_id
1141 assert contexts_by_id[child_id]["parent_context_id"] == parent_id
1142 assert child_id not in task_ids
1143 finally:
1144 AgentContext.remove(parent_id)
1145 AgentContext.remove(child_id)
1146
1147
1148 def test_chats_sidebar_projects_parallel_children_as_indented_accordion() -> None:
1149 store = (PROJECT_ROOT / "webui/components/sidebar/chats/chats-store.js").read_text(
1150 encoding="utf-8"
1151 )
1152 html = (PROJECT_ROOT / "webui/components/sidebar/chats/chats-list.html").read_text(
1153 encoding="utf-8"
1154 )
1155
1156 assert "parent_context_id" in store
1157 assert "this.expandedParents[selectedId] === undefined" in store
1158 assert "...this.expandedParents," in store
1159 assert "[selectedId]: true," in store
1160 assert "topLevelContexts()" in html
1161 assert "childContexts(context.id)" in html
1162 assert "chat-child-container" in html
1163 assert "keyboard_arrow_up" in html
1164 assert "keyboard_arrow_down" in html
1165 assert ".chats-config-list .chat-tree-item" in html
1166 assert ".chats-config-list .chat-child-list > li" in html
1167 assert 'x-show="$store.chats.hasChildren(context.id)"' in html
1168 assert "'chat-has-children': $store.chats.hasChildren(context.id)" in html
1169 assert ".chat-container.chat-has-children .chat-list-button" in html
1170 assert "left: 2px" in html
1171 assert "padding-left: 24px" in html
1172 assert "color: var(--color-text-muted)" in html
1173
1174
1175 def test_parallel_result_json_is_compact() -> None:
1176 result = parallel_tools.format_parallel_results(
1177 [{"job_id": "wait-1", "tool_name": "wait", "state": "success"}]
1178 )
1179
1180 assert result == '{"status":"success","jobs":[{"job_id":"wait-1","tool_name":"wait","state":"success"}]}'