cleveragents-core/robot/helper_autonomy_guardrails.py

"""Helper script for Robot Framework autonomy guardrails tests."""

from __future__ import annotations

import sys
from datetime import UTC, datetime, timedelta
from pathlib import Path
from typing import Any

sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))

from cleveragents.application.services.autonomy_guardrail_service import (
    AutonomyGuardrailService,
)
from cleveragents.domain.models.core.autonomy_guardrails import (
    ActorLimits,
    AutonomyGuardrails,
    GuardrailAuditEntry,
    GuardrailAuditTrail,
    GuardrailEventType,
    GuardrailResult,
)


def _test_step_allow() -> None:
    """Step limit allows within bounds."""
    g = AutonomyGuardrails(max_steps=10)
    g.step_count = 5
    allowed, _ = g.check_step_limit()
    assert allowed is True
    print("step-allow-ok")


def _test_step_block() -> None:
    """Step limit blocks at max."""
    g = AutonomyGuardrails(max_steps=5)
    g.step_count = 5
    allowed, reason = g.check_step_limit()
    assert allowed is False
    assert reason is not None
    assert "Step limit" in reason
    print("step-block-ok")


def _test_budget_allow() -> None:
    """Budget allows within limit."""
    g = AutonomyGuardrails(tool_budget=100.0)
    allowed, _ = g.check_tool_budget(50.0)
    assert allowed is True
    print("budget-allow-ok")


def _test_budget_block() -> None:
    """Budget blocks when exceeded."""
    g = AutonomyGuardrails(tool_budget=10.0, budget_spent=8.0)
    allowed, reason = g.check_tool_budget(5.0)
    assert allowed is False
    assert reason is not None
    assert "Budget exceeded" in reason
    print("budget-block-ok")


def _test_confirm_required() -> None:
    """Confirmation required for listed operation."""
    g = AutonomyGuardrails(required_confirmations=["deploy", "delete"])
    required, reason = g.check_confirmation_required("deploy")
    assert required is True
    assert reason is not None
    not_required, _ = g.check_confirmation_required("read")
    assert not_required is False
    print("confirm-required-ok")


def _test_audit_trail() -> None:
    """Audit trail records enforcement events."""
    svc = AutonomyGuardrailService()
    svc.configure_guardrails(
        "test-plan",
        AutonomyGuardrails(max_steps=5, tool_budget=50.0),
    )
    svc.check_step_limit("test-plan", 3)
    svc.check_tool_budget("test-plan", 10.0)
    trail = svc.get_audit_trail("test-plan")
    assert len(trail.entries) == 2
    assert trail.entries[0].guard_name == "step_limit"
    assert trail.entries[1].guard_name == "tool_budget"
    print("audit-trail-ok")


def _test_metadata() -> None:
    """Service can serialize and restore from metadata."""
    svc = AutonomyGuardrailService()
    svc.configure_guardrails(
        "meta-plan",
        AutonomyGuardrails(max_steps=10, tool_budget=25.0),
    )
    svc.check_step_limit("meta-plan", 2)

    metadata = svc.get_audit_trail_as_metadata("meta-plan")
    assert "guardrail_audit_trail" in metadata
    assert "autonomy_guardrails" in metadata

    svc2 = AutonomyGuardrailService()
    svc2.load_from_metadata("meta-plan", metadata)
    g = svc2.get_guardrails("meta-plan")
    assert g is not None
    assert g.max_steps == 10
    print("metadata-ok")


def _test_summary() -> None:
    """Summary of autonomy guardrail features."""
    g = AutonomyGuardrails(
        max_steps=100,
        tool_budget=500.0,
        required_confirmations=["deploy", "delete"],
    )
    assert g.max_steps == 100
    assert g.tool_budget == 500.0
    assert len(g.required_confirmations) == 2
    g.increment_step()
    assert g.step_count == 1
    g.record_cost(10.0)
    assert g.budget_spent == 10.0
    trail = GuardrailAuditTrail()
    entry = GuardrailAuditEntry(
        event_type=GuardrailEventType.STEP_ALLOWED,
        guard_name="step_limit",
        result=GuardrailResult.ALLOWED,
    )
    trail.add_entry(entry)
    assert len(trail.entries) == 1
    print("summary-ok")


def _test_wall_clock_allow() -> None:
    """Wall-clock allows when within limit."""
    g = AutonomyGuardrails(max_wall_clock_seconds=600.0)
    g.mark_started()
    allowed, _ = g.check_wall_clock()
    assert allowed is True
    print("wall-clock-allow-ok")


def _test_wall_clock_block() -> None:
    """Wall-clock blocks when expired."""
    g = AutonomyGuardrails(max_wall_clock_seconds=0.001)
    g.start_time = (datetime.now(tz=UTC) - timedelta(seconds=10)).isoformat()
    allowed, reason = g.check_wall_clock()
    assert allowed is False
    assert reason is not None
    assert "Wall-clock" in reason
    print("wall-clock-block-ok")


def _test_actor_limit() -> None:
    """Actor tool-call limit enforcement."""
    g = AutonomyGuardrails(
        actor_limits=ActorLimits(max_tool_calls_per_invocation=5),
    )
    allowed, _ = g.check_actor_tool_calls(3)
    assert allowed is True
    blocked, reason = g.check_actor_tool_calls(5)
    assert blocked is False
    assert reason is not None
    assert "Actor" in reason
    print("actor-limit-ok")


def _test_eviction() -> None:
    """Audit trail evicts oldest entries when full."""
    trail = GuardrailAuditTrail(max_entries=3)
    for i in range(5):
        entry = GuardrailAuditEntry(
            event_type=GuardrailEventType.STEP_ALLOWED,
            guard_name=f"step_{i}",
            result=GuardrailResult.ALLOWED,
        )
        trail.add_entry(entry)
    assert len(trail.entries) == 3
    # Oldest two (step_0, step_1) should be evicted
    assert trail.entries[0].guard_name == "step_2"
    assert trail.allowed_count == 3
    print("eviction-ok")


def _test_zero_budget() -> None:
    """Zero budget blocks any positive cost."""
    g = AutonomyGuardrails(tool_budget=0.0)
    allowed, _ = g.check_tool_budget(0.0)
    assert allowed is True
    blocked, reason = g.check_tool_budget(0.01)
    assert blocked is False
    assert reason is not None
    print("zero-budget-ok")


if __name__ == "__main__":
    cmd = sys.argv[1] if len(sys.argv) > 1 else "summary"
    dispatch: dict[str, Any] = {
        "step-allow": _test_step_allow,
        "step-block": _test_step_block,
        "budget-allow": _test_budget_allow,
        "budget-block": _test_budget_block,
        "confirm-required": _test_confirm_required,
        "audit-trail": _test_audit_trail,
        "metadata": _test_metadata,
        "wall-clock-allow": _test_wall_clock_allow,
        "wall-clock-block": _test_wall_clock_block,
        "actor-limit": _test_actor_limit,
        "eviction": _test_eviction,
        "zero-budget": _test_zero_budget,
        "summary": _test_summary,
    }
    fn = dispatch.get(cmd)
    if fn:
        fn()
    else:
        print(f"Unknown command: {cmd}", file=sys.stderr)
        sys.exit(1)