test: annotations + coverage for the supervise RPC seam
tracker-policy-pr / check-pr (pull_request) Successful in 6s
test / integration-docker (pull_request) Successful in 15s
test / unit (pull_request) Successful in 39s
test / integration-firecracker (pull_request) Successful in 3m20s
test / coverage (pull_request) Successful in 44s
test / publish-infra (pull_request) Has been skipped
tracker-policy-pr / check-pr (pull_request) Successful in 6s
test / integration-docker (pull_request) Successful in 15s
test / unit (pull_request) Successful in 39s
test / integration-firecracker (pull_request) Successful in 3m20s
test / coverage (pull_request) Successful in 44s
test / publish-infra (pull_request) Has been skipped
Add pyright-strict parameter/return annotations to the fake resolvers and test helpers, and cover the new control-plane validation/403 branches (/supervise/propose + /supervise/poll) plus the supervise-server and egress poll-error fail-closed paths, so the diff-coverage gate stays above 90%. No production behavior change. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
@@ -277,16 +277,18 @@ class _SuperviseRpcFake:
|
|||||||
|
|
||||||
supervise_status: "str | None" = None
|
supervise_status: "str | None" = None
|
||||||
propose_error: bool = False
|
propose_error: bool = False
|
||||||
|
poll_error: bool = False
|
||||||
|
|
||||||
@property
|
@property
|
||||||
def propose_calls(self) -> list:
|
def propose_calls(self) -> list[dict[str, str]]:
|
||||||
if not hasattr(self, "_propose_calls"):
|
if not hasattr(self, "_propose_calls"):
|
||||||
self._propose_calls: list = []
|
self._propose_calls: list[dict[str, str]] = []
|
||||||
return self._propose_calls
|
return self._propose_calls
|
||||||
|
|
||||||
def propose_supervise(
|
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
|
del proposed_file, justification
|
||||||
self.propose_calls.append(
|
self.propose_calls.append(
|
||||||
{"source_ip": source_ip, "identity_token": identity_token, "tool": tool}
|
{"source_ip": source_ip, "identity_token": identity_token, "tool": tool}
|
||||||
@@ -295,8 +297,12 @@ class _SuperviseRpcFake:
|
|||||||
raise PolicyResolveError("orchestrator down")
|
raise PolicyResolveError("orchestrator down")
|
||||||
return "prop-1"
|
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
|
del source_ip, identity_token, proposal_id
|
||||||
|
if self.poll_error:
|
||||||
|
raise PolicyResolveError("orchestrator down")
|
||||||
if self.supervise_status is None:
|
if self.supervise_status is None:
|
||||||
return {"status": "pending"}
|
return {"status": "pending"}
|
||||||
return {"status": self.supervise_status, "notes": "", "final_file": None}
|
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.assertEqual(403, flow.response.status_code)
|
||||||
self.assertIn("timed out", flow.response.get_text())
|
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
|
# Inbound DLP on responses
|
||||||
|
|||||||
@@ -21,7 +21,7 @@ from unittest.mock import patch
|
|||||||
|
|
||||||
from bot_bottle.orchestrator.broker import StubBroker
|
from bot_bottle.orchestrator.broker import StubBroker
|
||||||
from bot_bottle.orchestrator.control_plane import dispatch, make_server
|
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.orchestrator.service import Orchestrator
|
||||||
from bot_bottle.store_manager import StoreManager
|
from bot_bottle.store_manager import StoreManager
|
||||||
from bot_bottle.supervise import (
|
from bot_bottle.supervise import (
|
||||||
@@ -455,13 +455,13 @@ class TestDispatchSuperviseAgentRpc(unittest.TestCase):
|
|||||||
return self.store.register(
|
return self.store.register(
|
||||||
source_ip, metadata=json.dumps({"slug": slug}), policy="routes: []\n")
|
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({
|
return dispatch(self.orch, "POST", "/supervise/propose", _body({
|
||||||
"source_ip": rec.source_ip, "identity_token": rec.identity_token,
|
"source_ip": rec.source_ip, "identity_token": rec.identity_token,
|
||||||
"tool": TOOL_EGRESS_ALLOW, "proposed_file": proposed, "justification": "need it",
|
"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({
|
return dispatch(self.orch, "POST", "/supervise/poll", _body({
|
||||||
"source_ip": rec.source_ip, "identity_token": rec.identity_token,
|
"source_ip": rec.source_ip, "identity_token": rec.identity_token,
|
||||||
"proposal_id": proposal_id,
|
"proposal_id": proposal_id,
|
||||||
@@ -474,9 +474,11 @@ class TestDispatchSuperviseAgentRpc(unittest.TestCase):
|
|||||||
pid = payload["proposal_id"]
|
pid = payload["proposal_id"]
|
||||||
assert isinstance(pid, str) and pid
|
assert isinstance(pid, str) and pid
|
||||||
_, listing = dispatch(self.orch, "GET", "/supervise/proposals", b"")
|
_, 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.
|
# 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:
|
def test_propose_unattributed_is_403(self) -> None:
|
||||||
status, payload = dispatch(self.orch, "POST", "/supervise/propose", _body({
|
status, payload = dispatch(self.orch, "POST", "/supervise/propose", _body({
|
||||||
@@ -531,6 +533,41 @@ class TestDispatchSuperviseAgentRpc(unittest.TestCase):
|
|||||||
self.assertEqual(200, status)
|
self.assertEqual(200, status)
|
||||||
self.assertEqual("unknown", poll["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__":
|
if __name__ == "__main__":
|
||||||
unittest.main()
|
unittest.main()
|
||||||
|
|||||||
@@ -144,6 +144,7 @@ class TestPolicyResolver(unittest.TestCase):
|
|||||||
{"status": "approved", "notes": "ok", "final_file": None})
|
{"status": "approved", "notes": "ok", "final_file": None})
|
||||||
) as m:
|
) as m:
|
||||||
result = self.r.poll_supervise("10.243.0.7", "tok", "p-7")
|
result = self.r.poll_supervise("10.243.0.7", "tok", "p-7")
|
||||||
|
assert result is not None
|
||||||
self.assertEqual("approved", result["status"])
|
self.assertEqual("approved", result["status"])
|
||||||
req = m.call_args.args[0]
|
req = m.call_args.args[0]
|
||||||
self.assertTrue(req.full_url.endswith("/supervise/poll"))
|
self.assertTrue(req.full_url.endswith("/supervise/poll"))
|
||||||
|
|||||||
@@ -15,7 +15,6 @@ import time
|
|||||||
import types
|
import types
|
||||||
import unittest
|
import unittest
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from unittest.mock import patch
|
|
||||||
|
|
||||||
from tests.unit import use_bottle_root
|
from tests.unit import use_bottle_root
|
||||||
|
|
||||||
@@ -66,14 +65,21 @@ class _FakeSuperviseResolver:
|
|||||||
|
|
||||||
def __init__(
|
def __init__(
|
||||||
self, bottle_id: str | None = "dev", raises: bool = False, policy: str = "",
|
self, bottle_id: str | None = "dev", raises: bool = False, policy: str = "",
|
||||||
|
poll_raises: bool = False, poll_none: bool = False,
|
||||||
) -> None:
|
) -> None:
|
||||||
self.bottle_id = bottle_id
|
self.bottle_id = bottle_id
|
||||||
self.raises = raises
|
self.raises = raises
|
||||||
self._policy = policy
|
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(
|
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
|
del source_ip, identity_token
|
||||||
if self.raises:
|
if self.raises:
|
||||||
raise supervise_server.PolicyResolveError("orchestrator down")
|
raise supervise_server.PolicyResolveError("orchestrator down")
|
||||||
@@ -86,11 +92,13 @@ class _FakeSuperviseResolver:
|
|||||||
_sv.write_proposal(proposal)
|
_sv.write_proposal(proposal)
|
||||||
return proposal.id
|
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
|
del source_ip, identity_token
|
||||||
if self.raises:
|
if self.raises or self.poll_raises:
|
||||||
raise supervise_server.PolicyResolveError("orchestrator down")
|
raise supervise_server.PolicyResolveError("orchestrator down")
|
||||||
if self.bottle_id is None:
|
if self.bottle_id is None or self.poll_none:
|
||||||
return None
|
return None
|
||||||
try:
|
try:
|
||||||
response = _sv.read_response(self.bottle_id, proposal_id)
|
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.
|
# 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
|
del source_ip, identity_token
|
||||||
if self.raises:
|
if self.raises:
|
||||||
raise supervise_server.PolicyResolveError("orchestrator down")
|
raise supervise_server.PolicyResolveError("orchestrator down")
|
||||||
return self._policy, self.bottle_id, {}
|
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(
|
return handle_tools_call(
|
||||||
params, config or ServerConfig(),
|
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(
|
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
|
surfaces as ERR_INTERNAL — the daemon fails closed rather than leaking the
|
||||||
cause to the agent."""
|
cause to the agent."""
|
||||||
|
|
||||||
_ARGS = {
|
_ARGS: dict[str, object] = {
|
||||||
"name": _sv.TOOL_EGRESS_ALLOW,
|
"name": _sv.TOOL_EGRESS_ALLOW,
|
||||||
"arguments": {
|
"arguments": {
|
||||||
"routes_yaml": "routes:\n - host: example.com\n",
|
"routes_yaml": "routes:\n - host: example.com\n",
|
||||||
@@ -493,6 +506,25 @@ class TestHandleToolsCall(unittest.TestCase):
|
|||||||
self.assertIn("proposal remains queued", text)
|
self.assertIn("proposal remains queued", text)
|
||||||
self.assertEqual(1, len(_sv.list_pending_proposals("dev")))
|
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):
|
class TestResponseTimeoutEnv(unittest.TestCase):
|
||||||
def test_unset_uses_default(self):
|
def test_unset_uses_default(self):
|
||||||
@@ -776,6 +808,18 @@ class TestNonBlockingSupervise(unittest.TestCase):
|
|||||||
self.assertTrue(result["isError"])
|
self.assertTrue(result["isError"])
|
||||||
self.assertIn("status: unknown", result["content"][0]["text"]) # type: ignore[index]
|
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):
|
def test_check_missing_id_raises(self):
|
||||||
with self.assertRaises(_RpcClientError) as cm:
|
with self.assertRaises(_RpcClientError) as cm:
|
||||||
_check(self.resolver, {"arguments": {}})
|
_check(self.resolver, {"arguments": {}})
|
||||||
|
|||||||
Reference in New Issue
Block a user