grok-token-broker: hold a refreshable session credential behind a rotation-safety gate

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",

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.

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', '<not found>')}")
+    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())