Files
pi-extensions/extensions/aws-sso.ts
T
pi@mbp-m1-2020 32a397edb4 aws-sso: renew just after the browser sign-in ends, not before; auto-continue
AWS: CLI sessions refresh only while the user interactive session (browser
sign-in) is valid; a login inside it is silent but inherits its end (green
box 'can be used until your session expires', 2026-10-11). Renewing early
gains nothing, so renew 2 min after the estimated end, forcing a fresh IdP
sign-in while the last 1-hour token still works. Silent renewals keep the
anchor. On an expired-session request failure, renew then send 'continue'
(once per 10 min). New: /aws-sso continue on|off, signed-in HH:MM.
2026-10-11 17:40:40 +02:00

416 lines
21 KiB
TypeScript
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
/**
* 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.
*
* Timing (AWS IdC docs, "Understanding authentication sessions"): the CLI holds an
* *application session* whose 1-hour token the SDK refreshes by itself "as long as
* the underlying user interactive session remains valid". The interactive session
* is the browser's AWS sign-in, admin-capped (8 h measured). A login while it is
* still valid completes silently but INHERITS its end — renewing early gains
* nothing. A login just AFTER it ends makes the browser sign in again (through the
* IdP, silent while "stay signed in" holds) and starts a fresh interactive session.
* So the keeper renews a couple of minutes after the estimated end, while the CLI's
* last 1-hour token still works, and if a request fails anyway it renews and can
* send "continue" so an unattended run resumes.
* The session start is not stored anywhere: it is inferred from refresh-token
* changes (a change after the end = new session; before the end = inherited).
* 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;
renewAfterMinutes?: number;
autoContinue?: boolean;
/** since = start of the browser's interactive session (estimated). */
anchor?: { refreshFp: string; since: number };
/** A request failed on an expired session: the next new token starts a fresh session. */
sessionEnded?: boolean;
lastLogin?: { at: number; ok: boolean; mode: string; detail?: string };
}
export const DEFAULTS = { sessionHours: 8, renewAfterMinutes: 2 };
/**
* Where the session stands (endsAt = estimated end of the browser's sign-in).
* none – no cached token
* expired – past the end AND the CLI's 1-hour token has run out too
* due – past the end + renewAfterMinutes: renew now (a fresh sign-in)
* ok – before that
*/
export function assess(state: State, token: { refreshFp: string; accessExpiresAt: number } | null, now: number) {
const hours = state.sessionHours ?? DEFAULTS.sessionHours;
const after = (state.renewAfterMinutes ?? DEFAULTS.renewAfterMinutes) * 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 + after) return { status: "due" as const, endsAt };
return { status: "ok" as const, endsAt };
}
/**
* Anchor update when the refresh token changed (null = unchanged).
* Before the old session's end a new token is a silent renewal: it keeps the end.
* After the end (or after an observed expiry) it is a fresh sign-in: new start.
*/
export function reanchor(state: State, refreshFp: string, firstSeen: number, now: number): State["anchor"] | null {
const a = state.anchor;
if (a && a.refreshFp === refreshFp) return null;
if (!a) return { refreshFp, since: firstSeen };
const ended = state.sessionEnded || now >= a.since + (state.sessionHours ?? DEFAULTS.sessionHours) * 3600e3;
return { refreshFp, since: ended ? firstSeen : a.since };
}
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()), Date.now());
if (next) {
const inherited = !!state.anchor && next.since === state.anchor.since;
state.anchor = next; state.sessionEnded = false; saveState(state);
logEvent("anchor", { refreshFp: next.refreshFp, since: new Date(next.since).toISOString(), inherited });
}
}
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"
: o.a.status === "due" ? "AWS ⚠ renewing" : left === null ? "AWS ✓" : `AWS ✓ ${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: login complete." : `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;
// Just after the end: one attempt per session end, so an 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("session ended");
}
}
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: renew, then (optionally) resume the run.
let lastAutoContinue = 0;
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;
const st = loadState(); st.sessionEnded = true; saveState(st);
logEvent("request_failed_expired");
void login("a request failed: session expired").then((ok) => {
if (!ok) return;
// At most once per 10 min: if the renewed session fails at once, stop instead of looping.
if ((loadState().autoContinue ?? true) && Date.now() - lastAutoContinue > 10 * 60e3) {
lastAutoContinue = Date.now();
logEvent("auto_continue");
ui?.notify("AWS SSO renewed — resuming with 'continue'.", "info");
pi.sendUserMessage("continue", { deliverAs: "followUp" });
} else 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 | continue on|off | signed-in HH:MM | 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 "continue":
if (val !== "on" && val !== "off") { ctx.ui.notify("Usage: /aws-sso continue on|off", "error"); break; }
st.autoContinue = val === "on"; saveState(st); ctx.ui.notify(`AWS SSO auto-continue after a failed request: ${val}`, "info"); break;
case "signed-in": {
// Correct the estimate: when did the browser last sign in through the IdP?
const hm = String(val ?? "").match(/^(\d{1,2}):(\d{2})$/);
const o = observe();
if (!hm || !o.token) { ctx.ui.notify("Usage: /aws-sso signed-in HH:MM (local time, today; needs a cached token)", "error"); break; }
const d = new Date(); d.setHours(Number(hm[1]), Number(hm[2]), 0, 0);
if (d.getTime() > Date.now()) d.setDate(d.getDate() - 1);
const s2 = loadState(); s2.anchor = { refreshFp: o.token.refreshFp, since: d.getTime() }; s2.sessionEnded = false; saveState(s2);
logEvent("anchor_set", { since: d.toISOString() });
render(); ctx.ui.notify(`AWS SSO: session start set to ${d.toLocaleString()}.`, "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 · auto-continue: ${(st.autoContinue ?? true) ? "on" : "off"}`,
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 ? ` · browser sign-in ${new Date(o.a.endsAt - (st.sessionHours ?? DEFAULTS.sessionHours) * 3600e3).toLocaleTimeString()} → ends ≈ ${new Date(o.a.endsAt).toLocaleTimeString()} (${fmtLeft(o.a.endsAt - Date.now())} left) · renew ≈ ${new Date(o.a.endsAt + (st.renewAfterMinutes ?? DEFAULTS.renewAfterMinutes) * 60e3).toLocaleTimeString()}` : " · 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");
}
}
},
});
}