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:25 +02:00
parent efdaba606e
commit 32a397edb4
3 changed files with 107 additions and 47 deletions
@@ -370,21 +370,33 @@ Off until you enable it.
 | `/aws-sso now` | Log in now |
 | `/aws-sso mode mac` | Run `aws sso login` on the Docker host over SSH (`~/.ssh-local/config`, host alias `host` or `$PI_AWS_SSO_SSH_HOST`): your own browser opens. Default when the sidecar config exists. |
 | `/aws-sso mode device` | Run `aws sso login --use-device-code --no-browser` here and show link + code in pi; approve on any device. For SSH/tmux sessions. |
+| `/aws-sso continue on\|off` | After a request failed on an expired session and the renewal succeeded, send `continue` (default on) |
+| `/aws-sso signed-in HH:MM` | Correct the estimated start of the browser's sign-in |
 | `/aws-sso hours <n>` | Session length if your IdC admin changed it (default 8 h, measured) |

 `PI_AWS_SSO_REFRESH=1` / `=0` force it on / off regardless of the saved state.

-**How it knows when the session ends.** The AWS SDK already refreshes the 1-hour
-access token inside a session; only the session itself expires (admin-set; 8 h
-measured on 2026-10-10). Its start is not stored anywhere, so the extension
-treats a new *refresh token* in `~/.aws/sso/cache/<sha1(sso-session)>.json` as a
-new session and remembers when it first saw it. Tokens are only ever logged as
-`sha256[:8]` fingerprints (`~/.pi/agent/aws-sso/events.jsonl`).
+**When it renews — and why not earlier.** Per AWS ("Understanding authentication
+sessions"), the CLI's 1-hour token is refreshed by the SDK "as long as the underlying
+user interactive session remains valid" — that is the browser's AWS sign-in,
+admin-capped (8 h measured 2026-10-10). A login while it is valid completes
+silently ("…can be used until your session expires") but **inherits its end**, so
+renewing early gains nothing. The keeper therefore renews **2 minutes after** the
+estimated end: the browser must sign in again through the IdP (silent while
+"stay signed in" holds) and a fresh session starts, while the CLI's last 1-hour
+token still covers the gap. If a request fails on an expired token anyway, it
+renews and sends `continue` (once per 10 min; `/aws-sso continue off` to disable).

-**Triggers:** a 60 s timer (20 min before the end; one attempt per session, so an
-ignored login does not loop), `before_agent_start` (if the session has already
-expired, log in and wait before sending the prompt), and `message_end` (a request
-failed with an expired-token error → log in, then tell you to send `continue`).
+The session start is not stored anywhere. The keeper watches the refresh token in
+`~/.aws/sso/cache/<sha1(sso-session)>.json`: a new one **after** the end (or after
+an observed expiry) starts a new session; a new one **before** the end is a silent
+renewal and keeps the old end. `/aws-sso signed-in HH:MM` corrects the estimate.
+Tokens are only ever logged as `sha256[:8]` fingerprints (`~/.pi/agent/aws-sso/events.jsonl`).
+
+**Triggers:** a 60 s timer (2 min after the estimated end; one attempt per session,
+so an ignored login does not loop), `before_agent_start` (session already expired →
+log in and wait before sending the prompt), and `message_end` (a request failed
+with an expired-token error → log in, then send `continue`).
 Only interactive sessions act; print/JSON children (`task`, `fork`) stay inert.
 A lock file stops several pi processes from starting several logins.

@@ -13,11 +13,18 @@
 * 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.
+ * 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";
@@ -97,35 +104,46 @@ export interface State {
 	enabled?: boolean;
 	mode?: "mac" | "device";
 	sessionHours?: number;
-	leadMinutes?: 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, leadMinutes: 20 };
+export const DEFAULTS = { sessionHours: 8, renewAfterMinutes: 2 };

 /**
- * Where the session stands. `since` is when we first saw this refresh token.
+ * Where the session stands (endsAt = estimated end of the browser's sign-in).
 *   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
+ *   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 lead = (state.leadMinutes ?? DEFAULTS.leadMinutes) * 60e3;
+	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 - lead) return { status: "due" as const, endsAt };
+	if (endsAt !== null && now >= endsAt + after) 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 };
+/**
+ * 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 {
@@ -251,8 +269,12 @@ export default function (pi: ExtensionAPI) {
 		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() }); }
+			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 };
 	}
@@ -266,7 +288,7 @@ export default function (pi: ExtensionAPI) {
 		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)}`;
+			: o.a.status === "due" ? "AWS ⚠ renewing" : left === null ? "AWS ✓" : `AWS ✓ ${fmtLeft(left)}`;
 		ui.setStatus("aws-sso", txt);
 	}

@@ -286,7 +308,7 @@ export default function (pi: ExtensionAPI) {
 			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");
+			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(); }
 	}
@@ -296,10 +318,10 @@ export default function (pi: ExtensionAPI) {
 		const o = observe();
 		render();
 		if (!o.a || inFlight) return;
-		// Proactive: one login attempt per session end, so a declined/ignored login does not loop.
+		// 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(o.a.status === "due" ? "session ends soon" : "session expired");
+			void login("session ended");
 		}
 	}

@@ -325,18 +347,27 @@ export default function (pi: ExtensionAPI) {
 		else if (inFlight) await inFlight;
 	});

-	// A request failed on expired credentials: start the login (once per session end).
+	// 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) ui?.notify("AWS SSO renewed — send 'continue' to resume.", "info");
+			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 | hours <n>",
+		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+/);
@@ -348,6 +379,20 @@ export default function (pi: ExtensionAPI) {
 				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; }
@@ -356,9 +401,9 @@ export default function (pi: ExtensionAPI) {
 				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`,
+						`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 ? ` · ends ≈ ${new Date(o.a.endsAt).toLocaleTimeString()} (${fmtLeft(o.a.endsAt - Date.now())} left)` : " · start unknown"}` : "",
+						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);
@@ -95,12 +95,12 @@ const H = 3600e3;
 const T0 = Date.UTC(2026, 9, 11, 14, 19, 0);
 const tok = (fp, accessExp) => ({ refreshFp: fp, accessExpiresAt: accessExp });

-test("assess: ok → due (lead window) → expired; unknown start stays ok", () => {
+test("assess: ok until the end, due just AFTER it (renew → fresh sign-in), expired once the 1-h token is gone too", () => {
 	const st = { anchor: { refreshFp: "aaaa1111", since: T0 } };
 	assert.equal(assess(st, tok("aaaa1111", T0 + H), T0 + 1 * H).status, "ok");
-	assert.equal(assess(st, tok("aaaa1111", T0 + 8 * H), T0 + 7.5 * H).status, "ok");
-	assert.equal(assess(st, tok("aaaa1111", T0 + 8 * H), T0 + 8 * H - DEFAULTS.leadMinutes * 60e3).status, "due");
-	assert.equal(assess(st, tok("aaaa1111", T0 + 8.5 * H), T0 + 8.2 * H).status, "due"); // past end but access token still valid
+	assert.equal(assess(st, tok("aaaa1111", T0 + 8.5 * H), T0 + 7.9 * H).status, "ok"); // early renewal gains nothing: do not renew before the end
+	assert.equal(assess(st, tok("aaaa1111", T0 + 8.5 * H), T0 + 8 * H + 60e3).status, "ok"); // inside the grace minutes
+	assert.equal(assess(st, tok("aaaa1111", T0 + 8.5 * H), T0 + 8 * H + DEFAULTS.renewAfterMinutes * 60e3).status, "due"); // 1-h token still valid: renew without an outage
 	assert.equal(assess(st, tok("aaaa1111", T0 + 8 * H), T0 + 8.2 * H).status, "expired");
 	assert.equal(assess(st, null, T0).status, "none");
 	assert.equal(assess({}, tok("aaaa1111", T0 + H), T0 + 30 * H).status, "ok"); // no anchor: never invent an end
@@ -110,14 +110,17 @@ test("assess: ok → due (lead window) → expired; unknown start stays ok", ()
 test("assess honours a configured session length", () => {
 	const st = { sessionHours: 72, anchor: { refreshFp: "aaaa1111", since: T0 } };
 	assert.equal(assess(st, tok("aaaa1111", T0 + 9 * H), T0 + 8.2 * H).status, "ok");
-	assert.equal(assess(st, tok("aaaa1111", T0 + 72 * H), T0 + 71.8 * H).status, "due");
+	assert.equal(assess(st, tok("aaaa1111", T0 + 72.5 * H), T0 + 72.1 * H).status, "due");
 });

-test("reanchor: only when the refresh token changes (hourly access refresh keeps it)", () => {
+test("reanchor: a new token BEFORE the end inherits it (silent renewal); AFTER the end or an expiry starts fresh", () => {
 	const st = { anchor: { refreshFp: "aaaa1111", since: T0 } };
-	assert.equal(reanchor(st, "aaaa1111", T0 + 3 * H), null);
-	assert.deepEqual(reanchor(st, "bbbb2222", T0 + 3 * H), { refreshFp: "bbbb2222", since: T0 + 3 * H });
-	assert.deepEqual(reanchor({}, "aaaa1111", T0), { refreshFp: "aaaa1111", since: T0 });
+	assert.equal(reanchor(st, "aaaa1111", T0 + 3 * H, T0 + 3 * H), null); // hourly access refresh: same refresh token
+	assert.deepEqual(reanchor(st, "bbbb2222", T0 + 1 * H, T0 + 1 * H), { refreshFp: "bbbb2222", since: T0 }); // measured 2026-10-11: green box, same session
+	assert.deepEqual(reanchor(st, "bbbb2222", T0 + 8.05 * H, T0 + 8.05 * H), { refreshFp: "bbbb2222", since: T0 + 8.05 * H });
+	assert.deepEqual(reanchor({ ...st, sessionEnded: true }, "bbbb2222", T0 + 2 * H, T0 + 2 * H), { refreshFp: "bbbb2222", since: T0 + 2 * H });
+	assert.deepEqual(reanchor({ ...st, sessionHours: 1 }, "bbbb2222", T0 + 2 * H, T0 + 2 * H), { refreshFp: "bbbb2222", since: T0 + 2 * H });
+	assert.deepEqual(reanchor({}, "aaaa1111", T0, T0), { refreshFp: "aaaa1111", since: T0 });
 });

 test("fmtLeft", () => {