Import AITURK IDE 1.0.0-beta.1 from Hermes 63279301; preserve MIT license

agent/command_token_source.py

@@ -0,0 +1,190 @@
+"""Mint a provider API key by running a command (``key_cmd``).
+
+Static API keys are the exception at enterprise gateways: SSO/OIDC brokers,
+cloud IAM, and internal auth proxies all issue SHORT-LIVED bearers instead.
+A key copied into ``.env`` (``key_env``) is stale within the hour, so every
+request after that 401s and the user has to restart the session.
+
+``key_cmd`` names a command that PRINTS a token, so the credential is derived
+rather than stored::
+
+    providers:
+      my-gateway:
+        base_url: https://gateway.internal.example.com/v1
+        api_mode: chat_completions
+        key_cmd: my-auth-cli print-token --profile prod
+
+This is the established pattern for agent tooling — Claude Code's
+``apiKeyHelper``, the ``gcloud auth print-access-token`` / ``aws ecr
+get-login-password`` idiom, and vendor helpers such as ``databricks auth
+token`` all expose exactly this contract. Hermes already accepts a callable
+API key on both wire clients (the Entra ID / Azure identity path) and invokes
+it per request, so nothing downstream changes: the token is simply always
+fresh. It is cached until shortly before expiry, so the command runs about
+once per token lifetime rather than once per request.
+
+Output contract: print ONLY the token on stdout, either bare or as JSON with
+an ``access_token`` field (``expires_in`` is honoured when present) — the
+shape OAuth 2.0 token endpoints and the helpers above already emit.
+
+Precedence: an explicit ``--api-key`` still wins (the one-off recovery escape
+hatch); otherwise ``key_cmd`` is preferred over a static ``api_key`` /
+``key_env`` on the same entry.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import subprocess
+import threading
+import time
+from typing import Callable, Optional
+
+logger = logging.getLogger(__name__)
+
+# Treat a cached token as spent slightly before its stated expiry, so a request
+# can't be signed with a token that dies in flight. 60s matches the leeway used
+# by comparable OAuth token caches.
+_TOKEN_REFRESH_LEEWAY_SECONDS = 60.0
+# A token helper reads a local credential cache and should answer in
+# milliseconds; anything approaching this budget is hung, not slow.
+_MINT_TIMEOUT_SECONDS = 15
+# When a helper advertises NO expiry, the token cannot be cached for the life
+# of the process: nothing in the request path re-mints on 401 (the SDK retries
+# 429/5xx only), so an expired no-TTL token would 401 every request until
+# restart. Re-mint on a bounded window instead — the helper answers from a
+# local credential cache in milliseconds, so a periodic re-run is cheap, and a
+# helper that wants a longer cache can simply advertise its real expiry.
+_NO_TTL_REFRESH_SECONDS = 900.0
+
+
+class CommandTokenError(RuntimeError):
+    """A ``key_cmd`` failed to produce a usable token."""
+
+
+def _mint(command: str, label: str) -> tuple[str, Optional[float]]:
+    """Run *command*, returning ``(token, ttl_seconds_or_None)``."""
+    try:
+        completed = subprocess.run(
+            command,
+            shell=True,
+            capture_output=True,
+            text=True,
+            timeout=_MINT_TIMEOUT_SECONDS,
+        )
+    except subprocess.TimeoutExpired as exc:
+        raise CommandTokenError(
+            f"key_cmd for provider {label!r} timed out after "
+            f"{_MINT_TIMEOUT_SECONDS}s"
+        ) from exc
+    except OSError as exc:
+        raise CommandTokenError(
+            f"key_cmd for provider {label!r} could not be executed: {exc}"
+        ) from exc
+
+    if completed.returncode != 0:
+        # NEVER include stdout/stderr: a partially-successful auth helper can
+        # print a token or refresh secret there. The command STRING is also
+        # withheld — a key_cmd can legitimately embed a secret
+        # (`print-token --client-secret=…`), so echoing it back would leak the
+        # very credential this module exists to protect. Name the provider so
+        # the user knows which config entry to run by hand.
+        raise CommandTokenError(
+            f"key_cmd for provider {label!r} exited {completed.returncode}. "
+            f"Run that provider's key_cmd manually to see why "
+            f"(e.g. `databricks auth login` if its OAuth session expired)."
+        )
+
+    stdout = completed.stdout or ""
+    if not stdout.strip():
+        raise CommandTokenError(f"key_cmd for provider {label!r} produced no output")
+
+    # JSON payload — the shape `databricks auth token --output json` prints.
+    # Token extraction mirrors databricks/ucode's get_databricks_token:
+    #   json.loads(result.stdout or "{}").get("access_token", "")
+    if stdout.lstrip().startswith("{"):
+        try:
+            payload = json.loads(stdout)
+        except json.JSONDecodeError:
+            payload = None
+        if isinstance(payload, dict):
+            token = str(payload.get("access_token") or "").strip()
+            if not token:
+                raise CommandTokenError(
+                    f"key_cmd for provider {label!r} returned JSON without an "
+                    "'access_token' field"
+                )
+            ttl = payload.get("expires_in")
+            if isinstance(ttl, (int, float)) and ttl > 0:
+                return token, float(ttl)
+            # A relative lifetime is the OAuth 2.0 field, but CLI token helpers
+            # commonly print an absolute ISO 8601 deadline instead. Treating
+            # that as "no TTL advertised" caches the token for the life of the
+            # process, so every request 401s once the deadline passes.
+            # Imported lazily: hermes_cli.auth imports from agent.* at module
+            # level, so a top-level import here would risk a cycle.
+            from hermes_cli.auth import _parse_iso_timestamp
+
+            for field in ("expiry", "expiresOn"):
+                deadline = _parse_iso_timestamp(payload.get(field))
+                if deadline is not None:
+                    remaining = deadline - time.time()
+                    if remaining > 0:
+                        return token, remaining
+            return token, None
+
+    # Bare token. The contract every comparable helper documents is "stdout
+    # carries the token and nothing else" — extra output would be consumed as
+    # part of the credential. Strip surrounding whitespace and take the rest
+    # verbatim; do NOT silently keep one line of several, which converts a
+    # misconfigured helper (banner, warning, two tokens) into a corrupt-key 401
+    # that is far harder to diagnose than an explicit refusal.
+    token = stdout.strip()
+    if "\n" in token:
+        raise CommandTokenError(
+            f"key_cmd for provider {label!r} printed multiple lines; it must "
+            "print only the token (or JSON with an 'access_token' field)"
+        )
+    return token, None
+
+
+class CommandTokenSource:
+    """Callable returning a bearer token, cached until shortly before expiry."""
+
+    def __init__(self, command: str, label: str = "custom") -> None:
+        self._command = command
+        self._label = label or "custom"
+        self._lock = threading.Lock()
+        self._token = ""
+        self._expires_at: float = 0.0
+
+    def __call__(self) -> str:
+        with self._lock:
+            if self._token and time.monotonic() < self._expires_at:
+                return self._token
+            token, ttl = _mint(self._command, self._label)
+            self._token = token
+            self._expires_at = (
+                time.monotonic() + max(ttl - _TOKEN_REFRESH_LEEWAY_SECONDS, 5.0)
+                if ttl
+                # No advertised TTL: bounded cache (see _NO_TTL_REFRESH_SECONDS)
+                # — there is no 401-driven re-mint hook to fall back on.
+                else time.monotonic() + _NO_TTL_REFRESH_SECONDS
+            )
+            logger.debug(
+                "Minted key_cmd token for provider %s (ttl=%s)",
+                self._label, f"{int(ttl)}s" if ttl else "unknown",
+            )
+            return token
+
+
+def build_command_token_provider(
+    key_cmd: str,
+    provider_label: str = "custom",
+) -> Optional[Callable[[], str]]:
+    """A per-request token provider for *key_cmd*, or ``None`` when unset."""
+    command = str(key_cmd or "").strip()
+    if not command:
+        return None
+    return CommandTokenSource(command, provider_label)