agentcap-cli 0.1.0__py3-none-any.whl

agentcap/__init__.py

@@ -0,0 +1,7 @@
+"""AgentCap - local, serverless spend-cap proxy for AI agents."""
+
+from __future__ import annotations
+
+__version__ = "0.1.0"
+
+__all__ = ["__version__"]

agentcap/alerts.py

@@ -0,0 +1,312 @@
+"""Alert fan-out for spend and safety events.
+
+Every alert goes to the terminal (stderr, red).  Pro users can additionally
+receive webhook and Telegram notifications; those are delivered from a
+background thread so a slow endpoint never blocks the proxy.
+
+Payloads are built from a strict field whitelist and pass through
+:func:`redact` before leaving the process, so prompts, response text and
+credentials can never be included.
+"""
+
+from __future__ import annotations
+
+import os
+import re
+import sys
+import threading
+from datetime import datetime, timezone
+from typing import Any, Callable, Iterable, Mapping, TextIO
+
+import httpx
+
+RED = "\x1b[31m"
+YELLOW = "\x1b[33m"
+RESET = "\x1b[0m"
+
+TERMINAL_EVENTS = frozenset(
+    {
+        "cap_threshold",
+        "cap_exceeded",
+        "loop_detected",
+        "usage_missing",
+        "overshoot",
+        "output_clamped",
+        "request_too_large",
+    }
+)
+
+# Alerts that must never be repeated for one run (a run has one kill and one
+# budget-clamp notice, no matter how many later requests are rejected).
+ONCE_EVENTS = frozenset({"cap_exceeded", "loop_detected", "output_clamped"})
+
+# Only these fields may appear in an alert payload.  Anything else is dropped.
+ALLOWED_FIELDS = frozenset(
+    {
+        "type",
+        "reason",
+        "scope",
+        "cap_usd",
+        "spent_usd",
+        "needed_usd",
+        "cost_usd",
+        "reserved_usd",
+        "run_spent_usd",
+        "charged_usd",
+        "model",
+        "provider",
+        "threshold",
+        "window_s",
+        "request_hash",
+        "run_id",
+        "allowed_output_tokens",
+        "remaining_usd",
+    }
+)
+
+_SECRET_PATTERNS = (
+    re.compile(r"sk-[A-Za-z0-9_\-]{4,}"),
+    re.compile(r"(?i)bearer\s+[A-Za-z0-9._\-]+"),
+    re.compile(r"(?i)(x-goog-api-key|api[_-]?key|authorization)\s*[:=]\s*\S+"),
+)
+
+
+def redact(text: str) -> str:
+    """Scrub credential-shaped substrings from a string."""
+
+    if not isinstance(text, str):
+        return text
+    result = text
+    for pattern in _SECRET_PATTERNS:
+        result = pattern.sub("[redacted]", result)
+    return result
+
+
+def sanitize(payload: Mapping[str, Any]) -> dict[str, Any]:
+    """Keep only whitelisted fields and redact every string value."""
+
+    clean: dict[str, Any] = {}
+    for key, value in payload.items():
+        if key not in ALLOWED_FIELDS:
+            continue
+        if isinstance(value, str):
+            clean[key] = redact(value)
+        elif isinstance(value, (int, float, bool)) or value is None:
+            clean[key] = value
+        else:
+            clean[key] = redact(str(value))
+    return clean
+
+
+def format_alert(event: Mapping[str, Any]) -> str:
+    """One-line, human-readable terminal message for an event."""
+
+    kind = str(event.get("type", "event"))
+    run = event.get("run_id")
+    prefix = "agentcap" + (f" [{redact(str(run))}]" if run else "")
+
+    if kind == "cap_threshold":
+        threshold = event.get("threshold", "?")
+        return (
+            f"{prefix}: {threshold}% of the ${event.get('cap_usd', '?')} run cap "
+            f"used (${event.get('spent_usd', '?')})."
+        )
+    if kind == "cap_exceeded":
+        return (
+            f"{prefix}: RUN CAP HIT (scope={event.get('scope', 'run')}, "
+            f"cap=${event.get('cap_usd', '?')}, spent=${event.get('spent_usd', '?')}). "
+            "Stopping the agent."
+        )
+    if kind == "loop_detected":
+        return (
+            f"{prefix}: runaway loop detected (identical request "
+            f"{event.get('threshold', '?')} times within "
+            f"{event.get('window_s', '?')}s). Stopping the agent."
+        )
+    if kind == "usage_missing":
+        return (
+            f"{prefix}: upstream returned no usage; charged the full "
+            f"reservation (${event.get('charged_usd', '?')}) to stay conservative."
+        )
+    if kind == "overshoot":
+        return (
+            f"{prefix}: a request settled above its reservation "
+            f"(${event.get('cost_usd', '?')} > ${event.get('reserved_usd', '?')}); "
+            "the cap is still enforced on the next request."
+        )
+    if kind == "output_clamped":
+        return (
+            f"{prefix}: budget low, limiting output to "
+            f"{event.get('allowed_output_tokens', '?')} tokens per request"
+        )
+    if kind == "request_too_large":
+        return (
+            f"{prefix}: a request needs ${event.get('needed_usd', '?')} but only "
+            f"${event.get('remaining_usd', event.get('cap_usd', '?'))} of budget "
+            "remains; not sent. Raise the cap or enable output clamping."
+        )
+    return f"{prefix}: {kind}"
+
+
+class AlertDispatcher:
+    """Terminal + optional webhook/Telegram alert fan-out.
+
+    ``transport`` is an ``httpx`` transport, injected in tests so no real
+    network call is ever made.
+    """
+
+    def __init__(
+        self,
+        *,
+        features: Iterable[str] | None = None,
+        webhook_url: str | None = None,
+        telegram_token: str | None = None,
+        telegram_chat_id: str | None = None,
+        transport: httpx.BaseTransport | None = None,
+        stderr: TextIO | None = None,
+        clock: Callable[[], datetime] | None = None,
+        background: bool = True,
+        timeout_s: float = 5.0,
+    ) -> None:
+        self.features = frozenset(features or ())
+        self.webhook_url = webhook_url or None
+        self.telegram_token = telegram_token or None
+        self.telegram_chat_id = telegram_chat_id or None
+        self.transport = transport
+        self.stderr = stderr if stderr is not None else sys.stderr
+        self.clock = clock or (lambda: datetime.now(timezone.utc))
+        self.background = background
+        self.timeout_s = float(timeout_s)
+        self._threads: list[threading.Thread] = []
+        self._lock = threading.Lock()
+        self._seen_once: set[tuple[str, object]] = set()
+
+    # --------------------------------------------------------------- sending
+    def __call__(self, event: Mapping[str, Any]) -> None:
+        self.emit(event)
+
+    def emit(self, event: Mapping[str, Any]) -> None:
+        if not isinstance(event, Mapping):
+            return
+        kind = str(event.get("type", ""))
+        if kind not in TERMINAL_EVENTS:
+            # Per-request accounting events stay local; only alerts fan out.
+            return
+        if kind in ONCE_EVENTS:
+            key = (kind, event.get("run_id"))
+            with self._lock:
+                if key in self._seen_once:
+                    return
+                self._seen_once.add(key)
+        message = format_alert(event)
+        color = YELLOW if kind in {"output_clamped", "request_too_large"} else RED
+        self._terminal(message, color=color)
+
+        clean = sanitize(event)
+        clean.setdefault("run_id", None)
+        payload = {
+            "source": "agentcap",
+            "sent_at": self.clock().isoformat(),
+            "event": clean,
+            "message": redact(message),
+        }
+        if self.webhook_url and "webhook_alerts" in self.features:
+            self._deliver(self._post_webhook, payload)
+        if (
+            self.telegram_token
+            and self.telegram_chat_id
+            and "telegram_alerts" in self.features
+        ):
+            self._deliver(self._post_telegram, payload)
+
+    def _terminal(self, message: str, *, color: str = RED) -> None:
+        try:
+            print(f"{color}{redact(message)}{RESET}", file=self.stderr, flush=True)
+        except Exception:
+            pass
+
+    def _deliver(self, sender: Callable[[dict[str, Any]], None], payload: dict[str, Any]) -> None:
+        if not self.background:
+            sender(payload)
+            return
+        thread = threading.Thread(target=sender, args=(payload,), daemon=True)
+        with self._lock:
+            self._threads.append(thread)
+        thread.start()
+
+    def wait(self, timeout: float | None = 5.0) -> None:
+        """Join outstanding delivery threads (used by tests and shutdown)."""
+
+        with self._lock:
+            threads = list(self._threads)
+            self._threads.clear()
+        for thread in threads:
+            thread.join(timeout)
+
+    def _post_webhook(self, payload: dict[str, Any]) -> None:
+        try:
+            with httpx.Client(
+                transport=self.transport, timeout=self.timeout_s, trust_env=False
+            ) as client:
+                response = client.post(self.webhook_url, json=payload)
+                response.raise_for_status()
+        except Exception as exc:
+            # Never echo the URL or exception text: it may contain a token.
+            self._warn("webhook", exc)
+
+    def _post_telegram(self, payload: dict[str, Any]) -> None:
+        try:
+            text = payload.get("message") or "agentcap alert"
+            url = f"https://api.telegram.org/bot{self.telegram_token}/sendMessage"
+            with httpx.Client(
+                transport=self.transport, timeout=self.timeout_s, trust_env=False
+            ) as client:
+                response = client.post(
+                    url,
+                    json={"chat_id": self.telegram_chat_id, "text": text},
+                )
+                response.raise_for_status()
+        except Exception as exc:
+            self._warn("telegram", exc)
+
+    def _warn(self, channel: str, exc: BaseException) -> None:
+        try:
+            print(
+                f"agentcap: alert delivery failed ({channel}: "
+                f"{type(exc).__name__}); continuing.",
+                file=self.stderr,
+                flush=True,
+            )
+        except Exception:
+            pass
+
+
+def dispatcher_from_env(
+    *,
+    features: Iterable[str] | None = None,
+    transport: httpx.BaseTransport | None = None,
+    stderr: TextIO | None = None,
+    background: bool = True,
+) -> AlertDispatcher:
+    """Build a dispatcher from ``AGENTCAP_WEBHOOK_URL`` / Telegram env vars."""
+
+    return AlertDispatcher(
+        features=features,
+        webhook_url=os.environ.get("AGENTCAP_WEBHOOK_URL"),
+        telegram_token=os.environ.get("AGENTCAP_TELEGRAM_BOT_TOKEN"),
+        telegram_chat_id=os.environ.get("AGENTCAP_TELEGRAM_CHAT_ID"),
+        transport=transport,
+        stderr=stderr,
+        background=background,
+    )
+
+
+__all__ = [
+    "AlertDispatcher",
+    "ALLOWED_FIELDS",
+    "TERMINAL_EVENTS",
+    "dispatcher_from_env",
+    "format_alert",
+    "redact",
+    "sanitize",
+]

agentcap/clamp.py

@@ -0,0 +1,174 @@
+"""Budget-fit output clamping.
+
+Most agents never set an output limit, so a full worst-case reservation
+(``max_output_tokens`` of the model) is much larger than the actual request.
+With a small cap that reservation alone would cross the cap and kill a run that
+has spent nothing yet.
+
+This module computes how many output tokens actually fit in the remaining
+budget and rewrites *only* the request's output limit.  It never raises an
+existing limit, never touches prompts and never changes settlement (which still
+uses real usage).
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from decimal import Decimal, ROUND_CEILING, ROUND_FLOOR
+from typing import Any, Mapping
+
+from .estimate import candidate_count, explicit_output_tokens
+from .pricing import MICRO_USD, ONE_MILLION, ModelPrice, effective_rates
+
+#: Clamping below this many tokens is pointless: treat the budget as exhausted.
+MIN_CLAMP_OUTPUT = 256
+
+
+@dataclass(frozen=True)
+class OutputFit:
+    """A clamped output allowance and the exact reservation it implies."""
+
+    allowed_output: int
+    input_cost: Decimal
+    reservation: Decimal
+
+
+def _positive_int(value: Any) -> int | None:
+    if isinstance(value, bool) or value is None:
+        return None
+    if isinstance(value, int):
+        return value if value > 0 else None
+    if isinstance(value, float) and value.is_integer() and value > 0:
+        return int(value)
+    if isinstance(value, str) and value.strip().isdigit() and int(value.strip()) > 0:
+        return int(value.strip())
+    return None
+
+
+def fit_output(
+    model_price: ModelPrice,
+    *,
+    provider: str,
+    body: Mapping[str, Any],
+    input_tokens: int,
+    remaining: Decimal,
+    now: Any = None,
+    cache_write: bool = False,
+    tier_multiplier: Decimal = Decimal("1"),
+    minimum: int = MIN_CLAMP_OUTPUT,
+) -> OutputFit:
+    """Largest per-candidate output limit whose reservation fits ``remaining``.
+
+    The returned ``allowed_output`` can be below ``minimum`` (or zero); the
+    caller decides whether that is still worth sending.
+    """
+
+    if remaining <= 0:
+        return OutputFit(0, Decimal("0"), Decimal("0"))
+
+    input_rate, _cached_rate, output_rate = effective_rates(
+        model_price, input_tokens, now=now
+    )
+    if cache_write and model_price.cache_write_per_mtok is not None:
+        input_rate = max(input_rate, model_price.cache_write_per_mtok)
+
+    input_cost = (
+        Decimal(input_tokens) * input_rate / ONE_MILLION
+    ).quantize(MICRO_USD, rounding=ROUND_CEILING)
+    if tier_multiplier != 1:
+        input_cost = (input_cost * tier_multiplier).quantize(
+            MICRO_USD, rounding=ROUND_CEILING
+        )
+
+    output_per_token = output_rate * tier_multiplier / ONE_MILLION
+    candidates = candidate_count(provider, body)
+    if output_per_token <= 0:
+        return OutputFit(0, input_cost, input_cost)
+
+    budget_for_output = remaining - input_cost
+    if budget_for_output <= 0:
+        allowed = 0
+    else:
+        raw = budget_for_output / (output_per_token * Decimal(candidates))
+        allowed = int(raw.to_integral_value(rounding=ROUND_FLOOR))
+    if allowed < 0:
+        allowed = 0
+
+    existing = explicit_output_tokens(provider, body)
+    if existing is not None:
+        # Never raise a limit the client already set (math should prevent it).
+        allowed = min(allowed, existing)
+
+    def _reservation(value: int) -> Decimal:
+        total = input_cost + (
+            Decimal(value * candidates) * output_rate * tier_multiplier / ONE_MILLION
+        )
+        return total.quantize(MICRO_USD, rounding=ROUND_CEILING)
+
+    reservation = _reservation(allowed)
+    # Ceiling rounding can push the reservation one micro-dollar over; lower the
+    # allowance until it genuinely fits.
+    while reservation > remaining and allowed > 0:
+        allowed -= 1
+        reservation = _reservation(allowed)
+
+    return OutputFit(allowed, input_cost, reservation)
+
+
+def apply_output_clamp(
+    provider: str,
+    path: str,
+    body: dict[str, Any],
+    allowed_output: int,
+) -> None:
+    """Lower the request's output limit to ``allowed_output`` in place.
+
+    * OpenAI chat: ``max_completion_tokens`` (or the legacy ``max_tokens`` when
+      the body already used it; never both).
+    * OpenAI responses: ``max_output_tokens``.
+    * Gemini: ``generationConfig.maxOutputTokens`` (created when missing).
+    """
+
+    allowed = _positive_int(allowed_output)
+    if allowed is None:
+        return
+
+    if provider == "gemini":
+        config = body.get("generationConfig")
+        config = dict(config) if isinstance(config, Mapping) else {}
+        current = _positive_int(
+            config.get("maxOutputTokens", config.get("max_output_tokens"))
+        )
+        config["maxOutputTokens"] = min(allowed, current) if current else allowed
+        body["generationConfig"] = config
+        return
+
+    normalised = path.strip("/")
+    if normalised == "v1/responses":
+        current = _positive_int(body.get("max_output_tokens"))
+        body["max_output_tokens"] = min(allowed, current) if current else allowed
+        return
+
+    # Chat Completions (and the legacy /v1/completions shape).
+    if normalised == "v1/completions":
+        current = _positive_int(body.get("max_tokens"))
+        body["max_tokens"] = min(allowed, current) if current else allowed
+        return
+    legacy = _positive_int(body.get("max_tokens"))
+    modern = _positive_int(body.get("max_completion_tokens"))
+    if legacy is not None:
+        body["max_tokens"] = min(allowed, legacy)
+        if modern is not None:
+            # Both fields already existed; lower both so neither stays above
+            # the reservation.  A new second field is never added.
+            body["max_completion_tokens"] = min(allowed, modern)
+        return
+    body["max_completion_tokens"] = min(allowed, modern) if modern else allowed
+
+
+__all__ = [
+    "MIN_CLAMP_OUTPUT",
+    "OutputFit",
+    "apply_output_clamp",
+    "fit_output",
+]