mirror of
https://github.com/agent0ai/agent-zero.git
synced 2026-07-09 17:08:29 +00:00
78a50e1431
Skip persisted tool-result files for BACKGROUND contexts, remove ephemeral direct parallel worker chat folders during cleanup, and ignore chat directories without chat.json during startup loading. Add focused regressions for background output persistence, parallel worker cleanup, and orphan directory loading.
486 lines
16 KiB
Python
486 lines
16 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import time
|
|
import sys
|
|
from types import SimpleNamespace
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
PROJECT_ROOT = Path(__file__).resolve().parents[1]
|
|
if str(PROJECT_ROOT) not in sys.path:
|
|
sys.path.insert(0, str(PROJECT_ROOT))
|
|
|
|
from helpers import parallel_tools
|
|
from helpers.tool import Response
|
|
|
|
|
|
class _FakeLogItem:
|
|
def __init__(self, type_, heading="", content="", kvps=None, id_=None, **kwargs) -> None:
|
|
self.type = type_
|
|
self.heading = heading
|
|
self.content = content
|
|
self.kvps = dict(kvps or {})
|
|
self.kvps.update(kwargs)
|
|
self.id = id_
|
|
|
|
def update(self, content=None, kvps=None, **kwargs):
|
|
if content is not None:
|
|
self.content = content
|
|
if kvps:
|
|
self.kvps.update(kvps)
|
|
self.kvps.update(kwargs)
|
|
|
|
|
|
class _FakeLog:
|
|
def __init__(self) -> None:
|
|
self.items = []
|
|
|
|
def log(self, type, heading="", content="", kvps=None, id=None, **kwargs):
|
|
item = _FakeLogItem(type, heading, content, kvps, id, **kwargs)
|
|
self.items.append(item)
|
|
return item
|
|
|
|
|
|
class _FakeContext:
|
|
def __init__(self) -> None:
|
|
self.id = "ctx"
|
|
self.data = {}
|
|
self.log = _FakeLog()
|
|
|
|
def get_data(self, key: str, recursive: bool = True):
|
|
return self.data.get(key)
|
|
|
|
def set_data(self, key: str, value, recursive: bool = True):
|
|
self.data[key] = value
|
|
|
|
|
|
class _FakeAgent:
|
|
def __init__(self) -> None:
|
|
self.context = _FakeContext()
|
|
self.agent_name = "A0"
|
|
|
|
|
|
class _FakeDeferredTask:
|
|
def __init__(self, *, ready: bool = False, alive: bool = True, result=None) -> None:
|
|
self.ready = ready
|
|
self.alive = alive
|
|
self._result = result
|
|
self.killed = 0
|
|
|
|
def is_ready(self):
|
|
return self.ready
|
|
|
|
def is_alive(self):
|
|
return self.alive
|
|
|
|
async def result(self):
|
|
return self._result
|
|
|
|
def kill(self):
|
|
self.killed += 1
|
|
self.alive = False
|
|
|
|
|
|
def test_normalize_parallel_tool_calls_accepts_normal_tool_request_shapes() -> None:
|
|
calls = parallel_tools.normalize_parallel_tool_calls(
|
|
[
|
|
{
|
|
"tool_name": "text_editor:read",
|
|
"tool_args": {"path": "README.md"},
|
|
},
|
|
{
|
|
"tool": "scheduler",
|
|
"args": {"method": "list_tasks"},
|
|
},
|
|
]
|
|
)
|
|
|
|
assert calls[0].tool_name == "text_editor"
|
|
assert calls[0].tool_args == {"path": "README.md", "action": "read"}
|
|
assert calls[1].tool_name == "scheduler"
|
|
assert calls[1].tool_args == {"method": "list_tasks", "action": "list_tasks"}
|
|
|
|
|
|
def test_normalize_parallel_tool_calls_rejects_nested_parallel() -> None:
|
|
with pytest.raises(ValueError, match="cannot be nested"):
|
|
parallel_tools.normalize_parallel_tool_calls(
|
|
[{"tool_name": "parallel", "tool_args": {"tool_calls": []}}]
|
|
)
|
|
|
|
|
|
def test_normalize_parallel_tool_calls_rejects_document_query() -> None:
|
|
with pytest.raises(ValueError, match="document_query.*parallel"):
|
|
parallel_tools.normalize_parallel_tool_calls(
|
|
[
|
|
{
|
|
"tool_name": "document_query",
|
|
"tool_args": {"document": "/tmp/report.pdf"},
|
|
}
|
|
]
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_parallel_jobs_extras_lists_running_and_ready_jobs() -> None:
|
|
agent = _FakeAgent()
|
|
running = parallel_tools.ParallelJob(
|
|
id="search-1234abcd",
|
|
parent_context_id="ctx",
|
|
index=0,
|
|
tool_name="search_engine",
|
|
tool_args={"query": "Agent Zero"},
|
|
kind="tool",
|
|
state="running",
|
|
started_at=time.time() - 2,
|
|
)
|
|
ready = parallel_tools.ParallelJob(
|
|
id="callsubordin-5678efgh",
|
|
parent_context_id="ctx",
|
|
index=1,
|
|
tool_name="call_subordinate",
|
|
tool_args={"message": "Summarize"},
|
|
kind="subordinate",
|
|
state="success",
|
|
started_at=time.time() - 4,
|
|
completed_at=time.time() - 1,
|
|
result="done",
|
|
)
|
|
agent.context.set_data(
|
|
parallel_tools.PARALLEL_JOBS_KEY,
|
|
{running.id: running, ready.id: ready},
|
|
)
|
|
|
|
extras = await parallel_tools.build_parallel_jobs_extras(agent) # type: ignore[arg-type]
|
|
|
|
assert "search-1234abcd" in extras
|
|
assert "callsubordin-5678efgh" in extras
|
|
assert "ready to collect" in extras
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_parallel_await_timeout_keeps_running_jobs_awaitable(monkeypatch) -> None:
|
|
agent = _FakeAgent()
|
|
task = _FakeDeferredTask(alive=True)
|
|
job = parallel_tools.ParallelJob(
|
|
id="wait-1234abcd",
|
|
parent_context_id="ctx",
|
|
index=0,
|
|
tool_name="wait",
|
|
tool_args={"seconds": 60},
|
|
kind="tool",
|
|
state="running",
|
|
created_at=99.0,
|
|
started_at=99.0,
|
|
deferred_task=task, # type: ignore[arg-type]
|
|
)
|
|
agent.context.set_data(parallel_tools.PARALLEL_JOBS_KEY, {job.id: job})
|
|
times = iter([100.0, 102.0])
|
|
monkeypatch.setattr(
|
|
parallel_tools.time,
|
|
"time",
|
|
lambda: next(times, 102.0),
|
|
)
|
|
|
|
results = await parallel_tools.await_parallel_jobs( # type: ignore[arg-type]
|
|
agent,
|
|
[job.id],
|
|
timeout=1,
|
|
collect=True,
|
|
wait=True,
|
|
)
|
|
payload = json.loads(parallel_tools.format_parallel_results(results))
|
|
|
|
assert results[0]["state"] == "running"
|
|
assert results[0]["wait_timed_out"] is True
|
|
assert payload["status"] == "waiting"
|
|
assert payload["wait_timeout"] is True
|
|
assert "await" in payload["instruction"]
|
|
assert task.killed == 0
|
|
assert agent.context.get_data(parallel_tools.PARALLEL_JOBS_KEY)[job.id] is job
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_parallel_collect_returns_running_jobs_without_waiting_or_canceling() -> None:
|
|
from tools.parallel import ParallelTool
|
|
|
|
agent = _FakeAgent()
|
|
task = _FakeDeferredTask(alive=True)
|
|
job = parallel_tools.ParallelJob(
|
|
id="wait-collect",
|
|
parent_context_id="ctx",
|
|
index=0,
|
|
tool_name="wait",
|
|
tool_args={"seconds": 60},
|
|
kind="tool",
|
|
state="running",
|
|
deferred_task=task, # type: ignore[arg-type]
|
|
)
|
|
agent.context.set_data(parallel_tools.PARALLEL_JOBS_KEY, {job.id: job})
|
|
tool = ParallelTool(
|
|
agent, # type: ignore[arg-type]
|
|
"parallel",
|
|
None,
|
|
{"action": "collect", "job_ids": [job.id]},
|
|
"",
|
|
None,
|
|
)
|
|
|
|
response = await tool.execute(**tool.args)
|
|
payload = json.loads(response.message)
|
|
|
|
assert payload["status"] == "running"
|
|
assert payload["jobs"][0]["job_id"] == job.id
|
|
assert "wait_timeout" not in payload
|
|
assert task.killed == 0
|
|
assert agent.context.get_data(parallel_tools.PARALLEL_JOBS_KEY)[job.id] is job
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_parallel_cancel_still_stops_and_removes_running_jobs() -> None:
|
|
agent = _FakeAgent()
|
|
task = _FakeDeferredTask(alive=True)
|
|
job = parallel_tools.ParallelJob(
|
|
id="wait-cancel",
|
|
parent_context_id="ctx",
|
|
index=0,
|
|
tool_name="wait",
|
|
tool_args={"seconds": 60},
|
|
kind="tool",
|
|
state="running",
|
|
deferred_task=task, # type: ignore[arg-type]
|
|
)
|
|
agent.context.set_data(parallel_tools.PARALLEL_JOBS_KEY, {job.id: job})
|
|
|
|
results = await parallel_tools.cancel_parallel_jobs(agent, [job.id]) # type: ignore[arg-type]
|
|
payload = json.loads(parallel_tools.format_parallel_results(results))
|
|
|
|
assert results[0]["state"] == "cancelled"
|
|
assert payload["status"] == "cancelled"
|
|
assert task.killed == 1
|
|
assert job.id not in agent.context.get_data(parallel_tools.PARALLEL_JOBS_KEY)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_parallel_remove_context_deletes_persisted_worker_chat(monkeypatch) -> None:
|
|
removed = []
|
|
|
|
from helpers import persist_chat
|
|
|
|
monkeypatch.setattr(persist_chat, "remove_chat", removed.append)
|
|
|
|
await parallel_tools._remove_context("missing-worker")
|
|
|
|
assert removed == ["missing-worker"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_parallel_recursion_guard_allows_subordinate_children_but_blocks_tool_workers() -> None:
|
|
from extensions.python.tool_execute_before._20_block_parallel_recursion import (
|
|
BlockParallelRecursion,
|
|
)
|
|
from helpers.errors import RepairableException
|
|
|
|
agent = _FakeAgent()
|
|
agent.context.set_data(parallel_tools.PARALLEL_WORKER_JOB_KEY, "legacy-job")
|
|
assert parallel_tools.is_parallel_worker(agent) is True # type: ignore[arg-type]
|
|
|
|
agent.context.set_data(parallel_tools.PARALLEL_WORKER_KIND_KEY, "subordinate")
|
|
assert parallel_tools.is_parallel_worker(agent) is False # type: ignore[arg-type]
|
|
await BlockParallelRecursion(agent=agent).execute(tool_name="parallel") # type: ignore[arg-type]
|
|
|
|
agent.context.set_data(parallel_tools.PARALLEL_WORKER_KIND_KEY, "tool")
|
|
assert parallel_tools.is_parallel_worker(agent) is True # type: ignore[arg-type]
|
|
with pytest.raises(RepairableException, match="cannot be used inside a parallel worker"):
|
|
await BlockParallelRecursion(agent=agent).execute(tool_name="parallel") # type: ignore[arg-type]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_parallel_subordinate_jobs_are_visible_child_logs_not_scheduler_tasks(monkeypatch) -> None:
|
|
class FakeDeferredTask:
|
|
def __init__(self, thread_name=None) -> None:
|
|
self.thread_name = thread_name
|
|
self.started = None
|
|
|
|
def start_task(self, func, *args):
|
|
self.started = (func, args)
|
|
return self
|
|
|
|
def is_ready(self):
|
|
return False
|
|
|
|
def is_alive(self):
|
|
return True
|
|
|
|
def kill(self):
|
|
pass
|
|
|
|
monkeypatch.setattr(parallel_tools, "DeferredTask", FakeDeferredTask)
|
|
agent = _FakeAgent()
|
|
|
|
jobs = await parallel_tools.start_parallel_jobs(
|
|
agent, # type: ignore[arg-type]
|
|
[
|
|
parallel_tools.NormalizedToolCall(
|
|
index=0,
|
|
tool_name="call_subordinate",
|
|
tool_args={
|
|
"profile": "developer",
|
|
"message": "Return ALPHA=1",
|
|
"reset": True,
|
|
},
|
|
)
|
|
],
|
|
)
|
|
|
|
assert jobs[0].kind == "subordinate"
|
|
snapshot = parallel_tools._job_snapshot(jobs[0], include_result=False)
|
|
assert "scheduler_task_uuid" not in snapshot
|
|
assert agent.context.log.items[0].type == "subagent"
|
|
assert agent.context.log.items[0].kvps == {
|
|
"profile": "developer",
|
|
"message": "Return ALPHA=1",
|
|
"reset": True,
|
|
}
|
|
assert "id" not in agent.context.log.items[0].kvps
|
|
assert "tool_name" not in agent.context.log.items[0].kvps
|
|
assert "parallel_child" not in agent.context.log.items[0].kvps
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_parallel_direct_tool_jobs_log_normal_tool_metadata(monkeypatch) -> None:
|
|
class FakeDeferredTask:
|
|
def __init__(self, thread_name=None) -> None:
|
|
self.thread_name = thread_name
|
|
self.started = None
|
|
|
|
def start_task(self, func, *args):
|
|
self.started = (func, args)
|
|
return self
|
|
|
|
def is_ready(self):
|
|
return False
|
|
|
|
def is_alive(self):
|
|
return True
|
|
|
|
def kill(self):
|
|
pass
|
|
|
|
monkeypatch.setattr(parallel_tools, "DeferredTask", FakeDeferredTask)
|
|
agent = _FakeAgent()
|
|
|
|
jobs = await parallel_tools.start_parallel_jobs(
|
|
agent, # type: ignore[arg-type]
|
|
[
|
|
parallel_tools.NormalizedToolCall(
|
|
index=0,
|
|
tool_name="wait",
|
|
tool_args={"seconds": 1},
|
|
)
|
|
],
|
|
)
|
|
|
|
assert jobs[0].kind == "tool"
|
|
assert agent.context.log.items[0].type == "tool"
|
|
assert agent.context.log.items[0].kvps == {"seconds": 1, "_tool_name": "wait"}
|
|
|
|
parallel_tools._finish_job(jobs[0], "success", result="done")
|
|
|
|
assert agent.context.log.items[0].content == "done"
|
|
assert agent.context.log.items[0].kvps == {"seconds": 1, "_tool_name": "wait"}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_parallel_tool_keeps_wrapper_out_of_visible_log() -> None:
|
|
from tools.parallel import ParallelTool
|
|
|
|
class HistoryAgent(_FakeAgent):
|
|
def __init__(self) -> None:
|
|
super().__init__()
|
|
self.tool_results = []
|
|
|
|
def hist_add_tool_result(self, tool_name, tool_result, **kwargs):
|
|
self.tool_results.append((tool_name, tool_result, kwargs))
|
|
|
|
agent = HistoryAgent()
|
|
tool = ParallelTool(agent, "parallel", None, {}, "", None) # type: ignore[arg-type]
|
|
|
|
await tool.before_execution()
|
|
await tool.after_execution(Response(message="done", break_loop=False, additional={"extra": "value"}))
|
|
|
|
assert agent.context.log.items == []
|
|
assert agent.tool_results == [("parallel", "done", {"extra": "value"})]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_parallel_child_contexts_are_chats_not_tasks(monkeypatch) -> None:
|
|
from agent import AgentContext
|
|
from initialize import initialize_agent
|
|
from helpers import state_snapshot
|
|
|
|
class NoTaskScheduler:
|
|
def get_task_by_uuid(self, _task_id):
|
|
return None
|
|
|
|
monkeypatch.setattr(
|
|
state_snapshot,
|
|
"TaskScheduler",
|
|
SimpleNamespace(get=lambda: NoTaskScheduler()),
|
|
)
|
|
|
|
parent_id = "ctx-par-parent"
|
|
child_id = "ctx-par-child"
|
|
parent = AgentContext(config=initialize_agent(), id=parent_id, name="Parent", set_current=False)
|
|
child = AgentContext(config=initialize_agent(), id=child_id, name="Child", set_current=False)
|
|
try:
|
|
child.set_output_data(parallel_tools.CHILD_PARENT_CONTEXT_ID_KEY, parent.id)
|
|
child.set_output_data(parallel_tools.CHILD_PARENT_CONTEXT_KIND_KEY, "parallel")
|
|
child.set_output_data(parallel_tools.CHILD_PARENT_CONTEXT_LABEL_KEY, "Child task")
|
|
child.set_output_data(parallel_tools.CHILD_PARALLEL_JOB_ID_KEY, "job-123")
|
|
|
|
payload = await state_snapshot.build_snapshot(
|
|
context=parent.id,
|
|
log_from=0,
|
|
notifications_from=0,
|
|
timezone="UTC",
|
|
)
|
|
|
|
contexts_by_id = {ctx["id"]: ctx for ctx in payload["contexts"]}
|
|
task_ids = {task["id"] for task in payload["tasks"]}
|
|
assert parent_id in contexts_by_id
|
|
assert child_id in contexts_by_id
|
|
assert contexts_by_id[child_id]["parent_context_id"] == parent_id
|
|
assert child_id not in task_ids
|
|
finally:
|
|
AgentContext.remove(parent_id)
|
|
AgentContext.remove(child_id)
|
|
|
|
|
|
def test_chats_sidebar_projects_parallel_children_as_indented_accordion() -> None:
|
|
store = (PROJECT_ROOT / "webui/components/sidebar/chats/chats-store.js").read_text(
|
|
encoding="utf-8"
|
|
)
|
|
html = (PROJECT_ROOT / "webui/components/sidebar/chats/chats-list.html").read_text(
|
|
encoding="utf-8"
|
|
)
|
|
|
|
assert "parent_context_id" in store
|
|
assert "const nextExpandedParents = { ...this.expandedParents };" in store
|
|
assert "nextExpandedParents[selectedId] === undefined" in store
|
|
assert "nextExpandedParents[selectedId] = true;" in store
|
|
assert "topLevelContexts()" in html
|
|
assert "childContexts(context.id)" in html
|
|
assert "chat-child-container" in html
|
|
assert "keyboard_arrow_up" in html
|
|
assert "keyboard_arrow_down" in html
|
|
assert ".chats-config-list .chat-tree-item" in html
|
|
assert ".chats-config-list .chat-child-list > li" in html
|
|
assert 'x-show="$store.chats.hasChildren(context.id)"' in html
|
|
assert "'chat-has-children': $store.chats.hasChildren(context.id)" in html
|
|
assert ".chat-container.chat-has-children .chat-list-button" in html
|
|
assert "left: 2px" in html
|
|
assert "padding-left: 24px" in html
|
|
assert "color: var(--color-text-muted)" in html
|
|
assert "padding: 8px;" in html
|