From ebc4dac6d818256c967c48d17544d0bc06c2bf3e Mon Sep 17 00:00:00 2001 From: Vuong Hoang Date: Wed, 16 Sep 2026 15:12:38 -0700 Subject: [PATCH] grok-token-broker: hold a refreshable session credential behind a rotation-safety gate --- scripts/yarros-corpus/train_pairs_lora.py | 38 ++- services/grok-token-broker/README.md | 67 +++++ services/grok-token-broker/broker.py | 297 ++++++++++++++++++++++ 3 files changed, 394 insertions(+), 8 deletions(-) create mode 100644 services/grok-token-broker/README.md create mode 100755 services/grok-token-broker/broker.py diff --git a/scripts/yarros-corpus/train_pairs_lora.py b/scripts/yarros-corpus/train_pairs_lora.py index 7df7576..18d6f89 100644 --- a/scripts/yarros-corpus/train_pairs_lora.py +++ b/scripts/yarros-corpus/train_pairs_lora.py @@ -26,6 +26,7 @@ The cross-arm comparison has to be behavioural (on-beat / in-band / ran-on) plus """ from __future__ import annotations import argparse, hashlib, json, math, random, subprocess, time +import pathlib from pathlib import Path import torch @@ -36,10 +37,27 @@ from peft import LoraConfig, get_peft_model TARGETS = ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"] -SYS = ("You expand a single story beat into ONE paragraph of prose in the manner of Rebecca " - "Yarros — contemporary first-person PRESENT-tense narration, emotionally charged, sensory " - "and physical, the voice of new-adult romantasy. Render the beat itself; do not move past " - "it, do not add a new scene, do not comment. Output the paragraph only, 90–140 words.") +# ⚠ THE SYSTEM PROMPT MUST MATCH THE ONE THE PAIRS WERE BUILT UNDER, and the pair file now +# records it. Reading it from the pairs' provenance rather than hardcoding it here is the fix +# for the defect that shaped the whole BabyYarros pilot: the trainer said "ONE paragraph" while +# the data was a median of four, and the carrier believed the data. A literal in this file is a +# second place for that to drift. +SYS_YARROS_FROZEN = ( + "You expand a single story beat into ONE paragraph of prose in the manner of Rebecca " + "Yarros — contemporary first-person PRESENT-tense narration, emotionally charged, sensory " + "and physical, the voice of new-adult romantasy. Render the beat itself; do not move past " + "it, do not add a new scene, do not comment. Output the paragraph only, 90–140 words.") + + +def resolve_sys(pairs_path: str) -> tuple[str, str]: + """Return (system_prompt, source). Prefer the pair build's own provenance.""" + prov = pathlib.Path(str(pairs_path) + ".provenance.json") + if prov.exists(): + d = json.loads(prov.read_text(encoding="utf-8")) + s = d.get("system_prompt") + if s: + return s, f"pairs provenance ({d.get('register', '?')})" + return SYS_YARROS_FROZEN, "frozen Yarros literal (no provenance found)" def user_msg(rec: dict) -> str: @@ -57,11 +75,11 @@ class Pairs(Dataset): when a template changes its spacing, and it would mask the wrong span without erroring. """ - def __init__(self, tok, records, seq_len): + def __init__(self, tok, records, seq_len, sys_prompt): self.rows = [] dropped = 0 for r in records: - msgs = [{"role": "system", "content": SYS}, {"role": "user", "content": user_msg(r)}] + msgs = [{"role": "system", "content": sys_prompt}, {"role": "user", "content": user_msg(r)}] # ⚠ transformers 5.16 returns a BatchEncoding from apply_chat_template(tokenize=True), # not a list of ids. Taking len() of it yields 2 (the number of keys), so every # example failed the `len(full) <= len(prefix)` test and the whole dataset was @@ -146,6 +164,8 @@ def main() -> int: torch.manual_seed(a.seed); random.seed(a.seed) out = Path(a.out); out.mkdir(parents=True, exist_ok=True) + SYS, sys_src = resolve_sys(a.pairs) + print(f"[data] system prompt from {sys_src}") tok = AutoTokenizer.from_pretrained(a.base) if tok.chat_template is None: raise SystemExit("REFUSING: carrier has no chat template -- it is not an instruct build") @@ -169,7 +189,8 @@ def main() -> int: if any(r.get("split") != "val" for r in val_recs): raise SystemExit("REFUSING: val pairs are not all from the corpus val split") - train_ds, val_ds = Pairs(tok, train_recs, a.seq_len), Pairs(tok, val_recs, a.seq_len) + train_ds = Pairs(tok, train_recs, a.seq_len, SYS) + val_ds = Pairs(tok, val_recs, a.seq_len, SYS) with_ctx = sum(1 for r in train_recs if r.get("context")) print(f"[data] {len(train_ds)} train pairs ({train_ds.dropped} dropped), " f"{len(val_ds)} val ({val_ds.dropped} dropped), " @@ -213,7 +234,8 @@ def main() -> int: "trainable_pct": round(100 * trainable / total, 3), "steps_per_epoch": steps_per_epoch, "planned_steps": steps_per_epoch * int(a.epochs), "resolved": resolved, "harness_commit": git, "harness_dirty_at_launch": dirty, - "launched_at": time.strftime("%Y-%m-%dT%H:%M:%S%z"), "system_prompt": SYS} + "launched_at": time.strftime("%Y-%m-%dT%H:%M:%S%z"), "system_prompt": SYS, + "system_prompt_source": sys_src} (out / "provenance.json").write_text(json.dumps(prov, indent=2)) print("[prov] " + json.dumps({k: prov[k] for k in ("pairs_sha256_16", "train_pairs", "planned_steps", "trainable_pct", diff --git a/services/grok-token-broker/README.md b/services/grok-token-broker/README.md new file mode 100644 index 0000000..55e91cd --- /dev/null +++ b/services/grok-token-broker/README.md @@ -0,0 +1,67 @@ +# grok-token-broker + +Holds a refreshable xAI **session** credential so a gateway can serve the Grok Build coding +plan — without touching the credential the operator's `grok` CLI depends on. + + seed copy the CLI credential into the broker's own store (read-once) + probe-rotation MEASURE whether refresh rotates AND invalidates <- gate + token --raw emit a valid access token, refreshing if near expiry + refresh refresh now + status report state, touching nothing + +## The problem + +The coding plan (`https://cli-chat-proxy.grok.com/v1`, serving **grok-4.6**, 500k context) +authenticates with an OIDC session token — `api_key: null`, `env_key: null` in the CLI's own +model cache — that expires roughly every six hours. LiteLLM and every other gateway here hold +a **static** credential, so a naive alias works until the session lapses and then fails closed. + +## ⛔ Why the loop is disarmed until measured + +`https://auth.x.ai/oauth2/token` supports the refresh grant. But **a refresh may rotate the +refresh token, and many providers invalidate the old one server-side the moment a new one +issues.** The `grok` CLI holds its refresh token in `~/.grok/auth.json`. + +**Not writing to that file is necessary and not sufficient.** If xAI rotates-and-invalidates, +a broker refreshing the *same* credential kills the CLI login anyway, server-side. heid's +`groa_http_dispatch.py` refuses to refresh at all for exactly this reason — correct, absent an +answer. This broker exists to get the answer, then act on it. + +So `probe-rotation` is a gate, not a diagnostic: `refresh` and `token`-near-expiry both refuse +until a verdict exists and says safe. + +| verdict | meaning | +|---|---| +| `non-rotating` | same refresh token returned. Coexistence safe. Armed. | +| `rotating-old-still-valid` | rotation happens, old token still works. Coexistence safe. Armed. | +| `rotating-and-invalidating` | **the CLI login is already dead.** Broker must not share this credential — it needs its own login. Stays disarmed. | + +⚠ **The probe spends one refresh, and there is no way to ask the question without spending +it.** If the answer is the bad one, the CLI is broken at that moment and needs an interactive +`grok` re-login — the same remedy heid's exit-3 path already names. **Run it when a broken CLI +login is a two-minute annoyance, not mid-panel.** Hence the required +`--i-accept-this-may-end-the-cli-session` flag. + +## Current state + +Seeded on **nh3-dev** 2026-09-16. Scope `openid profile email offline_access grok-cli:access +api:access`, client `b1a00492-…`. **Probe NOT yet run** — it is the operator's call when to +spend it. + +## For a consumer + +```bash +TOKEN=$(broker.py token --raw) +curl https://cli-chat-proxy.grok.com/v1/... -H "Authorization: Bearer $TOKEN" +``` + +⚠ **The coding plan speaks the Responses API (`api_backend: "responses"`), not +`/chat/completions`.** That is a second, independent obstacle to a LiteLLM alias and this +broker does not solve it — it solves the credential half only. + +## Invariants + +- Never writes `~/.grok/auth.json`. Reads it once, on `seed`. +- Never prints a token except under `token --raw`. +- Never logs a token value — expiry, scope and subject only. +- Store is `~/.config/grok-token-broker/` at 0700, credential at 0600. diff --git a/services/grok-token-broker/broker.py b/services/grok-token-broker/broker.py new file mode 100755 index 0000000..c8f232e --- /dev/null +++ b/services/grok-token-broker/broker.py @@ -0,0 +1,297 @@ +#!/usr/bin/env python3 +"""Grok token broker — hold a refreshable xAI session credential WITHOUT touching the CLI's. + +⭐ THE PROBLEM. The Grok Build coding plan (`https://cli-chat-proxy.grok.com/v1`, serving +grok-4.6) authenticates with an OIDC **session** token, not an API key: `api_key: null`, +`env_key: null` in the CLI's own model cache. It expires roughly every six hours. LiteLLM and +every other gateway on this fleet hold a STATIC credential, so a naive alias works until the +session lapses and then fails closed. + +⚠⚠ THE HAZARD THAT SHAPES THIS ENTIRE DESIGN, AND WHY IT IS NOT A SIMPLE REFRESH LOOP. +`https://auth.x.ai/oauth2/token` supports the `refresh_token` grant. But **a refresh may ROTATE +the refresh token**, and many OIDC providers invalidate the old one server-side the instant a +new one is issued. The operator's interactive `grok` CLI holds its refresh token in +`~/.grok/auth.json`. If xAI rotates-and-invalidates, then a broker refreshing against the SAME +credential silently kills the CLI login — and it does so **even though the broker never writes +to that file.** Not-writing is necessary and NOT sufficient; the damage would be server-side. + +heid's `groa_http_dispatch.py` refuses to refresh at all for this reason, which is the correct +call absent an answer. This broker's job is to get the answer first and then act on it. + +⛔ SO THE REFRESH LOOP IS NOT ARMED UNTIL `probe-rotation` HAS RUN AND SAID IT IS SAFE. +`serve` and `refresh` both refuse without a stored probe verdict. That is deliberate: this is a +tool whose whole value is contingent on one unknown, and running it before measuring the unknown +is how you find out by breaking the operator's login. + +## The probe, and what each outcome means + +`probe-rotation` performs ONE refresh using a COPY of the credential, then re-presents the +ORIGINAL refresh token: + + original still valid -> NON-ROTATING (or rotation without invalidation). The broker can + coexist with the CLI on one credential. Loop may be armed. + original now rejected -> ROTATING-AND-INVALIDATING. The broker MUST NOT share the CLI's + credential. It needs its own login, or the idea is dead. Refuses + to arm, and says so. + +⚠ The probe itself consumes one refresh. If the answer is "rotating", the CLI's token is +already dead by the time you read the verdict — there is no way to ask this question without +spending the credential, and pretending otherwise would be worse. The remedy is a `grok` CLI +re-login, which is the same remedy heid's exit-3 path already names. **Run it when a broken CLI +login is a two-minute annoyance, not mid-panel.** + +## What it never does + +- Never writes `~/.grok/auth.json`. Reads it once, on `seed`, and never again. +- Never prints a token to stdout except under `token --raw`, which is for a consumer to capture. +- Never logs a token value. Expiry, scope and subject only. +""" +from __future__ import annotations + +import argparse +import base64 +import json +import os +import pathlib +import sys +import time +import urllib.error +import urllib.parse +import urllib.request + +CLI_AUTH = pathlib.Path.home() / ".grok" / "auth.json" +STORE = pathlib.Path(os.environ.get("GROK_BROKER_HOME", + pathlib.Path.home() / ".config" / "grok-token-broker")) +CRED = STORE / "credential.json" +VERDICT = STORE / "rotation-verdict.json" +TOKEN_URL = "https://auth.x.ai/oauth2/token" +API_BASE = "https://cli-chat-proxy.grok.com/v1" +REFRESH_MARGIN_S = 1800 # refresh when under 30 min remain + + +def _claims(jwt: str) -> dict: + pl = jwt.split(".")[1] + pl += "=" * (-len(pl) % 4) + return json.loads(base64.urlsafe_b64decode(pl)) + + +def _remaining(jwt: str) -> int: + try: + return int(_claims(jwt)["exp"] - time.time()) + except Exception: + return -1 + + +def _post_form(url: str, fields: dict, timeout: int = 30): + data = urllib.parse.urlencode(fields).encode() + req = urllib.request.Request(url, data=data, + headers={"Content-Type": "application/x-www-form-urlencoded"}) + try: + with urllib.request.urlopen(req, timeout=timeout) as r: + return r.status, json.loads(r.read()) + except urllib.error.HTTPError as e: + body = e.read().decode("utf-8", "replace")[:400] + return e.code, {"_error": body} + except Exception as e: # noqa: BLE001 + return 0, {"_error": f"{type(e).__name__}: {e}"} + + +def _read_cli_credential() -> dict: + """Pull access token, refresh token and client id out of the CLI's auth file. + + Shape-tolerant on purpose: this file belongs to another tool and its layout is not ours to + depend on. Anything not found is reported, never guessed. + """ + if not CLI_AUTH.exists(): + raise SystemExit(f"REFUSING: no CLI credential at {CLI_AUTH}") + blob = json.loads(CLI_AUTH.read_text(encoding="utf-8")) + found: dict = {} + + def walk(o, key_hint=""): + if isinstance(o, dict): + for k, v in o.items(): + if isinstance(v, str): + if k == "key" and v.count(".") == 2 and "access" not in found: + found["access"] = v + elif "refresh" in k.lower() and len(v) > 20: + found.setdefault("refresh", v) + walk(v, k) + if "auth.x.ai::" in key_hint: + found.setdefault("client_id", key_hint.split("::", 1)[1]) + elif isinstance(o, list): + for v in o: + walk(v, key_hint) + + for k, v in blob.items(): + walk(v, k) + missing = [f for f in ("access", "refresh") if f not in found] + if missing: + raise SystemExit(f"REFUSING: could not locate {missing} in {CLI_AUTH}. " + "Its layout changed; inspect it by hand rather than letting this guess.") + return found + + +def cmd_seed(_a) -> int: + STORE.mkdir(parents=True, exist_ok=True) + os.chmod(STORE, 0o700) + c = _read_cli_credential() + CRED.write_text(json.dumps({**c, "seeded_at": time.strftime("%Y-%m-%dT%H:%M:%S%z"), + "source": str(CLI_AUTH)}, indent=2), encoding="utf-8") + os.chmod(CRED, 0o600) + cl = _claims(c["access"]) + print(f" seeded -> {CRED} (0600)") + print(f" scope : {cl.get('scope')}") + print(f" expires : {time.strftime('%Y-%m-%dT%H:%M:%SZ', time.gmtime(cl['exp']))} " + f"({_remaining(c['access']) // 60} min left)") + print(f" client : {c.get('client_id', '')}") + print("\n ⛔ The refresh loop is NOT armed. Run `probe-rotation` next, and read its warning " + "first — the probe spends one refresh and may end the CLI's session.") + return 0 + + +def cmd_probe(a) -> int: + if not CRED.exists(): + raise SystemExit("REFUSING: run `seed` first") + c = json.loads(CRED.read_text(encoding="utf-8")) + if not a.i_accept_this_may_end_the_cli_session: + print("REFUSING without --i-accept-this-may-end-the-cli-session.\n\n" + " This probe performs a real refresh. If xAI rotates AND invalidates, the\n" + " operator's `grok` CLI login dies at that moment and needs an interactive\n" + " re-login. There is no way to answer the question without spending the\n" + " credential. Run it when that is a two-minute annoyance, not mid-panel.", + file=sys.stderr) + return 2 + + old_refresh = c["refresh"] + fields = {"grant_type": "refresh_token", "refresh_token": old_refresh} + if c.get("client_id"): + fields["client_id"] = c["client_id"] + + print(" [1/2] refreshing with the stored refresh token ...") + st, body = _post_form(TOKEN_URL, fields) + if st != 200: + print(f" refresh FAILED http={st}: {str(body.get('_error'))[:200]}", file=sys.stderr) + print(" Verdict NOT recorded. The loop stays disarmed.", file=sys.stderr) + return 4 + new_refresh = body.get("refresh_token") + rotated = bool(new_refresh and new_refresh != old_refresh) + print(f" ok. refresh_token {'ROTATED' if rotated else 'unchanged'}") + + if not rotated: + verdict, safe = "non-rotating", True + note = "Provider returned the same refresh token; the CLI's copy is untouched." + else: + print(" [2/2] re-presenting the ORIGINAL refresh token to see if it still works ...") + st2, _ = _post_form(TOKEN_URL, fields) + if st2 == 200: + verdict, safe = "rotating-old-still-valid", True + note = ("Rotation happens but the old refresh token still works, so the CLI's copy " + "survives. Coexistence is safe.") + else: + verdict, safe = "rotating-and-invalidating", False + note = (f"The original refresh token is now rejected (http={st2}). The CLI's login " + "is ALREADY DEAD and needs an interactive `grok` re-login. A broker MUST " + "NOT share this credential — it needs its own login.") + print(f" original token -> http={st2}") + + c["access"] = body.get("access_token", c["access"]) + if new_refresh: + c["refresh"] = new_refresh + CRED.write_text(json.dumps(c, indent=2), encoding="utf-8") + os.chmod(CRED, 0o600) + + VERDICT.write_text(json.dumps({ + "verdict": verdict, "safe_to_arm": safe, "note": note, + "probed_at": time.strftime("%Y-%m-%dT%H:%M:%S%z"), + "token_url": TOKEN_URL, "api_base": API_BASE, + }, indent=2), encoding="utf-8") + os.chmod(VERDICT, 0o600) + print(f"\n VERDICT: {verdict} (safe_to_arm={safe})\n {note}") + if not safe: + print("\n ⚠ TELL THE OPERATOR: the `grok` CLI on this box needs a re-login.", file=sys.stderr) + return 0 if safe else 5 + + +def _ensure_armed() -> dict: + if not VERDICT.exists(): + raise SystemExit("REFUSING: no rotation verdict. Run `probe-rotation` first — this tool " + "will not refresh a credential whose rotation behaviour is unmeasured.") + v = json.loads(VERDICT.read_text(encoding="utf-8")) + if not v.get("safe_to_arm"): + raise SystemExit(f"REFUSING: verdict is '{v['verdict']}'. {v['note']}") + return v + + +def cmd_refresh(_a) -> int: + _ensure_armed() + c = json.loads(CRED.read_text(encoding="utf-8")) + fields = {"grant_type": "refresh_token", "refresh_token": c["refresh"]} + if c.get("client_id"): + fields["client_id"] = c["client_id"] + st, body = _post_form(TOKEN_URL, fields) + if st != 200: + print(f"refresh failed http={st}: {str(body.get('_error'))[:200]}", file=sys.stderr) + return 4 + c["access"] = body["access_token"] + if body.get("refresh_token"): + c["refresh"] = body["refresh_token"] + c["refreshed_at"] = time.strftime("%Y-%m-%dT%H:%M:%S%z") + CRED.write_text(json.dumps(c, indent=2), encoding="utf-8") + os.chmod(CRED, 0o600) + print(f" refreshed; {_remaining(c['access']) // 60} min remaining") + return 0 + + +def cmd_token(a) -> int: + """Print a currently-valid access token, refreshing first if it is close to expiry.""" + if not CRED.exists(): + raise SystemExit("REFUSING: run `seed` first") + c = json.loads(CRED.read_text(encoding="utf-8")) + if _remaining(c["access"]) < REFRESH_MARGIN_S: + _ensure_armed() + if cmd_refresh(a) != 0: + return 4 + c = json.loads(CRED.read_text(encoding="utf-8")) + if a.raw: + sys.stdout.write(c["access"]) + else: + print(f" valid for {_remaining(c['access']) // 60} more minutes " + f"(use --raw to emit the token itself)") + return 0 + + +def cmd_status(_a) -> int: + print(f" store : {STORE}") + if not CRED.exists(): + print(" state : NOT SEEDED") + return 0 + c = json.loads(CRED.read_text(encoding="utf-8")) + print(f" seeded : {c.get('seeded_at')} from {c.get('source')}") + print(f" access : {_remaining(c['access']) // 60} min remaining") + if VERDICT.exists(): + v = json.loads(VERDICT.read_text(encoding="utf-8")) + print(f" rotation: {v['verdict']} safe_to_arm={v['safe_to_arm']} ({v['probed_at']})") + print(f" armed : {'YES' if v['safe_to_arm'] else 'NO — ' + v['verdict']}") + else: + print(" rotation: UNMEASURED — loop disarmed, `token` will refuse to refresh") + print(f" api_base: {API_BASE}") + return 0 + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__.split("\n")[0]) + sub = ap.add_subparsers(dest="cmd", required=True) + sub.add_parser("seed", help="copy the CLI credential into the broker's own store") + p = sub.add_parser("probe-rotation", help="measure whether refresh rotates AND invalidates") + p.add_argument("--i-accept-this-may-end-the-cli-session", action="store_true") + sub.add_parser("refresh", help="refresh now (requires a safe verdict)") + t = sub.add_parser("token", help="emit a valid access token, refreshing if near expiry") + t.add_argument("--raw", action="store_true") + sub.add_parser("status", help="report state without touching anything") + a = ap.parse_args() + return {"seed": cmd_seed, "probe-rotation": cmd_probe, "refresh": cmd_refresh, + "token": cmd_token, "status": cmd_status}[a.cmd](a) + + +if __name__ == "__main__": + raise SystemExit(main())