agentium-tracer 0.1.0__tar.gz

agentium_tracer-0.1.0/PKG-INFO

@@ -0,0 +1,6 @@
+Metadata-Version: 2.4
+Name: agentium-tracer
+Version: 0.1.0
+Summary: Python SDK for Agent Runtime Layer traces
+License: AGPL-3.0-or-later
+Requires-Python: >=3.10

agentium_tracer-0.1.0/agent_runtime_layer/__init__.py

@@ -0,0 +1,16 @@
+from agent_runtime_layer.capture import capture_command
+from agent_runtime_layer.client import AgentRuntimeClient
+from agent_runtime_layer.integrations.aider import capture_aider
+from agent_runtime_layer.trace import TraceEvent
+from agent_runtime_layer.tracer import AgentRuntimeTracer, context_hash, estimate_cost, prompt_hash
+
+__all__ = [
+    "AgentRuntimeClient",
+    "TraceEvent",
+    "capture_command",
+    "capture_aider",
+    "AgentRuntimeTracer",
+    "context_hash",
+    "estimate_cost",
+    "prompt_hash",
+]

agentium_tracer-0.1.0/agent_runtime_layer/budget.py

@@ -0,0 +1,195 @@
+"""Phase 1.8: Budget Governor — session cost and retry state tracker.
+
+Reads .agentium/config.yaml and enforces per-run budget caps and retry limits.
+Used by Claude Code and Codex hook handlers to block runaway runs.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import time
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any
+from uuid import uuid4
+
+
+DEFAULT_CONFIG = {
+    "max_cost_per_run": 0.10,
+    "max_retries_per_task": 5,
+    "alert_threshold": 0.05,
+    "token_limit_per_call": 200000,
+    "enabled": True,
+}
+
+
+@dataclass
+class BudgetState:
+    session_id: str
+    task_id: str | None = None
+    cost: float = 0.0
+    retries: int = 0
+    blocked: bool = False
+    block_reason: str = ""
+    started_at: float = field(default_factory=time.time)
+
+
+class BudgetGovernor:
+    """Session-scoped budget governor. One instance per agent session."""
+
+    def __init__(
+        self,
+        repo_path: Path | str = ".",
+        base_url: str = "http://localhost:8000/api",
+        session_id: str | None = None,
+    ) -> None:
+        self.repo_path = Path(repo_path)
+        self.base_url = base_url
+        self.session_id = session_id or f"bsess_{uuid4().hex[:12]}"
+        self._config = self._load_config()
+        self._state = BudgetState(session_id=self.session_id)
+
+    def _load_config(self) -> dict[str, Any]:
+        """Load .agentium/config.yaml or fall back to defaults."""
+        config_path = self.repo_path / ".agentium" / "config.yaml"
+        if config_path.exists():
+            try:
+                import yaml  # type: ignore[import]
+                with open(config_path, encoding="utf-8") as f:
+                    loaded = yaml.safe_load(f) or {}
+                return {**DEFAULT_CONFIG, **loaded}
+            except Exception:
+                pass
+        # Try JSON fallback
+        json_path = self.repo_path / ".agentium" / "config.json"
+        if json_path.exists():
+            try:
+                with open(json_path, encoding="utf-8") as f:
+                    loaded = json.load(f)
+                return {**DEFAULT_CONFIG, **loaded}
+            except Exception:
+                pass
+        return dict(DEFAULT_CONFIG)
+
+    @property
+    def enabled(self) -> bool:
+        return bool(self._config.get("enabled", True))
+
+    @property
+    def max_cost(self) -> float:
+        return float(self._config.get("max_cost_per_run", 0.10))
+
+    @property
+    def max_retries(self) -> int:
+        return int(self._config.get("max_retries_per_task", 5))
+
+    @property
+    def alert_threshold(self) -> float:
+        return float(self._config.get("alert_threshold", 0.05))
+
+    def set_task(self, task_id: str) -> None:
+        self._state.task_id = task_id
+
+    def add_cost(self, cost_dollars: float) -> None:
+        self._state.cost += cost_dollars
+
+    def increment_retry(self) -> None:
+        self._state.retries += 1
+
+    def check(self) -> tuple[bool, str]:
+        """Check if current session is within budget and retry limits.
+        Returns (allowed: bool, reason: str).
+        """
+        if not self.enabled:
+            return True, ""
+
+        if self._state.cost >= self.max_cost:
+            reason = (
+                f"Budget cap reached: ${self._state.cost:.4f} >= ${self.max_cost:.4f}. "
+                f"Run terminated by Agentium Budget Governor."
+            )
+            self._state.blocked = True
+            self._state.block_reason = reason
+            self._report_block("budget_cap", reason)
+            return False, reason
+
+        if self._state.retries >= self.max_retries:
+            reason = (
+                f"Retry limit reached: {self._state.retries} >= {self.max_retries} retries. "
+                f"Run terminated by Agentium Budget Governor."
+            )
+            self._state.blocked = True
+            self._state.block_reason = reason
+            self._report_block("retry_limit", reason)
+            return False, reason
+
+        if self._state.cost >= self.alert_threshold and self._state.cost < self.max_cost:
+            # Alert but don't block
+            pass
+
+        return True, ""
+
+    def _report_block(self, event_type: str, reason: str) -> None:
+        """Report a budget block event to the backend API."""
+        try:
+            import urllib.request
+            payload = json.dumps({
+                "event_id": f"bev_{uuid4().hex[:12]}",
+                "session_id": self.session_id,
+                "task_id": self._state.task_id,
+                "event_type": event_type,
+                "reason": reason,
+                "cost_at_block": self._state.cost,
+                "retries_at_block": self._state.retries,
+                "budget_limit": self.max_cost,
+                "retry_limit": self.max_retries,
+            }).encode("utf-8")
+            req = urllib.request.Request(
+                f"{self.base_url}/budget/events",
+                data=payload,
+                headers={"Content-Type": "application/json"},
+                method="POST",
+            )
+            urllib.request.urlopen(req, timeout=2)
+        except Exception:
+            pass  # Never block the agent because of reporting failures
+
+    @property
+    def current_cost(self) -> float:
+        return self._state.cost
+
+    @property
+    def current_retries(self) -> int:
+        return self._state.retries
+
+    @property
+    def summary(self) -> dict[str, Any]:
+        return {
+            "session_id": self.session_id,
+            "cost": self._state.cost,
+            "retries": self._state.retries,
+            "blocked": self._state.blocked,
+            "block_reason": self._state.block_reason,
+            "max_cost": self.max_cost,
+            "max_retries": self.max_retries,
+        }
+
+
+def write_default_config(repo_path: Path | str = ".") -> Path:
+    """Write a default .agentium/config.yaml to the given repo path."""
+    config_dir = Path(repo_path) / ".agentium"
+    config_dir.mkdir(parents=True, exist_ok=True)
+    config_path = config_dir / "config.yaml"
+    if not config_path.exists():
+        config_path.write_text(
+            "# Agentium Phase 1.8 Budget Governor configuration\n"
+            "# All settings apply per agent session.\n\n"
+            "max_cost_per_run: 0.10        # Stop run if cost exceeds this (dollars)\n"
+            "max_retries_per_task: 5       # Stop run after this many retries\n"
+            "alert_threshold: 0.05         # Log a warning at this cost (dollars)\n"
+            "token_limit_per_call: 200000  # Warn if a single call exceeds this token count\n"
+            "enabled: true                 # Set to false to disable all budget enforcement\n",
+            encoding="utf-8",
+        )
+    return config_path

agentium_tracer-0.1.0/agent_runtime_layer/capture.py

@@ -0,0 +1,246 @@
+import json
+import subprocess
+import time
+from dataclasses import dataclass
+from datetime import datetime, timezone
+from pathlib import Path
+from uuid import uuid4
+
+from agent_runtime_layer.client import AgentRuntimeClient
+from agent_runtime_layer.redaction import redact_text, redact_value
+
+
+DEFAULT_TRACE_DIR = Path(".agent-runtime") / "traces"
+SUMMARY_LIMIT = 4000
+
+
+def utc_now() -> str:
+    return datetime.now(timezone.utc).isoformat(timespec="milliseconds").replace("+00:00", "Z")
+
+
+def summarize_stream(value: str, limit: int = SUMMARY_LIMIT) -> str:
+    redacted = redact_text(value)
+    if len(redacted) <= limit:
+        return redacted
+    omitted = len(redacted) - limit
+    return f"{redacted[:limit]}\n...[truncated {omitted} chars]"
+
+
+def command_display(command: list[str]) -> str:
+    return " ".join(command)
+
+
+def safe_git_diff_summary() -> str | None:
+    try:
+        result = subprocess.run(
+            ["git", "diff", "--stat", "--no-ext-diff"],
+            capture_output=True,
+            text=True,
+            timeout=10,
+            check=False,
+        )
+    except (OSError, subprocess.SubprocessError):
+        return None
+    output = (result.stdout or result.stderr or "").strip()
+    return summarize_stream(output, 2000) if output else None
+
+
+@dataclass
+class TraceCaptureResult:
+    task_id: str
+    trace_path: Path
+    exit_code: int
+    uploaded: bool = False
+    upload_response: dict | None = None
+
+
+class LocalTraceCapture:
+    def __init__(
+        self,
+        name: str,
+        command: list[str],
+        project_id: str = "default",
+        trace_dir: Path = DEFAULT_TRACE_DIR,
+        capture_full_logs: bool = False,
+        capture_diff: bool = False,
+    ) -> None:
+        if not command:
+            raise ValueError("A command is required after --")
+        self.name = name
+        self.command = command
+        self.project_id = project_id
+        self.trace_dir = trace_dir
+        self.capture_full_logs = capture_full_logs
+        self.capture_diff = capture_diff
+        self.task_id = f"task_local_{uuid4().hex[:12]}"
+        self.task_span_id = f"span_task_{uuid4().hex[:8]}"
+        self.tool_span_id = f"span_tool_{uuid4().hex[:8]}"
+
+    def event(
+        self,
+        event_type: str,
+        span_id: str,
+        name: str,
+        attributes: dict | None = None,
+        payload: dict | None = None,
+        parent_span_id: str | None = None,
+        timestamp: str | None = None,
+    ) -> dict:
+        return {
+            "event_id": f"evt_{uuid4().hex[:12]}",
+            "task_id": self.task_id,
+            "timestamp": timestamp or utc_now(),
+            "event_type": event_type,
+            "span_id": span_id,
+            "parent_span_id": parent_span_id,
+            "name": name,
+            "attributes": redact_value(attributes or {}),
+            "payload": redact_value(payload or {}),
+        }
+
+    def run(self) -> tuple[dict, int]:
+        started_at = utc_now()
+        start_time = time.perf_counter()
+        display = command_display(self.command)
+        events = [
+            self.event(
+                "task_start",
+                self.task_span_id,
+                "task_start",
+                payload={
+                    "goal": self.name,
+                    "agent_type": "local_command",
+                    "latency_slo_seconds": None,
+                },
+                timestamp=started_at,
+            ),
+            self.event(
+                "tool_call_start",
+                self.tool_span_id,
+                "local_command",
+                attributes={
+                    "tool_name": "terminal",
+                    "command": display,
+                    "risk_level": "user_provided",
+                },
+                parent_span_id=self.task_span_id,
+                timestamp=started_at,
+            ),
+        ]
+
+        completed = subprocess.run(
+            self.command,
+            capture_output=True,
+            text=True,
+            check=False,
+        )
+        ended_at = utc_now()
+        duration_ms = int((time.perf_counter() - start_time) * 1000)
+        status = "success" if completed.returncode == 0 else "failed"
+        stdout_preview = summarize_stream(completed.stdout or "")
+        stderr_preview = summarize_stream(completed.stderr or "")
+        terminal_payload = {
+            "stdout_preview": stdout_preview,
+            "stderr_preview": stderr_preview,
+        }
+        if self.capture_full_logs:
+            terminal_payload["stdout"] = redact_text(completed.stdout or "")
+            terminal_payload["stderr"] = redact_text(completed.stderr or "")
+        if self.capture_diff:
+            diff_summary = safe_git_diff_summary()
+            if diff_summary:
+                terminal_payload["git_diff_summary"] = diff_summary
+
+        events.extend(
+            [
+                self.event(
+                    "terminal_event",
+                    f"span_terminal_{uuid4().hex[:8]}",
+                    "local_command_output",
+                    attributes={
+                        "command": display,
+                        "duration_ms": duration_ms,
+                        "exit_code": completed.returncode,
+                        "stdout_preview": stdout_preview,
+                        "stderr_preview": stderr_preview,
+                    },
+                    payload=terminal_payload,
+                    parent_span_id=self.tool_span_id,
+                    timestamp=ended_at,
+                ),
+                self.event(
+                    "tool_call_end",
+                    self.tool_span_id,
+                    "local_command_end",
+                    attributes={
+                        "latency_ms": duration_ms,
+                        "status": status,
+                        "exit_code": completed.returncode,
+                    },
+                    parent_span_id=self.task_span_id,
+                    timestamp=ended_at,
+                ),
+                self.event(
+                    "task_end",
+                    self.task_span_id,
+                    "task_end",
+                    payload={
+                        "status": "completed" if completed.returncode == 0 else "failed",
+                        "summary": f"Command exited with code {completed.returncode}.",
+                    },
+                    timestamp=ended_at,
+                ),
+            ]
+        )
+
+        trace = {
+            "project_id": self.project_id,
+            "task": {
+                "task_id": self.task_id,
+                "project_id": self.project_id,
+                "goal": self.name,
+                "agent_type": "local_command",
+                "budget_dollars": None,
+                "latency_slo_seconds": None,
+            },
+            "events": events,
+        }
+        return trace, completed.returncode
+
+    def write(self, trace: dict) -> Path:
+        self.trace_dir.mkdir(parents=True, exist_ok=True)
+        trace_path = self.trace_dir / f"{self.task_id}.json"
+        trace_path.write_text(json.dumps(trace, indent=2), encoding="utf-8")
+        return trace_path
+
+
+def capture_command(
+    name: str,
+    command: list[str],
+    project_id: str = "default",
+    trace_dir: Path = DEFAULT_TRACE_DIR,
+    capture_full_logs: bool = False,
+    capture_diff: bool = False,
+    upload: bool = False,
+    base_url: str = "http://localhost:8000/api",
+) -> TraceCaptureResult:
+    capture = LocalTraceCapture(
+        name=name,
+        command=command,
+        project_id=project_id,
+        trace_dir=trace_dir,
+        capture_full_logs=capture_full_logs,
+        capture_diff=capture_diff,
+    )
+    trace, exit_code = capture.run()
+    trace_path = capture.write(trace)
+    upload_response = None
+    if upload:
+        upload_response = AgentRuntimeClient(base_url).import_trace_file(trace_path)
+    return TraceCaptureResult(
+        task_id=capture.task_id,
+        trace_path=trace_path,
+        exit_code=exit_code,
+        uploaded=upload,
+        upload_response=upload_response,
+    )