@@ -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 –
* own browser opens. For when you sit at the laptop.
* device –
* 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 –
* expired –
* due –
* ok –
*/
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 ) * 60 e3 ;
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 * 3600 e3 ;
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 / 60 e3 ) ;
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 * 60 e3 ) : 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 : - / b i n / z s h } " - l c ' a w s s s o l o g i n - - s s o - s e s s i o n $ { s . s e s s i o n N a m e . r e p l a c e ( / [ ^ \ 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");
}
}
},
});
}
pi@mbp-m1-2020