agent_os_kernel 3.1.0__py3-none-any.whl

agent_os/policies/bridge.py

@@ -0,0 +1,169 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+"""
+Bridge between legacy GovernancePolicy and declarative PolicyDocument.
+
+Provides bidirectional conversion so existing code continues to work while
+new code can use the declarative policy format.
+"""
+
+from __future__ import annotations
+
+from ..integrations.base import GovernancePolicy, PatternType
+from .schema import (
+    PolicyAction,
+    PolicyCondition,
+    PolicyDefaults,
+    PolicyDocument,
+    PolicyOperator,
+    PolicyRule,
+)
+
+
+def governance_to_document(policy: GovernancePolicy) -> PolicyDocument:
+    """Convert an existing GovernancePolicy to a declarative PolicyDocument."""
+    rules: list[PolicyRule] = []
+    priority = 100
+
+    # Token limit rule
+    rules.append(
+        PolicyRule(
+            name="max_tokens",
+            condition=PolicyCondition(
+                field="token_count",
+                operator=PolicyOperator.GT,
+                value=policy.max_tokens,
+            ),
+            action=PolicyAction.DENY,
+            priority=priority,
+            message=f"Token count exceeds limit of {policy.max_tokens}",
+        )
+    )
+    priority -= 1
+
+    # Tool call limit rule
+    rules.append(
+        PolicyRule(
+            name="max_tool_calls",
+            condition=PolicyCondition(
+                field="tool_call_count",
+                operator=PolicyOperator.GT,
+                value=policy.max_tool_calls,
+            ),
+            action=PolicyAction.DENY,
+            priority=priority,
+            message=f"Tool call count exceeds limit of {policy.max_tool_calls}",
+        )
+    )
+    priority -= 1
+
+    # Allowed tools rule (if a tool is not in the allowed list, deny)
+    if policy.allowed_tools:
+        rules.append(
+            PolicyRule(
+                name="allowed_tools",
+                condition=PolicyCondition(
+                    field="tool_name",
+                    operator=PolicyOperator.IN,
+                    value=policy.allowed_tools,
+                ),
+                action=PolicyAction.ALLOW,
+                priority=priority,
+                message="Tool is in the allowed list",
+            )
+        )
+        priority -= 1
+
+    # Blocked patterns
+    for i, pattern in enumerate(policy.blocked_patterns):
+        if isinstance(pattern, str):
+            pat_str = pattern
+            operator = PolicyOperator.CONTAINS
+        else:
+            pat_str, pat_type = pattern
+            if pat_type == PatternType.REGEX:
+                operator = PolicyOperator.MATCHES
+            elif pat_type == PatternType.GLOB:
+                operator = PolicyOperator.MATCHES
+            else:
+                operator = PolicyOperator.CONTAINS
+
+        rules.append(
+            PolicyRule(
+                name=f"blocked_pattern_{i}",
+                condition=PolicyCondition(
+                    field="content",
+                    operator=operator,
+                    value=pat_str,
+                ),
+                action=PolicyAction.BLOCK,
+                priority=priority,
+                message=f"Content matches blocked pattern: {pat_str}",
+            )
+        )
+        priority -= 1
+
+    # Confidence threshold rule
+    rules.append(
+        PolicyRule(
+            name="confidence_threshold",
+            condition=PolicyCondition(
+                field="confidence",
+                operator=PolicyOperator.LT,
+                value=policy.confidence_threshold,
+            ),
+            action=PolicyAction.DENY,
+            priority=priority,
+            message=f"Confidence below threshold of {policy.confidence_threshold}",
+        )
+    )
+
+    return PolicyDocument(
+        version=policy.version,
+        name=policy.name,
+        description=f"Auto-converted from GovernancePolicy '{policy.name}'",
+        rules=rules,
+        defaults=PolicyDefaults(
+            action=PolicyAction.ALLOW,
+            max_tokens=policy.max_tokens,
+            max_tool_calls=policy.max_tool_calls,
+            confidence_threshold=policy.confidence_threshold,
+        ),
+    )
+
+
+def document_to_governance(doc: PolicyDocument) -> GovernancePolicy:
+    """Convert a declarative PolicyDocument back to a GovernancePolicy."""
+    max_tokens = doc.defaults.max_tokens
+    max_tool_calls = doc.defaults.max_tool_calls
+    confidence_threshold = doc.defaults.confidence_threshold
+    allowed_tools: list[str] = []
+    blocked_patterns: list[str | tuple[str, PatternType]] = []
+
+    for rule in doc.rules:
+        cond = rule.condition
+
+        if rule.name == "max_tokens" and cond.field == "token_count":
+            max_tokens = int(cond.value)
+        elif rule.name == "max_tool_calls" and cond.field == "tool_call_count":
+            max_tool_calls = int(cond.value)
+        elif rule.name == "allowed_tools" and cond.field == "tool_name":
+            if isinstance(cond.value, list):
+                allowed_tools = list(cond.value)
+        elif rule.name.startswith("blocked_pattern_") and cond.field == "content":
+            if cond.operator == PolicyOperator.MATCHES:
+                blocked_patterns.append((str(cond.value), PatternType.REGEX))
+            elif cond.operator == PolicyOperator.CONTAINS:
+                blocked_patterns.append(str(cond.value))
+        elif rule.name == "confidence_threshold" and cond.field == "confidence":
+            confidence_threshold = float(cond.value)
+
+    return GovernancePolicy(
+        name=doc.name,
+        max_tokens=max_tokens,
+        max_tool_calls=max_tool_calls,
+        allowed_tools=allowed_tools,
+        blocked_patterns=blocked_patterns,
+        confidence_threshold=confidence_threshold,
+        version=doc.version,
+    )

agent_os/policies/budget.py

@@ -0,0 +1,85 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+"""Budget policy rules for token, cost, and tool-call limits."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Optional
+
+
+@dataclass
+class BudgetPolicy:
+    """Resource consumption limits for a governed task."""
+
+    max_tokens: Optional[int] = None
+    max_tool_calls: Optional[int] = None
+    max_cost_usd: Optional[float] = None
+    max_duration_seconds: Optional[float] = None
+
+
+@dataclass
+class BudgetTracker:
+    """Tracks resource consumption against a BudgetPolicy.
+
+    Example::
+
+        policy = BudgetPolicy(max_tokens=8000, max_tool_calls=20)
+        tracker = BudgetTracker(policy)
+        tracker.record_tokens(1500)
+        tracker.record_tool_call()
+        if tracker.is_exceeded():
+            print(tracker.exceeded_reasons())
+    """
+
+    policy: BudgetPolicy
+    tokens_used: int = 0
+    tool_calls_used: int = 0
+    cost_usd_used: float = 0.0
+    duration_seconds_used: float = 0.0
+
+    def record_tokens(self, count: int) -> None:
+        self.tokens_used += count
+
+    def record_tool_call(self) -> None:
+        self.tool_calls_used += 1
+
+    def record_cost(self, amount: float) -> None:
+        self.cost_usd_used += amount
+
+    def record_duration(self, seconds: float) -> None:
+        self.duration_seconds_used += seconds
+
+    def is_exceeded(self) -> bool:
+        return len(self.exceeded_reasons()) > 0
+
+    def exceeded_reasons(self) -> list[str]:
+        reasons = []
+        p = self.policy
+        if p.max_tokens is not None and self.tokens_used > p.max_tokens:
+            reasons.append(f"tokens: {self.tokens_used}/{p.max_tokens}")
+        if p.max_tool_calls is not None and self.tool_calls_used > p.max_tool_calls:
+            reasons.append(f"tool_calls: {self.tool_calls_used}/{p.max_tool_calls}")
+        if p.max_cost_usd is not None and self.cost_usd_used > p.max_cost_usd:
+            reasons.append(f"cost: ${self.cost_usd_used:.2f}/${p.max_cost_usd:.2f}")
+        if p.max_duration_seconds is not None and self.duration_seconds_used > p.max_duration_seconds:
+            reasons.append(f"duration: {self.duration_seconds_used:.1f}s/{p.max_duration_seconds:.1f}s")
+        return reasons
+
+    def remaining(self) -> dict[str, float | int | None]:
+        p = self.policy
+        return {
+            "tokens": (p.max_tokens - self.tokens_used) if p.max_tokens else None,
+            "tool_calls": (p.max_tool_calls - self.tool_calls_used) if p.max_tool_calls else None,
+            "cost_usd": round(p.max_cost_usd - self.cost_usd_used, 4) if p.max_cost_usd else None,
+            "duration_seconds": round(p.max_duration_seconds - self.duration_seconds_used, 1) if p.max_duration_seconds else None,
+        }
+
+    def utilization(self) -> dict[str, float | None]:
+        p = self.policy
+        return {
+            "tokens": round(self.tokens_used / p.max_tokens, 3) if p.max_tokens else None,
+            "tool_calls": round(self.tool_calls_used / p.max_tool_calls, 3) if p.max_tool_calls else None,
+            "cost_usd": round(self.cost_usd_used / p.max_cost_usd, 3) if p.max_cost_usd else None,
+            "duration_seconds": round(self.duration_seconds_used / p.max_duration_seconds, 3) if p.max_duration_seconds else None,
+        }

agent_os/policies/cli.py

@@ -0,0 +1,294 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+"""
+Policy CLI - Manage and validate security policies.
+
+Commands:
+- validate: Check a policy file against the PolicyDocument schema
+- test: Run security scenarios against policy definitions
+- diff: Compare two policy versions structurally
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import re
+import sys
+from pathlib import Path
+from typing import Any, List
+
+import yaml
+from pydantic import ValidationError
+
+from agent_os.policies.schema import PolicyDocument
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def _load_file(path: Path) -> dict:
+    """Load a YAML or JSON file and return the parsed dict."""
+    text = path.read_text(encoding="utf-8")
+    if path.suffix == ".json":
+        return json.loads(text)
+    return yaml.safe_load(text)
+
+
+def _evaluate_condition(condition: dict, context: dict) -> bool:
+    """Evaluate a single policy condition against a context dict."""
+    field = condition.get("field", "")
+    operator = condition.get("operator", "eq")
+    value = condition.get("value")
+
+    if field not in context:
+        return False
+
+    ctx_value = context[field]
+
+    if operator == "eq":
+        return ctx_value == value
+    if operator == "ne":
+        return ctx_value != value
+    if operator == "gt":
+        return ctx_value > value
+    if operator == "lt":
+        return ctx_value < value
+    if operator == "gte":
+        return ctx_value >= value
+    if operator == "lte":
+        return ctx_value <= value
+    if operator == "in":
+        return ctx_value in value
+    if operator == "contains":
+        return value in ctx_value
+    if operator == "matches":
+        return bool(re.match(str(value), str(ctx_value)))
+
+    return False
+
+
+def _resolve_action(rules: list[dict], defaults: dict | None, context: dict) -> str:
+    """Find the first matching rule (highest priority) and return its action."""
+    sorted_rules = sorted(rules, key=lambda r: r.get("priority", 0), reverse=True)
+
+    for rule in sorted_rules:
+        condition = rule.get("condition", {})
+        if _evaluate_condition(condition, context):
+            return rule.get("action", "allow")
+
+    if defaults:
+        return defaults.get("action", "allow")
+    return "allow"
+
+
+# ---------------------------------------------------------------------------
+# Sub-commands
+# ---------------------------------------------------------------------------
+
+
+def _cmd_validate(args: argparse.Namespace) -> int:
+    """Validate a policy file against the PolicyDocument Pydantic schema."""
+    path = Path(args.file)
+
+    if not path.exists():
+        print("ERROR: File not found: " + str(path), file=sys.stderr)
+        return 2
+
+    try:
+        data = _load_file(path)
+    except Exception:
+        print("ERROR: Could not parse file: " + str(path), file=sys.stderr)
+        return 2
+
+    if data is None:
+        data = {}
+
+    try:
+        PolicyDocument.model_validate(data)
+    except ValidationError as exc:
+        print("FAIL: " + str(exc), file=sys.stderr)
+        return 1
+
+    print("OK")
+    return 0
+
+
+def _cmd_test(args: argparse.Namespace) -> int:
+    """Run test scenarios against a policy."""
+    policy_path = Path(args.policy)
+    scenarios_path = Path(args.scenarios)
+
+    if not policy_path.exists():
+        print("ERROR: Policy file not found: " + str(policy_path), file=sys.stderr)
+        return 2
+
+    if not scenarios_path.exists():
+        print("ERROR: Scenarios file not found: " + str(scenarios_path), file=sys.stderr)
+        return 2
+
+    try:
+        policy_data = _load_file(policy_path)
+        scenarios_data = _load_file(scenarios_path)
+    except Exception:
+        print("ERROR: Could not parse files", file=sys.stderr)
+        return 2
+
+    rules = (policy_data or {}).get("rules", [])
+    defaults = (policy_data or {}).get("defaults", {})
+    scenarios: list[dict[str, Any]] = (scenarios_data or {}).get("scenarios", [])
+
+    if not scenarios:
+        print("ERROR: No scenarios provided", file=sys.stderr)
+        return 2
+
+    passed = 0
+    total = len(scenarios)
+
+    for scenario in scenarios:
+        context = scenario.get("context", {})
+        expected_action = scenario.get("expected_action")
+        expected_allowed = scenario.get("expected_allowed")
+
+        actual_action = _resolve_action(rules, defaults, context)
+        actual_allowed = actual_action == "allow"
+
+        ok = True
+        if expected_action is not None and actual_action != expected_action:
+            ok = False
+        if expected_allowed is not None and actual_allowed != expected_allowed:
+            ok = False
+
+        if ok:
+            passed += 1
+        else:
+            print(
+                f"FAIL: {scenario.get('name', 'unnamed')}: "
+                f"expected {expected_action}, got {actual_action}",
+                file=sys.stderr,
+            )
+
+    print(f"{passed}/{total} scenarios passed")
+    return 0 if passed == total else 1
+
+
+def _cmd_diff(args: argparse.Namespace) -> int:
+    """Compare two policy files structurally."""
+    base_path = Path(args.base)
+    target_path = Path(args.target)
+
+    if not base_path.exists():
+        print("ERROR: File not found: " + str(base_path), file=sys.stderr)
+        return 2
+    if not target_path.exists():
+        print("ERROR: File not found: " + str(target_path), file=sys.stderr)
+        return 2
+
+    try:
+        base_data = _load_file(base_path)
+        target_data = _load_file(target_path)
+    except Exception:
+        print("ERROR: Could not parse files", file=sys.stderr)
+        return 2
+
+    if base_data == target_data:
+        print("No differences")
+        return 0
+
+    diffs: list[str] = []
+
+    # --- rules ---
+    base_rules = {r["name"]: r for r in (base_data.get("rules") or [])}
+    target_rules = {r["name"]: r for r in (target_data.get("rules") or [])}
+
+    for name in target_rules:
+        if name not in base_rules:
+            diffs.append(f"rule added: {name}")
+
+    for name in base_rules:
+        if name not in target_rules:
+            diffs.append(f"rule removed: {name}")
+
+    for name in base_rules:
+        if name in target_rules:
+            br = base_rules[name]
+            tr = target_rules[name]
+            all_keys = set(br.keys()) | set(tr.keys())
+            for key in sorted(all_keys):
+                if br.get(key) != tr.get(key):
+                    diffs.append(f"rule {name}: {key}: {br.get(key)} -> {tr.get(key)}")
+
+    # --- defaults ---
+    base_defaults = (base_data.get("defaults") or {})
+    target_defaults = (target_data.get("defaults") or {})
+    all_def_keys = set(base_defaults.keys()) | set(target_defaults.keys())
+    for key in sorted(all_def_keys):
+        if base_defaults.get(key) != target_defaults.get(key):
+            diffs.append(f"defaults: {key}: {base_defaults.get(key)} -> {target_defaults.get(key)}")
+
+    # --- other top-level fields ---
+    for key in sorted(set(base_data.keys()) | set(target_data.keys())):
+        if key in ("rules", "defaults"):
+            continue
+        if base_data.get(key) != target_data.get(key):
+            diffs.append(f"{key}: {base_data.get(key)} -> {target_data.get(key)}")
+
+    for d in diffs:
+        print(d)
+
+    return 1 if diffs else 0
+
+
+# ---------------------------------------------------------------------------
+# Parser / entry-point
+# ---------------------------------------------------------------------------
+
+
+def build_parser() -> argparse.ArgumentParser:
+    """Build the argument parser."""
+    parser = argparse.ArgumentParser(
+        prog="policies-cli",
+        description="Agent OS Policy Management Tools",
+    )
+    subparsers = parser.add_subparsers(dest="command", help="Command to run")
+
+    # -- validate -----------------------------------------------------------
+    val_parser = subparsers.add_parser("validate", help="Validate a policy file")
+    val_parser.add_argument("file", help="Policy file to validate")
+
+    # -- test ---------------------------------------------------------------
+    test_parser = subparsers.add_parser("test", help="Test policy against scenarios")
+    test_parser.add_argument("policy", help="Policy file to test")
+    test_parser.add_argument("scenarios", help="Scenarios file (YAML/JSON)")
+
+    # -- diff ---------------------------------------------------------------
+    diff_parser = subparsers.add_parser("diff", help="Compare two policy files")
+    diff_parser.add_argument("base", help="Base policy file")
+    diff_parser.add_argument("target", help="Target policy file")
+
+    return parser
+
+
+def main(argv: List[str] | None = None) -> int:
+    """CLI entry point."""
+    parser = build_parser()
+    args = parser.parse_args(argv)
+
+    if not args.command:
+        parser.print_help()
+        return 2
+
+    if args.command == "validate":
+        return _cmd_validate(args)
+    if args.command == "test":
+        return _cmd_test(args)
+    if args.command == "diff":
+        return _cmd_diff(args)
+
+    return 2
+
+
+if __name__ == "__main__":
+    sys.exit(main())