Files
pi-extensions/extensions/aws-sso.ts
T
pi@mbp-m1-2020 fdee3e42d6 aws-sso: detect the real session end by a refresh attempt, renew before requests fail
Measured: a session lasted <= 6 h (2026-10-11), not the 8 h assumed, so a
clock cannot predict the end. token-providers 3.1127.0 refreshes only inside
a 5-min window and, on refusal, silently returns the still-valid token. The
keeper now makes that attempt itself (pi's own SDK, same cache) when the
window opens, and if expiresAt did not advance, logs in while the old token
still works. Clock-based renewal, signed-in and hours commands removed.
2026-10-11 22:33:40 +02:00

452 lines
23 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.
*
* Detection (AWS IdC docs, "Understanding authentication sessions"): the CLI holds an
* *application session* whose 1-hour token is refreshed "as long as the underlying
* user interactive session remains valid" (the browser's AWS sign-in, admin-capped;
* measured both 8 h and <= 6 h, so its length is NOT predictable). When it ends,
* the next refresh is refused. Inside the SDK's 5-minute window a refused refresh is
* silent: the SDK keeps returning the still-valid token, so requests only fail once
* that token runs out. The keeper makes the refresh attempt itself as soon as the
* window opens (same SDK, same cache file) and checks whether the cache advanced.
* If not, the session has ended: it logs in right away, while the old token still
* covers requests. A login while the browser sign-in is still valid is silent and
* inherits its end, which is why nothing is renewed on a clock. If a request fails
* anyway, it renews and can send "continue" so an unattended run resumes.
* 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, realpathSync, renameSync, statSync, writeFileSync, openSync, closeSync, unlinkSync } from "node:fs";
import { createRequire } from "node:module";
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;
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 };
/** The AWS SDK's own refresh window (token-providers EXPIRE_WINDOW_MS). */
export const REFRESH_WINDOW_MS = 5 * 60e3;
/**
* Should the keeper make a refresh attempt now? Only inside the SDK's window (outside
* it the SDK does not refresh, so an attempt proves nothing). Every 60 s while the
* token is still valid; every 10 min once it has run out (laptop slept, login pending).
*/
export function probeDue(accessExpiresAt: number, now: number, lastProbeAt: number): boolean {
const left = accessExpiresAt - now;
if (left > REFRESH_WINDOW_MS) return false;
return now - lastProbeAt >= (left > 0 ? 60e3 : 10 * 60e3);
}
export type ProbeResult = { result: "fresh" | "refreshed" | "refused" | "unavailable"; before: number; after: number; error?: string };
/**
* One refresh attempt through the SDK's own token provider, judged by the cache:
* refreshed = expiresAt advanced; refused = it did not (session over, or offline).
* "unavailable" = the SDK could not be loaded: says nothing about the session.
* The token the provider returns is dropped unread.
*/
export async function probeWith(loadFromSso: () => (init: object) => () => Promise<unknown>, readExpiresAt: () => number, profile: string, now: () => number = Date.now): Promise<ProbeResult> {
const before = readExpiresAt();
if (before - now() > REFRESH_WINDOW_MS) return { result: "fresh", before, after: before };
let provider: (init: object) => () => Promise<unknown>;
try { provider = loadFromSso(); } catch (e) { return { result: "unavailable", before, after: before, error: (e as Error).message }; }
let error: string | undefined;
try { await provider({ profile, ignoreCache: true })(); } catch (e) { error = (e as Error).message; }
const after = readExpiresAt();
return { result: after > before ? "refreshed" : "refused", before, after, error };
}
/**
* 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; }
}
/** The SDK pi itself uses (resolved from pi's install, not from this file's location). */
let fromSsoCached: ((init: object) => () => Promise<unknown>) | null = null;
function loadFromSso() {
if (fromSsoCached) return fromSsoCached;
const bases = [process.argv[1], "/usr/lib/node_modules/@earendil-works/pi-coding-agent/package.json"].filter(Boolean) as string[];
let last: unknown;
for (const b of bases) {
try { fromSsoCached = createRequire(realpathSync(b))("@aws-sdk/token-providers").fromSso; if (fromSsoCached) return fromSsoCached; }
catch (e) { last = e; }
}
throw last instanceof Error ? last : new Error("@aws-sdk/token-providers not found");
}
/** 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 probing = false;
let lastProbeAt = 0;
let lastAutoLogin = 0;
let lastProbe: { at: number; result: string } | 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, token: null, state: loadState() };
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, token, state };
}
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) { ui.setStatus("aws-sso", `AWS ✗ ${o.error}`); return; }
if (!o.token) { ui.setStatus("aws-sso", "AWS ✗ not logged in"); return; }
if (o.state.sessionEnded) { ui.setStatus("aws-sso", "AWS ✗ session ended"); return; }
const since = o.state.anchor ? new Date(o.state.anchor.since).toLocaleTimeString([], { hour: "2-digit", minute: "2-digit" }) : null;
ui.setStatus("aws-sso", since ? `AWS ✓ since ${since}` : "AWS ✓");
}
/** Refresh attempt inside the SDK window; on refusal, log in while the old token still works. */
async function probe(why: string): Promise<ProbeResult | null> {
const o = observe();
if (!o.s || !o.token || probing) return null;
probing = true;
lastProbeAt = Date.now();
try {
const r = await probeWith(loadFromSso, () => readToken(o.s!.sessionName)?.accessExpiresAt ?? 0, o.s.chain[o.s.chain.length - 1]);
lastProbe = { at: lastProbeAt, result: r.result };
if (r.result !== "fresh") logEvent("probe", { why, result: r.result, left_s: Math.round((r.before - lastProbeAt) / 1000), after: r.after ? new Date(r.after).toISOString() : null, error: r.result === "refreshed" ? undefined : r.error?.slice(0, 160) });
if (r.result === "refused") { const st = loadState(); st.sessionEnded = true; saveState(st); }
return r;
} finally { probing = false; render(); }
}
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(); }
}
async function tick() {
if (!enabled() || inFlight || probing) { render(); return; }
const o = observe();
render();
if (!o.token) return;
if (!probeDue(o.token.accessExpiresAt, Date.now(), lastProbeAt)) return;
const r = await probe("refresh window");
// At most one unprompted login per 30 min, so a login waiting for you does not pile up tabs.
if (r?.result === "refused" && Date.now() - lastAutoLogin >= 30 * 60e3) {
lastAutoLogin = Date.now();
void login("session ended (refresh refused)");
}
}
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(() => void tick(), 20e3);
if (enabled()) void 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;
if (inFlight) { await inFlight; return; }
let o = observe();
if (!o.token) { await login("not logged in"); return; }
// Token in (or past) the refresh window: find out now whether the session still exists.
if (o.token.accessExpiresAt - Date.now() <= REFRESH_WINDOW_MS && !o.state.sessionEnded) await probe("before prompt");
o = observe();
if (o.state.sessionEnded) await login("session ended");
});
// 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",
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(() => void tick(), 20e3); void 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 "status": default: {
const o = observe();
const t = (ms: number) => new Date(ms).toLocaleTimeString([], { hour: "2-digit", minute: "2-digit" });
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()} · auto-continue: ${(st.autoContinue ?? true) ? "on" : "off"}`,
o.error ? `config: ✗ ${o.error}` : `sso-session: ${o.s!.sessionName} (profile chain ${o.s!.chain.join(" → ")})`,
o.token ? `session: ${o.state.sessionEnded ? "ENDED" : "alive"}${o.state.anchor ? ` · started ≈ ${t(o.state.anchor.since)}` : ""} · 1-h token until ${t(o.token.accessExpiresAt)} (next check from ${t(o.token.accessExpiresAt - REFRESH_WINDOW_MS)})` : "token: none",
lastProbe ? `last check: ${t(lastProbe.at)} · ${lastProbe.result}` : "last check: — (this pi process)",
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");
}
}
},
});
}