436 lines
15 KiB
Python
436 lines
15 KiB
Python
from __future__ import annotations
|
|
|
|
import base64
|
|
import hashlib
|
|
import json
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey
|
|
|
|
from aios_baseline.optimization import (
|
|
OptimizationController,
|
|
OptimizationEnvelope,
|
|
OptimizationError,
|
|
canonical_json_bytes,
|
|
validate_proposal,
|
|
)
|
|
|
|
|
|
def _signed_envelope(**overrides):
|
|
private_key = Ed25519PrivateKey.generate()
|
|
document = {
|
|
"schema": "aios.optimization-envelope.v1",
|
|
"envelope_id": "opt-envelope-1",
|
|
"tenant_id": "tenant-a",
|
|
"allowed_parameters": {
|
|
"retry.max_attempts": {"minimum": 1, "maximum": 5, "integer": True},
|
|
"timeout.operation_seconds": {"minimum": 5, "maximum": 60},
|
|
"evidence_reuse.max_age_seconds": {
|
|
"minimum": 0,
|
|
"maximum": 600,
|
|
"integer": True,
|
|
},
|
|
},
|
|
"allowed_components": ["task-worker"],
|
|
"allowed_environments": ["staging"],
|
|
"max_change_size": 2,
|
|
"max_compute_budget": 20,
|
|
"quality_floor": {
|
|
"success_rate": {"direction": "higher", "threshold": 0.95},
|
|
"error_rate": {"direction": "lower", "threshold": 0.05},
|
|
},
|
|
"rollback_requirement": True,
|
|
"promotion_rules": {
|
|
"objective_metric": "success_rate",
|
|
"minimum_improvement": 0.01,
|
|
"mandatory_tests": ["test_worker_safety"],
|
|
"mandatory_quality_gates": ["no_tenant_leak"],
|
|
"required_receipt_schema": "aios.governed-task-receipt.v1",
|
|
},
|
|
"source_config_digest": "sha256:" + "a" * 64,
|
|
"key_id": "optimization-config-key-1",
|
|
}
|
|
document.update(overrides)
|
|
signature = private_key.sign(canonical_json_bytes(document))
|
|
document["signature"] = base64.b64encode(signature).decode("ascii")
|
|
|
|
def verify(key_id, payload, encoded_signature):
|
|
assert key_id == "optimization-config-key-1"
|
|
try:
|
|
private_key.public_key().verify(base64.b64decode(encoded_signature), payload)
|
|
except Exception:
|
|
return False
|
|
return True
|
|
|
|
return document, verify
|
|
|
|
|
|
def _envelope():
|
|
document, verifier = _signed_envelope()
|
|
return OptimizationEnvelope.from_dict(document, verify_envelope=verifier)
|
|
|
|
|
|
def _context(**overrides):
|
|
value = {
|
|
"tenant_id": "tenant-a",
|
|
"subject_id": "optimizer-service",
|
|
"grant_id": "grant-123",
|
|
"mission_id": "mission-456",
|
|
"service_identity": "aios-root-service",
|
|
}
|
|
value.update(overrides)
|
|
return value
|
|
|
|
|
|
def _proposal(**overrides):
|
|
value = {
|
|
"schema": "aios.optimization-proposal.v1",
|
|
"proposal_id": "proposal-1",
|
|
"envelope_id": "opt-envelope-1",
|
|
"tenant_id": "tenant-a",
|
|
"component": "task-worker",
|
|
"environment": "staging",
|
|
"changes": [
|
|
{"path": "retry.max_attempts", "value": 4},
|
|
{"path": "timeout.operation_seconds", "value": 30},
|
|
],
|
|
"compute_budget": 10,
|
|
"tests": ["test_worker_safety", "test_candidate_latency"],
|
|
"quality_gates": ["no_tenant_leak"],
|
|
"baseline_metrics": {"success_rate": 0.95, "error_rate": 0.04},
|
|
"rollback_plan": {
|
|
"strategy": "restore_previous_values",
|
|
"reference": "baseline-config-17",
|
|
},
|
|
}
|
|
value.update(overrides)
|
|
return value
|
|
|
|
|
|
def _signed_receipt_verifier():
|
|
private_key = Ed25519PrivateKey.generate()
|
|
|
|
def sign(receipt):
|
|
signed = dict(receipt)
|
|
signed["key_id"] = "governed-receipt-key-1"
|
|
signed["signature"] = base64.b64encode(
|
|
private_key.sign(canonical_json_bytes(signed))
|
|
).decode("ascii")
|
|
return signed
|
|
|
|
def verify(receipt):
|
|
unsigned = dict(receipt)
|
|
signature = unsigned.pop("signature", "")
|
|
if unsigned.get("key_id") != "governed-receipt-key-1":
|
|
return False
|
|
try:
|
|
private_key.public_key().verify(
|
|
base64.b64decode(signature), canonical_json_bytes(unsigned)
|
|
)
|
|
except Exception:
|
|
return False
|
|
return True
|
|
|
|
return sign, verify
|
|
|
|
|
|
def _controller(tmp_path: Path):
|
|
submitted = []
|
|
sign, verify = _signed_receipt_verifier()
|
|
|
|
def submit_task(task):
|
|
submitted.append(task)
|
|
return {"task_id": "evaluation-task-1", "state": "QUEUED"}
|
|
|
|
controller = OptimizationController(
|
|
tmp_path / "optimization.sqlite3",
|
|
_envelope(),
|
|
submit_task,
|
|
verify_receipt=verify,
|
|
)
|
|
return controller, submitted, sign
|
|
|
|
|
|
def _evaluation(proposal_record, sign, **overrides):
|
|
result = {
|
|
"candidate_metrics": {"success_rate": 0.97, "error_rate": 0.03},
|
|
"compute_used": 8,
|
|
"tests_passed": ["test_worker_safety", "test_candidate_latency"],
|
|
"quality_gates_passed": ["no_tenant_leak"],
|
|
}
|
|
result.update(overrides)
|
|
receipt = sign(
|
|
{
|
|
"schema": "aios.governed-task-receipt.v1",
|
|
"task_id": proposal_record["evaluation_task_id"],
|
|
"proposal_id": proposal_record["proposal_id"],
|
|
"proposal_digest": proposal_record["proposal_digest"],
|
|
"action": "evaluate",
|
|
"success": True,
|
|
"result_digest": "sha256:"
|
|
+ hashlib.sha256(canonical_json_bytes(result)).hexdigest(),
|
|
}
|
|
)
|
|
return {**result, "receipt": receipt}
|
|
|
|
|
|
def test_envelope_requires_a_valid_signature_over_the_exact_payload():
|
|
document, verifier = _signed_envelope()
|
|
document["max_compute_budget"] = 10_000
|
|
|
|
with pytest.raises(OptimizationError, match="OPTIMIZATION_ENVELOPE_SIGNATURE_INVALID"):
|
|
OptimizationEnvelope.from_dict(document, verify_envelope=verifier)
|
|
|
|
with pytest.raises(OptimizationError, match="OPTIMIZATION_ENVELOPE_VERIFIER_REQUIRED"):
|
|
OptimizationEnvelope.from_dict(document, verify_envelope=None)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"path",
|
|
[
|
|
"authorization.accept_without_grant",
|
|
"trust_root.key",
|
|
"tenant_isolation.enabled",
|
|
"acceptance.threshold",
|
|
"destructive_capability.enabled",
|
|
"owner_grant.semantics",
|
|
"production_migration.auto_apply",
|
|
"code.patch",
|
|
],
|
|
)
|
|
def test_envelope_rejects_forbidden_or_generic_parameter_paths(path):
|
|
document, verifier = _signed_envelope(
|
|
allowed_parameters={path: {"minimum": 0, "maximum": 1}}
|
|
)
|
|
# Re-sign the changed trusted document so this reaches semantic validation.
|
|
private_key = Ed25519PrivateKey.generate()
|
|
unsigned = dict(document)
|
|
unsigned.pop("signature")
|
|
document["signature"] = base64.b64encode(
|
|
private_key.sign(canonical_json_bytes(unsigned))
|
|
).decode("ascii")
|
|
|
|
def changed_verifier(_key_id, payload, signature):
|
|
try:
|
|
private_key.public_key().verify(base64.b64decode(signature), payload)
|
|
except Exception:
|
|
return False
|
|
return True
|
|
|
|
with pytest.raises(OptimizationError, match="OPTIMIZATION_PARAMETER_PATH_FORBIDDEN"):
|
|
OptimizationEnvelope.from_dict(document, verify_envelope=changed_verifier)
|
|
|
|
|
|
def test_proposal_is_bounded_and_cannot_remove_mandatory_checks():
|
|
envelope = _envelope()
|
|
normalized = validate_proposal(envelope, _proposal())
|
|
assert normalized["changes"][0] == {"path": "retry.max_attempts", "value": 4}
|
|
|
|
with pytest.raises(OptimizationError, match="OPTIMIZATION_PARAMETER_OUT_OF_RANGE"):
|
|
validate_proposal(
|
|
envelope,
|
|
_proposal(changes=[{"path": "retry.max_attempts", "value": 99}]),
|
|
)
|
|
with pytest.raises(OptimizationError, match="OPTIMIZATION_MANDATORY_TEST_MISSING"):
|
|
validate_proposal(envelope, _proposal(tests=["test_candidate_latency"]))
|
|
with pytest.raises(OptimizationError, match="OPTIMIZATION_MANDATORY_GATE_MISSING"):
|
|
validate_proposal(envelope, _proposal(quality_gates=[]))
|
|
|
|
|
|
def test_evidence_reuse_cannot_widen_tenant_component_or_environment_scope():
|
|
envelope = _envelope()
|
|
proposal = _proposal(
|
|
changes=[{"path": "evidence_reuse.max_age_seconds", "value": 300}],
|
|
evidence_reuse={
|
|
"scope": {
|
|
"tenant_id": "tenant-b",
|
|
"component": "task-worker",
|
|
"environment": "staging",
|
|
},
|
|
"evidence_ids": ["evidence-1"],
|
|
},
|
|
)
|
|
|
|
with pytest.raises(OptimizationError, match="OPTIMIZATION_EVIDENCE_SCOPE_MISMATCH"):
|
|
validate_proposal(envelope, proposal)
|
|
|
|
|
|
def test_propose_only_persists_a_record_and_submits_an_identity_bound_task(tmp_path):
|
|
controller, submitted, _ = _controller(tmp_path)
|
|
record = controller.propose(_proposal(), authenticated_context=_context())
|
|
|
|
assert record["status"] == "EVALUATION_SUBMITTED"
|
|
assert record["evaluation_task_id"] == "evaluation-task-1"
|
|
assert submitted == [
|
|
{
|
|
"schema": "aios.optimization-task-request.v1",
|
|
"task_id": "optimization-evaluate-proposal-1",
|
|
"action": "evaluate",
|
|
"tenant_id": "tenant-a",
|
|
"subject_id": "optimizer-service",
|
|
"grant_id": "grant-123",
|
|
"mission_id": "mission-456",
|
|
"proposal_id": "proposal-1",
|
|
"proposal_digest": record["proposal_digest"],
|
|
"component": "task-worker",
|
|
"environment": "staging",
|
|
"compute_budget": 10.0,
|
|
}
|
|
]
|
|
|
|
|
|
def test_invalid_candidate_is_stopped_without_task_submission(tmp_path):
|
|
controller, submitted, _ = _controller(tmp_path)
|
|
with pytest.raises(OptimizationError, match="OPTIMIZATION_COMPONENT_FORBIDDEN"):
|
|
controller.propose(
|
|
_proposal(component="authorization-service"),
|
|
authenticated_context=_context(),
|
|
)
|
|
|
|
stopped = controller.get_proposal("proposal-1", authenticated_context=_context())
|
|
assert stopped["status"] == "STOPPED_INVALID"
|
|
assert submitted == []
|
|
|
|
|
|
def test_cross_tenant_or_subject_access_is_denied(tmp_path):
|
|
controller, _, _ = _controller(tmp_path)
|
|
controller.propose(_proposal(), authenticated_context=_context())
|
|
|
|
with pytest.raises(OptimizationError, match="OPTIMIZATION_CONTEXT_TENANT_MISMATCH"):
|
|
controller.get_proposal(
|
|
"proposal-1", authenticated_context=_context(tenant_id="tenant-b")
|
|
)
|
|
with pytest.raises(OptimizationError, match="OPTIMIZATION_CONTEXT_SUBJECT_MISMATCH"):
|
|
controller.get_proposal(
|
|
"proposal-1", authenticated_context=_context(subject_id="other-service")
|
|
)
|
|
|
|
|
|
def test_forged_or_wrong_task_receipt_is_rejected(tmp_path):
|
|
controller, _, sign = _controller(tmp_path)
|
|
proposal = controller.propose(_proposal(), authenticated_context=_context())
|
|
evaluation = _evaluation(proposal, sign)
|
|
evaluation["receipt"]["task_id"] = "some-other-task"
|
|
|
|
with pytest.raises(OptimizationError, match="OPTIMIZATION_RECEIPT_UNTRUSTED"):
|
|
controller.record_evaluation(
|
|
"proposal-1", evaluation, authenticated_context=_context()
|
|
)
|
|
|
|
validly_signed_wrong_binding = _evaluation(proposal, sign)
|
|
unsigned = dict(validly_signed_wrong_binding["receipt"])
|
|
unsigned["task_id"] = "some-other-task"
|
|
unsigned.pop("signature")
|
|
validly_signed_wrong_binding["receipt"] = sign(unsigned)
|
|
with pytest.raises(OptimizationError, match="OPTIMIZATION_RECEIPT_TASK_MISMATCH"):
|
|
controller.record_evaluation(
|
|
"proposal-1", validly_signed_wrong_binding, authenticated_context=_context()
|
|
)
|
|
|
|
|
|
def test_no_improvement_stops_promotion(tmp_path):
|
|
controller, submitted, sign = _controller(tmp_path)
|
|
proposal = controller.propose(_proposal(), authenticated_context=_context())
|
|
evaluation = _evaluation(
|
|
proposal,
|
|
sign,
|
|
candidate_metrics={"success_rate": 0.95, "error_rate": 0.04},
|
|
)
|
|
result = controller.record_evaluation(
|
|
"proposal-1", evaluation, authenticated_context=_context()
|
|
)
|
|
|
|
assert result["status"] == "STOPPED_NO_IMPROVEMENT"
|
|
assert result["effect_comparison"] == {
|
|
"error_rate": 0.0,
|
|
"success_rate": 0.0,
|
|
}
|
|
with pytest.raises(OptimizationError, match="OPTIMIZATION_NOT_PROMOTABLE"):
|
|
controller.promote(
|
|
"proposal-1",
|
|
{
|
|
"evaluation_task_id": "evaluation-task-1",
|
|
"proposal_digest": proposal["proposal_digest"],
|
|
},
|
|
authenticated_context=_context(),
|
|
)
|
|
assert len(submitted) == 1
|
|
|
|
|
|
def test_quality_floor_or_baseline_regression_stops_promotion(tmp_path):
|
|
controller, _, sign = _controller(tmp_path)
|
|
proposal = controller.propose(_proposal(), authenticated_context=_context())
|
|
evaluation = _evaluation(
|
|
proposal,
|
|
sign,
|
|
candidate_metrics={"success_rate": 0.97, "error_rate": 0.05},
|
|
)
|
|
result = controller.record_evaluation(
|
|
"proposal-1", evaluation, authenticated_context=_context()
|
|
)
|
|
assert result["status"] == "STOPPED_QUALITY_REGRESSION"
|
|
assert result["effect_comparison"]["error_rate"] == pytest.approx(-0.01)
|
|
|
|
improved = _evaluation(proposal, sign)
|
|
with pytest.raises(OptimizationError, match="OPTIMIZATION_EVALUATION_STATE_INVALID"):
|
|
controller.record_evaluation(
|
|
"proposal-1", improved, authenticated_context=_context()
|
|
)
|
|
assert controller.get_proposal(
|
|
"proposal-1", authenticated_context=_context()
|
|
)["status"] == "STOPPED_QUALITY_REGRESSION"
|
|
|
|
|
|
def test_eligible_candidate_promotion_is_another_governed_task_only(tmp_path):
|
|
controller, submitted, sign = _controller(tmp_path)
|
|
proposal = controller.propose(_proposal(), authenticated_context=_context())
|
|
result = controller.record_evaluation(
|
|
"proposal-1",
|
|
_evaluation(proposal, sign),
|
|
authenticated_context=_context(),
|
|
)
|
|
assert result["status"] == "ELIGIBLE"
|
|
|
|
promoted = controller.promote(
|
|
"proposal-1",
|
|
{
|
|
"evaluation_task_id": "evaluation-task-1",
|
|
"proposal_digest": proposal["proposal_digest"],
|
|
},
|
|
authenticated_context=_context(),
|
|
)
|
|
assert promoted["status"] == "PROMOTION_SUBMITTED"
|
|
assert len(submitted) == 2
|
|
promotion_task = submitted[-1]
|
|
assert promotion_task["action"] == "promote"
|
|
assert promotion_task["tenant_id"] == "tenant-a"
|
|
assert promotion_task["grant_id"] == "grant-123"
|
|
assert promotion_task["mission_id"] == "mission-456"
|
|
assert "changes" not in promotion_task
|
|
|
|
|
|
def test_restart_reads_persisted_baseline_candidate_and_effect(tmp_path):
|
|
path = tmp_path / "optimization.sqlite3"
|
|
submitted = []
|
|
sign, verify = _signed_receipt_verifier()
|
|
|
|
def submit_task(task):
|
|
submitted.append(task)
|
|
return {"task_id": "evaluation-task-1"}
|
|
|
|
first = OptimizationController(path, _envelope(), submit_task, verify_receipt=verify)
|
|
proposal = first.propose(_proposal(), authenticated_context=_context())
|
|
first.record_evaluation(
|
|
"proposal-1", _evaluation(proposal, sign), authenticated_context=_context()
|
|
)
|
|
|
|
restarted = OptimizationController(path, _envelope(), submit_task, verify_receipt=verify)
|
|
restored = restarted.get_proposal("proposal-1", authenticated_context=_context())
|
|
assert restored["baseline_metrics"] == {"error_rate": 0.04, "success_rate": 0.95}
|
|
assert restored["candidate_metrics"] == {"error_rate": 0.03, "success_rate": 0.97}
|
|
assert restored["effect_comparison"] == {
|
|
"error_rate": pytest.approx(0.01),
|
|
"success_rate": pytest.approx(0.02),
|
|
}
|