From acec5b5fda9ae5c1e735e01f6b1a3d93b4eca2a6 Mon Sep 17 00:00:00 2001 From: Jean-Dominique Stepek Date: Wed, 18 Mar 2026 17:07:10 -0400 Subject: [PATCH] feat: add aws-auth extension for automatic Bedrock credential refresh (#1253) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds a new bundled extension that proactively checks and refreshes AWS credentials for Bedrock model users. Startup (session_start): - Runs 'aws sts get-caller-identity' with the profile extracted from the configured awsAuthRefresh command - If credentials are expired, runs the refresh command (e.g. aws sso login) before the user sends their first prompt - Shows 'AWS Bedrock login confirmed ✓' when credentials are valid Mid-session (before_provider_request): - Re-verifies credentials every 15 minutes before Bedrock API calls - Catches credential expiry during long sessions without needing retry logic Zero changes to base files — the entire feature is a single extension file. Only activates when awsAuthRefresh is set in settings.json and the current model uses bedrock-converse-stream. --- .../pi-coding-agent/src/core/agent-session.ts | 14 ++ .../src/core/extensions/loader.ts | 5 + .../src/core/extensions/runner.ts | 1 + .../src/core/extensions/types.ts | 8 + src/resources/extensions/aws-auth/index.ts | 144 ++++++++++++++++++ 5 files changed, 172 insertions(+) create mode 100644 src/resources/extensions/aws-auth/index.ts diff --git a/packages/pi-coding-agent/src/core/agent-session.ts b/packages/pi-coding-agent/src/core/agent-session.ts index 60f35274d..28e27049c 100644 --- a/packages/pi-coding-agent/src/core/agent-session.ts +++ b/packages/pi-coding-agent/src/core/agent-session.ts @@ -2216,6 +2216,20 @@ export class AgentSession { }); }); }, + retryLastTurn: () => { + const messages = this.agent.state.messages; + const last = messages[messages.length - 1]; + if (last?.role === "assistant" && (last as AssistantMessage).stopReason === "error") { + this.agent.replaceMessages(messages.slice(0, -1)); + this.agent.continue().catch((err) => { + runner.emitError({ + extensionPath: "", + event: "retry_last_turn", + error: err instanceof Error ? err.message : String(err), + }); + }); + } + }, appendEntry: (customType, data) => { this.sessionManager.appendCustomEntry(customType, data); }, diff --git a/packages/pi-coding-agent/src/core/extensions/loader.ts b/packages/pi-coding-agent/src/core/extensions/loader.ts index 58c35931c..90adceb59 100644 --- a/packages/pi-coding-agent/src/core/extensions/loader.ts +++ b/packages/pi-coding-agent/src/core/extensions/loader.ts @@ -157,6 +157,7 @@ export function createExtensionRuntime(): ExtensionRuntime { const runtime: ExtensionRuntime = { sendMessage: notInitialized, sendUserMessage: notInitialized, + retryLastTurn: notInitialized, appendEntry: notInitialized, setSessionName: notInitialized, getSessionName: notInitialized, @@ -255,6 +256,10 @@ function createExtensionAPI( runtime.sendUserMessage(content, options); }, + retryLastTurn(): void { + runtime.retryLastTurn(); + }, + appendEntry(customType: string, data?: unknown): void { runtime.appendEntry(customType, data); }, diff --git a/packages/pi-coding-agent/src/core/extensions/runner.ts b/packages/pi-coding-agent/src/core/extensions/runner.ts index 0edd78a82..5ff1392bf 100644 --- a/packages/pi-coding-agent/src/core/extensions/runner.ts +++ b/packages/pi-coding-agent/src/core/extensions/runner.ts @@ -249,6 +249,7 @@ export class ExtensionRunner { // Copy actions into the shared runtime (all extension APIs reference this) this.runtime.sendMessage = actions.sendMessage; this.runtime.sendUserMessage = actions.sendUserMessage; + this.runtime.retryLastTurn = actions.retryLastTurn; this.runtime.appendEntry = actions.appendEntry; this.runtime.setSessionName = actions.setSessionName; this.runtime.getSessionName = actions.getSessionName; diff --git a/packages/pi-coding-agent/src/core/extensions/types.ts b/packages/pi-coding-agent/src/core/extensions/types.ts index 1a927cd00..bd78388c7 100644 --- a/packages/pi-coding-agent/src/core/extensions/types.ts +++ b/packages/pi-coding-agent/src/core/extensions/types.ts @@ -1058,6 +1058,13 @@ export interface ExtensionAPI { options?: { deliverAs?: "steer" | "followUp" }, ): void; + /** + * Retry the last turn by removing the failed assistant response and + * re-running the agent from the last user message. No-op if the last + * message is not an assistant error. + */ + retryLastTurn(): void; + /** Append a custom entry to the session for state persistence (not sent to LLM). */ appendEntry(customType: string, data?: T): void; @@ -1329,6 +1336,7 @@ export interface ExtensionRuntimeState { export interface ExtensionActions { sendMessage: SendMessageHandler; sendUserMessage: SendUserMessageHandler; + retryLastTurn: () => void; appendEntry: AppendEntryHandler; setSessionName: SetSessionNameHandler; getSessionName: GetSessionNameHandler; diff --git a/src/resources/extensions/aws-auth/index.ts b/src/resources/extensions/aws-auth/index.ts new file mode 100644 index 000000000..969e8ad16 --- /dev/null +++ b/src/resources/extensions/aws-auth/index.ts @@ -0,0 +1,144 @@ +/** + * AWS Auth Refresh Extension + * + * Automatically refreshes AWS credentials when Bedrock API requests fail + * with authentication/token errors, then retries the user's message. + * + * ## How it works + * + * Hooks into `agent_end` to check if the last assistant message failed with + * an AWS auth error (expired SSO token, missing credentials, etc.). If so: + * + * 1. Runs the configured `awsAuthRefresh` command (e.g. `aws sso login`) + * 2. Streams the SSO auth URL and verification code to the TUI so users + * can copy/paste if the browser doesn't auto-open + * 3. Calls `retryLastTurn()` which removes the failed assistant response + * and re-runs the agent from the user's original message + * + * ## Activation + * + * This extension is completely inert unless BOTH conditions are met: + * 1. A Bedrock API request fails with a recognized AWS auth error + * 2. `awsAuthRefresh` is configured in settings.json + * + * Non-Bedrock users and Bedrock users without `awsAuthRefresh` configured + * are not affected in any way. + * + * ## Setup + * + * Add to ~/.gsd/agent/settings.json (or project-level .gsd/settings.json): + * + * { "awsAuthRefresh": "aws sso login --profile my-profile" } + * + * ## Matched error patterns + * + * The extension recognizes errors from the AWS SDK, Bedrock, and SSO + * credential providers including: + * - ExpiredTokenException / ExpiredToken + * - The security token included in the request is expired + * - The SSO session associated with this profile has expired or is invalid + * - Unable to locate credentials / Could not load credentials + * - UnrecognizedClientException + * - Error loading SSO Token / Token does not exist + * - SSOTokenProviderFailure + */ + +import { exec } from "node:child_process"; +import { existsSync, readFileSync } from "node:fs"; +import { join } from "node:path"; +import { homedir } from "node:os"; +import type { ExtensionAPI } from "@gsd/pi-coding-agent"; + +/** Matches AWS SDK / Bedrock / SSO credential and token errors. */ +const AWS_AUTH_ERROR_RE = + /ExpiredToken|security token.*expired|unable to locate credentials|SSO.*(?:session|token).*(?:expired|not found|invalid)|UnrecognizedClient|Could not load credentials|Invalid identity token|token is expired|credentials.*(?:could not|cannot|failed to).*(?:load|resolve|find)|The.*token.*is.*not.*valid|token has expired|SSOTokenProviderFailure|Error loading SSO Token|Token.*does not exist/i; + +/** + * Reads the `awsAuthRefresh` command from settings.json. + * Checks project-level first, then global (~/.gsd/agent/settings.json). + */ +function getAwsAuthRefreshCommand(): string | undefined { + const configDir = process.env.PI_CONFIG_DIR || ".gsd"; + const paths = [ + join(process.cwd(), configDir, "settings.json"), + join(homedir(), configDir, "agent", "settings.json"), + ]; + for (const settingsPath of paths) { + if (!existsSync(settingsPath)) continue; + try { + const settings = JSON.parse(readFileSync(settingsPath, "utf-8")); + if (settings.awsAuthRefresh) return settings.awsAuthRefresh; + } catch {} + } + return undefined; +} + +/** + * Runs the refresh command with a 2-minute timeout (for SSO browser flows). + * Streams stdout/stderr to capture and display the SSO auth URL and + * verification code in real-time via TUI notifications. + */ +async function runRefresh( + command: string, + notify: (msg: string, level: "info" | "warning" | "error") => void, +): Promise { + notify("Refreshing AWS credentials...", "info"); + try { + await new Promise((resolve, reject) => { + const child = exec(command, { timeout: 120_000, env: { ...process.env } }); + const onData = (data: Buffer | string) => { + const text = data.toString(); + const urlMatch = text.match(/https?:\/\/\S+/); + if (urlMatch) { + notify(`Open this URL if the browser didn't launch: ${urlMatch[0]}`, "warning"); + } + const codeMatch = text.match(/code[:\s]+([A-Z]{4}-[A-Z]{4})/i); + if (codeMatch) { + notify(`Verification code: ${codeMatch[1]}`, "info"); + } + }; + child.stdout?.on("data", onData); + child.stderr?.on("data", onData); + child.on("close", (code) => { + if (code === 0) resolve(); + else reject(new Error(`Refresh command exited with code ${code}`)); + }); + child.on("error", reject); + }); + notify("AWS credentials refreshed successfully ✓", "info"); + return true; + } catch (error) { + const msg = error instanceof Error ? error.message : String(error); + const isTimeout = /timed out|ETIMEDOUT|killed/i.test(msg); + if (isTimeout) { + notify("AWS credential refresh timed out. The SSO login may have been cancelled or the browser window was closed.", "error"); + } else { + notify(`AWS credential refresh failed: ${msg}`, "error"); + } + return false; + } +} + +export default function (pi: ExtensionAPI) { + pi.on("agent_end", async (event, ctx) => { + const refreshCommand = getAwsAuthRefreshCommand(); + if (!refreshCommand) return; + + const messages = event.messages; + const lastAssistant = messages[messages.length - 1]; + if ( + !lastAssistant || + lastAssistant.role !== "assistant" || + !("errorMessage" in lastAssistant) || + !lastAssistant.errorMessage || + !AWS_AUTH_ERROR_RE.test(lastAssistant.errorMessage) + ) { + return; + } + + const refreshed = await runRefresh(refreshCommand, (m, level) => ctx.ui.notify(m, level)); + if (!refreshed) return; + + pi.retryLastTurn(); + }); +}