import type { ExtensionContext } from "@earendil-works/pi-coding-agent"; import { PROVIDER_ID, USER_AGENT } from "./config.ts"; import { creditsUrl } from "./urls.ts"; const CREDITS_FETCH_TIMEOUT_MS = 30_000; const CREDITS_RETRY_DELAY_MS_INITIAL = 5_000; const CREDITS_RETRY_DELAY_MS_MAX = 5 * 60_000; const CREDITS_RETRY_EXPONENT_MAX = 6; export function parseCreditsRemaining(payload: unknown): number { const record = asRecord(payload); if (!record) { throw new Error("Experiential Labs credits response is invalid"); } const data = asRecord(record.data) ?? record; const remaining = asNumber(data.remaining); if (remaining !== undefined) { return remaining; } const total = asNumber(data.total_credits); if (total !== undefined) { return total - (asNumber(data.total_usage) ?? 0); } const balance = asNumber(data.balance) ?? asNumber(data.credits); if (balance !== undefined) { return balance; } throw new Error("Experiential Labs credits response is missing a balance"); } export function formatCreditsUsd(amount: number): string { return amount.toLocaleString("en-US", { style: "currency", currency: "USD", minimumFractionDigits: 2, maximumFractionDigits: 2, }); } function asRecord(value: unknown): Record | undefined { return value !== null && typeof value === "object" && !Array.isArray(value) ? (value as Record) : undefined; } function asNumber(value: unknown): number | undefined { return typeof value === "number" && Number.isFinite(value) ? value : undefined; } function isExplabsModel(model: ExtensionContext["model"]): model is NonNullable { return model?.provider === PROVIDER_ID; } function isTransientCreditError(error: unknown): boolean { if (error instanceof DOMException && error.name === "TimeoutError") { return true; } if (error instanceof TypeError) { return true; } const status = (error as { status?: number } | undefined)?.status; return status === 408 || status === 429 || (typeof status === "number" && status >= 500 && status <= 599); } async function fetchRemaining(baseUrl: string, apiKey: string, signal: AbortSignal): Promise { const timeout = AbortSignal.timeout(CREDITS_FETCH_TIMEOUT_MS); const response = await fetch(creditsUrl(baseUrl), { headers: { Accept: "application/json", Authorization: `Bearer ${apiKey}`, "User-Agent": USER_AGENT, }, signal: AbortSignal.any([signal, timeout]), }); if (!response.ok) { const error = new Error(`Experiential Labs credits failed: HTTP ${response.status}`) as Error & { status: number; }; error.status = response.status; throw error; } return parseCreditsRemaining(await response.json()); } export interface CreditStatusRuntime { refresh(ctx: ExtensionContext, selectedModel: ExtensionContext["model"]): Promise; dispose(): void; } export function createCreditStatusRuntime(baseUrl: string): CreditStatusRuntime { let invocationSequence = 0; let committedInvocation = 0; let credentialEpoch = 0; let consecutiveFailures = 0; let retryAtMs = 0; let currentApiKey: string | undefined; let cachedRemaining: number | undefined; let inFlight: | { apiKey: string; credentialEpoch: number; controller: AbortController; operation: Promise } | undefined; let disposed = false; type CredentialLease = { apiKey: string; credentialEpoch: number }; function clearCreditState(): void { cachedRemaining = undefined; consecutiveFailures = 0; retryAtMs = 0; } function invalidateCredential(): void { credentialEpoch += 1; inFlight?.controller.abort(); } function renderStatus(ctx: ExtensionContext): void { if (cachedRemaining === undefined) { ctx.ui.setStatus(PROVIDER_ID, undefined); return; } ctx.ui.setStatus(PROVIDER_ID, formatCreditsUsd(cachedRemaining)); } function ownsCredential(lease: CredentialLease): boolean { return lease.apiKey === currentApiKey && lease.credentialEpoch === credentialEpoch; } async function fetchAndCache(lease: CredentialLease, signal: AbortSignal): Promise { try { const remaining = await fetchRemaining(baseUrl, lease.apiKey, signal); if (disposed || !ownsCredential(lease)) { return; } cachedRemaining = remaining; consecutiveFailures = 0; retryAtMs = 0; } catch (error) { if (disposed || !ownsCredential(lease)) { throw error; } if (!isTransientCreditError(error)) { clearCreditState(); throw error; } consecutiveFailures += 1; retryAtMs = Date.now() + Math.min( CREDITS_RETRY_DELAY_MS_INITIAL * 2 ** Math.min(consecutiveFailures - 1, CREDITS_RETRY_EXPONENT_MAX), CREDITS_RETRY_DELAY_MS_MAX, ); throw error; } } async function shareFetch(apiKey: string, expectedCredentialEpoch: number): Promise { const active = inFlight; if (active?.apiKey === apiKey && active.credentialEpoch === expectedCredentialEpoch) { await active.operation; return; } const controller = new AbortController(); const operation = fetchAndCache({ apiKey, credentialEpoch: expectedCredentialEpoch }, controller.signal); const started = { apiKey, credentialEpoch: expectedCredentialEpoch, controller, operation }; inFlight = started; try { await operation; } finally { if (inFlight === started) { inFlight = undefined; } } } return { async refresh(ctx, selectedModel) { if (disposed || !ctx.hasUI) { return; } const invocation = invocationSequence + 1; invocationSequence = invocation; if (!isExplabsModel(selectedModel)) { committedInvocation = invocation; invalidateCredential(); ctx.ui.setStatus(PROVIDER_ID, undefined); return; } renderStatus(ctx); try { const auth = await ctx.modelRegistry.getApiKeyAndHeaders(selectedModel); if (disposed || invocation < committedInvocation) { return; } if (!auth.ok || !auth.apiKey) { committedInvocation = invocation; invalidateCredential(); currentApiKey = undefined; clearCreditState(); renderStatus(ctx); return; } committedInvocation = invocation; if (currentApiKey === auth.apiKey && Date.now() < retryAtMs) { return; } if (currentApiKey !== auth.apiKey) { currentApiKey = auth.apiKey; invalidateCredential(); clearCreditState(); renderStatus(ctx); } const lease = { apiKey: auth.apiKey, credentialEpoch }; await shareFetch(auth.apiKey, lease.credentialEpoch); if (disposed || !ownsCredential(lease) || invocation !== invocationSequence) { return; } renderStatus(ctx); } catch { if (disposed || invocation !== invocationSequence) { return; } renderStatus(ctx); } }, dispose() { if (disposed) { return; } disposed = true; committedInvocation = ++invocationSequence; invalidateCredential(); }, }; }