Adam Moussa
b6c66877f5
Security follow-ups from the per-PR review (non-blocking MEDIUMs): - Eager _get_token() at make_app build time so a missing AGENT_TEAM_API_TOKEN fails fast instead of serving requests first (matches the docstring contract). - Disable /docs, /redoc, /openapi.json (no auth dependency in FastAPI) — the API is VPN-only/127.0.0.1 and should not expose its schema unauthenticated. - Scrub raw exception text and subprocess stderr from 500 response bodies; log server-side instead (avoid internal-path/state disclosure). - Bound /orchestrator/invoke concurrency with a semaphore (429 over the cap) so an authenticated caller cannot exhaust the box via many 600s subprocesses. Also pin fastapi/uvicorn in requirements.txt (WS1 dep). With fastapi now installed in CI, the previously skip-guarded TestClient tests run for real; the importorskip guard stays as a no-op safety net. Tests: 23 pass (adds docs-disabled + concurrency-429 cases).
433 lines
15 KiB
Python
433 lines
15 KiB
Python
"""WS1 tests: invoker_multi dispatch + api.py HTTP API seam.
|
|
|
|
Covers:
|
|
* invoker_multi.multi_invoke dispatches to the correct model factory (mocked)
|
|
* invoker_multi.make_fast_coder_invoker / make_scanner_invoker return callables
|
|
* invoker_multi.bind_multi_invoker wires the review loop seam
|
|
* api.py: 401 without bearer token; 200 with valid token (endpoints stubbed)
|
|
* api.py: make_app returns a FastAPI app; GET /tasks 404 when no state
|
|
* review_loop_llm default_plan_reviewer is now make_cross_reviewer_invoker (in-process)
|
|
* builders_llm subprocess_build is kept as opt-in; make_fast_coder_invoker callable
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import importlib.util
|
|
import sys
|
|
from pathlib import Path
|
|
from typing import Any
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import pytest
|
|
|
|
_AGENT_TEAM_DIR = Path(__file__).resolve().parents[1]
|
|
if str(_AGENT_TEAM_DIR) not in sys.path:
|
|
sys.path.insert(0, str(_AGENT_TEAM_DIR))
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# invoker_multi: basic dispatch via mocked models
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _make_fake_model(return_text: str = "model output") -> MagicMock:
|
|
"""Return a mock model whose .invoke() yields a fake response."""
|
|
mock = MagicMock()
|
|
result = MagicMock()
|
|
result.content = return_text
|
|
mock.invoke.return_value = result
|
|
return mock
|
|
|
|
|
|
def test_multi_invoke_cross_reviewer(monkeypatch: Any) -> None:
|
|
from agent_team import invoker_multi
|
|
|
|
fake_model = _make_fake_model("APPROVE verdict")
|
|
fake_models = MagicMock()
|
|
fake_models.get_cross_reviewer.return_value = fake_model
|
|
fake_models.get_fast_coder.return_value = MagicMock()
|
|
fake_models.get_scanner.return_value = MagicMock()
|
|
|
|
with patch.dict(sys.modules, {"models": fake_models}):
|
|
result = invoker_multi.multi_invoke("review this plan", model="cross_reviewer")
|
|
|
|
assert result == "APPROVE verdict"
|
|
fake_model.invoke.assert_called_once_with("review this plan")
|
|
|
|
|
|
def test_multi_invoke_fast_coder(monkeypatch: Any) -> None:
|
|
from agent_team import invoker_multi
|
|
|
|
fake_model = _make_fake_model("--- a/file.py\n+++ b/file.py\n@@")
|
|
fake_models = MagicMock()
|
|
fake_models.get_fast_coder.return_value = fake_model
|
|
fake_models.get_cross_reviewer.return_value = MagicMock()
|
|
fake_models.get_scanner.return_value = MagicMock()
|
|
|
|
with patch.dict(sys.modules, {"models": fake_models}):
|
|
result = invoker_multi.multi_invoke("add a login form", model="fast_coder")
|
|
|
|
assert "@@" in result
|
|
|
|
|
|
def test_multi_invoke_scanner(monkeypatch: Any) -> None:
|
|
from agent_team import invoker_multi
|
|
|
|
fake_model = _make_fake_model("no issues found")
|
|
fake_models = MagicMock()
|
|
fake_models.get_scanner.return_value = fake_model
|
|
fake_models.get_cross_reviewer.return_value = MagicMock()
|
|
fake_models.get_fast_coder.return_value = MagicMock()
|
|
|
|
with patch.dict(sys.modules, {"models": fake_models}):
|
|
result = invoker_multi.multi_invoke("scan the diff", model="scanner")
|
|
|
|
assert result == "no issues found"
|
|
|
|
|
|
def test_multi_invoke_unknown_model() -> None:
|
|
from agent_team import invoker_multi
|
|
|
|
fake_models = MagicMock()
|
|
with patch.dict(sys.modules, {"models": fake_models}):
|
|
with pytest.raises(ValueError, match="unknown model key"):
|
|
invoker_multi.multi_invoke("hi", model="gpt_4_turbo")
|
|
|
|
|
|
def test_multi_invoke_models_unavailable() -> None:
|
|
from agent_team import invoker_multi
|
|
|
|
saved = sys.modules.pop("models", None)
|
|
try:
|
|
# Ensure "models" is NOT in sys.modules so the import inside multi_invoke fails.
|
|
# We also need to make sure it's not importable from the path.
|
|
with pytest.raises(RuntimeError, match="Cannot import orchestrator"):
|
|
# patch.dict with a None value keeps it out of sys.modules
|
|
with patch.dict(sys.modules, {"models": None}): # type: ignore[dict-item]
|
|
invoker_multi.multi_invoke("hi", model="cross_reviewer")
|
|
finally:
|
|
if saved is not None:
|
|
sys.modules["models"] = saved
|
|
|
|
|
|
def test_make_fast_coder_invoker_returns_callable() -> None:
|
|
from agent_team.invoker_multi import make_fast_coder_invoker
|
|
|
|
invoker = make_fast_coder_invoker()
|
|
assert callable(invoker)
|
|
|
|
|
|
def test_make_scanner_invoker_returns_callable() -> None:
|
|
from agent_team.invoker_multi import make_scanner_invoker
|
|
|
|
invoker = make_scanner_invoker()
|
|
assert callable(invoker)
|
|
|
|
|
|
def test_make_fast_coder_invoker_calls_multi_invoke() -> None:
|
|
from agent_team import invoker_multi
|
|
|
|
fake_model = _make_fake_model("the diff")
|
|
fake_models = MagicMock()
|
|
fake_models.get_fast_coder.return_value = fake_model
|
|
fake_models.get_cross_reviewer.return_value = MagicMock()
|
|
fake_models.get_scanner.return_value = MagicMock()
|
|
|
|
invoker = invoker_multi.make_fast_coder_invoker()
|
|
with patch.dict(sys.modules, {"models": fake_models}):
|
|
result = invoker("build instruction")
|
|
|
|
assert result == "the diff"
|
|
|
|
|
|
def test_bind_multi_invoker_sets_review_seam() -> None:
|
|
from agent_team import invoker_multi
|
|
from agent_team.nodes import review_loop
|
|
|
|
original = review_loop._review_invoker # save
|
|
try:
|
|
fake_model = _make_fake_model("APPROVE")
|
|
fake_models = MagicMock()
|
|
fake_models.get_cross_reviewer.return_value = fake_model
|
|
fake_models.get_fast_coder.return_value = MagicMock()
|
|
fake_models.get_scanner.return_value = MagicMock()
|
|
|
|
with patch.dict(sys.modules, {"models": fake_models}):
|
|
invoker_multi.bind_multi_invoker()
|
|
|
|
# The review loop seam should now be the cross-reviewer invoker.
|
|
assert review_loop._review_invoker is not original
|
|
finally:
|
|
review_loop._review_invoker = original # restore
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# review_loop_llm: default is now in-process, subprocess kept as opt-in
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_review_loop_llm_default_is_cross_reviewer_invoker() -> None:
|
|
from agent_team.nodes import review_loop_llm
|
|
|
|
# The module-level default should be the closure from make_cross_reviewer_invoker,
|
|
# NOT the subprocess-based make_run_py_invoker closure. We test by name.
|
|
# Both are closures, so we check the closure cell names via __code__.
|
|
fn = review_loop_llm.default_plan_reviewer
|
|
assert callable(fn)
|
|
# make_cross_reviewer_invoker creates a _invoke that uses get_cross_reviewer
|
|
# internally; make_run_py_invoker creates one that uses subprocess.run.
|
|
# We verify by checking that the default raises RuntimeError (models absent)
|
|
# NOT FileNotFoundError (run.py absent), confirming it's the in-process path.
|
|
saved = sys.modules.pop("models", None)
|
|
try:
|
|
with patch.dict(sys.modules, {"models": None}): # type: ignore[dict-item]
|
|
with pytest.raises(
|
|
(RuntimeError, AttributeError, TypeError, ModuleNotFoundError)
|
|
):
|
|
fn("any prompt")
|
|
finally:
|
|
if saved is not None:
|
|
sys.modules["models"] = saved
|
|
|
|
|
|
def test_review_loop_llm_make_run_py_invoker_still_available() -> None:
|
|
from agent_team.nodes.review_loop_llm import make_run_py_invoker
|
|
|
|
fn = make_run_py_invoker()
|
|
assert callable(fn)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# builders_llm: subprocess_build opt-in; make_fast_coder_invoker available
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_builders_llm_subprocess_build_available() -> None:
|
|
from agent_team.nodes.builders_llm import subprocess_build
|
|
|
|
assert callable(subprocess_build)
|
|
|
|
|
|
def test_builders_llm_make_fast_coder_invoker_available() -> None:
|
|
from agent_team.nodes.builders_llm import make_fast_coder_invoker
|
|
|
|
fn = make_fast_coder_invoker()
|
|
assert callable(fn)
|
|
|
|
|
|
def test_builders_llm_default_build_uses_in_process() -> None:
|
|
from agent_team.nodes import builders_llm
|
|
|
|
fake_model = _make_fake_model("--- a/x.py\n+++ b/x.py\n@@ -1 +1 @@\n-old\n+new")
|
|
fake_models = MagicMock()
|
|
fake_models.get_fast_coder.return_value = fake_model
|
|
fake_models.get_cross_reviewer.return_value = MagicMock()
|
|
fake_models.get_scanner.return_value = MagicMock()
|
|
|
|
with patch.dict(sys.modules, {"models": fake_models}):
|
|
result = builders_llm.default_build("do the thing")
|
|
|
|
assert "@@" in result
|
|
fake_model.invoke.assert_called_once()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# api.py: bearer-token auth + endpoint smoke tests via FastAPI TestClient
|
|
# ---------------------------------------------------------------------------
|
|
|
|
# fastapi is a box-only dependency (installed on the R720 during attended
|
|
# deploy, never in CI or the root requirements). api.py imports it lazily, so
|
|
# these TestClient-based smoke tests are the only thing that hard-requires it.
|
|
# Skip them when it is absent rather than failing collection in CI.
|
|
requires_fastapi = pytest.mark.skipif(
|
|
importlib.util.find_spec("fastapi") is None,
|
|
reason="fastapi not installed (box-only dependency)",
|
|
)
|
|
|
|
|
|
@pytest.fixture()
|
|
def api_token(monkeypatch: Any) -> str:
|
|
token = "test-secret-token-for-ws1"
|
|
monkeypatch.setenv("AGENT_TEAM_API_TOKEN", token)
|
|
return token
|
|
|
|
|
|
def _make_stub_coordinator(thread_id: str = "t-test-123") -> MagicMock:
|
|
"""Return a coordinator stub that start_task returns a fixed thread_id."""
|
|
coord = MagicMock()
|
|
coord.start_task.return_value = thread_id
|
|
coord.graph = None # no graph for simple smoke tests
|
|
coord.setup = MagicMock()
|
|
return coord
|
|
|
|
|
|
@requires_fastapi
|
|
def test_api_401_without_token(api_token: str) -> None:
|
|
from fastapi.testclient import TestClient
|
|
|
|
from agent_team.api import make_app
|
|
|
|
app = make_app(coordinator=_make_stub_coordinator())
|
|
client = TestClient(app, raise_server_exceptions=False)
|
|
resp = client.post("/tasks", json={"task": "add login"})
|
|
assert resp.status_code == 401
|
|
|
|
|
|
@requires_fastapi
|
|
def test_api_401_wrong_token(api_token: str) -> None:
|
|
from fastapi.testclient import TestClient
|
|
|
|
from agent_team.api import make_app
|
|
|
|
app = make_app(coordinator=_make_stub_coordinator())
|
|
client = TestClient(app, raise_server_exceptions=False)
|
|
resp = client.post(
|
|
"/tasks",
|
|
json={"task": "add login"},
|
|
headers={"Authorization": "Bearer wrong-token"},
|
|
)
|
|
assert resp.status_code == 401
|
|
|
|
|
|
@requires_fastapi
|
|
def test_api_post_tasks_200(api_token: str) -> None:
|
|
from fastapi.testclient import TestClient
|
|
|
|
from agent_team.api import make_app
|
|
|
|
stub = _make_stub_coordinator("thread-abc")
|
|
app = make_app(coordinator=stub)
|
|
client = TestClient(app)
|
|
resp = client.post(
|
|
"/tasks",
|
|
json={"task": "add a login form"},
|
|
headers={"Authorization": f"Bearer {api_token}"},
|
|
)
|
|
assert resp.status_code == 200, resp.text
|
|
assert resp.json()["thread_id"] == "thread-abc"
|
|
stub.start_task.assert_called_once()
|
|
|
|
|
|
@requires_fastapi
|
|
def test_api_get_task_404_no_state(api_token: str) -> None:
|
|
from fastapi.testclient import TestClient
|
|
|
|
from agent_team.api import make_app
|
|
|
|
stub = _make_stub_coordinator()
|
|
app = make_app(coordinator=stub)
|
|
client = TestClient(app, raise_server_exceptions=False)
|
|
# graph is None -> 503 not ready
|
|
resp = client.get(
|
|
"/tasks/t-missing",
|
|
headers={"Authorization": f"Bearer {api_token}"},
|
|
)
|
|
assert resp.status_code == 503
|
|
|
|
|
|
@requires_fastapi
|
|
def test_api_get_task_with_graph_state(api_token: str) -> None:
|
|
from fastapi.testclient import TestClient
|
|
|
|
from agent_team.api import make_app
|
|
|
|
stub = _make_stub_coordinator()
|
|
# Attach a mock graph that returns a snapshot with values
|
|
snapshot = MagicMock()
|
|
snapshot.values = {"task": "add login", "current_phase": "CLARIFY"}
|
|
stub.graph = MagicMock()
|
|
stub.graph.get_state.return_value = snapshot
|
|
app = make_app(coordinator=stub)
|
|
client = TestClient(app)
|
|
resp = client.get(
|
|
"/tasks/t-abc",
|
|
headers={"Authorization": f"Bearer {api_token}"},
|
|
)
|
|
assert resp.status_code == 200
|
|
data = resp.json()
|
|
assert data["thread_id"] == "t-abc"
|
|
assert data["state"]["task"] == "add login"
|
|
|
|
|
|
@requires_fastapi
|
|
def test_api_orchestrator_invoke_calls_subprocess(
|
|
api_token: str, tmp_path: Any
|
|
) -> None:
|
|
from fastapi.testclient import TestClient
|
|
|
|
from agent_team.api import make_app
|
|
|
|
# Write a fake run.py that echos the first argv
|
|
fake_run = tmp_path / "run.py"
|
|
fake_run.write_text("import sys; print(sys.argv[1])")
|
|
|
|
stub = _make_stub_coordinator()
|
|
app = make_app(coordinator=stub)
|
|
|
|
# Patch _resolve_run_py to return our fake script
|
|
with patch("agent_team.api._resolve_run_py", return_value=fake_run):
|
|
client = TestClient(app)
|
|
resp = client.post(
|
|
"/orchestrator/invoke",
|
|
json={"prompt": "hello world"},
|
|
headers={"Authorization": f"Bearer {api_token}"},
|
|
)
|
|
|
|
assert resp.status_code == 200
|
|
assert "hello world" in resp.json()["text"]
|
|
|
|
|
|
@requires_fastapi
|
|
def test_api_no_token_env_raises_at_build(monkeypatch: Any) -> None:
|
|
from agent_team.api import make_app
|
|
|
|
monkeypatch.delenv("AGENT_TEAM_API_TOKEN", raising=False)
|
|
# Fail-fast: an unconfigured token must raise at app-build time, not serve a
|
|
# request first (the eager _get_token() in make_app).
|
|
with pytest.raises(RuntimeError, match="AGENT_TEAM_API_TOKEN"):
|
|
make_app(coordinator=_make_stub_coordinator())
|
|
|
|
|
|
@requires_fastapi
|
|
def test_api_docs_and_openapi_disabled(api_token: str) -> None:
|
|
from fastapi.testclient import TestClient
|
|
|
|
from agent_team.api import make_app
|
|
|
|
app = make_app(coordinator=_make_stub_coordinator())
|
|
client = TestClient(app, raise_server_exceptions=False)
|
|
# The unauthenticated docs/schema routes must be disabled (VPN-only API).
|
|
for path in ("/docs", "/redoc", "/openapi.json"):
|
|
assert client.get(path).status_code == 404, path
|
|
|
|
|
|
@requires_fastapi
|
|
def test_api_invoke_concurrency_cap_returns_429(api_token: str, tmp_path: Any) -> None:
|
|
import agent_team.api as api_mod
|
|
from fastapi.testclient import TestClient
|
|
|
|
from agent_team.api import make_app
|
|
|
|
fake_run = tmp_path / "run.py"
|
|
fake_run.write_text("import sys; print(sys.argv[1])")
|
|
app = make_app(coordinator=_make_stub_coordinator())
|
|
|
|
# Exhaust the bounded semaphore so the request sees no free slot → 429.
|
|
acquired = [
|
|
api_mod._invoke_semaphore.acquire(blocking=False)
|
|
for _ in range(api_mod._MAX_CONCURRENT_INVOKES)
|
|
]
|
|
try:
|
|
assert all(acquired)
|
|
with patch("agent_team.api._resolve_run_py", return_value=fake_run):
|
|
client = TestClient(app, raise_server_exceptions=False)
|
|
resp = client.post(
|
|
"/orchestrator/invoke",
|
|
json={"prompt": "hello"},
|
|
headers={"Authorization": f"Bearer {api_token}"},
|
|
)
|
|
assert resp.status_code == 429
|
|
finally:
|
|
for ok in acquired:
|
|
if ok:
|
|
api_mod._invoke_semaphore.release()
|