From 89973e7bcecb6809508dabca75c700570554faaf Mon Sep 17 00:00:00 2001 From: Jiri Kuncar Date: Thu, 14 May 2026 17:56:40 +0200 Subject: [PATCH 1/2] Introduce OAuth authentication --- .claude-plugin/marketplace.json | 2 +- .claude-plugin/plugin.json | 2 +- CLAUDE.md | 12 +- README.md | 65 +++++- bin/logfire-auth | 10 + commands/login.md | 21 ++ commands/logout.md | 10 + commands/refresh.md | 12 + commands/status.md | 11 + pyproject.toml | 2 +- scripts/auth.py | 377 +++++++++++++++++++++++++++++++ scripts/log-event.py | 43 ++-- scripts/oauth_token.py | 381 ++++++++++++++++++++++++++++++++ uv.lock | 2 +- 14 files changed, 924 insertions(+), 26 deletions(-) create mode 100755 bin/logfire-auth create mode 100644 commands/login.md create mode 100644 commands/logout.md create mode 100644 commands/refresh.md create mode 100644 commands/status.md create mode 100644 scripts/auth.py create mode 100644 scripts/oauth_token.py diff --git a/.claude-plugin/marketplace.json b/.claude-plugin/marketplace.json index 0fd4bf0..2564c08 100644 --- a/.claude-plugin/marketplace.json +++ b/.claude-plugin/marketplace.json @@ -12,7 +12,7 @@ "name": "logfire-session-capture", "source": "./", "description": "Captures Claude Code sessions and exports pydantic-ai compatible OTel traces to Pydantic Logfire, with local JSONL fallback.", - "version": "0.4.6", + "version": "0.5.0", "author": { "name": "Pydantic" }, diff --git a/.claude-plugin/plugin.json b/.claude-plugin/plugin.json index 0c0b44d..47a0859 100644 --- a/.claude-plugin/plugin.json +++ b/.claude-plugin/plugin.json @@ -1,5 +1,5 @@ { "name": "logfire-session-capture", - "version": "0.4.6", + "version": "0.5.0", "description": "Captures Claude Code sessions and exports pydantic-ai compatible OTel traces to Pydantic Logfire, with local JSONL fallback." } diff --git a/CLAUDE.md b/CLAUDE.md index b691e99..c6c167f 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -8,9 +8,17 @@ A Claude Code plugin that captures sessions and exports pydantic-ai compatible O ## Architecture -The entire plugin is a single Python script (`scripts/log-event.py`) invoked by every hook event defined in `hooks/hooks.json`. The plugin manifest lives at `.claude-plugin/plugin.json`. +The plugin has three Python scripts (all stdlib only) and an `hooks/hooks.json`: -**Data flow:** Claude Code hook fires -> stdin JSON piped to `log-event.py` -> appends JSONL locally -> if `LOGFIRE_TOKEN` set and event is SessionStart/Stop/SubagentStop/SessionEnd, builds OTLP/HTTP JSON payload and sends via `urllib`. +- `scripts/log-event.py` — invoked by every hook event; builds and ships OTLP spans. +- `scripts/oauth_token.py` — shared module: load/save/refresh OAuth token bundle at `~/.logfire/claude-code-logfire-plugin.json`. Used by both the hook and the CLI. +- `scripts/auth.py` — user-facing CLI for the OAuth Device Authorization Grant (RFC 8628) login / logout / status. Invoked by the `/logfire-session-capture:login`, `/logfire-session-capture:logout`, `/logfire-session-capture:status` slash commands under `commands/`. + +The plugin manifest lives at `.claude-plugin/plugin.json`. + +**Data flow:** Claude Code hook fires -> stdin JSON piped to `log-event.py` -> appends JSONL locally -> if a token is available (either `LOGFIRE_TOKEN` env var or a stored OAuth bundle that gets auto-refreshed inline) and event is SessionStart/Stop/SubagentStop/SessionEnd, builds OTLP/HTTP JSON payload and sends via `urllib`. + +**Auth resolution order:** `LOGFIRE_TOKEN` env var (highest precedence, backwards-compatible; uses `$LOGFIRE_BASE_URL` for the OTLP endpoint) → stored OAuth bundle (a single bundle at `~/.logfire/claude-code-logfire-plugin.json` recording its own `base_url`; refreshed if within 60s of expiry; refresh is serialised via an `os.mkdir` lock so concurrent sessions don't race refresh-token rotation). Only one OAuth bundle is stored at a time — re-running login overwrites the previous one. If neither auth source is available, the OTel export path is skipped silently. **Session state:** A temp file (`$TMPDIR/claude-logfire-{session_id}.json`) persists the root span ID, start time, transcript line offset, accumulated messages, usage totals, and cost details between hook invocations. Created on `SessionStart`, deleted on `SessionEnd`. diff --git a/README.md b/README.md index a71c80c..68fb1a5 100644 --- a/README.md +++ b/README.md @@ -29,7 +29,13 @@ From within Claude Code, run: /plugin update logfire-session-capture@pydantic-claude-code-logfire-plugin ``` -### Set your Logfire token +### Authenticate to Logfire + +The plugin supports two authentication modes. Pick whichever fits your setup — +they're checked in order, so you can have both configured and `LOGFIRE_TOKEN` +will win. + +#### Option A: Fixed write token (simplest) ```bash export LOGFIRE_TOKEN="your-logfire-write-token" @@ -43,10 +49,54 @@ For the EU region: export LOGFIRE_BASE_URL="https://logfire-eu.pydantic.dev" ``` +#### Option B: OAuth with automatic refresh (no long-lived secret in your shell) + +Run the device flow once — the plugin then refreshes the access token silently +on every hook invocation until you log out: + +```bash +# From inside a Claude Code session: +/logfire-session-capture:login + +# Or directly from your shell: +python3 ~/.claude/plugins/.../scripts/auth.py login +``` + +You'll see a one-time user code and a browser opens to authorize the plugin +against your Logfire org / project (RFC 8628 Device Authorization Grant with +PKCE; the access token carries the `project:write_otlp` scope and is bound to +the Fusionfire OTLP intake via the RFC 8707 `resource` parameter). The +`client_id` is a Client ID Metadata Document URL — by default +`https://logfire.pydantic.dev/clients/claude-code-logfire.json` for production +hosts and `https://logfire.pydantic.info/...` for staging — so the +authorization server fetches the canonical client metadata directly and no +per-install registration is needed. On success the access + refresh tokens +are written to `~/.logfire/claude-code-logfire-plugin.json` (mode 0600), along +with the chosen `base_url`. The plugin reads the bundle on every hook event and +exchanges the refresh token whenever the access token is within 60s of expiry; +the new bundle is written back atomically and shared across all your Claude +Code sessions. + +Only **one** bundle is stored at a time — re-running `/logfire-session-capture:login` +overwrites it. `logout` / `status` / `refresh` operate on whatever's stored, so +they don't need a `--base-url` argument. + +Four slash commands ship with the plugin: + +- `/logfire-session-capture:login` — start (or re-run) the device flow (pass `--base-url ...` for EU / self-hosted) +- `/logfire-session-capture:status` — print expiry, scope, and client info for the stored bundle +- `/logfire-session-capture:refresh` — force-exchange the refresh token (debugging / manual rotation; hooks refresh lazily already) +- `/logfire-session-capture:logout` — delete the stored bundle + +Unset `LOGFIRE_TOKEN` to make the plugin use the OAuth bundle. If both are +present, `LOGFIRE_TOKEN` wins (no surprise migrations). + +#### Configuration reference + | Variable | Required | Default | Description | |---|---|---|---| -| `LOGFIRE_TOKEN` | Yes | _(none)_ | Logfire write token | -| `LOGFIRE_BASE_URL` | No | `https://logfire-us.pydantic.dev` | Logfire ingest endpoint | +| `LOGFIRE_TOKEN` | One of token/OAuth | _(none)_ | Logfire write token. Takes precedence over the OAuth bundle when both are present. | +| `LOGFIRE_BASE_URL` | No | `https://logfire-us.pydantic.dev` | Default Logfire ingest endpoint when using `LOGFIRE_TOKEN`; default `--base-url` for `/logfire-session-capture:login`. In OAuth mode the stored bundle's `base_url` wins over this variable. | | `LOGFIRE_LOCAL_LOG` | No | `false` | Set to `true` to write JSONL event logs locally | | `LOGFIRE_DIAGNOSTICS` | No | `false` | Set to `true` to write diagnostic logs (enabled automatically when `LOGFIRE_LOCAL_LOG` is set) | | `LOGFIRE_ENVIRONMENT` | No | _(none)_ | Sets `deployment.environment.name` on every trace (e.g. `production`, `dev`) | @@ -54,7 +104,9 @@ export LOGFIRE_BASE_URL="https://logfire-eu.pydantic.dev" | `OTEL_SERVICE_NAME` | No | `claude-code-plugin` | Overrides the `service.name` resource attribute | | `OTEL_RESOURCE_ATTRIBUTES` | No | _(none)_ | Standard OTel env var for additional resource attributes, e.g. `deployment.environment.name=prod,service.instance.id=worker-1` | -Without `LOGFIRE_TOKEN`, no traces are sent. The plugin does nothing unless at least one of `LOGFIRE_TOKEN` or `LOGFIRE_LOCAL_LOG` is set. +Without either `LOGFIRE_TOKEN` or a stored OAuth bundle, no traces are sent. +The plugin does nothing unless at least one auth mode or `LOGFIRE_LOCAL_LOG` +is configured. ## What you get @@ -118,8 +170,9 @@ Diagnostic logs are written to `.claude/logs/diagnostics.jsonl` in the project d **Common issues:** -- **No traces appearing in Logfire** -- Check that `LOGFIRE_TOKEN` is set and valid. Enable diagnostics to see if OTLP exports are failing. -- **Export errors (HTTP 401/403)** -- Your Logfire token may be invalid or expired. Generate a new write token in the Logfire console. +- **No traces appearing in Logfire** -- Check that `LOGFIRE_TOKEN` is set and valid, or that `/logfire-session-capture:status` reports a stored OAuth bundle. Enable diagnostics to see if OTLP exports are failing. +- **Export errors (HTTP 401/403)** -- Your Logfire token may be invalid or expired. Generate a new write token in the Logfire console, or run `/logfire-session-capture:login` to refresh the OAuth bundle. +- **OAuth refresh keeps failing** -- The refresh token may have been revoked or expired. Run `/logfire-session-capture:logout` then `/logfire-session-capture:login` to start a fresh flow. - **Export errors (HTTP 4xx/5xx)** -- Check `LOGFIRE_BASE_URL` if using a non-default region. The plugin logs HTTP status codes to stderr and diagnostics. ## Development diff --git a/bin/logfire-auth b/bin/logfire-auth new file mode 100755 index 0000000..3edd13f --- /dev/null +++ b/bin/logfire-auth @@ -0,0 +1,10 @@ +#!/usr/bin/env bash +# Plugin-PATH wrapper for scripts/auth.py. +# +# Claude Code adds this plugin's bin/ to PATH, so slash commands can invoke +# `logfire-auth ...` directly without an explicit $CLAUDE_PLUGIN_ROOT expansion +# (which the slash-command permission checker rejects as "simple_expansion"). +# $CLAUDE_PLUGIN_ROOT isn't exported into the spawned subshell, so we locate +# the plugin root from the wrapper's own path: bin/'s parent dir. +PLUGIN_ROOT=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/.." &> /dev/null && pwd) +exec python3 "$PLUGIN_ROOT/scripts/auth.py" "$@" diff --git a/commands/login.md b/commands/login.md new file mode 100644 index 0000000..89c4338 --- /dev/null +++ b/commands/login.md @@ -0,0 +1,21 @@ +--- +description: Log in to Logfire via OAuth device flow (alternative to LOGFIRE_TOKEN) +allowed-tools: Bash(logfire-auth *) +--- + +Run the OAuth device authorization flow against Logfire and persist a token +bundle the plugin will auto-refresh on every Claude Code session. + +Use this instead of setting a fixed `LOGFIRE_TOKEN`. Once logged in, the plugin +silently refreshes the access token whenever it's near expiry (no manual +re-login until the refresh token itself is revoked). + +The script reads `$LOGFIRE_BASE_URL` for the default region (US if unset). Any +slash-command arguments are forwarded, so you can override with +`--base-url http://localhost:3000` (or pass `--no-browser`, `--client-id ...`, +etc). The chosen base URL is recorded inside the stored bundle, so subsequent +`/logfire-session-capture:logout`, `:status`, `:refresh` don't need any +arguments — only one base URL is stored at a time, and re-running login +overwrites the previous bundle. + +!`logfire-auth login $ARGUMENTS` diff --git a/commands/logout.md b/commands/logout.md new file mode 100644 index 0000000..5b25fbe --- /dev/null +++ b/commands/logout.md @@ -0,0 +1,10 @@ +--- +description: Remove the stored Logfire OAuth token +allowed-tools: Bash(logfire-auth *) +--- + +Delete the stored OAuth token bundle at `~/.logfire/claude-code-logfire-plugin.json`. +The bundle's own base URL is used — no extra argument needed. After logout the +plugin falls back to `LOGFIRE_TOKEN` (or stays silent if no token is configured). + +!`logfire-auth logout` diff --git a/commands/refresh.md b/commands/refresh.md new file mode 100644 index 0000000..1f620dd --- /dev/null +++ b/commands/refresh.md @@ -0,0 +1,12 @@ +--- +description: Force-refresh the stored Logfire OAuth access token +allowed-tools: Bash(logfire-auth *) +--- + +Force an OAuth refresh against the Logfire authorization server using the +stored refresh token, even if the access token isn't near expiry. Operates on +the stored bundle's own base URL. The plugin's hooks already refresh lazily on +every invocation, so this is mostly useful for debugging or to deliberately +cycle the access token without re-running the device flow. + +!`logfire-auth refresh` diff --git a/commands/status.md b/commands/status.md new file mode 100644 index 0000000..409765e --- /dev/null +++ b/commands/status.md @@ -0,0 +1,11 @@ +--- +description: Show stored Logfire OAuth token status (expiry, scope, etc.) +allowed-tools: Bash(logfire-auth *) +--- + +Print metadata about the stored OAuth token bundle: the base URL, scope, +client_id, time until expiry, and whether a refresh token is present. Does not +display the access token itself. Operates on whatever's in the store — only one +bundle is kept at a time. + +!`logfire-auth status` diff --git a/pyproject.toml b/pyproject.toml index c17b93c..531ddab 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "claude-code-logfire-plugin" -version = "0.4.1" +version = "0.5.0" requires-python = ">=3.7" [dependency-groups] diff --git a/scripts/auth.py b/scripts/auth.py new file mode 100644 index 0000000..2a330d0 --- /dev/null +++ b/scripts/auth.py @@ -0,0 +1,377 @@ +#!/usr/bin/env python3 +"""OAuth 2.0 Device Authorization Grant CLI for the Logfire plugin. + +Lets a user log in to Logfire interactively and persist an access / +refresh-token bundle that ``log-event.py`` then uses on every hook +invocation instead of a fixed ``LOGFIRE_TOKEN``. + +Subcommands +----------- + login Run the device flow against ``LOGFIRE_BASE_URL`` (or --base-url) + and save the resulting tokens under + ``~/.logfire/claude-code-logfire-plugin.json``. Overwrites any + previously stored bundle — only one base URL is active at a + time. + logout Remove the stored token bundle. + status Show whether a token is stored, when it expires, and which + scopes it carries. + refresh Force-exchange the stored refresh_token for a fresh access token + (the hook hot path already refreshes lazily; this is for + debugging or for cycling tokens on demand). + +``logout`` / ``status`` / ``refresh`` operate on whatever's in the store — +they don't take ``--base-url``. The bundle records its own ``base_url``, +so switching regions is done by re-running ``login --base-url ...``. + +The flow uses a Client ID Metadata Document (CIMD): the ``client_id`` is the +URL of a JSON document that the authorization server fetches to learn the +client's grants, scopes, and redirect URIs. + +The flow is RFC 8628 (Device Authorization Grant) with PKCE (RFC 7636); the +``project:write_otlp`` scope is requested by default — the same scope the +Fusionfire intake checks on ``/v1/traces``. + +Usage +----- + python3 scripts/auth.py login + python3 scripts/auth.py login --base-url https://logfire-eu.pydantic.dev + python3 scripts/auth.py status + python3 scripts/auth.py logout + +Stdlib only — runs on the same Python 3.7+ baseline as the rest of the plugin. +""" + +from __future__ import annotations + +import argparse +import base64 +import hashlib +import json +import os +import secrets +import sys +import time +import urllib.error +import webbrowser +from pathlib import Path + +# Allow ``python3 scripts/auth.py ...`` to import ``oauth_token`` without +# needing the directory on PYTHONPATH. +sys.path.insert(0, str(Path(__file__).resolve().parent)) + +from oauth_token import ( # noqa: E402 + DEFAULT_BASE_URL, + DEFAULT_SCOPE, + _bundle_from_token_response, + _http_post_form, + delete_bundle, + discover_metadata, + discover_resource, + force_refresh, + load_bundle, + save_bundle, +) + +DEVICE_CODE_GRANT = "urn:ietf:params:oauth:grant-type:device_code" +CIMD = "https://logfire.pydantic.dev/clients/claude-code-logfire.json" + + +# --------------------------------------------------------------------------- +# PKCE +# --------------------------------------------------------------------------- + + +def _pkce_pair() -> tuple[str, str]: + verifier = base64.urlsafe_b64encode(secrets.token_bytes(32)).decode().rstrip("=") + challenge = base64.urlsafe_b64encode(hashlib.sha256(verifier.encode()).digest()).decode().rstrip("=") + return verifier, challenge + + +# --------------------------------------------------------------------------- +# Network helpers (stdlib + readable error surfacing) +# --------------------------------------------------------------------------- + + +def _extract_oauth_error(body: bytes) -> tuple[str, str]: + """Pull a (code, description) tuple out of an OAuth error response. + + Logfire wraps OAuth errors in a FastAPI ``detail`` envelope; spec-pure + servers put the fields at the top level. Handle both. + """ + try: + data = json.loads(body.decode()) + except (ValueError, UnicodeDecodeError): + return "", body[:200].decode(errors="replace") + if isinstance(data, dict): + detail = data.get("detail") if isinstance(data.get("detail"), dict) else None + source = detail or data + return source.get("error", ""), source.get("error_description", "") + return "", "" + + +def _post_form(url: str, fields: dict, timeout: float = 30.0) -> dict: + return _http_post_form(url, fields, timeout=timeout) + + +# --------------------------------------------------------------------------- +# Device flow +# --------------------------------------------------------------------------- + + +def request_device_code( + metadata: dict, + *, + client_id: str, + code_challenge: str, + scope: str, + resource: str, +) -> dict: + fields = { + "client_id": client_id, + "scope": scope, + "code_challenge": code_challenge, + "code_challenge_method": "S256", + } + if resource: + fields["resource"] = resource + return _post_form(metadata["device_authorization_endpoint"], fields) + + +def poll_for_token( + metadata: dict, + *, + client_id: str, + device_code: str, + code_verifier: str, + resource: str, + interval: int, + expires_in: int, +) -> dict: + deadline = time.time() + expires_in + delay = max(interval, 1) + while time.time() < deadline: + time.sleep(delay) + fields = { + "grant_type": DEVICE_CODE_GRANT, + "device_code": device_code, + "client_id": client_id, + "code_verifier": code_verifier, + } + if resource: + fields["resource"] = resource + try: + return _post_form(metadata["token_endpoint"], fields) + except urllib.error.HTTPError as e: + code, desc = _extract_oauth_error(e.read()) + if code == "authorization_pending": + continue + if code == "slow_down": + delay += 5 + continue + if code == "access_denied": + raise SystemExit("Authorization denied") from None + if code == "expired_token": + raise SystemExit("Device code expired before authorization completed") from None + raise SystemExit(f"Device flow failed: {code or e.code} {desc}".rstrip()) from None + raise SystemExit("Device code expired before authorization completed") + + +# --------------------------------------------------------------------------- +# Subcommands +# --------------------------------------------------------------------------- + + +def cmd_login(args: argparse.Namespace) -> int: + base_url = args.base_url.rstrip("/") + scope = args.scope + existing = load_bundle() + if existing and existing.get("base_url") and existing["base_url"] != base_url: + print(f"Overwriting existing token for {existing['base_url']} (only one base URL is stored at a time).") + print(f"Discovering OAuth metadata for {base_url} ...") + metadata = discover_metadata(base_url) + resource = "" if args.no_resource else discover_resource(base_url) + + if not args.client_id: + raise SystemExit( + f"No CIMD client_id is known for {base_url}. " + f"Pass --client-id explicitly (e.g. the URL of a Client ID " + f"Metadata Document hosted by your Logfire deployment)." + ) + print(f"Using CIMD client_id={args.client_id}") + + verifier, challenge = _pkce_pair() + device = request_device_code( + metadata, + client_id=args.client_id, + code_challenge=challenge, + scope=scope, + resource=resource, + ) + + verification_url = device.get("verification_uri_complete") or device["verification_uri"] + user_code = device.get("user_code", "") + expires_in = int(device.get("expires_in", 600)) + interval = int(device.get("interval", 5)) + + print() + print(f"User code: {user_code}") + print(f"Open this URL to authorize: {verification_url}") + print(f"Code expires in {expires_in}s") + print() + + if not args.no_browser: + try: + webbrowser.open(verification_url) + except webbrowser.Error: + pass + + print("Waiting for browser authorization (Ctrl+C to cancel) ...") + token = poll_for_token( + metadata, + client_id=args.client_id, + device_code=device["device_code"], + code_verifier=verifier, + resource=resource, + interval=interval, + expires_in=expires_in, + ) + + bundle = _bundle_from_token_response( + token, + base_url=base_url, + client_id=args.client_id, + resource=resource, + fallback_scope=scope, + ) + save_bundle(bundle) + print() + print(f"Logged in to {base_url}") + print(f" scope: {bundle.get('scope', '')}") + print(f" expires in: {int(bundle['expires_at'] - time.time())}s") + print(" stored at: ~/.logfire/claude-code-logfire-plugin.json") + print() + print("Unset LOGFIRE_TOKEN in your shell to make the plugin use this OAuth bundle.") + return 0 + + +def cmd_logout(args: argparse.Namespace) -> int: + bundle = load_bundle() + if not bundle: + print("No stored token.") + return 0 + delete_bundle() + print(f"Removed stored token for {bundle.get('base_url', '(unknown)')}.") + return 0 + + +def cmd_refresh(args: argparse.Namespace) -> int: + try: + refreshed = force_refresh() + except LookupError as exc: + raise SystemExit(str(exc)) from None + except urllib.error.HTTPError as exc: + body = exc.read() if hasattr(exc, "read") else b"" + code, desc = _extract_oauth_error(body) + raise SystemExit(f"Refresh failed: {code or exc.code} {desc}".rstrip()) from None + except urllib.error.URLError as exc: + raise SystemExit(f"Refresh failed: {exc.reason}") from None + remaining = int(float(refreshed.get("expires_at", 0)) - time.time()) + print(f"Refreshed token for {refreshed.get('base_url', '')}") + print(f" scope: {refreshed.get('scope', '')}") + print(f" expires in: {remaining}s") + return 0 + + +def cmd_status(args: argparse.Namespace) -> int: + bundle = load_bundle() + if not bundle: + print("Not logged in. Run `login` to authenticate.") + return 0 + remaining = int(float(bundle.get("expires_at", 0)) - time.time()) + state = "valid" if remaining > 0 else "expired (refresh on next use)" + print(f"Base URL: {bundle.get('base_url', '')}") + print(f"Client ID: {bundle.get('client_id', '')}") + print(f"Scope: {bundle.get('scope', '')}") + print(f"Resource: {bundle.get('resource', '') or '(none)'}") + print(f"Expires in: {remaining}s ({state})") + print(f"Has refresh token: {'yes' if bundle.get('refresh_token') else 'no'}") + return 0 + + +# --------------------------------------------------------------------------- +# Entry point +# --------------------------------------------------------------------------- + + +def _default_base_url() -> str: + return os.environ.get("LOGFIRE_BASE_URL", DEFAULT_BASE_URL) + + +def main() -> int: + parser = argparse.ArgumentParser( + description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + sub = parser.add_subparsers(dest="cmd") + + login_p = sub.add_parser("login", help="Run device flow and store a token bundle") + login_p.add_argument( + "--base-url", + default=_default_base_url(), + help="Logfire base URL (default: $LOGFIRE_BASE_URL or %(default)s)", + ) + login_p.add_argument( + "--client-id", + default=CIMD, + help=("OAuth client_id (a Client ID Metadata Document URL)"), + ) + login_p.add_argument( + "--scope", + default=DEFAULT_SCOPE, + help=f"OAuth scope to request (default: {DEFAULT_SCOPE})", + ) + login_p.add_argument( + "--no-browser", + action="store_true", + help="Don't try to open the verification URL in a browser", + ) + login_p.add_argument( + "--no-resource", + action="store_true", + help="Skip RFC 8707 resource discovery (use only for non-standard backends)", + ) + + sub.add_parser("logout", help="Remove the stored token bundle") + sub.add_parser("status", help="Print info about the stored token bundle") + sub.add_parser("refresh", help="Force-refresh the stored token using its refresh_token") + + args = parser.parse_args() + if args.cmd == "login": + return cmd_login(args) + if args.cmd == "logout": + return cmd_logout(args) + if args.cmd == "status": + return cmd_status(args) + if args.cmd == "refresh": + return cmd_refresh(args) + parser.print_help() + return 0 + + +if __name__ == "__main__": + try: + sys.exit(main()) + except KeyboardInterrupt: + print("\nCancelled", file=sys.stderr) + sys.exit(130) + except urllib.error.HTTPError as exc: + body = exc.read() if hasattr(exc, "read") else b"" + code, desc = _extract_oauth_error(body) + if code: + print(f"HTTP {exc.code}: {code} {desc}".rstrip(), file=sys.stderr) + else: + print(f"HTTP {exc.code}: {body[:500].decode(errors='replace')}", file=sys.stderr) + sys.exit(1) + except urllib.error.URLError as exc: + print(f"Network error: {exc.reason}", file=sys.stderr) + sys.exit(1) diff --git a/scripts/log-event.py b/scripts/log-event.py index 567871d..7300ce8 100755 --- a/scripts/log-event.py +++ b/scripts/log-event.py @@ -26,7 +26,15 @@ from urllib.parse import unquote from urllib.request import Request, urlopen -VERSION = "0.4.6" +# ``oauth_token`` reads VERSION from .claude-plugin/plugin.json so we don't +# have to keep it in sync by hand. Adjust sys.path so the import works +# regardless of where this script is invoked from. +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +try: + from oauth_token import USER_AGENT, VERSION +except ImportError: + VERSION = "unknown" + USER_AGENT = "claude-code-logfire-plugin" OTLP_EVENTS = {"SessionStart", "Stop", "SubagentStop", "SessionEnd"} @@ -291,7 +299,7 @@ def send_otlp(payload: dict, endpoint: str, token: str) -> None: headers={ "Content-Type": "application/json", "Authorization": f"Bearer {token}", - "User-Agent": "claude-code-logfire-plugin", + "User-Agent": USER_AGENT, }, method="POST", ) @@ -701,11 +709,7 @@ def _extract_tools_from_messages(messages: list[dict]) -> tuple[list[str], str | def _has_thinking(messages: list[dict]) -> bool: """Check if any message contains a thinking block.""" - return any( - p.get("type") == "thinking" - for msg in messages - for p in msg.get("parts", []) - ) + return any(p.get("type") == "thinking" for msg in messages for p in msg.get("parts", [])) def _extract_user_snippet(input_messages: list[dict], max_len: int = 60) -> str | None: @@ -739,9 +743,7 @@ def _describe_call(input_messages: list[dict], output_messages: list[dict]) -> s user_snippet = _extract_user_snippet(input_messages) has_tool_results = any( - p.get("type") == "tool_call_response" - for msg in input_messages - for p in msg.get("parts", []) + p.get("type") == "tool_call_response" for msg in input_messages for p in msg.get("parts", []) ) thinking_prefix = "Thinking + " if thinking else "" @@ -1259,12 +1261,25 @@ def main() -> None: except OSError: log_diag("error", "Failed to write JSONL log entry") - # OTel export requires token + # OTel export requires a token. Prefer the explicit LOGFIRE_TOKEN env var + # for backwards compatibility; otherwise fall back to a stored OAuth + # bundle (auto-refreshed) populated by ``scripts/auth.py login``. The + # OAuth bundle records its own base_url and overrides $LOGFIRE_BASE_URL, + # since the access token is bound to that resource. logfire_token = os.environ.get("LOGFIRE_TOKEN", "") - if not logfire_token: - return + if logfire_token: + base_url = os.environ.get("LOGFIRE_BASE_URL", "https://logfire-us.pydantic.dev").rstrip("/") + else: + try: + from oauth_token import get_access_token - base_url = os.environ.get("LOGFIRE_BASE_URL", "https://logfire-us.pydantic.dev").rstrip("/") + result = get_access_token() + except (ImportError, OSError, ValueError) as exc: + log_diag("warn", "OAuth token lookup failed", str(exc)) + result = None + if not result: + return + logfire_token, base_url = result otlp_endpoint = f"{base_url}/v1/traces" # Only process OTLP events diff --git a/scripts/oauth_token.py b/scripts/oauth_token.py new file mode 100644 index 0000000..e2dc9bc --- /dev/null +++ b/scripts/oauth_token.py @@ -0,0 +1,381 @@ +"""OAuth token storage and refresh for the Logfire plugin. + +Stdlib-only helpers shared by ``auth.py`` (interactive device flow login) and +``log-event.py`` (hot-path hook handler that needs a valid access token). + +A single token "bundle" is persisted at +``~/.logfire/claude-code-logfire-plugin.json`` with mode 0600: + +```json +{ + "base_url": "https://logfire-us.pydantic.dev", + "access_token": "...", + "refresh_token": "...", + "expires_at": 1234567890.0, + "scope": "project:write_otlp", + "client_id": "https://logfire.pydantic.dev/clients/claude-code-logfire.json", + "resource": "https://logfire-us.pydantic.dev/v1/traces" +} +``` + +Only one bundle is ever stored — logging in to a different ``base_url`` +overwrites the previous one. ``logout`` / ``status`` / ``refresh`` therefore +don't need a ``--base-url`` argument; they always operate on the stored +bundle's own ``base_url``. + +The file is rewritten atomically (``mkstemp`` + ``os.replace``) and refresh is +serialised across sessions via an ``os.mkdir`` lock to keep refresh-token +rotation from racing when several Claude Code sessions start simultaneously. +""" + +from __future__ import annotations + +import contextlib +import json +import os +import tempfile +import time +import urllib.error +import urllib.parse +import urllib.request +from pathlib import Path + +DEFAULT_BASE_URL = "https://logfire-us.pydantic.dev" +DEFAULT_SCOPE = "project:write_otlp" + + +def _read_plugin_version() -> str: + """Read the plugin version from ``.claude-plugin/plugin.json``. + + Single source of truth shared with the plugin manifest. Falls back to + ``"unknown"`` if the file is missing or malformed (e.g. when the script + is invoked outside the plugin directory). + """ + manifest = Path(__file__).resolve().parent.parent / ".claude-plugin" / "plugin.json" + try: + with open(manifest) as f: + data = json.load(f) + version = data.get("version") + if isinstance(version, str) and version: + return version + except (OSError, ValueError): + pass + return "unknown" + + +VERSION = _read_plugin_version() + +# Identify ourselves on every outbound HTTP call. Python's default +# ``Python-urllib/X.Y`` User-Agent is blocked by Cloudflare's bot signature +# rules on logfire.pydantic.* (Error 1010), so we always send our own. +USER_AGENT = f"claude-code-logfire-plugin/{VERSION}" + +# Refresh the access token when this many seconds (or fewer) remain before +# expiry. Matches the buffer used by ``scripts/oauth_intake_example.py`` in +# the platform repo. +REFRESH_BUFFER_SECONDS = 60 + +# Cap how long we wait for the cross-session refresh lock. Hooks have a hard +# timeout (10s for Stop/PreToolUse, 30s for SessionEnd), so we must give up +# well before that to leave time for the OTLP send itself. +LOCK_MAX_WAIT_SECONDS = 5.0 +LOCK_POLL_INTERVAL = 0.1 +LOCK_STALE_SECONDS = 30 + +TOKEN_DIR = Path.home() / ".logfire" +TOKEN_FILE = TOKEN_DIR / "claude-code-logfire-plugin.json" +LOCK_DIR = TOKEN_DIR / ".claude-code-logfire-plugin.lock" + + +# --------------------------------------------------------------------------- +# File I/O +# --------------------------------------------------------------------------- + + +def _read_file() -> dict | None: + """Return the stored bundle, or ``None`` if no usable bundle exists.""" + try: + with open(TOKEN_FILE) as f: + data = json.load(f) + except (OSError, ValueError): + return None + if not isinstance(data, dict): + return None + # Bundle must have at least a base_url + access_token to be useful. + if not isinstance(data.get("base_url"), str) or not data.get("access_token"): + return None + return data + + +def _write_file(bundle: dict) -> None: + """Atomically write ``bundle`` to ``TOKEN_FILE``. Raises ``OSError`` on + failure so callers (notably ``cmd_login``) can surface the error rather + than silently telling the user the token was saved.""" + TOKEN_DIR.mkdir(parents=True, exist_ok=True) + fd, tmp = tempfile.mkstemp(prefix="claude-code-logfire-plugin.", dir=str(TOKEN_DIR)) + try: + with os.fdopen(fd, "w") as f: + json.dump(bundle, f, indent=2) + os.replace(tmp, str(TOKEN_FILE)) + except BaseException: + with contextlib.suppress(OSError): + os.unlink(tmp) + raise + # mkstemp already creates the file with mode 0600; the chmod is + # belt-and-suspenders in case a non-POSIX platform widens it. + with contextlib.suppress(OSError): + os.chmod(str(TOKEN_FILE), 0o600) + + +# --------------------------------------------------------------------------- +# Cross-process lock (matches the pattern in log-event.py) +# --------------------------------------------------------------------------- + + +@contextlib.contextmanager +def _lock(): + """Best-effort cross-process lock around the token store. + + Yields ``True`` if the lock was acquired, ``False`` if we timed out. + Callers that mutate state should still proceed on ``False`` + (last-writer-wins beats dropping a freshly-issued token on the floor); + callers that race refresh-token rotation should bail on ``False``. + + The release is gated on actual acquisition so a nested ``_lock()`` — + e.g. a refresh path that re-enters via ``save_bundle`` — can't tear + down its caller's lock. + """ + acquired = False + deadline = time.time() + LOCK_MAX_WAIT_SECONDS + TOKEN_DIR.mkdir(parents=True, exist_ok=True) + while True: + try: + os.mkdir(LOCK_DIR) + acquired = True + break + except FileExistsError: + try: + if time.time() - os.path.getmtime(LOCK_DIR) > LOCK_STALE_SECONDS: + with contextlib.suppress(OSError): + os.rmdir(LOCK_DIR) + continue + except OSError: + pass + if time.time() >= deadline: + break + time.sleep(LOCK_POLL_INTERVAL) + try: + yield acquired + finally: + if acquired: + with contextlib.suppress(OSError): + os.rmdir(LOCK_DIR) + + +# --------------------------------------------------------------------------- +# Bundle accessors +# --------------------------------------------------------------------------- + + +def load_bundle() -> dict | None: + """Return the stored bundle (with ``base_url`` field) or ``None``.""" + return _read_file() + + +def _save_bundle_unlocked(bundle: dict) -> None: + """Write ``bundle`` without taking the lock. Callers must already hold + ``_lock()`` (used by the refresh path, which can't re-enter the lock + it already owns).""" + if "base_url" not in bundle: + raise ValueError("bundle is missing required 'base_url' field") + _write_file(bundle) + + +def save_bundle(bundle: dict) -> None: + with _lock(): + _save_bundle_unlocked(bundle) + + +def delete_bundle() -> None: + with _lock(), contextlib.suppress(FileNotFoundError): + os.unlink(TOKEN_FILE) + + +# --------------------------------------------------------------------------- +# HTTP helpers (stdlib only) +# --------------------------------------------------------------------------- + + +def _http_get_json(url: str, timeout: float = 10.0) -> dict: + req = urllib.request.Request( + url, + headers={"Accept": "application/json", "User-Agent": USER_AGENT}, + ) + with urllib.request.urlopen(req, timeout=timeout) as r: + return json.load(r) + + +def _http_post_form(url: str, fields: dict, timeout: float = 30.0) -> dict: + body = urllib.parse.urlencode(fields).encode() + req = urllib.request.Request( + url, + data=body, + headers={ + "Content-Type": "application/x-www-form-urlencoded", + "Accept": "application/json", + "User-Agent": USER_AGENT, + }, + method="POST", + ) + with urllib.request.urlopen(req, timeout=timeout) as r: + return json.load(r) + + +def discover_metadata(base_url: str) -> dict: + """RFC 8414 OAuth Authorization Server Metadata.""" + return _http_get_json(base_url.rstrip("/") + "/.well-known/oauth-authorization-server") + + +def discover_resource(base_url: str) -> str: + """RFC 9728 protected-resource metadata — the JWT ``aud`` value the + intake expects. Falls back to ``base_url`` if the endpoint is unavailable + (e.g. self-hosted deployments without the metadata route).""" + try: + data = _http_get_json(base_url.rstrip("/") + "/.well-known/oauth-protected-resource/v1") + except (urllib.error.URLError, OSError, ValueError): + return base_url.rstrip("/") + resource = data.get("resource") + if isinstance(resource, str) and resource: + return resource + return base_url.rstrip("/") + + +# --------------------------------------------------------------------------- +# Refresh +# --------------------------------------------------------------------------- + + +def _bundle_from_token_response( + token: dict, + *, + base_url: str, + client_id: str, + resource: str, + fallback_refresh: str = "", + fallback_scope: str = "", +) -> dict: + return { + "base_url": base_url.rstrip("/"), + "access_token": token["access_token"], + "refresh_token": token.get("refresh_token") or fallback_refresh, + "expires_at": time.time() + int(token.get("expires_in", 0)), + "scope": token.get("scope") or fallback_scope, + "client_id": client_id, + "resource": resource, + } + + +def _refresh(bundle: dict) -> dict: + """Exchange the refresh_token for a new access token. Raises on failure.""" + base_url = bundle["base_url"] + metadata = discover_metadata(base_url) + resource = bundle.get("resource") or discover_resource(base_url) + fields = { + "grant_type": "refresh_token", + "refresh_token": bundle["refresh_token"], + "client_id": bundle.get("client_id", ""), + "scope": bundle.get("scope") or DEFAULT_SCOPE, + } + if resource: + # RFC 8707 — required when the bundle was originally issued with + # one, since the AS verifies the resource matches across refreshes. + fields["resource"] = resource + token = _http_post_form(metadata["token_endpoint"], fields) + new_bundle = _bundle_from_token_response( + token, + base_url=base_url, + client_id=bundle.get("client_id", ""), + resource=resource, + fallback_refresh=bundle["refresh_token"], + fallback_scope=bundle.get("scope", ""), + ) + # We're called from inside ``_lock()`` (via ``get_access_token`` / + # ``force_refresh``), so use the unlocked save to avoid the lock's + # 5-second timeout waiting on ourselves. + _save_bundle_unlocked(new_bundle) + return new_bundle + + +def force_refresh() -> dict: + """Force a refresh of the stored bundle and return it. + + Raises ``LookupError`` if there's no stored bundle (or no refresh token + to exchange), ``RuntimeError`` if the cross-process lock can't be + acquired, and the underlying ``urllib`` error if the AS rejects the + refresh. Intended for the ``/logfire-session-capture:refresh`` slash + command — ``get_access_token`` already refreshes lazily for the hook + hot path. + """ + with _lock() as acquired: + if not acquired: + raise RuntimeError("Could not acquire token-store lock") + bundle = load_bundle() + if not bundle: + raise LookupError("No stored token; run `login` first") + if not bundle.get("refresh_token"): + raise LookupError("Stored token has no refresh_token; run `login` again") + return _refresh(bundle) + + +def get_access_token() -> tuple[str, str] | None: + """Return ``(access_token, base_url)`` for the stored bundle, refreshing + if needed. + + Returns ``None`` if no bundle is stored, or if the token has expired and + cannot be refreshed. Designed to be called from the hook hot path so it + must never raise. + """ + bundle = load_bundle() + if not bundle: + return None + base_url = bundle["base_url"] + + now = time.time() + expires_at = float(bundle.get("expires_at", 0)) + if expires_at > now + REFRESH_BUFFER_SECONDS: + token = bundle.get("access_token", "") + return (token, base_url) if token else None + + if not bundle.get("refresh_token"): + # No refresh capability — return what we have if it's still valid, + # else give up. + if expires_at > now: + token = bundle.get("access_token", "") + return (token, base_url) if token else None + return None + + # Serialise refresh across sessions so refresh-token rotation doesn't + # race itself into a 4xx. + with _lock() as acquired: + if not acquired: + # Couldn't lock — fall back to the existing access token if it + # still has any life in it. + if expires_at > now: + token = bundle.get("access_token", "") + return (token, base_url) if token else None + return None + # Another session may have already refreshed while we waited. + latest = load_bundle() or bundle + latest_expires = float(latest.get("expires_at", 0)) + if latest_expires > now + REFRESH_BUFFER_SECONDS: + token = latest.get("access_token", "") + return (token, base_url) if token else None + try: + refreshed = _refresh(latest) + token = refreshed.get("access_token", "") + return (token, base_url) if token else None + except (urllib.error.HTTPError, urllib.error.URLError, OSError, ValueError, KeyError): + if latest_expires > now: + token = latest.get("access_token", "") + return (token, base_url) if token else None + return None diff --git a/uv.lock b/uv.lock index 40af9bb..984e924 100644 --- a/uv.lock +++ b/uv.lock @@ -4,7 +4,7 @@ requires-python = ">=3.7" [[package]] name = "claude-code-logfire-plugin" -version = "0.4.1" +version = "0.5.0" source = { virtual = "." } [package.dev-dependencies] From f6b10517e7099683e630f4136dacfc8b2b1e663b Mon Sep 17 00:00:00 2001 From: Jiri Kuncar Date: Tue, 26 May 2026 17:28:02 +0200 Subject: [PATCH 2/2] Keep SessionEnd hook ahead of the /exit SIGTERM MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The eager root-span send added in 0f1b117 had to ship before any blocking work, but the OAuth resolution path in main() could synchronously refresh a near-expiry token (lock + discovery + token POST, up to ~45s) before the eager send ever ran — defeating the guard the SessionEnd hook needs against /exit killing the subprocess mid-finalisation. Two fixes: - Skip the lazy refresh on SessionEnd if the access token still has any validity left (`skip_refresh_if_valid` on `get_access_token`). The trace ships first; the next session's SessionStart handles the refresh. - Cache `token_endpoint` in the stored bundle at login (and on the first refresh of a legacy bundle), so the refresh path skips the RFC 8414 discovery GET. Login already has the metadata in hand. Also fix a stale docstring example: the RFC 9728 `resource` value is `https://logfire-us.pydantic.dev/v1`, not the `/v1/traces` endpoint URL. Co-Authored-By: Claude Opus 4.7 (1M context) --- scripts/auth.py | 1 + scripts/log-event.py | 7 ++++- scripts/oauth_token.py | 69 +++++++++++++++++++++++++++++------------- 3 files changed, 55 insertions(+), 22 deletions(-) diff --git a/scripts/auth.py b/scripts/auth.py index 2a330d0..0274f11 100644 --- a/scripts/auth.py +++ b/scripts/auth.py @@ -241,6 +241,7 @@ def cmd_login(args: argparse.Namespace) -> int: base_url=base_url, client_id=args.client_id, resource=resource, + token_endpoint=metadata.get("token_endpoint", ""), fallback_scope=scope, ) save_bundle(bundle) diff --git a/scripts/log-event.py b/scripts/log-event.py index 7300ce8..380aca2 100755 --- a/scripts/log-event.py +++ b/scripts/log-event.py @@ -1273,7 +1273,12 @@ def main() -> None: try: from oauth_token import get_access_token - result = get_access_token() + # On SessionEnd we're racing a /exit-driven SIGTERM against a 30s + # hook budget, and the eager root-span send (added in 0f1b117) + # has to ship before any blocking work. Skip the lazy refresh if + # the access token still has any life left — the next session's + # SessionStart will refresh. + result = get_access_token(skip_refresh_if_valid=(_hook_event == "SessionEnd")) except (ImportError, OSError, ValueError) as exc: log_diag("warn", "OAuth token lookup failed", str(exc)) result = None diff --git a/scripts/oauth_token.py b/scripts/oauth_token.py index e2dc9bc..007ee23 100644 --- a/scripts/oauth_token.py +++ b/scripts/oauth_token.py @@ -14,10 +14,15 @@ "expires_at": 1234567890.0, "scope": "project:write_otlp", "client_id": "https://logfire.pydantic.dev/clients/claude-code-logfire.json", - "resource": "https://logfire-us.pydantic.dev/v1/traces" + "resource": "https://logfire-us.pydantic.dev/v1", + "token_endpoint": "https://logfire-us.pydantic.dev/oauth/token" } ``` +``token_endpoint`` is cached from the RFC 8414 metadata at login (or +re-populated on the first refresh of a legacy bundle) so the hot-path +refresh doesn't have to make a discovery roundtrip before the token POST. + Only one bundle is ever stored — logging in to a different ``base_url`` overwrites the previous one. ``logout`` / ``status`` / ``refresh`` therefore don't need a ``--base-url`` argument; they always operate on the stored @@ -39,6 +44,7 @@ import urllib.parse import urllib.request from pathlib import Path +from typing import Any DEFAULT_BASE_URL = "https://logfire-us.pydantic.dev" DEFAULT_SCOPE = "project:write_otlp" @@ -92,7 +98,7 @@ def _read_plugin_version() -> str: # --------------------------------------------------------------------------- -def _read_file() -> dict | None: +def _read_file() -> dict[str, Any] | None: """Return the stored bundle, or ``None`` if no usable bundle exists.""" try: with open(TOKEN_FILE) as f: @@ -107,7 +113,7 @@ def _read_file() -> dict | None: return data -def _write_file(bundle: dict) -> None: +def _write_file(bundle: dict[str, Any]) -> None: """Atomically write ``bundle`` to ``TOKEN_FILE``. Raises ``OSError`` on failure so callers (notably ``cmd_login``) can surface the error rather than silently telling the user the token was saved.""" @@ -177,12 +183,12 @@ def _lock(): # --------------------------------------------------------------------------- -def load_bundle() -> dict | None: +def load_bundle() -> dict[str, Any] | None: """Return the stored bundle (with ``base_url`` field) or ``None``.""" return _read_file() -def _save_bundle_unlocked(bundle: dict) -> None: +def _save_bundle_unlocked(bundle: dict[str, Any]) -> None: """Write ``bundle`` without taking the lock. Callers must already hold ``_lock()`` (used by the refresh path, which can't re-enter the lock it already owns).""" @@ -191,7 +197,7 @@ def _save_bundle_unlocked(bundle: dict) -> None: _write_file(bundle) -def save_bundle(bundle: dict) -> None: +def save_bundle(bundle: dict[str, Any]) -> None: with _lock(): _save_bundle_unlocked(bundle) @@ -206,7 +212,7 @@ def delete_bundle() -> None: # --------------------------------------------------------------------------- -def _http_get_json(url: str, timeout: float = 10.0) -> dict: +def _http_get_json(url: str, timeout: float = 10.0) -> dict[str, Any]: req = urllib.request.Request( url, headers={"Accept": "application/json", "User-Agent": USER_AGENT}, @@ -215,7 +221,7 @@ def _http_get_json(url: str, timeout: float = 10.0) -> dict: return json.load(r) -def _http_post_form(url: str, fields: dict, timeout: float = 30.0) -> dict: +def _http_post_form(url: str, fields: dict[str, Any], timeout: float = 30.0) -> dict[str, Any]: body = urllib.parse.urlencode(fields).encode() req = urllib.request.Request( url, @@ -231,7 +237,7 @@ def _http_post_form(url: str, fields: dict, timeout: float = 30.0) -> dict: return json.load(r) -def discover_metadata(base_url: str) -> dict: +def discover_metadata(base_url: str) -> dict[str, Any]: """RFC 8414 OAuth Authorization Server Metadata.""" return _http_get_json(base_url.rstrip("/") + "/.well-known/oauth-authorization-server") @@ -256,15 +262,16 @@ def discover_resource(base_url: str) -> str: def _bundle_from_token_response( - token: dict, + token: dict[str, Any], *, base_url: str, client_id: str, resource: str, + token_endpoint: str = "", fallback_refresh: str = "", fallback_scope: str = "", -) -> dict: - return { +) -> dict[str, Any]: + bundle: dict[str, Any] = { "base_url": base_url.rstrip("/"), "access_token": token["access_token"], "refresh_token": token.get("refresh_token") or fallback_refresh, @@ -273,14 +280,22 @@ def _bundle_from_token_response( "client_id": client_id, "resource": resource, } + if token_endpoint: + bundle["token_endpoint"] = token_endpoint + return bundle -def _refresh(bundle: dict) -> dict: +def _refresh(bundle: dict[str, Any]) -> dict[str, Any]: """Exchange the refresh_token for a new access token. Raises on failure.""" - base_url = bundle["base_url"] - metadata = discover_metadata(base_url) - resource = bundle.get("resource") or discover_resource(base_url) - fields = { + base_url: str = bundle["base_url"] + # Prefer the cached endpoint to skip the discovery roundtrip on the hot + # path; fall back to discovery for legacy bundles written before this + # field existed (the new bundle we save below will include it). + token_endpoint: str = bundle.get("token_endpoint") or "" + if not token_endpoint: + token_endpoint = discover_metadata(base_url)["token_endpoint"] + resource: str = bundle.get("resource") or discover_resource(base_url) + fields: dict[str, Any] = { "grant_type": "refresh_token", "refresh_token": bundle["refresh_token"], "client_id": bundle.get("client_id", ""), @@ -290,12 +305,13 @@ def _refresh(bundle: dict) -> dict: # RFC 8707 — required when the bundle was originally issued with # one, since the AS verifies the resource matches across refreshes. fields["resource"] = resource - token = _http_post_form(metadata["token_endpoint"], fields) + token = _http_post_form(token_endpoint, fields) new_bundle = _bundle_from_token_response( token, base_url=base_url, client_id=bundle.get("client_id", ""), resource=resource, + token_endpoint=token_endpoint, fallback_refresh=bundle["refresh_token"], fallback_scope=bundle.get("scope", ""), ) @@ -306,7 +322,7 @@ def _refresh(bundle: dict) -> dict: return new_bundle -def force_refresh() -> dict: +def force_refresh() -> dict[str, Any]: """Force a refresh of the stored bundle and return it. Raises ``LookupError`` if there's no stored bundle (or no refresh token @@ -327,18 +343,25 @@ def force_refresh() -> dict: return _refresh(bundle) -def get_access_token() -> tuple[str, str] | None: +def get_access_token(*, skip_refresh_if_valid: bool = False) -> tuple[str, str] | None: """Return ``(access_token, base_url)`` for the stored bundle, refreshing if needed. Returns ``None`` if no bundle is stored, or if the token has expired and cannot be refreshed. Designed to be called from the hook hot path so it must never raise. + + If ``skip_refresh_if_valid`` is set, return the stored access token as + long as it has any validity left (i.e. ``expires_at > now``), without + consulting the AS. Used by the SessionEnd hot path, which has a 30s + budget and races a ``/exit``-triggered SIGTERM — a near-expiry refresh + can blow that budget and leave the trace unclosed. The next session's + SessionStart picks up the refresh. """ bundle = load_bundle() if not bundle: return None - base_url = bundle["base_url"] + base_url: str = bundle["base_url"] now = time.time() expires_at = float(bundle.get("expires_at", 0)) @@ -346,6 +369,10 @@ def get_access_token() -> tuple[str, str] | None: token = bundle.get("access_token", "") return (token, base_url) if token else None + if skip_refresh_if_valid and expires_at > now: + token = bundle.get("access_token", "") + return (token, base_url) if token else None + if not bundle.get("refresh_token"): # No refresh capability — return what we have if it's still valid, # else give up.