mirror of https://github.com/aliasrobotics/cai.git
142 lines
4.2 KiB
Python
142 lines
4.2 KiB
Python
"""Tests for orchestration MAS heuristic nudge."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from types import SimpleNamespace
|
|
import pytest
|
|
|
|
from cai.config import reset_config
|
|
from cai.sdk.agents import orchestration_mas_hint as omh
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _reset_cfg():
|
|
reset_config()
|
|
yield
|
|
reset_config()
|
|
|
|
|
|
def test_user_message_suggests_multi_front_parallel_word() -> None:
|
|
text = (
|
|
"Please scan the web app and the API in parallel for obvious issues; "
|
|
"then summarize risks."
|
|
)
|
|
assert omh.user_message_suggests_multi_front(text) is True
|
|
|
|
|
|
def test_user_message_suggests_multi_front_numbered_list() -> None:
|
|
text = "Do the following:\n1. Enumerate hosts\n2. Check TLS\n3. Draft report"
|
|
assert omh.user_message_suggests_multi_front(text) is True
|
|
|
|
|
|
def test_user_message_suggests_multi_front_short_false() -> None:
|
|
assert omh.user_message_suggests_multi_front("run nmap on 10.0.0.1") is False
|
|
|
|
|
|
@pytest.mark.parametrize("enabled", [True, False])
|
|
def test_maybe_inject_respects_config(monkeypatch: pytest.MonkeyPatch, enabled: bool) -> None:
|
|
monkeypatch.setenv("CAI_ORCHESTRATION_MAS_HINT", "true" if enabled else "false")
|
|
reset_config()
|
|
|
|
calls: list[dict] = []
|
|
|
|
def add(msg: dict) -> None:
|
|
calls.append(msg)
|
|
|
|
agent = SimpleNamespace(
|
|
model=SimpleNamespace(
|
|
agent_type="orchestration_agent",
|
|
add_to_message_history=add,
|
|
)
|
|
)
|
|
tool = SimpleNamespace(name="run_specialist")
|
|
fr = SimpleNamespace(tool=tool)
|
|
rc = SimpleNamespace(trace_metadata={})
|
|
|
|
omh.maybe_inject_orchestration_mas_hint_after_tools(
|
|
agent=agent,
|
|
original_input=(
|
|
"Investigate in parallel: DNS records, HTTP headers, and TLS chain for example.com. "
|
|
"Be thorough."
|
|
),
|
|
function_results=[fr],
|
|
run_config=rc,
|
|
)
|
|
|
|
if enabled:
|
|
assert len(calls) == 1
|
|
assert calls[0]["role"] == "user"
|
|
assert "run_parallel_specialists" in calls[0]["content"]
|
|
assert rc.trace_metadata.get("_cai_mas_hint_injected") is True
|
|
else:
|
|
assert calls == []
|
|
|
|
|
|
def test_maybe_inject_skips_non_orchestration_agent(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
monkeypatch.setenv("CAI_ORCHESTRATION_MAS_HINT", "true")
|
|
reset_config()
|
|
|
|
calls: list[dict] = []
|
|
|
|
agent = SimpleNamespace(
|
|
model=SimpleNamespace(
|
|
agent_type="redteam_agent",
|
|
add_to_message_history=lambda m: calls.append(m),
|
|
)
|
|
)
|
|
tool = SimpleNamespace(name="run_specialist")
|
|
fr = SimpleNamespace(tool=tool)
|
|
rc = SimpleNamespace(trace_metadata={})
|
|
|
|
omh.maybe_inject_orchestration_mas_hint_after_tools(
|
|
agent=agent,
|
|
original_input="Do A and B in parallel for the target host.",
|
|
function_results=[fr],
|
|
run_config=rc,
|
|
)
|
|
assert calls == []
|
|
|
|
|
|
def test_maybe_inject_once_per_run_config() -> None:
|
|
calls: list[dict] = []
|
|
|
|
agent = SimpleNamespace(
|
|
model=SimpleNamespace(
|
|
agent_type="orchestration_agent",
|
|
add_to_message_history=lambda m: calls.append(m),
|
|
)
|
|
)
|
|
tool = SimpleNamespace(name="run_specialist")
|
|
fr = SimpleNamespace(tool=tool)
|
|
rc = SimpleNamespace(trace_metadata={"_cai_mas_hint_injected": True})
|
|
|
|
omh.maybe_inject_orchestration_mas_hint_after_tools(
|
|
agent=agent,
|
|
original_input="Run three checks in parallel on the subnet.",
|
|
function_results=[fr],
|
|
run_config=rc,
|
|
)
|
|
assert calls == []
|
|
|
|
|
|
def test_maybe_inject_when_parallel_already_used() -> None:
|
|
calls: list[dict] = []
|
|
|
|
agent = SimpleNamespace(
|
|
model=SimpleNamespace(
|
|
agent_type="orchestration_agent",
|
|
add_to_message_history=lambda m: calls.append(m),
|
|
)
|
|
)
|
|
rs = SimpleNamespace(name="run_specialist")
|
|
rp = SimpleNamespace(name="run_parallel_specialists")
|
|
rc = SimpleNamespace(trace_metadata={})
|
|
|
|
omh.maybe_inject_orchestration_mas_hint_after_tools(
|
|
agent=agent,
|
|
original_input="Do A and B in parallel for the target host.",
|
|
function_results=[SimpleNamespace(tool=rs), SimpleNamespace(tool=rp)],
|
|
run_config=rc,
|
|
)
|
|
assert calls == []
|