multi-forge 0.2.0__py3-none-any.whl

forge/core/reactive/cost_tracking.py

@@ -0,0 +1,300 @@
+"""Verb-level cost attribution via proxy metric snapshot deltas.
+
+Wraps subprocess invocations (panel, supervisor, handoff, etc.) to
+measure cost by snapshotting proxy metrics before and after execution.
+Results are logged to PID-sharded verb JSONL files.
+
+All verb costs are marked ``estimated`` because concurrent proxy traffic
+(e.g., the main interactive session) may share the same proxy during
+the measurement window.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import os
+import threading
+import time
+import urllib.request
+from contextlib import contextmanager
+from dataclasses import asdict, dataclass, field
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any
+
+from forge.core.paths import get_forge_home
+
+logger = logging.getLogger(__name__)
+
+_verb_lock = threading.Lock()
+
+
+@dataclass
+class ProxyCostDelta:
+    """Cost delta for a single proxy between two snapshots."""
+
+    base_url: str
+    cost_micros: int = 0
+    input_tokens: int = 0
+    output_tokens: int = 0
+    cached_tokens: int = 0
+    request_count: int = 0
+
+
+@dataclass
+class VerbCostResult:
+    """Aggregated cost attribution for one verb invocation."""
+
+    verb: str
+    total_cost_micros: int = 0
+    input_tokens: int = 0
+    output_tokens: int = 0
+    cached_tokens: int = 0
+    request_count: int = 0
+    duration_ms: float = 0.0
+    estimated: bool = True
+    per_proxy: list[ProxyCostDelta] = field(default_factory=list)
+
+
+def _fetch_snapshot(base_url: str, timeout: float = 2.0) -> dict[str, Any] | None:
+    """Fetch proxy metrics via GET /. Returns None on failure."""
+    try:
+        normalized = base_url if "://" in base_url else f"http://{base_url}"
+        url = normalized.rstrip("/") + "/"
+        with urllib.request.urlopen(url, timeout=timeout) as resp:
+            data = json.loads(resp.read())
+        if data.get("is_proxy") and "metrics" in data:
+            return data["metrics"]
+    except Exception as e:
+        logger.debug("Failed to fetch proxy snapshot from %s: %s", base_url, e)
+    return None
+
+
+def _compute_delta(before: dict[str, Any], after: dict[str, Any], base_url: str) -> ProxyCostDelta:
+    """Compute the difference between two proxy metric snapshots."""
+    b_tokens = before.get("tokens", {})
+    a_tokens = after.get("tokens", {})
+    b_costs = before.get("costs", {})
+    a_costs = after.get("costs", {})
+
+    return ProxyCostDelta(
+        base_url=base_url,
+        cost_micros=a_costs.get("total_micros", 0) - b_costs.get("total_micros", 0),
+        input_tokens=a_tokens.get("input", 0) - b_tokens.get("input", 0),
+        output_tokens=a_tokens.get("output", 0) - b_tokens.get("output", 0),
+        cached_tokens=a_tokens.get("cached", 0) - b_tokens.get("cached", 0),
+        request_count=after.get("total_requests", 0) - before.get("total_requests", 0),
+    )
+
+
+def _verb_log_dir() -> Path:
+    return get_forge_home() / "costs" / "verbs"
+
+
+def _log_verb_cost(result: VerbCostResult) -> None:
+    """Append a verb cost record to the PID-sharded JSONL log."""
+    record: dict[str, Any] = {
+        "ts": datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ"),
+        "verb": result.verb,
+        "total_cost_micros": result.total_cost_micros,
+        "estimated": result.estimated,
+        "input_tokens": result.input_tokens,
+        "output_tokens": result.output_tokens,
+        "cached_tokens": result.cached_tokens,
+        "request_count": result.request_count,
+        "duration_ms": round(result.duration_ms, 1),
+        "per_proxy": [asdict(p) for p in result.per_proxy],
+    }
+
+    try:
+        from forge.core.state import open_secure_append
+
+        log_dir = _verb_log_dir()
+        log_dir.mkdir(parents=True, exist_ok=True)
+        month = datetime.now(timezone.utc).strftime("%Y-%m")
+        path = log_dir / f"{month}_{os.getpid()}.jsonl"
+
+        with _verb_lock:
+            with open_secure_append(path) as f:
+                f.write(json.dumps(record, separators=(",", ":")) + "\n")
+    except Exception as e:
+        logger.warning("Failed to write verb cost log: %s", e)
+
+
+def resolve_subprocess_proxy_url() -> str | None:
+    """Resolve the current FORGE_SUBPROCESS_PROXY to a base URL, if configured."""
+    from forge.core.reactive.env import (
+        FORGE_SUBPROCESS_BASE_URL_VAR,
+        FORGE_SUBPROCESS_PROXY_VAR,
+    )
+    from forge.core.reactive.proxy import lookup_proxy_base_url
+
+    injected_url = os.environ.get(FORGE_SUBPROCESS_BASE_URL_VAR)
+    if injected_url:
+        return injected_url
+
+    proxy = os.environ.get(FORGE_SUBPROCESS_PROXY_VAR)
+    if not proxy:
+        return None
+
+    try:
+        return lookup_proxy_base_url(proxy)
+    except Exception:
+        return None
+
+
+def resolve_proxy_urls(specs: list[Any]) -> list[str]:
+    """Extract unique proxy base URLs from a list of ModelSpecs.
+
+    For specs with no explicit proxy, falls back to FORGE_SUBPROCESS_PROXY
+    when configured.
+    Deduplicates by resolved URL.
+    """
+    from forge.core.reactive.env import (
+        FORGE_SUBPROCESS_BASE_URL_VAR,
+        FORGE_SUBPROCESS_PROXY_VAR,
+    )
+    from forge.core.reactive.proxy import lookup_proxy_base_url
+
+    subprocess_proxy = os.environ.get(FORGE_SUBPROCESS_PROXY_VAR)
+    subprocess_base_url = os.environ.get(FORGE_SUBPROCESS_BASE_URL_VAR)
+    seen: set[str] = set()
+    urls: list[str] = []
+    for spec in specs:
+        proxy = getattr(spec, "preferred_proxy", None) or getattr(spec, "proxy", None) or subprocess_proxy
+        if not proxy:
+            continue
+        try:
+            url: str | None
+            if subprocess_base_url and proxy == subprocess_proxy:
+                url = subprocess_base_url
+            else:
+                url = lookup_proxy_base_url(proxy)
+            if url and url not in seen:
+                seen.add(url)
+                urls.append(url)
+        except Exception:
+            pass
+    return urls
+
+
+def resolve_proxy_urls_from_plan(plan: Any) -> list[str]:
+    """Extract unique proxy base URLs from a WorkerRoutingPlan.
+
+    Uses actual routing decisions (correct for --proxy, subprocess proxy,
+    route scan, and session proxy fallback).
+    """
+    seen: set[str] = set()
+    urls: list[str] = []
+    for result in plan.routes:
+        url = result.base_url
+        if url and url not in seen:
+            seen.add(url)
+            urls.append(url)
+    return urls
+
+
+@contextmanager
+def track_verb_cost(verb: str, proxy_base_urls: list[str]):
+    """Snapshot proxy metrics across all proxies before/after a verb invocation.
+
+    Args:
+        verb: Origin label ("panel", "supervisor", "handoff", etc.)
+        proxy_base_urls: ALL proxy base URLs this verb will use.
+            Direct workers (no proxy) are excluded — only proxied
+            requests have cost data at the proxy level.
+
+    Yields control to the caller. On exit, computes snapshot deltas,
+    logs the verb cost record, and discards. The caller does not
+    receive the result (it's fire-and-forget for the log).
+    """
+    unique_urls = list(dict.fromkeys(u for u in proxy_base_urls if u))
+
+    if not unique_urls:
+        yield
+        return
+
+    snapshots_before: dict[str, dict[str, Any]] = {}
+    for url in unique_urls:
+        snap = _fetch_snapshot(url)
+        if snap is not None:
+            snapshots_before[url] = snap
+
+    start = time.monotonic()
+    try:
+        yield
+    finally:
+        elapsed = (time.monotonic() - start) * 1000
+
+        try:
+            deltas: list[ProxyCostDelta] = []
+            for url in unique_urls:
+                if url not in snapshots_before:
+                    continue
+                after = _fetch_snapshot(url)
+                if after is None:
+                    continue
+                deltas.append(_compute_delta(snapshots_before[url], after, url))
+
+            total_cost = sum(d.cost_micros for d in deltas)
+            total_input = sum(d.input_tokens for d in deltas)
+            total_output = sum(d.output_tokens for d in deltas)
+            total_cached = sum(d.cached_tokens for d in deltas)
+            total_requests = sum(d.request_count for d in deltas)
+
+            result = VerbCostResult(
+                verb=verb,
+                total_cost_micros=total_cost,
+                input_tokens=total_input,
+                output_tokens=total_output,
+                cached_tokens=total_cached,
+                request_count=total_requests,
+                duration_ms=elapsed,
+                estimated=True,
+                per_proxy=deltas,
+            )
+            _log_verb_cost(result)
+        except Exception as e:
+            logger.warning("Failed to track verb cost for %s: %s", verb, e)
+
+
+def read_verb_logs(
+    period_start: datetime | None = None,
+    period_end: datetime | None = None,
+) -> list[dict[str, Any]]:
+    """Read and aggregate verb cost records from all PID shards."""
+    log_dir = _verb_log_dir()
+    if not log_dir.is_dir():
+        return []
+
+    records: list[dict[str, Any]] = []
+    for path in sorted(log_dir.glob("*.jsonl")):
+        try:
+            with open(path) as f:
+                for line in f:
+                    line = line.strip()
+                    if not line:
+                        continue
+                    try:
+                        record = json.loads(line)
+                    except json.JSONDecodeError:
+                        continue
+
+                    if period_start or period_end:
+                        ts_str = record.get("ts", "")
+                        try:
+                            ts = datetime.fromisoformat(ts_str.rstrip("Z").removesuffix("+00:00") + "+00:00")
+                        except (ValueError, TypeError):
+                            continue
+                        if period_start and ts < period_start:
+                            continue
+                        if period_end and ts >= period_end:
+                            continue
+
+                    records.append(record)
+        except OSError as e:
+            logger.warning("Failed to read verb log %s: %s", path, e)
+
+    records.sort(key=lambda r: r.get("ts", ""))
+    return records

forge/core/reactive/env.py

@@ -0,0 +1,180 @@
+"""Environment builder for Claude subprocess invocation.
+
+Provides ``build_claude_env()`` for constructing subprocess environments,
+and ``FORGE_DEPTH`` helpers for recursion-guarding hook → subprocess chains.
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+from collections.abc import Mapping
+
+logger = logging.getLogger(__name__)
+
+# Defense-in-depth: --bare prevents hook recursion in child processes,
+# but FORGE_DEPTH still guards against subprocess spawning at depth >= 2.
+FORGE_DEPTH_VAR = "FORGE_DEPTH"
+FORGE_MAX_DEPTH = 2
+
+FORGE_SUBPROCESS_PROXY_VAR = "FORGE_SUBPROCESS_PROXY"
+FORGE_SUBPROCESS_BASE_URL_VAR = "FORGE_SUBPROCESS_BASE_URL"
+FORGE_SUBPROCESS_PROXY_ID_VAR = "FORGE_SUBPROCESS_PROXY_ID"
+FORGE_SUBPROCESS_TEMPLATE_VAR = "FORGE_SUBPROCESS_TEMPLATE"
+FORGE_SIDECAR_VAR = "FORGE_SIDECAR"
+FORGE_LAUNCH_MODE_VAR = "FORGE_LAUNCH_MODE"
+
+# --bare (Claude Code >= 2.1.81) disables OAuth/keychain auth, requiring
+# ANTHROPIC_API_KEY in the environment. Only safe when the key is present.
+_BARE_AUTH_KEY = "ANTHROPIC_API_KEY"
+
+
+def can_use_bare(env: Mapping[str, str] | None = None) -> bool:
+    """True if ``--bare`` is safe for headless subprocesses.
+
+    ``--bare`` disables OAuth/keychain auth, so it requires
+    ANTHROPIC_API_KEY. When an explicit ``env`` dict is given, checks
+    only that dict (caller owns the env). When using os.environ
+    (default), also falls back to the credential file via
+    ``resolve_env_or_credential`` (which respects ``auth_ignore_env``).
+    """
+    if env is not None:
+        return bool(env.get(_BARE_AUTH_KEY))
+
+    from forge.core.auth.template_secrets import resolve_env_or_credential
+
+    return bool(resolve_env_or_credential(_BARE_AUTH_KEY))
+
+
+def get_forge_depth(env: Mapping[str, str] | None = None) -> int:
+    """Read current FORGE_DEPTH from the given env (or os.environ).
+
+    Invalid or missing values are treated as 0 (fail-open).
+    """
+    source = env if env is not None else os.environ
+    raw = source.get(FORGE_DEPTH_VAR, "0")
+    try:
+        return max(0, int(raw))
+    except (ValueError, TypeError):
+        return 0
+
+
+def should_spawn_subprocesses(env: Mapping[str, str] | None = None) -> bool:
+    """True if current depth allows spawning ``claude -p`` subprocesses.
+
+    Returns False when depth >= FORGE_MAX_DEPTH, meaning hooks should skip
+    subprocess-spawning work (supervisor, handoff agent, etc.) to prevent
+    runaway recursion.
+    """
+    return get_forge_depth(env) < FORGE_MAX_DEPTH
+
+
+def build_claude_env(
+    base_url: str | None = None,
+    extra_vars: dict[str, str] | None = None,
+    direct: bool = False,
+) -> dict[str, str]:
+    """Build environment dict for a Claude subprocess.
+
+    Starts with the current process environment. Sets ANTHROPIC_BASE_URL
+    if ``base_url`` is provided. When ``direct`` is True, removes any
+    inherited ANTHROPIC_BASE_URL and subprocess proxy so the child hits
+    Anthropic directly.
+    Applies ``extra_vars`` before routing and depth handling so explicit
+    function arguments remain authoritative.
+
+    Hydrates ANTHROPIC_API_KEY from the credential file when it's not in
+    the env (or when ``auth_ignore_env`` overrides it). This ensures
+    ``can_use_bare(env)`` and the subprocess both see the resolved key.
+
+    Args:
+        base_url: Proxy URL to route Claude requests through.
+        extra_vars: Additional environment variables to set/override.
+        direct: Force direct Anthropic routing (unset inherited proxy URL).
+
+    Returns:
+        Complete environment dict ready for ``subprocess.run(env=...)``.
+    """
+    env = os.environ.copy()
+    _hydrate_credentials(env)
+
+    # Apply extra_vars AFTER hydration so explicit caller overrides
+    # take precedence over credential-file values.
+    if extra_vars:
+        env.update(extra_vars)
+
+    if base_url:
+        env["ANTHROPIC_BASE_URL"] = base_url
+    elif direct:
+        env.pop("ANTHROPIC_BASE_URL", None)
+        env.pop(FORGE_SUBPROCESS_PROXY_VAR, None)
+        env.pop(FORGE_SUBPROCESS_BASE_URL_VAR, None)
+        env.pop(FORGE_SUBPROCESS_PROXY_ID_VAR, None)
+        env.pop(FORGE_SUBPROCESS_TEMPLATE_VAR, None)
+    else:
+        # No explicit base_url and not forced direct: check subprocess proxy fallback.
+        # FORGE_SUBPROCESS_PROXY is set by `forge session start --subprocess-proxy`
+        # and inherited by all child processes.
+        injected_subprocess_base_url = env.get(FORGE_SUBPROCESS_BASE_URL_VAR)
+        if injected_subprocess_base_url:
+            env["ANTHROPIC_BASE_URL"] = injected_subprocess_base_url
+        elif subprocess_proxy := env.get(FORGE_SUBPROCESS_PROXY_VAR):
+            resolved = _resolve_subprocess_proxy(subprocess_proxy)
+            if resolved:
+                env["ANTHROPIC_BASE_URL"] = resolved
+            else:
+                env.pop("ANTHROPIC_BASE_URL", None)
+
+    # Increment FORGE_DEPTH so child subprocesses know their nesting level
+    current_depth = get_forge_depth(env)
+    env[FORGE_DEPTH_VAR] = str(current_depth + 1)
+
+    return env
+
+
+def _hydrate_credentials(env: dict[str, str]) -> None:
+    """Ensure resolved credentials are in the subprocess env dict.
+
+    When ``auth_ignore_env`` is active, removes the inherited env value
+    for ANTHROPIC_API_KEY and injects the credential-file value instead.
+    When inactive, injects the credential-file value only if the env
+    var is absent (so ``can_use_bare(env)`` and the subprocess agree).
+    """
+    from forge.core.auth.template_secrets import resolve_env_or_credential
+
+    resolved = resolve_env_or_credential(_BARE_AUTH_KEY)
+
+    try:
+        from forge.runtime_config import get_runtime_config
+
+        ignore_env = get_runtime_config().auth_ignore_env
+    except Exception as e:
+        logger.debug("Could not read auth_ignore_env; using environment credentials: %s", e)
+        ignore_env = False
+
+    if ignore_env:
+        if resolved:
+            env[_BARE_AUTH_KEY] = resolved
+        else:
+            env.pop(_BARE_AUTH_KEY, None)
+    elif resolved and not env.get(_BARE_AUTH_KEY):
+        env[_BARE_AUTH_KEY] = resolved
+
+
+def _resolve_subprocess_proxy(proxy_id: str) -> str | None:
+    """Resolve subprocess proxy to a base URL, or None if unavailable.
+
+    Direct URL lookup only (not resolve_subprocess_routing). build_claude_env()
+    sets env vars for child processes and only needs a URL. Model compatibility
+    validation happens at workflow routing time (resolve_invocation_routing).
+    """
+    try:
+        from forge.core.reactive.proxy import lookup_proxy_base_url
+
+        url = lookup_proxy_base_url(proxy_id)
+        if url:
+            logger.debug("Subprocess proxy %r resolved to %s", proxy_id, url)
+        return url
+    except Exception as e:
+        logger.warning("Subprocess proxy %r unavailable: %s", proxy_id, e)
+        return None

forge/core/reactive/proxy.py

@@ -0,0 +1,78 @@
+"""Proxy registry lookup utility.
+
+Provides a shared function for looking up proxy base URLs from the
+registry. Used by the semantic supervisor, handoff agent, and review engine.
+
+Note: This module is intentionally NOT re-exported from __init__.py
+because it lazy-imports forge.proxy.proxies (a top-level component).
+Consumers import directly: ``from forge.core.reactive.proxy import lookup_proxy_base_url``.
+"""
+
+from __future__ import annotations
+
+
+def lookup_proxy_base_url(proxy: str | None) -> str | None:
+    """Look up base_url from the proxy registry by proxy_id or template name.
+
+    Internal boundary: if ``proxy`` is provided but resolution fails, the
+    exception propagates (ProxyResolutionError, AmbiguousProxyError, or
+    ProxyRegistryCorruptedError). Callers decide how to handle the failure.
+
+    Args:
+        proxy: Proxy identifier or template name to look up. If None, returns None.
+
+    Returns:
+        The proxy's base_url if found, None if proxy is None/empty.
+
+    Raises:
+        ProxyResolutionError: If proxy name is provided but cannot be resolved.
+        ProxyRegistryCorruptedError: If the registry file is unreadable.
+    """
+    if not proxy:
+        return None
+
+    from forge.proxy.proxies import ProxyRegistryStore, resolve_proxy
+
+    registry = ProxyRegistryStore().read()
+    entry = resolve_proxy(registry, proxy)
+    return entry.base_url
+
+
+def check_proxy_reachable(
+    proxy: str,
+    timeout_s: float = 1.0,
+) -> tuple[bool, str, str | None]:
+    """Check if a named proxy is locally routable.
+
+    Resolves via registry, then HTTP health-checks the endpoint.
+    "Ready" means the proxy responds at its base_url with valid
+    proxy metadata -- not just that it exists in the registry.
+
+    Returns:
+        (reachable, reason, base_url):
+        - reachable: True if proxy resolves AND health check passes
+        - reason: empty if reachable, human-readable otherwise
+        - base_url: proxy URL if resolved, None otherwise
+    """
+    from forge.proxy.proxies import ProxyRegistryStore, resolve_proxy
+    from forge.proxy.proxy_orchestrator import check_proxy_health
+
+    try:
+        registry = ProxyRegistryStore().read()
+    except Exception as e:
+        return (False, f"Registry error: {e}", None)
+
+    try:
+        entry = resolve_proxy(registry, proxy)
+    except Exception as e:
+        return (False, str(e), None)
+
+    if not check_proxy_health(
+        base_url=entry.base_url,
+        expected_template=entry.template,
+        expected_proxy_id=entry.proxy_id,
+        timeout_s=timeout_s,
+    ):
+        return (False, f"Proxy '{proxy}' not responding at {entry.base_url}", entry.base_url)
+
+    return (True, "", entry.base_url)