609 lines
18 KiB
Python
609 lines
18 KiB
Python
# -*- coding: utf-8 -*-
|
|
"""Tests for bounded multi-agent startup scheduling."""
|
|
# pylint: disable=protected-access
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import os
|
|
import subprocess
|
|
import sys
|
|
from types import SimpleNamespace
|
|
from unittest.mock import AsyncMock, MagicMock
|
|
|
|
import pytest
|
|
|
|
import pineagents.app.multi_agent_manager as multi_agent_manager_module
|
|
import pineagents.constant as constants
|
|
from pineagents.app.agent_startup import AgentStartupStatus
|
|
from pineagents.app.multi_agent_manager import MultiAgentManager
|
|
from pineagents.app.task_tracker import REPLAY_END_SSE, TaskTracker
|
|
from pineagents.constant import BUILTIN_QA_AGENT_ID
|
|
|
|
|
|
def _config(*agent_ids: str):
|
|
profiles = {
|
|
agent_id: SimpleNamespace(
|
|
id=agent_id,
|
|
workspace_dir=f"/tmp/{agent_id}",
|
|
enabled=True,
|
|
)
|
|
for agent_id in agent_ids
|
|
}
|
|
return SimpleNamespace(
|
|
agents=SimpleNamespace(profiles=profiles),
|
|
)
|
|
|
|
|
|
class _ReloadServiceManager:
|
|
def __init__(self) -> None:
|
|
self.services = {}
|
|
|
|
def get_reusable_services(self) -> dict:
|
|
return {}
|
|
|
|
|
|
class _ReloadWorkspace:
|
|
def __init__(self, agent_id: str) -> None:
|
|
self.agent_id = agent_id
|
|
self.task_tracker = TaskTracker()
|
|
self._service_manager = _ReloadServiceManager()
|
|
self.started = False
|
|
self.stopped = False
|
|
self.manager = None
|
|
|
|
def set_task_tracker(self, task_tracker: TaskTracker) -> None:
|
|
assert not self.started
|
|
self.task_tracker = task_tracker
|
|
|
|
async def set_reusable_components(self, _components: dict) -> None:
|
|
return None
|
|
|
|
async def start(self) -> None:
|
|
self.started = True
|
|
|
|
def set_manager(self, manager: MultiAgentManager) -> None:
|
|
self.manager = manager
|
|
|
|
async def stop(self, final: bool = True) -> None:
|
|
del final
|
|
self.stopped = True
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_reload_reuses_tracker_for_active_stream_reconnect(
|
|
monkeypatch,
|
|
) -> None:
|
|
manager = MultiAgentManager()
|
|
config = _config("agent-1")
|
|
monkeypatch.setattr(
|
|
"pineagents.app.multi_agent_manager.load_config",
|
|
lambda: config,
|
|
)
|
|
old_workspace = _ReloadWorkspace("agent-1")
|
|
new_workspace = _ReloadWorkspace("agent-1")
|
|
manager.agents["agent-1"] = old_workspace
|
|
manager._create_workspace = MagicMock(return_value=new_workspace)
|
|
release = asyncio.Event()
|
|
emitted = asyncio.Event()
|
|
|
|
async def producer(_payload):
|
|
yield "data: replayed\n\n"
|
|
emitted.set()
|
|
await release.wait()
|
|
yield "data: live\n\n"
|
|
|
|
original_queue, _ = await old_workspace.task_tracker.attach_or_start(
|
|
"chat-1",
|
|
None,
|
|
producer,
|
|
owner=old_workspace,
|
|
)
|
|
await asyncio.wait_for(emitted.wait(), timeout=1)
|
|
await asyncio.sleep(0)
|
|
|
|
assert await manager.reload_agent("agent-1") is True
|
|
assert manager.agents["agent-1"] is new_workspace
|
|
assert new_workspace.task_tracker is old_workspace.task_tracker
|
|
|
|
reconnect_queue = await new_workspace.task_tracker.attach("chat-1")
|
|
assert reconnect_queue is not None
|
|
assert await reconnect_queue.get() == "data: replayed\n\n"
|
|
assert await reconnect_queue.get() == REPLAY_END_SSE
|
|
|
|
cleanup_tasks = list(manager._cleanup_tasks)
|
|
assert cleanup_tasks
|
|
release.set()
|
|
assert await reconnect_queue.get() == "data: live\n\n"
|
|
assert await reconnect_queue.get() is None
|
|
async for _ in old_workspace.task_tracker.stream_from_queue(
|
|
original_queue,
|
|
"chat-1",
|
|
):
|
|
pass
|
|
await asyncio.wait_for(
|
|
asyncio.gather(*cleanup_tasks),
|
|
timeout=1,
|
|
)
|
|
assert old_workspace.stopped is True
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cleanup_keeps_old_workspace_alive_after_wait_timeout() -> None:
|
|
manager = MultiAgentManager()
|
|
old_workspace = _ReloadWorkspace("agent-1")
|
|
task = asyncio.Future()
|
|
old_workspace.task_tracker.wait_tasks_done = AsyncMock(
|
|
side_effect=[False, True],
|
|
)
|
|
|
|
await manager._graceful_stop_old_instance(
|
|
old_workspace,
|
|
"agent-1",
|
|
active_tasks={"chat-1": task},
|
|
)
|
|
cleanup_tasks = list(manager._cleanup_tasks)
|
|
assert cleanup_tasks
|
|
await asyncio.wait_for(
|
|
asyncio.gather(*cleanup_tasks),
|
|
timeout=1,
|
|
)
|
|
|
|
assert old_workspace.task_tracker.wait_tasks_done.await_count == 2
|
|
assert old_workspace.stopped is True
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cleanup_forces_stop_after_maximum_wait_rounds(
|
|
monkeypatch,
|
|
) -> None:
|
|
manager = MultiAgentManager()
|
|
old_workspace = _ReloadWorkspace("agent-1")
|
|
task = asyncio.Future()
|
|
old_workspace.task_tracker.wait_tasks_done = AsyncMock(
|
|
return_value=False,
|
|
)
|
|
monkeypatch.setattr(
|
|
multi_agent_manager_module,
|
|
"_OLD_WORKSPACE_TASK_MAX_WAIT_ROUNDS",
|
|
2,
|
|
)
|
|
|
|
await manager._graceful_stop_old_instance(
|
|
old_workspace,
|
|
"agent-1",
|
|
active_tasks={"chat-1": task},
|
|
)
|
|
cleanup_tasks = list(manager._cleanup_tasks)
|
|
assert cleanup_tasks
|
|
await asyncio.wait_for(
|
|
asyncio.gather(*cleanup_tasks),
|
|
timeout=1,
|
|
)
|
|
|
|
assert old_workspace.task_tracker.wait_tasks_done.await_count == 2
|
|
assert old_workspace.stopped is True
|
|
|
|
|
|
def _read_custom_startup_concurrency(
|
|
value: str | None = None,
|
|
legacy_value: str | None = None,
|
|
) -> int:
|
|
"""Read the import-time setting in an isolated interpreter."""
|
|
env = os.environ.copy()
|
|
env.pop(constants.CUSTOM_AGENT_STARTUP_CONCURRENCY_ENV, None)
|
|
legacy_env = "COPAW_CUSTOM_AGENT_STARTUP_CONCURRENCY"
|
|
env.pop(legacy_env, None)
|
|
if value is not None:
|
|
env[constants.CUSTOM_AGENT_STARTUP_CONCURRENCY_ENV] = value
|
|
if legacy_value is not None:
|
|
env[legacy_env] = legacy_value
|
|
|
|
code = (
|
|
"from pineagents.constant import "
|
|
"CUSTOM_AGENT_STARTUP_CONCURRENCY; "
|
|
"print(CUSTOM_AGENT_STARTUP_CONCURRENCY)"
|
|
)
|
|
completed = subprocess.run(
|
|
[sys.executable, "-c", code],
|
|
env=env,
|
|
check=True,
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=10,
|
|
)
|
|
return int(completed.stdout.strip())
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_disabled_agent_is_not_started_or_mutated(monkeypatch) -> None:
|
|
"""Startup must preserve and skip an explicitly disabled profile."""
|
|
manager = MultiAgentManager()
|
|
config = _config("default", "disabled")
|
|
config.agents.profiles["disabled"].enabled = False
|
|
monkeypatch.setattr(
|
|
"pineagents.app.multi_agent_manager.load_config",
|
|
lambda: config,
|
|
)
|
|
manager.get_agent = AsyncMock(return_value=SimpleNamespace())
|
|
|
|
result = await manager.start_all_configured_agents()
|
|
|
|
assert result == {"default": True}
|
|
manager.get_agent.assert_awaited_once_with("default")
|
|
assert config.agents.profiles["disabled"].enabled is False
|
|
assert (
|
|
manager.get_agent_startup_status(
|
|
"disabled",
|
|
enabled=False,
|
|
)
|
|
== AgentStartupStatus.DISABLED
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("value", "expected"),
|
|
[(None, 5), ("invalid", 5), ("0", 1), ("4", 4)],
|
|
)
|
|
def test_custom_startup_concurrency_parsing(
|
|
value: str | None,
|
|
expected: int,
|
|
) -> None:
|
|
assert _read_custom_startup_concurrency(value=value) == expected
|
|
|
|
|
|
def test_custom_startup_concurrency_supports_legacy_env() -> None:
|
|
"""The legacy COPAW-prefixed environment variable remains supported."""
|
|
assert _read_custom_startup_concurrency(legacy_value="3") == 3
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_core_agents_overlap_before_custom_agents(
|
|
monkeypatch,
|
|
) -> None:
|
|
manager = MultiAgentManager()
|
|
config = _config("default", BUILTIN_QA_AGENT_ID, "custom")
|
|
monkeypatch.setattr(
|
|
"pineagents.app.multi_agent_manager.load_config",
|
|
lambda: config,
|
|
)
|
|
|
|
core_started = set()
|
|
both_core_started = asyncio.Event()
|
|
release_core = asyncio.Event()
|
|
custom_started = asyncio.Event()
|
|
|
|
async def get_agent(agent_id: str):
|
|
if agent_id in {"default", BUILTIN_QA_AGENT_ID}:
|
|
core_started.add(agent_id)
|
|
if len(core_started) == 2:
|
|
both_core_started.set()
|
|
await release_core.wait()
|
|
else:
|
|
custom_started.set()
|
|
return SimpleNamespace()
|
|
|
|
manager.get_agent = AsyncMock(side_effect=get_agent)
|
|
callback = MagicMock()
|
|
task = asyncio.create_task(
|
|
manager.start_all_configured_agents(
|
|
on_core_ready=callback,
|
|
),
|
|
)
|
|
|
|
await asyncio.wait_for(both_core_started.wait(), timeout=1)
|
|
assert not custom_started.is_set()
|
|
release_core.set()
|
|
result = await asyncio.wait_for(task, timeout=1)
|
|
|
|
assert result == {
|
|
"default": True,
|
|
BUILTIN_QA_AGENT_ID: True,
|
|
"custom": True,
|
|
}
|
|
callback.assert_called_once()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_core_ready_waits_for_enabled_qa(monkeypatch) -> None:
|
|
"""Ready is published only after both enabled core agents finish."""
|
|
manager = MultiAgentManager()
|
|
config = _config("default", BUILTIN_QA_AGENT_ID)
|
|
monkeypatch.setattr(
|
|
"pineagents.app.multi_agent_manager.load_config",
|
|
lambda: config,
|
|
)
|
|
default_done = asyncio.Event()
|
|
qa_started = asyncio.Event()
|
|
release_qa = asyncio.Event()
|
|
|
|
async def get_agent(agent_id: str):
|
|
if agent_id == "default":
|
|
default_done.set()
|
|
else:
|
|
qa_started.set()
|
|
await release_qa.wait()
|
|
return SimpleNamespace()
|
|
|
|
manager.get_agent = AsyncMock(side_effect=get_agent)
|
|
callback = MagicMock()
|
|
task = asyncio.create_task(
|
|
manager.start_all_configured_agents(on_core_ready=callback),
|
|
)
|
|
|
|
await asyncio.wait_for(default_done.wait(), timeout=1)
|
|
await asyncio.wait_for(qa_started.wait(), timeout=1)
|
|
callback.assert_not_called()
|
|
|
|
release_qa.set()
|
|
await asyncio.wait_for(task, timeout=1)
|
|
callback.assert_called_once()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_core_ready_does_not_wait_for_disabled_qa(monkeypatch) -> None:
|
|
"""A disabled QA agent is excluded from the core readiness phase."""
|
|
manager = MultiAgentManager()
|
|
config = _config("default", BUILTIN_QA_AGENT_ID)
|
|
config.agents.profiles[BUILTIN_QA_AGENT_ID].enabled = False
|
|
monkeypatch.setattr(
|
|
"pineagents.app.multi_agent_manager.load_config",
|
|
lambda: config,
|
|
)
|
|
manager.get_agent = AsyncMock(return_value=SimpleNamespace())
|
|
callback = MagicMock()
|
|
|
|
result = await manager.start_all_configured_agents(
|
|
on_core_ready=callback,
|
|
)
|
|
|
|
assert result == {"default": True}
|
|
manager.get_agent.assert_awaited_once_with("default")
|
|
callback.assert_called_once_with({"default": True})
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_startup_preserves_loaded_agent_status_during_core_phase(
|
|
monkeypatch,
|
|
) -> None:
|
|
"""A lazily loaded agent remains running while core agents start."""
|
|
manager = MultiAgentManager()
|
|
config = _config("default", "custom")
|
|
monkeypatch.setattr(
|
|
"pineagents.app.multi_agent_manager.load_config",
|
|
lambda: config,
|
|
)
|
|
manager.agents["custom"] = SimpleNamespace()
|
|
core_started = asyncio.Event()
|
|
release_core = asyncio.Event()
|
|
|
|
async def get_agent(agent_id: str):
|
|
if agent_id == "default":
|
|
core_started.set()
|
|
await release_core.wait()
|
|
return manager.agents.get(agent_id, SimpleNamespace())
|
|
|
|
manager.get_agent = AsyncMock(side_effect=get_agent)
|
|
task = asyncio.create_task(manager.start_all_configured_agents())
|
|
|
|
await asyncio.wait_for(core_started.wait(), timeout=1)
|
|
assert manager.get_agent_startup_status("custom") == (
|
|
AgentStartupStatus.RUNNING
|
|
)
|
|
assert not manager.is_agent_startup_in_progress("custom")
|
|
|
|
release_core.set()
|
|
result = await asyncio.wait_for(task, timeout=1)
|
|
assert result == {"default": True, "custom": True}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_custom_agent_startup_respects_concurrency(
|
|
monkeypatch,
|
|
) -> None:
|
|
custom_ids = [f"custom-{index}" for index in range(6)]
|
|
config = _config("default", BUILTIN_QA_AGENT_ID, *custom_ids)
|
|
monkeypatch.setattr(
|
|
"pineagents.app.multi_agent_manager.load_config",
|
|
lambda: config,
|
|
)
|
|
monkeypatch.setenv(
|
|
"QWENPAW_CUSTOM_AGENT_STARTUP_CONCURRENCY",
|
|
"2",
|
|
)
|
|
monkeypatch.setattr(
|
|
multi_agent_manager_module,
|
|
"CUSTOM_AGENT_STARTUP_CONCURRENCY",
|
|
2,
|
|
)
|
|
manager = MultiAgentManager()
|
|
|
|
active_custom = 0
|
|
peak_custom = 0
|
|
|
|
async def get_agent(agent_id: str):
|
|
nonlocal active_custom, peak_custom
|
|
if agent_id in custom_ids:
|
|
active_custom += 1
|
|
peak_custom = max(peak_custom, active_custom)
|
|
await asyncio.sleep(0.01)
|
|
active_custom -= 1
|
|
return SimpleNamespace()
|
|
|
|
manager.get_agent = AsyncMock(side_effect=get_agent)
|
|
startup_display = MagicMock()
|
|
result = await manager.start_all_configured_agents(
|
|
startup_display=startup_display,
|
|
)
|
|
|
|
assert all(result.values())
|
|
assert peak_custom == 2
|
|
startup_display.start_custom_agents.assert_called_once_with(6)
|
|
assert startup_display.advance.call_count == 6
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_runtime_startups_share_concurrency_and_pending_state(
|
|
monkeypatch,
|
|
) -> None:
|
|
"""Runtime-created agents use the same bounded startup scheduler."""
|
|
monkeypatch.setattr(
|
|
multi_agent_manager_module,
|
|
"CUSTOM_AGENT_STARTUP_CONCURRENCY",
|
|
1,
|
|
)
|
|
manager = MultiAgentManager()
|
|
config = _config("alpha", "beta")
|
|
monkeypatch.setattr(
|
|
"pineagents.app.multi_agent_manager.load_config",
|
|
lambda: config,
|
|
)
|
|
alpha_started = asyncio.Event()
|
|
release_alpha = asyncio.Event()
|
|
beta_started = asyncio.Event()
|
|
|
|
async def get_agent(agent_id: str):
|
|
if agent_id == "alpha":
|
|
alpha_started.set()
|
|
await release_alpha.wait()
|
|
else:
|
|
beta_started.set()
|
|
return SimpleNamespace()
|
|
|
|
manager.get_agent = AsyncMock(side_effect=get_agent)
|
|
|
|
alpha_task = manager.schedule_agent_startup("alpha")
|
|
beta_task = manager.schedule_agent_startup("beta")
|
|
await asyncio.wait_for(alpha_started.wait(), timeout=1)
|
|
|
|
assert manager.get_agent_startup_status("beta") == (
|
|
AgentStartupStatus.PENDING
|
|
)
|
|
assert manager.is_agent_startup_in_progress("beta")
|
|
assert not beta_started.is_set()
|
|
|
|
release_alpha.set()
|
|
await asyncio.wait_for(beta_started.wait(), timeout=1)
|
|
assert await asyncio.gather(alpha_task, beta_task) == [True, True]
|
|
await asyncio.sleep(0)
|
|
assert not manager._agent_startup_tasks
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_startup_display_skips_empty_custom_phase(monkeypatch) -> None:
|
|
manager = MultiAgentManager()
|
|
config = _config("default", BUILTIN_QA_AGENT_ID)
|
|
monkeypatch.setattr(
|
|
"pineagents.app.multi_agent_manager.load_config",
|
|
lambda: config,
|
|
)
|
|
manager.get_agent = AsyncMock(return_value=SimpleNamespace())
|
|
startup_display = MagicMock()
|
|
|
|
result = await manager.start_all_configured_agents(
|
|
startup_display=startup_display,
|
|
)
|
|
|
|
assert all(result.values())
|
|
startup_display.start_custom_agents.assert_not_called()
|
|
startup_display.advance.assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_default_failure_skips_custom_agent_phase(monkeypatch) -> None:
|
|
"""Custom agents must not start when the Default core agent fails."""
|
|
manager = MultiAgentManager()
|
|
config = _config("default", "custom")
|
|
monkeypatch.setattr(
|
|
"pineagents.app.multi_agent_manager.load_config",
|
|
lambda: config,
|
|
)
|
|
|
|
async def get_agent(agent_id: str):
|
|
if agent_id == "default":
|
|
raise RuntimeError("invalid default config")
|
|
return SimpleNamespace()
|
|
|
|
manager.get_agent = AsyncMock(side_effect=get_agent)
|
|
startup_display = MagicMock()
|
|
|
|
result = await manager.start_all_configured_agents(
|
|
startup_display=startup_display,
|
|
)
|
|
|
|
assert result == {"default": False, "custom": False}
|
|
manager.get_agent.assert_awaited_once_with("default")
|
|
startup_display.start_custom_agents.assert_not_called()
|
|
startup_display.advance.assert_not_called()
|
|
|
|
|
|
class _WorkspaceStub:
|
|
def __init__(self, start_event: asyncio.Event, release: asyncio.Event):
|
|
self._start_event = start_event
|
|
self._release = release
|
|
|
|
async def start(self) -> None:
|
|
self._start_event.set()
|
|
await self._release.wait()
|
|
|
|
def set_manager(self, _manager) -> None:
|
|
return None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_agent_updates_runtime_status(monkeypatch) -> None:
|
|
manager = MultiAgentManager()
|
|
config = _config("custom")
|
|
monkeypatch.setattr(
|
|
"pineagents.app.multi_agent_manager.load_config",
|
|
lambda: config,
|
|
)
|
|
started = asyncio.Event()
|
|
release = asyncio.Event()
|
|
workspace = _WorkspaceStub(started, release)
|
|
monkeypatch.setattr(
|
|
manager,
|
|
"_create_workspace",
|
|
lambda **_kwargs: workspace,
|
|
)
|
|
|
|
task = asyncio.create_task(manager.get_agent("custom"))
|
|
await asyncio.wait_for(started.wait(), timeout=1)
|
|
assert manager.get_agent_startup_status("custom") == (
|
|
AgentStartupStatus.STARTING
|
|
)
|
|
|
|
release.set()
|
|
assert await asyncio.wait_for(task, timeout=1) is workspace
|
|
assert manager.get_agent_startup_status("custom") == (
|
|
AgentStartupStatus.RUNNING
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cancelled_start_cleans_pending_state(monkeypatch) -> None:
|
|
manager = MultiAgentManager()
|
|
config = _config("custom")
|
|
monkeypatch.setattr(
|
|
"pineagents.app.multi_agent_manager.load_config",
|
|
lambda: config,
|
|
)
|
|
started = asyncio.Event()
|
|
never_release = asyncio.Event()
|
|
workspace = _WorkspaceStub(started, never_release)
|
|
monkeypatch.setattr(
|
|
manager,
|
|
"_create_workspace",
|
|
lambda **_kwargs: workspace,
|
|
)
|
|
|
|
task = asyncio.create_task(manager.get_agent("custom"))
|
|
await asyncio.wait_for(started.wait(), timeout=1)
|
|
task.cancel()
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await task
|
|
|
|
assert "custom" not in manager._pending_starts
|
|
assert manager.get_agent_startup_status("custom") == (
|
|
AgentStartupStatus.FAILED
|
|
)
|