diff --git a/tests/unit/test_egress_addon_request_flow.py b/tests/unit/test_egress_addon_request_flow.py index 3de75b9..6896e9d 100644 --- a/tests/unit/test_egress_addon_request_flow.py +++ b/tests/unit/test_egress_addon_request_flow.py @@ -277,16 +277,18 @@ class _SuperviseRpcFake: supervise_status: "str | None" = None propose_error: bool = False + poll_error: bool = False @property - def propose_calls(self) -> list: + def propose_calls(self) -> list[dict[str, str]]: if not hasattr(self, "_propose_calls"): - self._propose_calls: list = [] + self._propose_calls: list[dict[str, str]] = [] return self._propose_calls def propose_supervise( - self, source_ip, identity_token, *, tool, proposed_file, justification, - ): + self, source_ip: str, identity_token: str, *, + tool: str, proposed_file: str, justification: str, + ) -> str: del proposed_file, justification self.propose_calls.append( {"source_ip": source_ip, "identity_token": identity_token, "tool": tool} @@ -295,8 +297,12 @@ class _SuperviseRpcFake: raise PolicyResolveError("orchestrator down") return "prop-1" - def poll_supervise(self, source_ip, identity_token, proposal_id): + def poll_supervise( + self, source_ip: str, identity_token: str, proposal_id: str, + ) -> dict[str, object]: del source_ip, identity_token, proposal_id + if self.poll_error: + raise PolicyResolveError("orchestrator down") if self.supervise_status is None: return {"status": "pending"} return {"status": self.supervise_status, "notes": "", "final_file": None} @@ -583,6 +589,16 @@ class TestSuperviseBranch(unittest.TestCase): self.assertEqual(403, flow.response.status_code) self.assertIn("timed out", flow.response.get_text()) + def test_poll_error_during_wait_times_out_and_blocks(self) -> None: + # A transient orchestrator error on each poll is retried until the + # deadline, then fails closed (blocked) — never forwarded unsupervised. + addon = self._supervised_addon("approved") + cast(Any, addon._resolver).poll_error = True + flow = _Flow(_Request(host="api.example.com", method="POST", body=f"k={_OPENAI_KEY}")) + _run_request(addon, flow) + assert flow.response is not None + self.assertEqual(403, flow.response.status_code) + # --------------------------------------------------------------------------- # Inbound DLP on responses diff --git a/tests/unit/test_orchestrator_control_plane.py b/tests/unit/test_orchestrator_control_plane.py index fe445b2..c30222f 100644 --- a/tests/unit/test_orchestrator_control_plane.py +++ b/tests/unit/test_orchestrator_control_plane.py @@ -21,7 +21,7 @@ from unittest.mock import patch from bot_bottle.orchestrator.broker import StubBroker from bot_bottle.orchestrator.control_plane import dispatch, make_server -from bot_bottle.orchestrator.registry import RegistryStore +from bot_bottle.orchestrator.registry import BottleRecord, RegistryStore from bot_bottle.orchestrator.service import Orchestrator from bot_bottle.store_manager import StoreManager from bot_bottle.supervise import ( @@ -455,13 +455,13 @@ class TestDispatchSuperviseAgentRpc(unittest.TestCase): return self.store.register( source_ip, metadata=json.dumps({"slug": slug}), policy="routes: []\n") - def _propose(self, rec, proposed: str = "routes:\n - host: g.com\n"): + def _propose(self, rec: BottleRecord, proposed: str = "routes:\n - host: g.com\n"): return dispatch(self.orch, "POST", "/supervise/propose", _body({ "source_ip": rec.source_ip, "identity_token": rec.identity_token, "tool": TOOL_EGRESS_ALLOW, "proposed_file": proposed, "justification": "need it", })) - def _poll(self, rec, proposal_id: str): + def _poll(self, rec: BottleRecord, proposal_id: str): return dispatch(self.orch, "POST", "/supervise/poll", _body({ "source_ip": rec.source_ip, "identity_token": rec.identity_token, "proposal_id": proposal_id, @@ -474,9 +474,11 @@ class TestDispatchSuperviseAgentRpc(unittest.TestCase): pid = payload["proposal_id"] assert isinstance(pid, str) and pid _, listing = dispatch(self.orch, "GET", "/supervise/proposals", b"") - self.assertEqual(pid, listing["proposals"][0]["id"]) + proposals = listing["proposals"] + assert isinstance(proposals, list) + self.assertEqual(pid, proposals[0]["id"]) # Queued under the orchestrator-resolved bottle id, never a caller slug. - self.assertEqual(rec.bottle_id, listing["proposals"][0]["bottle_slug"]) + self.assertEqual(rec.bottle_id, proposals[0]["bottle_slug"]) def test_propose_unattributed_is_403(self) -> None: status, payload = dispatch(self.orch, "POST", "/supervise/propose", _body({ @@ -531,6 +533,41 @@ class TestDispatchSuperviseAgentRpc(unittest.TestCase): self.assertEqual(200, status) self.assertEqual("unknown", poll["status"]) + # --- request validation (400) + fail-closed (403) ---------------------- + + def test_propose_invalid_json_is_400(self) -> None: + status, _ = dispatch(self.orch, "POST", "/supervise/propose", b"{not json") + self.assertEqual(400, status) + + def test_propose_missing_fields_are_400(self) -> None: + rec = self._register() + base = { + "source_ip": rec.source_ip, "identity_token": rec.identity_token, + "tool": TOOL_EGRESS_ALLOW, "proposed_file": "routes:\n", "justification": "j", + } + for drop in ("source_ip", "proposed_file", "justification"): + body = {k: v for k, v in base.items() if k != drop} + status, _ = dispatch(self.orch, "POST", "/supervise/propose", _body(body)) + self.assertEqual(400, status, drop) + + def test_poll_invalid_json_is_400(self) -> None: + status, _ = dispatch(self.orch, "POST", "/supervise/poll", b"{not json") + self.assertEqual(400, status) + + def test_poll_missing_fields_are_400(self) -> None: + rec = self._register() + for body in ({"identity_token": rec.identity_token, "proposal_id": "p"}, + {"source_ip": rec.source_ip, "identity_token": rec.identity_token}): + status, _ = dispatch(self.orch, "POST", "/supervise/poll", _body(body)) + self.assertEqual(400, status) + + def test_poll_unattributed_is_403(self) -> None: + status, payload = dispatch(self.orch, "POST", "/supervise/poll", _body({ + "source_ip": "10.9.9.9", "identity_token": "wrong", "proposal_id": "p", + })) + self.assertEqual(403, status) + self.assertIn("unattributed", str(payload["error"])) + if __name__ == "__main__": unittest.main() diff --git a/tests/unit/test_policy_resolver.py b/tests/unit/test_policy_resolver.py index 6ab5d1c..cc5bde0 100644 --- a/tests/unit/test_policy_resolver.py +++ b/tests/unit/test_policy_resolver.py @@ -144,6 +144,7 @@ class TestPolicyResolver(unittest.TestCase): {"status": "approved", "notes": "ok", "final_file": None}) ) as m: result = self.r.poll_supervise("10.243.0.7", "tok", "p-7") + assert result is not None self.assertEqual("approved", result["status"]) req = m.call_args.args[0] self.assertTrue(req.full_url.endswith("/supervise/poll")) diff --git a/tests/unit/test_supervise_server.py b/tests/unit/test_supervise_server.py index 725943f..a3f956f 100644 --- a/tests/unit/test_supervise_server.py +++ b/tests/unit/test_supervise_server.py @@ -15,7 +15,6 @@ import time import types import unittest from pathlib import Path -from unittest.mock import patch from tests.unit import use_bottle_root @@ -66,14 +65,21 @@ class _FakeSuperviseResolver: def __init__( self, bottle_id: str | None = "dev", raises: bool = False, policy: str = "", + poll_raises: bool = False, poll_none: bool = False, ) -> None: self.bottle_id = bottle_id self.raises = raises self._policy = policy + # propose succeeds, but the later poll fails — models an orchestrator + # that becomes unreachable (poll_raises) or unattributes mid-flight + # (poll_none) after the proposal was queued. + self.poll_raises = poll_raises + self.poll_none = poll_none def propose_supervise( - self, source_ip, identity_token, *, tool, proposed_file, justification, - ): + self, source_ip: str, identity_token: str, *, + tool: str, proposed_file: str, justification: str, + ) -> str | None: del source_ip, identity_token if self.raises: raise supervise_server.PolicyResolveError("orchestrator down") @@ -86,11 +92,13 @@ class _FakeSuperviseResolver: _sv.write_proposal(proposal) return proposal.id - def poll_supervise(self, source_ip, identity_token, proposal_id): + def poll_supervise( + self, source_ip: str, identity_token: str, proposal_id: str, + ) -> dict[str, object] | None: del source_ip, identity_token - if self.raises: + if self.raises or self.poll_raises: raise supervise_server.PolicyResolveError("orchestrator down") - if self.bottle_id is None: + if self.bottle_id is None or self.poll_none: return None try: response = _sv.read_response(self.bottle_id, proposal_id) @@ -107,23 +115,28 @@ class _FakeSuperviseResolver: } # Used by list-egress-routes (`_resolved_routes_payload`), unchanged path. - def resolve_policy_and_bottle_id(self, source_ip, identity_token=""): + def resolve_policy_and_bottle_id( + self, source_ip: str, identity_token: str = "", + ) -> tuple[str | None, str | None, dict[str, str]]: del source_ip, identity_token if self.raises: raise supervise_server.PolicyResolveError("orchestrator down") return self._policy, self.bottle_id, {} -def _tools_call(resolver, params, config=None): +def _tools_call( + resolver: object, params: dict[str, object], + config: "ServerConfig | None" = None, +) -> dict[str, object]: return handle_tools_call( params, config or ServerConfig(), - resolver=resolver, source_ip=_SRC, identity_token=_TOK, + resolver=resolver, source_ip=_SRC, identity_token=_TOK, # type: ignore[arg-type] ) -def _check(resolver, params): +def _check(resolver: object, params: dict[str, object]) -> dict[str, object]: return handle_check_proposal( - params, resolver=resolver, source_ip=_SRC, identity_token=_TOK, + params, resolver=resolver, source_ip=_SRC, identity_token=_TOK, # type: ignore[arg-type] ) @@ -203,7 +216,7 @@ class TestRpcInternalErrorOnRpcFailure(unittest.TestCase): surfaces as ERR_INTERNAL — the daemon fails closed rather than leaking the cause to the agent.""" - _ARGS = { + _ARGS: dict[str, object] = { "name": _sv.TOOL_EGRESS_ALLOW, "arguments": { "routes_yaml": "routes:\n - host: example.com\n", @@ -493,6 +506,25 @@ class TestHandleToolsCall(unittest.TestCase): self.assertIn("proposal remains queued", text) self.assertEqual(1, len(_sv.list_pending_proposals("dev"))) + _ALLOW: dict[str, object] = { + "name": _sv.TOOL_EGRESS_ALLOW, + "arguments": {"routes_yaml": "routes:\n - host: example.com\n", "justification": "x"}, + } + + def test_poll_unreachable_after_queue_raises_internal(self): + # Proposal queues, then the orchestrator becomes unreachable on poll. + self.resolver.poll_raises = True + with self.assertRaises(_RpcInternalError) as cm: + _tools_call(self.resolver, self._ALLOW, ServerConfig(response_timeout_seconds=5)) + self.assertEqual(ERR_INTERNAL, cm.exception.code) + + def test_poll_unattributed_after_queue_raises_internal(self): + # Proposal queues, then poll comes back unattributed (fail-closed). + self.resolver.poll_none = True + with self.assertRaises(_RpcInternalError) as cm: + _tools_call(self.resolver, self._ALLOW, ServerConfig(response_timeout_seconds=5)) + self.assertEqual(ERR_INTERNAL, cm.exception.code) + class TestResponseTimeoutEnv(unittest.TestCase): def test_unset_uses_default(self): @@ -776,6 +808,18 @@ class TestNonBlockingSupervise(unittest.TestCase): self.assertTrue(result["isError"]) self.assertIn("status: unknown", result["content"][0]["text"]) # type: ignore[index] + def test_check_poll_unreachable_raises_internal(self): + self.resolver.poll_raises = True + with self.assertRaises(_RpcInternalError) as cm: + _check(self.resolver, {"arguments": {"proposal_id": "p"}}) + self.assertEqual(ERR_INTERNAL, cm.exception.code) + + def test_check_poll_unattributed_raises_internal(self): + self.resolver.poll_none = True + with self.assertRaises(_RpcInternalError) as cm: + _check(self.resolver, {"arguments": {"proposal_id": "p"}}) + self.assertEqual(ERR_INTERNAL, cm.exception.code) + def test_check_missing_id_raises(self): with self.assertRaises(_RpcClientError) as cm: _check(self.resolver, {"arguments": {}})