Files
bot-bottle/tests/unit/test_orchestrator_service.py
T
didericis-claude 7ab85e9ea6
prd-number-check / require-numbered-prds (pull_request) Failing after 12s
test / integration-docker (pull_request) Successful in 19s
tracker-policy-pr / check-pr (pull_request) Successful in 9s
test / unit (pull_request) Failing after 52s
lint / lint (push) Successful in 1m1s
test / coverage (pull_request) Has been skipped
fix(orchestrator): address review on host control server transport (#468)
Codex review on #496:

- **High — ambiguous delivery no longer orphans a launched bottle.** A
  timeout / dropped response from the host controller is now the ambiguous
  BrokerUnavailableError (distinct from the definite BrokerAuthError /
  BrokerClientError). OrchestratorCore.launch_bottle keeps the registry
  row on the ambiguous case instead of deregistering — deregistering would
  orphan a running container with no record (reconcile reaps rows, never
  containers). The row is left for reconcile to reap iff the bottle is not
  actually live. Definite failures still roll back, so a real failure
  leaves no orphan row.
- **Medium — the privileged endpoint bounds request bodies.** The host
  server rejects an oversized Content-Length with 413 before reading it,
  and sets a per-request socket timeout, so a caller that can merely reach
  the socket (no signed token) can't exhaust memory or a handler thread.

Tests: ambiguous-keep vs definite-rollback in the launch path; the
BrokerUnavailableError/BrokerClientError split in BrokerClient; the 413
body cap + handler error paths (driven in-thread, since daemon request
threads lose coverage); and the __main__ entrypoint broker selection.
Diff-coverage 98%; pyright clean; pylint 9.88.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
2026-07-26 08:43:00 +00:00

429 lines
20 KiB
Python

"""Unit tests for the OrchestratorCore launch lifecycle (PRD 0070)."""
from __future__ import annotations
import json
import secrets
import sqlite3
import tempfile
import unittest
from contextlib import closing
from pathlib import Path
from unittest.mock import patch
from bot_bottle.orchestrator.broker import (
BrokerUnavailableError,
LaunchBroker,
LaunchRequest,
StubBroker,
)
from bot_bottle.orchestrator.store.registry_store import RegistryStore
from bot_bottle.orchestrator.service import OrchestratorCore
from bot_bottle.orchestrator.store.secret_store import new_env_var_secret
from bot_bottle.orchestrator.store.store_manager import StoreManager
from bot_bottle.orchestrator.supervisor import (
Proposal,
STATUS_APPROVED,
TOOL_EGRESS_ALLOW,
Supervisor,
)
class _FailingBroker(LaunchBroker):
"""Verifies the token like any broker, then fails the launch *definitely* —
to exercise the orchestrator's registry rollback."""
def _launch(self, req: LaunchRequest) -> None:
raise RuntimeError("launch failed")
def _teardown(self, req: LaunchRequest) -> None:
pass
class _UnavailableBroker(LaunchBroker):
"""Verifies the token, then raises the *ambiguous* BrokerUnavailableError —
the host may already have launched — so the orchestrator must KEEP the
registry row rather than orphan a running container."""
def _launch(self, req: LaunchRequest) -> None:
raise BrokerUnavailableError("delivery dropped after send")
def _teardown(self, req: LaunchRequest) -> None:
pass
class TestOrchestrator(unittest.TestCase):
def setUp(self) -> None:
self._tmp = tempfile.TemporaryDirectory()
self.secret = secrets.token_bytes(16)
self.store = RegistryStore(Path(self._tmp.name) / "r.db")
self.store.migrate()
self.broker = StubBroker(self.secret)
self.orch = OrchestratorCore(self.store, self.broker, self.secret)
def tearDown(self) -> None:
self._tmp.cleanup()
def test_launch_registers_and_brokers_signed_request(self) -> None:
rec = self.orch.launch_bottle("10.243.0.1", image_ref="sha256:abc", slot=2)
self.assertIsNotNone(self.store.get(rec.bottle_id))
self.assertEqual(1, len(self.broker.launched))
req = self.broker.launched[0]
self.assertEqual("launch", req.op)
self.assertEqual(rec.bottle_id, req.bottle_id)
self.assertEqual("10.243.0.1", req.source_ip)
self.assertEqual("sha256:abc", req.image_ref)
self.assertEqual(2, req.slot)
def test_launch_then_attribute(self) -> None:
rec = self.orch.launch_bottle("10.243.0.3")
got = self.orch.attribute("10.243.0.3", rec.identity_token)
assert got is not None
self.assertEqual(rec.bottle_id, got.bottle_id)
self.assertIsNone(self.orch.attribute("10.243.0.3", "wrong-token"))
def test_teardown_brokers_and_deregisters(self) -> None:
rec = self.orch.launch_bottle("10.243.0.1")
self.assertTrue(self.orch.teardown_bottle(rec.bottle_id))
self.assertIsNone(self.store.get(rec.bottle_id))
self.assertEqual(["teardown"], [r.op for r in self.broker.torn_down])
def test_teardown_unknown_is_false(self) -> None:
self.assertFalse(self.orch.teardown_bottle("ghost"))
def test_launch_with_policy_is_resolvable(self) -> None:
rec = self.orch.launch_bottle("10.243.0.1", policy='{"routes":[]}')
got = self.orch.attribute("10.243.0.1", rec.identity_token)
assert got is not None
self.assertEqual('{"routes":[]}', got.policy)
def test_tokens_held_in_memory_and_cleared_on_teardown(self) -> None:
rec = self.orch.launch_bottle("10.243.0.5", tokens={"EGRESS_TOKEN_0": "sk"})
self.assertEqual({"EGRESS_TOKEN_0": "sk"}, self.orch.tokens_for(rec.bottle_id))
# not persisted in the registry (redacted record has no tokens field)
self.assertNotIn("tokens", rec.redacted())
self.orch.teardown_bottle(rec.bottle_id)
self.assertEqual({}, self.orch.tokens_for(rec.bottle_id))
def test_tokens_default_empty(self) -> None:
rec = self.orch.launch_bottle("10.243.0.6")
self.assertEqual({}, self.orch.tokens_for(rec.bottle_id))
def test_encrypted_tokens_can_be_reprovisioned_after_memory_loss(self) -> None:
key = new_env_var_secret()
rec = self.orch.launch_bottle(
"10.243.0.12", tokens={"EGRESS_TOKEN_0": "secret"},
env_var_secret=key,
)
self.assertNotEqual({}, self.store.get_agent_secrets(rec.bottle_id))
self.orch._tokens.clear()
self.assertTrue(self.orch.reprovision_from_secret(rec.bottle_id, key))
self.assertEqual({"EGRESS_TOKEN_0": "secret"}, self.orch.tokens_for(rec.bottle_id))
def test_reprovision_rejects_missing_rows_and_wrong_key(self) -> None:
self.assertFalse(self.orch.reprovision_from_secret("missing", new_env_var_secret()))
rec = self.orch.launch_bottle(
"10.243.0.13", tokens={"K": "value"},
env_var_secret=new_env_var_secret(),
)
self.assertFalse(self.orch.reprovision_from_secret(rec.bottle_id, new_env_var_secret()))
def test_set_policy_live_reload(self) -> None:
rec = self.orch.launch_bottle("10.243.0.3")
self.assertTrue(self.orch.set_policy(rec.bottle_id, '{"x":1}'))
got = self.orch.attribute("10.243.0.3", rec.identity_token)
assert got is not None
self.assertEqual('{"x":1}', got.policy)
def test_set_policy_unknown_is_false(self) -> None:
self.assertFalse(self.orch.set_policy("ghost", "{}"))
def test_resolve_requires_matching_token(self) -> None:
# Mandatory (source_ip, token) pair — no source-IP-only fallback.
rec = self.orch.launch_bottle("10.243.0.1", policy="P")
self.assertIsNone(self.orch.resolve("10.243.0.1", "")) # empty token denies
self.assertIsNone(self.orch.resolve("10.243.0.1", "wrong")) # mismatch denies
got = self.orch.resolve("10.243.0.1", rec.identity_token) # exact pair
assert got is not None
self.assertEqual(rec.bottle_id, got.bottle_id)
self.assertEqual("P", got.policy)
def test_resolve_with_token_stays_strict(self) -> None:
rec = self.orch.launch_bottle("10.243.0.3")
self.assertIsNotNone(self.orch.resolve("10.243.0.3", rec.identity_token))
self.assertIsNone(self.orch.resolve("10.243.0.3", "wrong-token"))
def test_launch_rolls_back_registry_on_definite_broker_failure(self) -> None:
orch = OrchestratorCore(self.store, _FailingBroker(self.secret), self.secret)
with self.assertRaises(RuntimeError):
orch.launch_bottle("10.243.0.9")
self.assertEqual([], self.store.all()) # no orphan row
def test_launch_keeps_registry_on_ambiguous_broker_failure(self) -> None:
# The host may already have launched the bottle before the response was
# lost, so deregistering would orphan a running container with no row.
# The row is kept for reconcile to reap iff the bottle is not live.
orch = OrchestratorCore(self.store, _UnavailableBroker(self.secret), self.secret)
with self.assertRaises(BrokerUnavailableError):
orch.launch_bottle("10.243.0.9")
self.assertEqual(1, len(self.store.all())) # row survives — no orphan container
def test_gateway_status_reports_unconfigured(self) -> None:
# The orchestrator no longer owns a standalone gateway lifecycle; the
# consolidated flow runs the gateway data plane in the per-host infra
# container/VM. `GET /gateway` therefore always reports unconfigured.
self.assertEqual({"configured": False}, self.orch.gateway_status())
class TestOrchestratorSupervise(unittest.TestCase):
"""Operator-approval flow: the orchestrator applies the decision
server-side against the single DB (queue + policy + audit)."""
def setUp(self) -> None:
self._tmp = tempfile.TemporaryDirectory()
root = Path(self._tmp.name)
db = root / "db" / "bot-bottle.db"
db.parent.mkdir(parents=True)
# One DB for registry + supervise queue + audit (as in the VM).
self._env = patch.dict("os.environ", {
"BOT_BOTTLE_ROOT": str(root),
"SUPERVISE_DB_PATH": str(db),
})
self._env.start()
self.store = RegistryStore(db)
self.store.migrate()
StoreManager(db).migrate()
secret = secrets.token_bytes(16)
self.orch = OrchestratorCore(self.store, StubBroker(secret), secret)
def tearDown(self) -> None:
self._env.stop()
self._tmp.cleanup()
def _register(self, slug: str, policy: str) -> str:
rec = self.store.register(
"10.243.0.1", metadata=json.dumps({"slug": slug}), policy=policy)
return rec.bottle_id
def _queue(self, slug: str, proposed: str) -> str:
p = Proposal.new(
bottle_slug=slug, tool=TOOL_EGRESS_ALLOW, proposed_file=proposed,
justification="need it")
Supervisor().write_proposal(p)
return p.id
def test_pending_lists_queued_proposal(self) -> None:
self._register("demo", "routes: []\n")
pid = self._queue("demo", "routes:\n - host: google.com\n")
pending = self.orch.supervise_pending()
self.assertEqual(1, len(pending))
self.assertEqual(pid, pending[0]["id"])
self.assertEqual("demo", pending[0]["bottle_slug"])
def test_approve_applies_policy_writes_response_and_clears_pending(self) -> None:
bottle_id = self._register("demo", "routes:\n - host: existing.com\n")
new_routes = "routes:\n - host: google.com\n"
pid = self._queue("demo", new_routes)
ok, err = self.orch.supervise_respond(
pid, bottle_slug="demo", decision="approve")
self.assertTrue(ok, err)
# policy live-applied so /resolve serves the new routes
rec = self.store.get(bottle_id)
assert rec is not None
self.assertEqual(new_routes, rec.policy)
# response written -> agent unblocks, proposal no longer pending
self.assertEqual(STATUS_APPROVED, Supervisor().read_response("demo", pid).status)
self.assertEqual([], self.orch.supervise_pending())
def test_pending_carries_human_label(self) -> None:
# The proposal is keyed by bottle_id, but pending dicts also expose the
# bottle's human slug so the operator sees a name, not a hex id.
bottle_id = self._register("codex-dev-a1b2c", "routes: []\n")
self._queue(bottle_id, "routes:\n - host: google.com\n")
pending = self.orch.supervise_pending()
self.assertEqual(bottle_id, pending[0]["bottle_slug"])
self.assertEqual("codex-dev-a1b2c", pending[0]["bottle_label"])
def test_pending_label_falls_back_to_slug_when_bottle_gone(self) -> None:
# No registry record (torn down): label is the id, never empty.
self._queue("ghost-id", "routes:\n - host: google.com\n")
pending = self.orch.supervise_pending()
self.assertEqual("ghost-id", pending[0]["bottle_label"])
def test_approve_by_bottle_id_applies_policy(self) -> None:
# Consolidated reality: the supervise server keys each proposal by the
# orchestrator-assigned bottle_id, not the human slug. Approval must
# resolve the record by that id and apply the policy (regression for
# the "bottle <id> is no longer registered" 409).
bottle_id = self._register("codex-dev-a1b2c", "routes: []\n")
new_routes = "routes:\n - host: google.com\n"
pid = self._queue(bottle_id, new_routes)
ok, err = self.orch.supervise_respond(
pid, bottle_slug=bottle_id, decision="approve")
self.assertTrue(ok, err)
rec = self.store.get(bottle_id)
assert rec is not None
self.assertEqual(new_routes, rec.policy)
def test_modify_applies_final_file_not_proposed(self) -> None:
bottle_id = self._register("demo", "routes: []\n")
pid = self._queue("demo", "routes:\n - host: google.com\n")
edited = "routes:\n - host: example.com\n"
ok, _ = self.orch.supervise_respond(
pid, bottle_slug="demo", decision="modify", final_file=edited)
self.assertTrue(ok)
rec = self.store.get(bottle_id)
assert rec is not None
self.assertEqual(edited, rec.policy)
def test_reject_leaves_policy_unchanged(self) -> None:
bottle_id = self._register("demo", "routes:\n - host: existing.com\n")
pid = self._queue("demo", "routes:\n - host: google.com\n")
ok, _ = self.orch.supervise_respond(
pid, bottle_slug="demo", decision="reject", notes="no")
self.assertTrue(ok)
rec = self.store.get(bottle_id)
assert rec is not None
self.assertEqual("routes:\n - host: existing.com\n", rec.policy)
self.assertEqual("rejected", Supervisor().read_response("demo", pid).status)
def test_unknown_proposal_is_error(self) -> None:
ok, err = self.orch.supervise_respond(
"ghost", bottle_slug="demo", decision="approve")
self.assertFalse(ok)
self.assertIn("no such proposal", err)
def test_unknown_decision_is_error(self) -> None:
self._register("demo", "routes: []\n")
pid = self._queue("demo", "routes:\n - host: google.com\n")
ok, err = self.orch.supervise_respond(
pid, bottle_slug="demo", decision="bogus")
self.assertFalse(ok)
self.assertIn("unknown decision", err)
def test_approve_when_bottle_gone_cannot_apply(self) -> None:
# Proposal queued but the bottle was torn down before the operator
# acted: an egress apply has no target, so respond fails closed.
pid = self._queue("ghost-bottle", "routes:\n - host: google.com\n")
ok, err = self.orch.supervise_respond(
pid, bottle_slug="ghost-bottle", decision="approve")
self.assertFalse(ok)
self.assertIn("no longer registered", err)
# --- agent half: queue + poll (issue #469) -----------------------------
def test_queue_proposal_then_poll_pending(self) -> None:
bottle_id = self._register("demo", "routes: []\n")
pid = self.orch.supervise_queue_proposal(
bottle_id, tool=TOOL_EGRESS_ALLOW,
proposed_file="routes:\n - host: google.com\n", justification="need it")
# Visible to the operator, keyed by the bottle id.
pending = self.orch.supervise_pending()
self.assertEqual([pid], [p["id"] for p in pending])
self.assertEqual(
{"status": "pending"}, self.orch.supervise_poll_response(bottle_id, pid))
def test_poll_is_idempotent_and_leaves_no_pending(self) -> None:
bottle_id = self._register("demo", "routes: []\n")
pid = self.orch.supervise_queue_proposal(
bottle_id, tool=TOOL_EGRESS_ALLOW,
proposed_file="routes:\n - host: google.com\n", justification="need it")
self.orch.supervise_respond(
pid, bottle_slug=bottle_id, decision="approve", notes="ok")
decided = self.orch.supervise_poll_response(bottle_id, pid)
self.assertEqual("approved", decided["status"])
self.assertEqual("ok", decided["notes"])
# Decided proposal drops off the operator's pending list (a response row
# exists), but poll does NOT archive — a re-poll returns the same
# decision so a dropped connection can't lose it (issue #469 review).
self.assertEqual([], self.orch.supervise_pending())
self.assertEqual(decided, self.orch.supervise_poll_response(bottle_id, pid))
def test_teardown_reaps_the_bottles_proposals(self) -> None:
rec = self.store.register("10.243.0.20", policy="routes: []\n")
pid = self.orch.supervise_queue_proposal(
rec.bottle_id, tool=TOOL_EGRESS_ALLOW,
proposed_file="routes:\n - host: google.com\n", justification="j")
self.orch.supervise_respond(
pid, bottle_slug=rec.bottle_id, decision="approve", notes="ok")
self.assertTrue(self.orch.teardown_bottle(rec.bottle_id))
# The gone bottle's decided-but-unconsumed proposal is archived, so a
# late poll returns 'unknown' rather than lingering forever.
self.assertEqual(
{"status": "unknown"}, self.orch.supervise_poll_response(rec.bottle_id, pid))
def test_reconcile_reaps_the_bottles_proposals(self) -> None:
rec = self.store.register("10.243.0.21", policy="routes: []\n")
pid = self.orch.supervise_queue_proposal(
rec.bottle_id, tool=TOOL_EGRESS_ALLOW,
proposed_file="routes:\n - host: google.com\n", justification="j")
# No live source IPs -> the bottle is reaped (grace 0 so it's immediate).
self.assertEqual([rec.bottle_id], self.orch.reconcile([], grace_seconds=0))
self.assertEqual(
{"status": "unknown"}, self.orch.supervise_poll_response(rec.bottle_id, pid))
def test_poll_unknown_for_other_bottle(self) -> None:
bottle_id = self._register("demo", "routes: []\n")
pid = self.orch.supervise_queue_proposal(
bottle_id, tool=TOOL_EGRESS_ALLOW,
proposed_file="routes:\n - host: google.com\n", justification="j")
# A different bottle id can't read demo's proposal (scoped by queue key).
self.assertEqual(
{"status": "unknown"}, self.orch.supervise_poll_response("other-bottle", pid))
if __name__ == "__main__":
unittest.main()
class TestOrchestratorReconcile(unittest.TestCase):
"""`reconcile` — drop rows for bottles that are no longer running."""
def setUp(self) -> None:
self._tmp = tempfile.TemporaryDirectory()
self.secret = secrets.token_bytes(16)
self.db = Path(self._tmp.name) / "r.db"
self.store = RegistryStore(self.db)
self.store.migrate()
self.broker = StubBroker(self.secret)
self.orch = OrchestratorCore(self.store, self.broker, self.secret)
def tearDown(self) -> None:
self._tmp.cleanup()
def _age_all(self, seconds: float) -> None:
"""Backdate every row past the reap grace window."""
with closing(sqlite3.connect(self.db)) as conn:
conn.execute(
"UPDATE orchestrator_bottles SET created_at = created_at - ?", (seconds,))
conn.commit()
def test_reaps_dead_bottle_and_forgets_its_tokens(self) -> None:
dead = self.orch.launch_bottle("10.243.0.1", tokens={"EGRESS_TOKEN_0": "s3cret"})
live = self.orch.launch_bottle("10.243.0.2", tokens={"EGRESS_TOKEN_0": "keep"})
self._age_all(600)
self.assertEqual([dead.bottle_id], self.orch.reconcile(["10.243.0.2"]))
self.assertIsNone(self.store.get(dead.bottle_id))
self.assertIsNotNone(self.store.get(live.bottle_id))
# The in-memory egress credential goes with the row.
self.assertEqual({}, self.orch.tokens_for(dead.bottle_id))
self.assertEqual({"EGRESS_TOKEN_0": "keep"}, self.orch.tokens_for(live.bottle_id))
def test_reconcile_does_not_broker_a_teardown(self) -> None:
"""The container is already gone — there is nothing to stop, and a
broker error must not stop the sweep clearing the row."""
self.orch.launch_bottle("10.243.0.1")
self._age_all(600)
self.broker.launched.clear()
self.orch.reconcile([])
self.assertEqual([], self.broker.torn_down)
def test_reconcile_keeps_everything_when_all_are_live(self) -> None:
a = self.orch.launch_bottle("10.243.0.1")
b = self.orch.launch_bottle("10.243.0.2")
self._age_all(600)
self.assertEqual([], self.orch.reconcile(["10.243.0.1", "10.243.0.2"]))
self.assertIsNotNone(self.store.get(a.bottle_id))
self.assertIsNotNone(self.store.get(b.bottle_id))