aws-sso: keep the Bedrock SSO session alive by running aws sso login before it ends
Like Claude Code's awsAuthRefresh: the user still approves each login. mode mac runs it on the Docker host over SSH (own browser); mode device shows the device code in pi for SSH/tmux use. Off by default (/aws-sso on). Session start inferred from refresh-token changes; 8 h default, measured.
This commit is contained in:
pi@mbp-m1-2020
@@ -0,0 +1,370 @@
|
||||
/**
|
||||
* aws-sso.ts — keep the Bedrock AWS SSO session alive, the way Claude Code's
|
||||
* `awsAuthRefresh` does: when the session nears its end (or has expired), run
|
||||
* `aws sso login` for you. You still approve the login yourself.
|
||||
*
|
||||
* Modes (`/aws-sso mode …`):
|
||||
* mac – run `aws sso login` on the Docker host over SSH (pi-devbox), so your
|
||||
* own browser opens. For when you sit at the laptop.
|
||||
* device – run `aws sso login --use-device-code --no-browser` here and show the
|
||||
* link + code in pi. Approve on any device (phone, other computer).
|
||||
* For SSH/tmux sessions.
|
||||
*
|
||||
* Inert until enabled: `/aws-sso on` (persists in ~/.pi/agent/aws-sso/state.json),
|
||||
* or env PI_AWS_SSO_REFRESH=1. PI_AWS_SSO_REFRESH=0 is a hard off switch.
|
||||
*
|
||||
* Why a timer is needed at all: within one IAM Identity Center session the AWS
|
||||
* SDK already refreshes the 1-hour access token by itself. The session itself has
|
||||
* a fixed lifetime set by the IdC admin (8 h measured here); after that only a new
|
||||
* `aws sso login` helps. The session start is not recorded anywhere, so it is
|
||||
* inferred from the refresh token: a new refresh token = a new session.
|
||||
* Tokens never leave this process except as sha256[:8] fingerprints.
|
||||
*/
|
||||
import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
|
||||
import { spawn } from "node:child_process";
|
||||
import { createHash } from "node:crypto";
|
||||
import { existsSync, mkdirSync, readFileSync, renameSync, statSync, writeFileSync, openSync, closeSync, unlinkSync } from "node:fs";
|
||||
import { homedir } from "node:os";
|
||||
import { join } from "node:path";
|
||||
|
||||
// ── pure helpers (tested in test/aws-sso.test.mjs) ─────────────────────────────
|
||||
|
||||
export function fingerprint(value: string | undefined | null): string {
|
||||
if (!value) return "-";
|
||||
return createHash("sha256").update(String(value)).digest("hex").slice(0, 8);
|
||||
}
|
||||
|
||||
/** botocore's cache file name for an sso-session: sha1(session name). */
|
||||
export function cacheKeyFor(sessionName: string): string {
|
||||
return createHash("sha1").update(sessionName).digest("hex");
|
||||
}
|
||||
|
||||
export function parseIni(text: string): Record<string, Record<string, string>> {
|
||||
const out: Record<string, Record<string, string>> = {};
|
||||
let cur: string | null = null;
|
||||
for (const raw of String(text).split(/\r?\n/)) {
|
||||
const line = raw.trim();
|
||||
if (!line || line.startsWith("#") || line.startsWith(";")) continue;
|
||||
const sec = line.match(/^\[\s*(.+?)\s*\]$/);
|
||||
if (sec) { cur = sec[1].replace(/\s+/g, " "); out[cur] ??= {}; continue; }
|
||||
const kv = line.match(/^([^=]+?)\s*=\s*(.*)$/);
|
||||
if (kv && cur) out[cur][kv[1].trim()] = kv[2].trim();
|
||||
}
|
||||
return out;
|
||||
}
|
||||
|
||||
export interface SsoSession { sessionName: string; startUrl: string; region: string | null; chain: string[] }
|
||||
|
||||
/** Follow source_profile from `profile` to the profile that names an sso_session. */
|
||||
export function resolveSsoSession(configText: string, profile: string): SsoSession {
|
||||
const ini = parseIni(configText);
|
||||
const chain: string[] = [];
|
||||
let name = profile;
|
||||
for (let hop = 0; hop < 10; hop++) {
|
||||
if (chain.includes(name)) throw new Error(`source_profile loop: ${[...chain, name].join(" -> ")}`);
|
||||
chain.push(name);
|
||||
const p = name === "default" ? (ini["default"] ?? ini["profile default"]) : ini[`profile ${name}`];
|
||||
if (!p) throw new Error(`profile '${name}' not found in ~/.aws/config`);
|
||||
if (p.sso_session) {
|
||||
const s = ini[`sso-session ${p.sso_session}`];
|
||||
if (!s) throw new Error(`[sso-session ${p.sso_session}] not found (referenced by profile '${name}')`);
|
||||
if (!s.sso_start_url) throw new Error(`[sso-session ${p.sso_session}] has no sso_start_url`);
|
||||
return { sessionName: p.sso_session, startUrl: s.sso_start_url, region: s.sso_region ?? null, chain };
|
||||
}
|
||||
if (p.sso_start_url) throw new Error(`profile '${name}' uses legacy sso_start_url; convert it to an [sso-session] block (legacy sessions cannot refresh)`);
|
||||
if (!p.source_profile) throw new Error(`profile '${name}' has neither sso_session nor source_profile`);
|
||||
name = p.source_profile;
|
||||
}
|
||||
throw new Error(`source_profile chain too long: ${chain.join(" -> ")}`);
|
||||
}
|
||||
|
||||
/** Device-code output → { code, url }. The URL must be on the configured start URL's host. */
|
||||
export function parseDeviceOutput(text: string, startUrl: string): { code: string; url: string } | null {
|
||||
const want = new URL(startUrl).host;
|
||||
const m = String(text).match(/https:\/\/\S*user_code=([A-Z0-9-]+)/);
|
||||
if (!m) return null;
|
||||
let url: URL;
|
||||
try { url = new URL(m[0]); } catch { return null; }
|
||||
if (url.host !== want || !/^[A-Z0-9]{4}-[A-Z0-9]{4}$/.test(m[1])) return null;
|
||||
return { code: m[1], url: url.toString() };
|
||||
}
|
||||
|
||||
export function isSsoExpiredError(message: string): boolean {
|
||||
return /token (is )?expired|session (has )?expired|invalid_grant|ExpiredToken|The SSO session associated with this profile|run 'aws sso login'/i.test(String(message));
|
||||
}
|
||||
|
||||
export interface State {
|
||||
enabled?: boolean;
|
||||
mode?: "mac" | "device";
|
||||
sessionHours?: number;
|
||||
leadMinutes?: number;
|
||||
anchor?: { refreshFp: string; since: number };
|
||||
lastLogin?: { at: number; ok: boolean; mode: string; detail?: string };
|
||||
}
|
||||
|
||||
export const DEFAULTS = { sessionHours: 8, leadMinutes: 20 };
|
||||
|
||||
/**
|
||||
* Where the session stands. `since` is when we first saw this refresh token.
|
||||
* none – no cached token
|
||||
* expired – the 1-hour access token has expired AND the session is past its end
|
||||
* due – inside the lead window before the session's end
|
||||
* ok – healthy
|
||||
*/
|
||||
export function assess(state: State, token: { refreshFp: string; accessExpiresAt: number } | null, now: number) {
|
||||
const hours = state.sessionHours ?? DEFAULTS.sessionHours;
|
||||
const lead = (state.leadMinutes ?? DEFAULTS.leadMinutes) * 60e3;
|
||||
if (!token) return { status: "none" as const, endsAt: null as number | null };
|
||||
const since = state.anchor && state.anchor.refreshFp === token.refreshFp ? state.anchor.since : null;
|
||||
const endsAt = since === null ? null : since + hours * 3600e3;
|
||||
if (endsAt !== null && now >= endsAt && now >= token.accessExpiresAt) return { status: "expired" as const, endsAt };
|
||||
if (endsAt !== null && now >= endsAt - lead) return { status: "due" as const, endsAt };
|
||||
return { status: "ok" as const, endsAt };
|
||||
}
|
||||
|
||||
/** New anchor when the refresh token changed; null when unchanged. */
|
||||
export function reanchor(state: State, refreshFp: string, firstSeen: number): State["anchor"] | null {
|
||||
if (state.anchor && state.anchor.refreshFp === refreshFp) return null;
|
||||
return { refreshFp, since: firstSeen };
|
||||
}
|
||||
|
||||
export function fmtLeft(ms: number): string {
|
||||
if (ms <= 0) return "0m";
|
||||
const m = Math.round(ms / 60e3);
|
||||
return m >= 60 ? `${Math.floor(m / 60)}h${String(m % 60).padStart(2, "0")}m` : `${m}m`;
|
||||
}
|
||||
|
||||
// ── I/O ───────────────────────────────────────────────────────────────────────
|
||||
|
||||
const DIR = join(homedir(), ".pi", "agent", "aws-sso");
|
||||
const STATE = join(DIR, "state.json");
|
||||
const LOCK = join(DIR, "login.lock");
|
||||
const LOG = join(DIR, "events.jsonl");
|
||||
const SSH_CONFIG = join(homedir(), ".ssh-local", "config");
|
||||
|
||||
function loadState(): State {
|
||||
try { return JSON.parse(readFileSync(STATE, "utf8")); } catch { return {}; }
|
||||
}
|
||||
function saveState(s: State) {
|
||||
mkdirSync(DIR, { recursive: true, mode: 0o700 });
|
||||
const tmp = `${STATE}.${process.pid}.tmp`;
|
||||
writeFileSync(tmp, JSON.stringify(s, null, 2));
|
||||
renameSync(tmp, STATE);
|
||||
}
|
||||
function logEvent(event: string, extra: Record<string, unknown> = {}) {
|
||||
try {
|
||||
mkdirSync(DIR, { recursive: true, mode: 0o700 });
|
||||
writeFileSync(LOG, JSON.stringify({ ts: new Date().toISOString(), pid: process.pid, event, ...extra }) + "\n", { flag: "a" });
|
||||
} catch { /* logging must never break pi */ }
|
||||
}
|
||||
|
||||
function ssoSession(): SsoSession | { error: string } {
|
||||
const profile = process.env.AWS_PROFILE;
|
||||
if (!profile) return { error: "AWS_PROFILE is not set" };
|
||||
try { return resolveSsoSession(readFileSync(join(homedir(), ".aws", "config"), "utf8"), profile); }
|
||||
catch (e) { return { error: (e as Error).message }; }
|
||||
}
|
||||
|
||||
function readToken(sessionName: string): { refreshFp: string; accessExpiresAt: number; mtime: number } | null {
|
||||
const f = join(homedir(), ".aws", "sso", "cache", `${cacheKeyFor(sessionName)}.json`);
|
||||
try {
|
||||
const d = JSON.parse(readFileSync(f, "utf8"));
|
||||
return { refreshFp: fingerprint(d.refreshToken), accessExpiresAt: Date.parse(d.expiresAt), mtime: statSync(f).mtimeMs };
|
||||
} catch { return null; }
|
||||
}
|
||||
|
||||
/** Cross-process lock so several pi processes never start several logins. */
|
||||
function tryLock(): boolean {
|
||||
mkdirSync(DIR, { recursive: true, mode: 0o700 });
|
||||
try { closeSync(openSync(LOCK, "wx")); writeFileSync(LOCK, String(process.pid)); return true; }
|
||||
catch {
|
||||
try {
|
||||
const pid = Number(readFileSync(LOCK, "utf8"));
|
||||
if (pid && pid !== process.pid) { try { process.kill(pid, 0); return false; } catch { /* stale */ } }
|
||||
unlinkSync(LOCK);
|
||||
closeSync(openSync(LOCK, "wx")); writeFileSync(LOCK, String(process.pid)); return true;
|
||||
} catch { return false; }
|
||||
}
|
||||
}
|
||||
function unlock() { try { if (Number(readFileSync(LOCK, "utf8")) === process.pid) unlinkSync(LOCK); } catch { /* gone */ } }
|
||||
|
||||
function defaultMode(): "mac" | "device" { return existsSync(SSH_CONFIG) ? "mac" : "device"; }
|
||||
|
||||
type Ui = { notify(m: string, t?: "info" | "warning" | "error"): void; setStatus(k: string, t: string | undefined): void };
|
||||
|
||||
/** Run one login. Resolves true when the cached refresh token changed. */
|
||||
function runLogin(mode: "mac" | "device", s: SsoSession, ui: Ui | null, timeoutMs = 10 * 60e3): Promise<{ ok: boolean; detail: string }> {
|
||||
const before = readToken(s.sessionName)?.refreshFp ?? "-";
|
||||
const args = mode === "mac"
|
||||
// Login shell on the host: a non-interactive ssh PATH may not include aws.
|
||||
? ["-F", SSH_CONFIG, "-o", "BatchMode=yes", "-o", "ConnectTimeout=8", process.env.PI_AWS_SSO_SSH_HOST || "host",
|
||||
`exec "\${SHELL:-/bin/zsh}" -lc 'aws sso login --sso-session ${s.sessionName.replace(/[^\w.-]/g, "")}'`]
|
||||
: ["sso", "login", "--sso-session", s.sessionName, "--use-device-code", "--no-browser"];
|
||||
const cmd = mode === "mac" ? "ssh" : "aws";
|
||||
logEvent("login_start", { mode, before });
|
||||
return new Promise((resolve) => {
|
||||
let out = "";
|
||||
let shown = false;
|
||||
const child = spawn(cmd, args, { env: { ...process.env, AWS_PAGER: "" }, stdio: ["ignore", "pipe", "pipe"] });
|
||||
const onData = (b: Buffer) => {
|
||||
out += b.toString();
|
||||
if (mode === "device" && !shown) {
|
||||
const d = parseDeviceOutput(out, s.startUrl);
|
||||
if (d) {
|
||||
shown = true;
|
||||
ui?.notify(`AWS SSO: approve code ${d.code} at\n${d.url}`, "warning");
|
||||
ui?.setStatus("aws-sso", `AWS ⏳ approve ${d.code}`);
|
||||
logEvent("device_code_shown", { code: d.code });
|
||||
}
|
||||
}
|
||||
};
|
||||
child.stdout.on("data", onData);
|
||||
child.stderr.on("data", onData);
|
||||
const timer = setTimeout(() => child.kill("SIGTERM"), timeoutMs);
|
||||
child.on("close", (rc) => {
|
||||
clearTimeout(timer);
|
||||
const after = readToken(s.sessionName)?.refreshFp ?? "-";
|
||||
const ok = rc === 0 && after !== before;
|
||||
const detail = ok ? "new session" : (out.trim().split("\n").filter(Boolean).pop() ?? `exit ${rc}`).slice(0, 200);
|
||||
logEvent("login_end", { mode, rc, before, after, ok, detail: ok ? undefined : detail });
|
||||
resolve({ ok, detail });
|
||||
});
|
||||
child.on("error", (e) => { clearTimeout(timer); resolve({ ok: false, detail: e.message }); });
|
||||
});
|
||||
}
|
||||
|
||||
// ── extension ─────────────────────────────────────────────────────────────────
|
||||
|
||||
export default function (pi: ExtensionAPI) {
|
||||
let timer: ReturnType<typeof setInterval> | null = null;
|
||||
let inFlight: Promise<{ ok: boolean; detail: string }> | null = null;
|
||||
let ui: Ui | null = null;
|
||||
let warnedEnd: number | null = null;
|
||||
|
||||
const enabled = () => process.env.PI_AWS_SSO_REFRESH === "0" ? false : process.env.PI_AWS_SSO_REFRESH === "1" ? true : !!loadState().enabled;
|
||||
|
||||
/** Observe the token, keep the anchor current, return the assessment. */
|
||||
function observe() {
|
||||
const s = ssoSession();
|
||||
if ("error" in s) return { s: null, error: s.error, a: null as ReturnType<typeof assess> | null, token: null };
|
||||
const token = readToken(s.sessionName);
|
||||
const state = loadState();
|
||||
if (token) {
|
||||
// First sight of a refresh token: its file mtime is when it was written (≈ login time).
|
||||
const next = reanchor(state, token.refreshFp, Math.min(token.mtime, Date.now()));
|
||||
if (next) { state.anchor = next; saveState(state); logEvent("anchor", { refreshFp: next.refreshFp, since: new Date(next.since).toISOString() }); }
|
||||
}
|
||||
return { s, error: null, a: assess(state, token, Date.now()), token };
|
||||
}
|
||||
|
||||
function render() {
|
||||
if (!ui) return;
|
||||
if (!enabled()) { ui.setStatus("aws-sso", undefined); return; }
|
||||
if (inFlight) return; // the login path owns the status while it runs
|
||||
const o = observe();
|
||||
if (o.error || !o.a) { ui.setStatus("aws-sso", `AWS ✗ ${o.error}`); return; }
|
||||
const left = o.a.endsAt === null ? null : o.a.endsAt - Date.now();
|
||||
const txt = o.a.status === "none" ? "AWS ✗ not logged in"
|
||||
: o.a.status === "expired" ? "AWS ✗ session expired"
|
||||
: left === null ? "AWS ✓" : `AWS ${o.a.status === "due" ? "⚠" : "✓"} ${fmtLeft(left)}`;
|
||||
ui.setStatus("aws-sso", txt);
|
||||
}
|
||||
|
||||
async function login(reason: string): Promise<boolean> {
|
||||
if (inFlight) return (await inFlight).ok;
|
||||
const o = observe();
|
||||
if (!o.s) { ui?.notify(`AWS SSO: ${o.error}`, "error"); return false; }
|
||||
if (!tryLock()) { ui?.notify("AWS SSO: another pi process is already running the login — approve it there.", "info"); return false; }
|
||||
const state = loadState();
|
||||
const mode = state.mode ?? defaultMode();
|
||||
ui?.setStatus("aws-sso", mode === "mac" ? "AWS ⏳ approve in your browser" : "AWS ⏳ starting device login…");
|
||||
if (mode === "mac") ui?.notify(`AWS SSO (${reason}): a login page is opening in your browser on the host — approve it there.`, "warning");
|
||||
logEvent("login_trigger", { reason, mode });
|
||||
inFlight = runLogin(mode, o.s, ui);
|
||||
try {
|
||||
const r = await inFlight;
|
||||
const st = loadState();
|
||||
st.lastLogin = { at: Date.now(), ok: r.ok, mode, detail: r.ok ? undefined : r.detail };
|
||||
saveState(st);
|
||||
ui?.notify(r.ok ? "AWS SSO: new session — good for the next ~" + (st.sessionHours ?? DEFAULTS.sessionHours) + " h." : `AWS SSO login did not complete: ${r.detail}`, r.ok ? "info" : "error");
|
||||
return r.ok;
|
||||
} finally { inFlight = null; unlock(); render(); }
|
||||
}
|
||||
|
||||
function tick() {
|
||||
if (!enabled()) { render(); return; }
|
||||
const o = observe();
|
||||
render();
|
||||
if (!o.a || inFlight) return;
|
||||
// Proactive: one login attempt per session end, so a declined/ignored login does not loop.
|
||||
if ((o.a.status === "due" || o.a.status === "expired") && o.a.endsAt !== null && warnedEnd !== o.a.endsAt) {
|
||||
warnedEnd = o.a.endsAt;
|
||||
void login(o.a.status === "due" ? "session ends soon" : "session expired");
|
||||
}
|
||||
}
|
||||
|
||||
pi.on("session_start", async (_e, ctx) => {
|
||||
ui = ctx.hasUI ? ctx.ui : null;
|
||||
if (!ui) return; // print/json children: never start logins or timers
|
||||
render();
|
||||
timer = setInterval(tick, 60e3);
|
||||
if (enabled()) tick();
|
||||
});
|
||||
|
||||
pi.on("session_shutdown", async () => {
|
||||
if (timer) clearInterval(timer);
|
||||
timer = null;
|
||||
unlock();
|
||||
});
|
||||
|
||||
// Before a prompt goes out: if the session is already gone, log in first and wait.
|
||||
pi.on("before_agent_start", async () => {
|
||||
if (!ui || !enabled()) return;
|
||||
const o = observe();
|
||||
if (o.a?.status === "expired" || o.a?.status === "none") await login("session expired");
|
||||
else if (inFlight) await inFlight;
|
||||
});
|
||||
|
||||
// A request failed on expired credentials: start the login (once per session end).
|
||||
pi.on("message_end", async (e) => {
|
||||
const m = e.message as { role?: string; errorMessage?: string };
|
||||
if (!ui || !enabled() || m.role !== "assistant" || !m.errorMessage || !isSsoExpiredError(m.errorMessage)) return;
|
||||
logEvent("request_failed_expired");
|
||||
void login("a request failed: session expired").then((ok) => {
|
||||
if (ok) ui?.notify("AWS SSO renewed — send 'continue' to resume.", "info");
|
||||
});
|
||||
});
|
||||
|
||||
pi.registerCommand("aws-sso", {
|
||||
description: "AWS SSO session keeper: on | off | status | now | mode mac|device | hours <n>",
|
||||
handler: async (args, ctx) => {
|
||||
ui = ctx.ui;
|
||||
const [cmd, val] = String(args ?? "").trim().split(/\s+/);
|
||||
const st = loadState();
|
||||
switch (cmd || "status") {
|
||||
case "on": st.enabled = true; saveState(st); if (!timer) timer = setInterval(tick, 60e3); tick(); ctx.ui.notify(`AWS SSO keeper ON (mode ${st.mode ?? defaultMode()}).`, "info"); break;
|
||||
case "off": st.enabled = false; saveState(st); render(); ctx.ui.notify("AWS SSO keeper OFF.", "info"); break;
|
||||
case "now": await login("requested"); break;
|
||||
case "mode":
|
||||
if (val !== "mac" && val !== "device") { ctx.ui.notify("Usage: /aws-sso mode mac|device", "error"); break; }
|
||||
st.mode = val; saveState(st); ctx.ui.notify(`AWS SSO mode: ${val}`, "info"); break;
|
||||
case "hours": {
|
||||
const n = Number(val);
|
||||
if (!(n > 0 && n <= 24 * 90)) { ctx.ui.notify("Usage: /aws-sso hours <session length in hours>", "error"); break; }
|
||||
st.sessionHours = n; saveState(st); render(); ctx.ui.notify(`AWS SSO session length set to ${n} h.`, "info"); break;
|
||||
}
|
||||
case "status": default: {
|
||||
const o = observe();
|
||||
const lines = [
|
||||
`keeper: ${enabled() ? "ON" : "off"}${process.env.PI_AWS_SSO_REFRESH ? ` (env PI_AWS_SSO_REFRESH=${process.env.PI_AWS_SSO_REFRESH})` : ""} · mode: ${st.mode ?? defaultMode()} · session length: ${st.sessionHours ?? DEFAULTS.sessionHours} h`,
|
||||
o.error ? `config: ✗ ${o.error}` : `sso-session: ${o.s!.sessionName} (profile chain ${o.s!.chain.join(" → ")})`,
|
||||
o.a ? `session: ${o.a.status}${o.a.endsAt ? ` · ends ≈ ${new Date(o.a.endsAt).toLocaleTimeString()} (${fmtLeft(o.a.endsAt - Date.now())} left)` : " · start unknown"}` : "",
|
||||
o.token ? `token: refresh ${o.token.refreshFp} · access valid until ${new Date(o.token.accessExpiresAt).toLocaleTimeString()}` : "token: none",
|
||||
st.lastLogin ? `last login: ${new Date(st.lastLogin.at).toLocaleString()} · ${st.lastLogin.ok ? "ok" : "FAILED: " + st.lastLogin.detail} (${st.lastLogin.mode})` : "last login: —",
|
||||
].filter(Boolean);
|
||||
ctx.ui.notify(lines.join("\n"), "info");
|
||||
}
|
||||
}
|
||||
},
|
||||
});
|
||||
}
|
||||
Reference in New Issue
Block a user