@clawdbuilder/jig-bedrock 0.0.0-stage → 0.1.1

package/src/provider.ts

@@ -0,0 +1,344 @@
+// The Bedrock model provider: the team's application inference profiles as Pi models, AWS SSO recovery, a monthly spend figure and
+// the commands around them. Its settings are the `bedrock` section when Jig is present, or the extension's own settings (see config.ts); it also
+// publishes the optional display events in events.ts.
+//
+// The factory registers and reads settings, and no more: AWS, usage program and the browser are touched in `session_start`, in the
+// background, or by a command the user types.
+import { readFileSync } from "node:fs";
+import { homedir } from "node:os";
+import path from "node:path";
+import type { Provider } from "@earendil-works/pi-ai";
+import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";
+import type { CatalogSource, Status } from "./aws-discovery.ts";
+import { createAwsDiscovery, discoveryHelp } from "./aws-discovery.ts";
+import { createSignedFetch } from "./bedrock-fetch.ts";
+import type { BudgetSource } from "./budget.ts";
+import { createBudget } from "./budget.ts";
+import type { TeamBedrockModel } from "./catalog.ts";
+import {
+    discoveredProfiles,
+    registerBedrockSelectionGuard,
+    registerTeamBedrock,
+    updateModels,
+    usePiRegistry,
+} from "./catalog.ts";
+import { isObject } from "./check.ts";
+import { isSubagentChild } from "./child.ts";
+import { registerStatisticsCommands } from "./commands.ts";
+import type { ConfigResult, ResolvedConfig } from "./config.ts";
+import { loadConfig } from "./config.ts";
+import { MODEL_NAMES_EVENT, modelNames, SEGMENT_EVENT, spendSegment } from "./events.ts";
+import { isJigBedrock } from "./identity.ts";
+import type { JigApi } from "./jig/types.ts";
+import type { RecoveryOptions } from "./login.ts";
+import { createRecovery } from "./login.ts";
+import { findPiHost, loadCredentialProvider } from "./pi-host.ts";
+import type { AwsCredentials } from "./sigv4.ts";
+import type { Recover, RecoveredBedrock } from "./sso.ts";
+import { noticeFor, registerBedrockSsoRecovery, shareRecovery } from "./sso.ts";
+import type { SsoState } from "./sso-check.ts";
+import { checkSso, systemSsoDeps } from "./sso-check.ts";
+import { registerSignInAtStart } from "./start-sign-in.ts";
+import type { StatisticsSource } from "./statistics.ts";
+import { createStatistics } from "./statistics.ts";
+import { registerSwitchNotice } from "./switch-cost.ts";
+
+export type { BudgetSource, CatalogSource, Recover, ResolvedConfig, StatisticsSource };
+
+export interface EventNames {
+    provider: string;
+    usage: string;
+}
+
+export interface BedrockOptions {
+    catalog?: CatalogSource;
+    /** The settings, already resolved. Without them they are read from the `bedrock` section through `api`. */
+    config?: ResolvedConfig;
+    /** jig-core's API, for reading the settings. */
+    api?: JigApi;
+    ssoRecover?: Recover;
+    /** The parts of the default sign-in a test replaces: the program that starts the login, its lock, its environment, WSL. */
+    login?: Omit<RecoveryOptions, "say">;
+    /** Whether the profile's SSO session is still good, for the sign-in at start. Without it the AWS SDK in Pi's install answers. */
+    ssoCheck?: (profile: string, signal: AbortSignal) => Promise<SsoState>;
+    budget?: BudgetSource;
+    statistics?: StatisticsSource;
+    events?: EventNames;
+    selection?: { agentDir?: string; explicitSelection?: () => boolean };
+    transformModels?: (models: TeamBedrockModel[]) => TeamBedrockModel[];
+    backgroundRefresh?: boolean;
+    usageEvents?: boolean;
+    preserveNativeAuth?: boolean;
+    /** The AWS credentials of the Messages route; the SDK's default provider for the profile unless a test supplies them. */
+    credentials?: () => Promise<AwsCredentials>;
+    /** The fetch signed requests go through, so a test can watch them. */
+    fetch?: typeof globalThis.fetch;
+    env?: NodeJS.ProcessEnv;
+}
+
+export interface BedrockHandle {
+    models: TeamBedrockModel[];
+    refresh: (force?: boolean) => Promise<Status>;
+    failed: boolean;
+}
+
+/** Pi's agent folder, from PI_CODING_AGENT_DIR or HOME, shared by settings and notices. */
+export function agentFolder(env: NodeJS.ProcessEnv = process.env): string {
+    return env.PI_CODING_AGENT_DIR === undefined || env.PI_CODING_AGENT_DIR === ""
+        ? path.join(env.HOME ?? homedir(), ".pi", "agent")
+        : env.PI_CODING_AGENT_DIR;
+}
+
+/** The user's own bedrock settings file; `bedrock` or the file itself may hold the keys. */
+function userSettings(env: NodeJS.ProcessEnv): Record<string, unknown> | undefined {
+    try {
+        const raw = JSON.parse(readFileSync(path.join(agentFolder(env), "bedrock.json"), "utf8")) as unknown;
+        const value = isObject(raw) && isObject(raw.bedrock) ? raw.bedrock : raw;
+        return isObject(value) ? value : undefined;
+    } catch {
+        return undefined;
+    }
+}
+
+/** Jig's section without the two keys that start programs. */
+function safeSection(api: JigApi | undefined): { section?: Record<string, unknown>; dropped: string[] } {
+    const raw = api?.config.section("bedrock");
+    if (!isObject(raw)) return { dropped: [] };
+    const section = { ...raw };
+    const dropped = ["ssoLoginCommand", "usageDir", "browser"].filter((key) => key in section);
+    delete section.ssoLoginCommand;
+    delete section.usageDir;
+    delete section.browser;
+    return { section, dropped };
+}
+
+/** The user's bedrock.json first; jig-core's section on top, except the two program keys. */
+export function settingsFrom(
+    api: JigApi | undefined,
+    env: NodeJS.ProcessEnv = process.env,
+): Record<string, unknown> | undefined {
+    const { section } = safeSection(api);
+    const user = userSettings(env);
+    if (user === undefined) return section;
+    return section === undefined ? user : { ...user, ...section };
+}
+
+export function createJigBedrock(
+    options: BedrockOptions = {},
+): (pi: ExtensionAPI) => Promise<BedrockHandle | undefined> {
+    return async (pi) => {
+        const resolved: ConfigResult = options.config
+            ? { config: options.config }
+            : loadConfig({
+                  raw: settingsFrom(options.api, options.env),
+                  ...(options.env === undefined ? {} : { env: options.env }),
+              });
+        const config = resolved.config;
+        const api = config?.api ?? "messages";
+        const events = options.events ?? { provider: "jig-bedrock:provider", usage: "jig-bedrock:usage" };
+        const models: TeamBedrockModel[] = [];
+        const guard = registerBedrockSelectionGuard(pi, models, options.selection ?? { agentDir: config?.agentDir });
+        {
+            const dropped = safeSection(options.api).dropped;
+            if (dropped.length > 0)
+                pi.on("session_start", (_event, ctx) =>
+                    ctx.ui.notify(
+                        `Bedrock ignores ${dropped.join(" and ")} from a project's jig.json; set them in ${path.join(agentFolder(options.env), "bedrock.json")} or the environment.`,
+                        "warning",
+                    ),
+                );
+        }
+        if (!config) {
+            const file = path.join(agentFolder(options.env), "bedrock.json");
+            const where =
+                options.api === undefined
+                    ? `Add "team" to ${file}, then /reload.`
+                    : `Set it in bedrock in jig.json or in ${file}, then /reload.`;
+            pi.on("session_start", (_event, ctx) =>
+                ctx.ui.notify(`Bedrock provider disabled: ${resolved.error} ${where}`, "warning"),
+            );
+            return undefined;
+        }
+        const catalog = options.catalog ?? createAwsDiscovery(config);
+        const statistics = options.statistics ?? createStatistics(config);
+        const budget = options.budget ?? createBudget(config, undefined, undefined, undefined, statistics);
+        registerSwitchNotice(pi, { env: options.env });
+        const lifetime = new AbortController();
+        // A subagent child never refreshes the catalog, polls usage program or publishes usage: the parent does, and nothing there shows it.
+        const child = isSubagentChild(options.env);
+        const usageEvents = options.usageEvents !== false && !child;
+        const backgroundRefresh = options.backgroundRefresh !== false && !child;
+        let ctx: ExtensionContext | undefined;
+        let native: Provider | undefined;
+        let messages: RecoveredBedrock["messages"];
+        let failed = false;
+        let disposed = false;
+        let timer: ReturnType<typeof setInterval> | undefined;
+        const readModels = () => {
+            const result = discoveredProfiles(catalog.load(), config.usageDir, api);
+            if (options.transformModels) result.available = options.transformModels(result.available);
+            return result;
+        };
+        const publish = (): void => {
+            if (!ctx) return;
+            const spent = budget.load();
+            pi.events?.emit(events.usage, { sessionId: ctx.sessionManager.getSessionId(), budget: spent });
+            pi.events?.emit(SEGMENT_EVENT, spendSegment(spent, Date.now()));
+        };
+        const provider = (): void => {
+            pi.events?.emit(events.provider, {
+                sessionId: ctx?.sessionManager.getSessionId(),
+                budget: budget.load(),
+                models,
+                unavailable: readModels().unavailable,
+                status: failed ? "failed" : models.length ? "ready" : "absent",
+            });
+            if (usageEvents) pi.events?.emit(MODEL_NAMES_EVENT, modelNames(models));
+        };
+        const authProfile = options.preserveNativeAuth ? undefined : config.profile;
+        // A repair the sign-in makes for WSL is told in one line, the way the notices of a waiting request are.
+        const say = (line: string): void => {
+            if (ctx) noticeFor(ctx)(line, "info");
+        };
+        const recovery = options.ssoRecover ?? createRecovery(config, { ...options.login, say });
+        // The Messages route signs with the profile's SSO credentials; the credential chain is loaded from Pi's install on the
+        // first request, so nothing here reaches AWS at registration. Tests supply both.
+        const signedFetch =
+            api === "messages"
+                ? createSignedFetch({
+                      credentials:
+                          options.credentials ??
+                          (() => {
+                              // One chain, made on the first request: it caches the SSO role's credentials and refreshes them.
+                              let chain: (() => Promise<AwsCredentials>) | undefined;
+                              return async () => {
+                                  chain ??= loadCredentialProvider(findPiHost(options.env))({
+                                      profile: authProfile ?? config.profile,
+                                  }) as () => Promise<AwsCredentials>;
+                                  return chain();
+                              };
+                          })(),
+                      ...(options.fetch === undefined ? {} : { fetch: options.fetch }),
+                  })
+                : undefined;
+        try {
+            const recovered = await registerBedrockSsoRecovery(pi, recovery, authProfile);
+            native = recovered.provider;
+            messages = recovered.messages;
+            usePiRegistry(native);
+            models.push(...readModels().available);
+            await registerTeamBedrock(
+                pi,
+                config.agentDir || "",
+                native,
+                models,
+                catalog.load,
+                authProfile,
+                messages,
+                signedFetch,
+            );
+        } catch {
+            failed = true;
+        }
+        const refresh = async (force = true): Promise<Status> => {
+            const status = await catalog.refresh(force);
+            if (disposed) return status;
+            const updated = readModels().available;
+            const selected = ctx?.model;
+            if (isJigBedrock(selected)) {
+                const previous = models.find((m) => m.id === selected?.id);
+                const replacement = updated.find((m) => m.id === selected?.id);
+                if (
+                    !replacement ||
+                    (previous &&
+                        (previous.arn !== replacement.arn ||
+                            previous.foundationModelId !== replacement.foundationModelId))
+                ) {
+                    guard.blockUntilChoice();
+                }
+            }
+            const change = updateModels(models, updated);
+            if (ctx && (change.added || change.removed)) {
+                ctx.ui.notify(
+                    `Bedrock catalog updated: ${change.added} added, ${change.removed} removed. Choose /model if your selection was removed.`,
+                    "info",
+                );
+            }
+            if (change.changed && native) {
+                await registerTeamBedrock(
+                    pi,
+                    config.agentDir || "",
+                    native,
+                    models,
+                    catalog.load,
+                    authProfile,
+                    messages,
+                    signedFetch,
+                );
+            }
+            provider();
+            return status;
+        };
+        // Sign in when Pi starts and the session has expired, through the recovery the commands and requests share. Not in a child:
+        // the parent does it, and a second window from a child would be one nobody expects.
+        if (config.signInAtStart && !child && authProfile !== undefined) {
+            registerSignInAtStart(pi, {
+                profile: authProfile,
+                recover: shareRecovery(recovery),
+                check: options.ssoCheck ?? ((profile, signal) => checkSso(profile, systemSsoDeps(options.env), signal)),
+                refresh: () => refresh(false),
+            });
+        }
+        pi.registerCommand("bedrock-refresh", {
+            description: "Refresh team Bedrock models",
+            handler: async (_args, context) => {
+                const status = await refresh();
+                context.ui.notify(
+                    status === "ok"
+                        ? `Bedrock catalog: ${models.length} models. Choose /model.`
+                        : `Bedrock refresh ${status}. ${discoveryHelp(status)}`,
+                    status === "ok" ? "info" : "warning",
+                );
+            },
+        });
+        registerStatisticsCommands(pi, { config, statistics, catalog, models: () => models, budget, publish });
+        pi.on("session_start", (_event, context) => {
+            disposed = false;
+            ctx = context;
+            provider();
+            if (usageEvents) publish();
+            const unavailable = readModels().unavailable;
+            if (failed) context.ui.notify("Bedrock unavailable; update Pi and /reload.", "warning");
+            else if (unavailable.length > 0)
+                context.ui.notify(
+                    `${unavailable.length} Bedrock profiles unavailable: ${unavailable[0]?.reason}`,
+                    "warning",
+                );
+            else if (!models.length)
+                context.ui.notify("Bedrock models not loaded yet; run /aws-login, then /bedrock-refresh.", "warning");
+            // Start in the background; session startup never awaits AWS or usage program.
+            if (backgroundRefresh) void refresh(false).catch(() => {});
+            if (usageEvents && context.mode === "tui") {
+                const tick = (): void => {
+                    if (isJigBedrock(ctx?.model)) {
+                        void budget
+                            .refresh({ signal: lifetime.signal })
+                            .then(publish)
+                            .catch(() => {});
+                    }
+                };
+                tick();
+                timer = setInterval(tick, 10_000);
+                timer.unref();
+            }
+        });
+        if (usageEvents) pi.on("model_select", publish);
+        pi.on("session_shutdown", () => {
+            lifetime.abort();
+            disposed = true;
+            ctx = undefined;
+            if (timer) clearInterval(timer);
+        });
+        return { models, refresh, failed };
+    };
+}

package/src/sigv4.ts

@@ -0,0 +1,92 @@
+// AWS Signature Version 4 for the one request shape this provider sends: a POST to Bedrock's Anthropic route with a JSON body
+// and the credentials the AWS SDK resolves for the configured profile. Kept here, not in an SDK import, so it can be tested
+// against AWS's published example and so nothing reaches for a credential store on its own.
+import { createHash, createHmac } from "node:crypto";
+
+export interface AwsCredentials {
+    accessKeyId: string;
+    secretAccessKey: string;
+    sessionToken?: string;
+}
+
+export interface SignInput {
+    method: string;
+    url: string;
+    /** The request's headers as sent: `content-type`, `anthropic-version`, and so on. Auth headers are dropped here. */
+    headers: Readonly<Record<string, string>>;
+    body: string | Uint8Array;
+    region: string;
+    service?: string;
+    credentials: AwsCredentials;
+    now?: Date;
+}
+
+/** The `YYYYMMDDTHHMMSSZ` SigV4 date. */
+export function amzDate(now: Date): string {
+    return now.toISOString().replace(/[:-]|\.\d{3}/g, "");
+}
+
+const sha256 = (data: string | Uint8Array): string => createHash("sha256").update(data).digest("hex");
+const hmac = (key: string | Buffer, data: string): Buffer => createHmac("sha256", key).update(data).digest();
+
+/** RFC 3986, which is what SigV4 canonicalizes with: encodeURIComponent leaves `!'()*` alone and this does not. */
+function encode(value: string): string {
+    return encodeURIComponent(value).replace(/[!'()*]/g, (c) => `%${c.charCodeAt(0).toString(16).toUpperCase()}`);
+}
+
+function canonicalQuery(search: string): string {
+    const pairs: Array<[string, string]> = [];
+    for (const [name, value] of new URLSearchParams(search)) pairs.push([name, value]);
+    return pairs
+        .sort(([a, av], [b, bv]) => (a === b ? (av < bv ? -1 : av > bv ? 1 : 0) : a < b ? -1 : 1))
+        .map(([name, value]) => `${encode(name)}=${encode(value)}`)
+        .join("&");
+}
+
+/**
+ * The headers that carry the signature: `authorization`, `x-amz-date`, and the session token when the credentials have one.
+ * Every header given is signed, apart from the auth headers this replaces; `host` is taken from the URL.
+ */
+export function signRequest(input: SignInput): Record<string, string> {
+    const url = new URL(input.url);
+    const now = input.now ?? new Date();
+    const timestamp = amzDate(now);
+    const day = timestamp.slice(0, 8);
+    const service = input.service ?? "bedrock";
+    const payload = sha256(input.body);
+
+    const headers = new Map<string, string>();
+    for (const [name, value] of Object.entries(input.headers)) {
+        const key = name.toLowerCase();
+        if (key === "authorization" || key === "x-api-key" || key === "x-amz-date") continue;
+        headers.set(key, value.trim().replace(/\s+/g, " "));
+    }
+    headers.set("host", url.host);
+    headers.set("x-amz-date", timestamp);
+    if (input.credentials.sessionToken) headers.set("x-amz-security-token", input.credentials.sessionToken);
+
+    const names = [...headers.keys()].sort();
+    const signedHeaders = names.join(";");
+    const canonicalHeaders = names.map((name) => `${name}:${headers.get(name)}\n`).join("");
+    const canonicalRequest = [
+        input.method.toUpperCase(),
+        url.pathname === "" ? "/" : url.pathname,
+        canonicalQuery(url.search),
+        canonicalHeaders,
+        signedHeaders,
+        payload,
+    ].join("\n");
+    const scope = `${day}/${input.region}/${service}/aws4_request`;
+    const toSign = ["AWS4-HMAC-SHA256", timestamp, scope, sha256(canonicalRequest)].join("\n");
+    const kDate = hmac(`AWS4${input.credentials.secretAccessKey}`, day);
+    const kRegion = hmac(kDate, input.region);
+    const kService = hmac(kRegion, service);
+    const kSigning = hmac(kService, "aws4_request");
+    const signature = createHmac("sha256", kSigning).update(toSign).digest("hex");
+
+    return {
+        authorization: `AWS4-HMAC-SHA256 Credential=${input.credentials.accessKeyId}/${scope}, SignedHeaders=${signedHeaders}, Signature=${signature}`,
+        "x-amz-date": timestamp,
+        ...(input.credentials.sessionToken ? { "x-amz-security-token": input.credentials.sessionToken } : {}),
+    };
+}

package/src/sso-check.ts

@@ -0,0 +1,111 @@
+// Is the profile's AWS SSO session still good? Answered the way Pi's Bedrock provider will resolve it: through the AWS SDK's node
+// credential chain, not the AWS CLI, whose cached role credentials outlive an expired SSO session. Based on the older
+// bin/aws-sso-check.ts. The answer is one word and never a credential: the SDK's result is thrown away, and the token cache is
+// read for its expiry time alone.
+import { createHash } from "node:crypto";
+import { readFileSync } from "node:fs";
+import { homedir } from "node:os";
+import path from "node:path";
+import type { CredentialProvider } from "./pi-host.ts";
+import { findPiHost, loadCredentialProvider } from "./pi-host.ts";
+
+/** `ok`: the credentials resolve. `sign-in`: the SSO session has expired. `error`: the check itself could not tell. */
+export type SsoState = "ok" | "sign-in" | "error";
+
+/** The SDK's credential errors that mean "run aws sso login": its SSO token and credential providers all end theirs with that hint. */
+export function classify(error: unknown): SsoState {
+    const text = error instanceof Error ? error.message : String(error);
+    return /aws sso login/i.test(text) || /^Token is expired/.test(text) ? "sign-in" : "error";
+}
+
+/** The SSO session name (or legacy start URL) a profile signs in with, read from the AWS config text. */
+export function ssoKey(config: string, profile: string): string | undefined {
+    let section = "";
+    let key: string | undefined;
+    for (const line of config.split(/\r?\n/)) {
+        const header = /^\s*\[(.+)\]\s*$/.exec(line);
+        if (header) {
+            section = (header[1] ?? "").trim();
+            continue;
+        }
+        if (section !== `profile ${profile}` && !(profile === "default" && section === "default")) continue;
+        const pair = /^\s*(sso_session|sso_start_url)\s*=\s*(\S+)/.exec(line);
+        if (pair && (pair[1] === "sso_session" || key === undefined)) key = pair[2];
+    }
+    return key;
+}
+
+/** True when the cached SSO token stays valid for `marginMs` more: then the SDK needs no sign-in and no network call. */
+export function tokenFresh(cache: string, now: number, marginMs = 15 * 60_000): boolean {
+    try {
+        const expires = Date.parse((JSON.parse(cache) as { expiresAt?: string }).expiresAt ?? "");
+        return Number.isFinite(expires) && expires - now > marginMs;
+    } catch {
+        return false;
+    }
+}
+
+/** The parts of the machine the check reads, which a test replaces. */
+export interface SsoCheckDeps {
+    env: NodeJS.ProcessEnv;
+    home: string;
+    /** The text of a file; throws when it cannot be read. */
+    readText(file: string): string;
+    now(): number;
+    /** The SDK's `defaultProvider`, from Pi's install. Throws when it cannot be loaded. */
+    credentialProvider(): CredentialProvider;
+    /** How long the SDK may take before the answer is `error`. */
+    timeoutMs: number;
+}
+
+export function systemSsoDeps(env: NodeJS.ProcessEnv = process.env): SsoCheckDeps {
+    return {
+        env,
+        home: homedir(),
+        readText: (file) => readFileSync(file, "utf8"),
+        now: Date.now,
+        credentialProvider: () => loadCredentialProvider(findPiHost(env)),
+        timeoutMs: 10_000,
+    };
+}
+
+/** A fresh cached token means the profile is signed in, found by two small file reads and no SDK. */
+function cachedTokenFresh(profile: string, deps: SsoCheckDeps): boolean {
+    try {
+        const config = deps.readText(deps.env.AWS_CONFIG_FILE || path.join(deps.home, ".aws/config"));
+        const key = ssoKey(config, profile);
+        if (!key) return false;
+        const file = path.join(deps.home, ".aws/sso/cache", `${createHash("sha1").update(key).digest("hex")}.json`);
+        return tokenFresh(deps.readText(file), deps.now());
+    } catch {
+        return false;
+    }
+}
+
+/** `work`, or a rejection as soon as `signal` fires, whether or not `work` listens to it. `work`'s own failure is always observed. */
+function until<T>(work: Promise<T>, signal: AbortSignal): Promise<T> {
+    return new Promise((resolve, reject) => {
+        const stop = () => reject(new Error("timed out or cancelled"));
+        if (signal.aborted) stop();
+        signal.addEventListener("abort", stop, { once: true });
+        work.then(resolve, reject).finally(() => signal.removeEventListener("abort", stop));
+    });
+}
+
+/** Checks `profile`. Never throws, and never returns or keeps anything of a credential. */
+export async function checkSso(
+    profile: string,
+    deps: SsoCheckDeps = systemSsoDeps(),
+    signal?: AbortSignal,
+): Promise<SsoState> {
+    // The fast path: a cached token that stays valid needs no network check and no SDK.
+    if (cachedTokenFresh(profile, deps)) return "ok";
+    const limit = AbortSignal.timeout(deps.timeoutMs);
+    try {
+        // The SDK's answer is dropped: only whether it resolved matters.
+        await until(deps.credentialProvider()({ profile })(), signal ? AbortSignal.any([signal, limit]) : limit);
+        return "ok";
+    } catch (error) {
+        return classify(error);
+    }
+}