@centerforagenticai/pi-multi-account 0.1.1

package/src/credential-refresh.ts

@@ -0,0 +1,248 @@
+/**
+ * Forces one credential refresh when a provider says 401 while the local
+ * credential still looks valid.
+ *
+ * Pi refreshes an OAuth credential only after LOCAL expiry. A server-side early
+ * revocation therefore leaves a token that looks fine locally and 401s on every
+ * request, forever, with nothing triggering a refresh. The warmer only warms
+ * near-expiry credentials, so it never touches this case either.
+ *
+ * Ported from the Sarrius reference (index.ts:2576-2646), with one structural
+ * improvement: persistence goes through `AuthStorage.modify()`, whose file
+ * backend does a locked read-modify-write, so the cross-process race Sarrius
+ * handles by hand is handled by the host. Sarrius's sharpest check is kept: a
+ * refresh that returns a DIFFERENT account identity is terminal, because the
+ * slot is now a different account and continuing would silently use credentials
+ * the operator did not intend.
+ *
+ * CREDENTIAL DISCIPLINE. Credential values transit the `modify` callback and
+ * the family refresher, and nowhere else: they never enter fields on this
+ * class, never reach diagnostics, and never appear in a return value. Callers
+ * receive only a bounded outcome string.
+ */
+
+import type { DiagnosticLog } from "./diagnostics.js";
+import type { AllowedFamily } from "./config.js";
+
+/**
+ * The slice of a stored OAuth credential this module touches. Matches the
+ * host's stored shape; extra fields pass through `modify` untouched.
+ */
+export interface RefreshableCredential {
+	readonly type?: string;
+	readonly access?: string;
+	readonly refresh?: string;
+	readonly expires?: number;
+	readonly accountId?: string;
+	readonly [key: string]: unknown;
+}
+
+/** Complete input required by both maintained family refresh callbacks. */
+export interface OAuthRefreshCredential extends RefreshableCredential {
+	readonly type: "oauth";
+	readonly access: string;
+	readonly refresh: string;
+	readonly expires: number;
+}
+
+/** Refreshes one family's credential; the captured upstream oauth surface. */
+export type FamilyRefresher = (
+	credentials: OAuthRefreshCredential,
+) => Promise<RefreshableCredential>;
+
+/**
+ * `CredentialStore.modify` shape, injected so tests never construct a real
+ * store and live wiring can bind the host's already-open locked store.
+ */
+export type AuthModify = (
+	provider: string,
+	fn: (
+		current: RefreshableCredential | undefined,
+	) => Promise<RefreshableCredential | undefined>,
+) => Promise<RefreshableCredential | undefined>;
+
+function isRecord(value: unknown): value is Record<string, unknown> {
+	return typeof value === "object" && value !== null;
+}
+
+function isOAuthRefreshCredential(
+	value: RefreshableCredential | undefined,
+): value is OAuthRefreshCredential {
+	return (
+		value?.type === "oauth" &&
+		typeof value.access === "string" &&
+		value.access.length > 0 &&
+		typeof value.refresh === "string" &&
+		value.refresh.length > 0 &&
+		typeof value.expires === "number" &&
+		Number.isFinite(value.expires)
+	);
+}
+
+function familyRefresher(oauth: unknown): FamilyRefresher | undefined {
+	const refreshToken = isRecord(oauth) ? oauth.refreshToken : undefined;
+	if (typeof refreshToken !== "function") return undefined;
+	return async (credential) => {
+		const refreshed: unknown = await refreshToken.call(oauth, credential);
+		return isRecord(refreshed) ? refreshed : {};
+	};
+}
+
+export interface HostForcedCredentialRefresherOptions {
+	readonly modelRegistry: unknown;
+	readonly oauthByFamily: Partial<Record<AllowedFamily, unknown>>;
+	readonly diagnostics: Pick<DiagnosticLog, "record">;
+}
+
+/**
+ * Binds the maintained OAuth callbacks to Pi's already-open credential store.
+ *
+ * `ModelRegistry.runtime` and `ModelRuntime.credentials` are private in the
+ * locked Pi 0.84.4 declaration files, while that exact installed runtime owns
+ * both fields. `RuntimeCredentials` forwards `modify` to `AuthStorage.modify`,
+ * whose file backend holds a cross-process lock for the whole async
+ * read-modify-write. If that verified host shape changes, this returns
+ * undefined and 401 handling fails closed into the existing
+ * invalidate-and-route path.
+ */
+export function createHostForcedCredentialRefresher(
+	options: HostForcedCredentialRefresherOptions,
+): ForcedCredentialRefresher | undefined {
+	const credentials = (
+		options.modelRegistry as
+			| {
+					runtime?: {
+						credentials?: { modify?: unknown };
+					};
+			  }
+			| undefined
+	)?.runtime?.credentials;
+	if (!credentials || typeof credentials.modify !== "function") return undefined;
+
+	const refreshers: Partial<Record<AllowedFamily, FamilyRefresher>> = {};
+	for (const family of [
+		"anthropic",
+		"openai-codex",
+		"google-antigravity",
+	] as const) {
+		const refresh = familyRefresher(options.oauthByFamily[family]);
+		if (refresh !== undefined) refreshers[family] = refresh;
+	}
+	return new ForcedCredentialRefresher({
+		modify: credentials.modify.bind(credentials) as AuthModify,
+		refreshers,
+		diagnostics: options.diagnostics,
+	});
+}
+
+export type ForcedRefreshOutcome =
+	/** A refreshed credential was persisted; the account deserves a fresh chance. */
+	| "refreshed"
+	/** The refresh identified a DIFFERENT account; invalidation must stand. */
+	| "identity-changed"
+	/** The refresh failed; nothing was persisted. */
+	| "failed"
+	/** No stored oauth credential with a refresh token; nothing to force. */
+	| "unsupported"
+	/** This provider already got its one forced attempt this session. */
+	| "already-attempted";
+
+function recordOutcome(
+	diagnostics: Pick<DiagnosticLog, "record">,
+	providerId: string,
+	outcome: ForcedRefreshOutcome,
+): ForcedRefreshOutcome {
+	diagnostics.record(
+		outcome === "refreshed" ? "info" : "warning",
+		"credential.force-refresh",
+		outcome === "refreshed"
+			? "Forced one credential refresh after an explicit provider 401."
+			: "A forced credential refresh did not recover the provider; normal invalidation remains active.",
+		{ providerId, outcome },
+	);
+	return outcome;
+}
+
+export interface ForcedCredentialRefresherOptions {
+	readonly modify: AuthModify;
+	readonly refreshers: Partial<Record<AllowedFamily, FamilyRefresher>>;
+	readonly diagnostics: Pick<DiagnosticLog, "record">;
+}
+
+export class ForcedCredentialRefresher {
+	readonly #modify: AuthModify;
+	readonly #refreshers: Partial<Record<AllowedFamily, FamilyRefresher>>;
+	readonly #diagnostics: Pick<DiagnosticLog, "record">;
+	/**
+	 * Exactly one forced attempt per provider per session, whatever the outcome.
+	 * A 401 that survives a fresh token is not a staleness problem, and retrying
+	 * the refresh on every exhausted turn would hammer the provider's token
+	 * endpoint with requests that cannot succeed.
+	 */
+	readonly #attempted = new Set<string>();
+
+	constructor(options: ForcedCredentialRefresherOptions) {
+		this.#modify = options.modify;
+		this.#refreshers = options.refreshers;
+		this.#diagnostics = options.diagnostics;
+	}
+
+	async attempt(
+		providerId: string,
+		family: AllowedFamily,
+	): Promise<ForcedRefreshOutcome> {
+		if (this.#attempted.has(providerId)) return "already-attempted";
+		this.#attempted.add(providerId);
+
+		const refresher = this.#refreshers[family];
+		if (refresher === undefined) return "unsupported";
+
+		let outcome: ForcedRefreshOutcome = "failed";
+		try {
+			await this.#modify(providerId, async (current) => {
+				if (!isOAuthRefreshCredential(current)) {
+					outcome = "unsupported";
+					return undefined;
+				}
+				let refreshed: RefreshableCredential;
+				try {
+					refreshed = await refresher(current);
+				} catch {
+					// The refresh itself failed. Persist nothing; the 401 keeps its
+					// terminal classification. Never record the provider's error body:
+					// OAuth endpoints can echo credential material.
+					outcome = "failed";
+					return undefined;
+				}
+				if (
+					typeof current.accountId === "string" &&
+					typeof refreshed.accountId === "string" &&
+					current.accountId !== refreshed.accountId
+				) {
+					// The slot now resolves to a different account. Persisting would
+					// silently authenticate as an account the operator did not choose --
+					// the exact concurrency failure this extension exists to avoid.
+					outcome = "identity-changed";
+					return undefined;
+				}
+				const merged: RefreshableCredential = {
+					...current,
+					...refreshed,
+					type: "oauth",
+				};
+				if (!isOAuthRefreshCredential(merged)) {
+					outcome = "failed";
+					return undefined;
+				}
+				outcome = "refreshed";
+				return merged;
+			});
+		} catch {
+			// Record only bounded metadata. The thrown value can originate in an OAuth
+			// implementation and is therefore not safe diagnostic material.
+			outcome = "failed";
+		}
+
+		return recordOutcome(this.#diagnostics, providerId, outcome);
+	}
+}

package/src/declaration-notice-marker.ts

@@ -0,0 +1,238 @@
+import { randomUUID } from "node:crypto";
+import {
+	chmodSync,
+	closeSync,
+	fsyncSync,
+	mkdirSync,
+	openSync,
+	readFileSync,
+	renameSync,
+	statSync,
+	unlinkSync,
+	writeFileSync,
+} from "node:fs";
+import { dirname, join } from "node:path";
+import { acquireMachineLease } from "./machine-lease.js";
+import { getPeriodBounds } from "./period-boundaries.js";
+
+const SCHEMA_VERSION = 1;
+export const MAX_MARKER_BYTES = 1_024;
+const LEASE_TTL_MS = 5_000;
+
+export type DeclarationNoticeCondition = "stale" | "not-installed";
+
+/**
+ * The only condition values that may ever reach the persisted marker.
+ *
+ * TypeScript's `DeclarationNoticeCondition` union is erased at runtime, so a
+ * JavaScript or casted caller could otherwise pass an arbitrary string that
+ * would be serialized into `declaration-notice.json`. This allow-list is the
+ * runtime boundary that keeps a caller from causing a path, identifier, or
+ * credential-like value to be retained (AC-6): anything outside it is rejected
+ * before the lease is acquired or any byte is written.
+ */
+const ALLOWED_CONDITIONS: ReadonlySet<DeclarationNoticeCondition> = new Set([
+	"stale",
+	"not-installed",
+]);
+
+function isAllowedCondition(
+	value: unknown,
+): value is DeclarationNoticeCondition {
+	return ALLOWED_CONDITIONS.has(value as DeclarationNoticeCondition);
+}
+
+
+interface DeclarationNoticeState {
+	readonly schemaVersion: 1;
+	readonly dayStartMs: number;
+	readonly condition: DeclarationNoticeCondition;
+}
+
+type MarkerReadResult =
+	| { readonly status: "valid"; readonly marker: DeclarationNoticeState }
+	| { readonly status: "missing" }
+	| { readonly status: "malformed" }
+	| { readonly status: "unavailable" };
+
+function defaultPaths(): { readonly path: string; readonly lockPath: string } {
+	const agentDir =
+		process.env.PI_CODING_AGENT_DIR ??
+		join(process.env.HOME ?? ".", ".pi", "agent");
+	// Package storage identity stays decoupled from the logical provider ID.
+	const directory = join(agentDir, "pi-multi-account");
+	return {
+		path: join(directory, "declaration-notice.json"),
+		lockPath: join(directory, "declaration-notice.lock"),
+	};
+}
+
+function parseMarker(value: unknown): DeclarationNoticeState | undefined {
+	if (typeof value !== "object" || value === null || Array.isArray(value)) {
+		return undefined;
+	}
+	const candidate = value as Record<string, unknown>;
+	if (
+		candidate.schemaVersion !== SCHEMA_VERSION ||
+		!Number.isSafeInteger(candidate.dayStartMs) ||
+		(candidate.dayStartMs as number) < 0 ||
+		(candidate.condition !== "stale" && candidate.condition !== "not-installed") ||
+		Object.keys(candidate).some(
+			(key) =>
+				key !== "schemaVersion" && key !== "dayStartMs" && key !== "condition",
+		)
+	) {
+		return undefined;
+	}
+	return {
+		schemaVersion: SCHEMA_VERSION,
+		dayStartMs: candidate.dayStartMs as number,
+		condition: candidate.condition,
+	};
+}
+
+/** Persistent once-per-UTC-day marker for the session-start declaration notice. */
+export class DeclarationNoticeMarker {
+	readonly #path: string;
+	readonly #lockPath: string;
+	readonly #now: () => number;
+
+	constructor(options: {
+		readonly path?: string;
+		readonly lockPath?: string;
+		readonly now?: () => number;
+	} = {}) {
+		const defaults = defaultPaths();
+		this.#path = options.path ?? defaults.path;
+		this.#lockPath = options.lockPath ?? defaults.lockPath;
+		this.#now = options.now ?? Date.now;
+	}
+
+	/**
+	 * Returns true when the caller should notify. Contention and lease/read failures
+	 * suppress; a post-decision marker write failure still returns true.
+	 */
+	shouldNotify(condition: DeclarationNoticeCondition): boolean {
+		// Reject any runtime value outside the allowed literals BEFORE acquiring
+		// the lease or writing, so a caller can never cause an arbitrary string to
+		// be persisted in the marker (AC-6).
+		if (!isAllowedCondition(condition)) return false;
+		let lease: ReturnType<typeof acquireMachineLease>;
+		try {
+			lease = acquireMachineLease({
+				lockPath: this.#lockPath,
+				ttlMs: LEASE_TTL_MS,
+				now: this.#now,
+				reclaimMalformed: true,
+			});
+		} catch {
+			return false;
+		}
+		if (lease === undefined) return false;
+		try {
+			const todayStartMs = getPeriodBounds(this.#now(), "day").startMs;
+			const existing = this.#readMarker();
+			if (existing.status === "unavailable") return false;
+			if (
+				existing.status === "valid" &&
+				existing.marker.dayStartMs === todayStartMs
+			) {
+				return false;
+			}
+			this.#writeMarker({
+				schemaVersion: SCHEMA_VERSION,
+				dayStartMs: todayStartMs,
+				condition,
+			});
+			return true;
+		} catch {
+			return false;
+		} finally {
+			try {
+				lease.release();
+			} catch {
+				// A release failure cannot block session startup.
+			}
+		}
+	}
+
+	/** A matched session removes the prior notice marker so later drift can notify. */
+	clear(): void {
+		let lease: ReturnType<typeof acquireMachineLease>;
+		try {
+			lease = acquireMachineLease({
+				lockPath: this.#lockPath,
+				ttlMs: LEASE_TTL_MS,
+				now: this.#now,
+				reclaimMalformed: true,
+			});
+		} catch {
+			return;
+		}
+		if (lease === undefined) return;
+		try {
+			unlinkSync(this.#path);
+		} catch {
+			// Marker cleanup failures cannot block session startup.
+		} finally {
+			try {
+				lease.release();
+			} catch {
+				// A release failure cannot block session startup.
+			}
+		}
+	}
+
+	#readMarker(): MarkerReadResult {
+		try {
+			const stats = statSync(this.#path);
+			if (!stats.isFile() || stats.size > MAX_MARKER_BYTES) {
+				return { status: "malformed" };
+			}
+			const parsed = parseMarker(JSON.parse(readFileSync(this.#path, "utf8")));
+			return parsed === undefined
+				? { status: "malformed" }
+				: { status: "valid", marker: parsed };
+		} catch (error) {
+			return (error as NodeJS.ErrnoException).code === "ENOENT"
+				? { status: "missing" }
+				: error instanceof SyntaxError
+					? { status: "malformed" }
+					: { status: "unavailable" };
+		}
+	}
+
+	#writeMarker(marker: DeclarationNoticeState): boolean {
+		const encoded = `${JSON.stringify(marker)}\n`;
+		if (Buffer.byteLength(encoded, "utf8") > MAX_MARKER_BYTES) return false;
+		const directory = dirname(this.#path);
+		const temporaryPath = `${this.#path}.${process.pid}.${randomUUID()}.tmp`;
+		let descriptor: number | undefined;
+		try {
+			mkdirSync(directory, { recursive: true, mode: 0o700 });
+			chmodSync(directory, 0o700);
+			descriptor = openSync(temporaryPath, "wx", 0o600);
+			writeFileSync(descriptor, encoded, { encoding: "utf8" });
+			fsyncSync(descriptor);
+			closeSync(descriptor);
+			descriptor = undefined;
+			renameSync(temporaryPath, this.#path);
+			chmodSync(this.#path, 0o600);
+			return true;
+		} catch {
+			if (descriptor !== undefined) {
+				try {
+					closeSync(descriptor);
+				} catch {
+					// Continue with temporary-file cleanup.
+				}
+			}
+			try {
+				unlinkSync(temporaryPath);
+			} catch {
+				// The notice decision stays fail-soft when persistence fails.
+			}
+			return false;
+		}
+	}
+}