import type { Api, Model } from "@earendil-works/pi-ai"; /** * The registry hands back models for whichever API their catalog entry uses. */ type AnyModel = Model; /** * A provider/model pair recorded on the session branch. */ export interface SessionModel { provider: string; modelId: string; } /** * Minimal view of the model registry needed to reconcile a session model. */ export interface SessionModelRegistry { find(provider: string, modelId: string): AnyModel | undefined; hasConfiguredAuth(model: AnyModel): boolean; refresh(provider: string): Promise; } export interface RestoreSessionModelOptions { providerId: string; entries: readonly unknown[]; activeModel: { provider: string; id: string } | undefined; registry: SessionModelRegistry; setModel(model: AnyModel): void; attempts?: number; retryDelayMs?: number; sleep?: (ms: number) => Promise; } export type RestoreSessionModelResult = "skipped" | "restored" | "unavailable"; export const DEFAULT_RESTORE_ATTEMPTS = 10; export const DEFAULT_RESTORE_RETRY_DELAY_MS = 50; function asRecord(value: unknown): Record | undefined { return value !== null && typeof value === "object" && !Array.isArray(value) ? (value as Record) : undefined; } function asString(value: unknown): string | undefined { return typeof value === "string" && value.length > 0 ? value : undefined; } /** * Recover the provider/model the session was last using. * * Pi folds `model_change` entries and the `provider`/`model` recorded on assistant * messages into the same "current model" notion, so both are considered here. */ export function recordedSessionModel(entries: readonly unknown[]): SessionModel | undefined { let recorded: SessionModel | undefined; for (const entry of entries) { const record = asRecord(entry); if (!record) { continue; } if (record.type === "model_change") { const provider = asString(record.provider); const modelId = asString(record.modelId); if (provider && modelId) { recorded = { provider, modelId }; } continue; } if (record.type === "message") { const message = asRecord(record.message); if (message?.role !== "assistant") { continue; } const provider = asString(message.provider); const modelId = asString(message.model); if (provider && modelId) { recorded = { provider, modelId }; } } } return recorded; } /** * Re-select the session's provider model when pi could not restore it on its own. * * Pi restores the session model during `createAgentSession` by requiring both the * catalog entry *and* `hasConfiguredAuth()`. The catalog is restored synchronously * from the model store, but the configured-provider snapshot is filled in by an * asynchronous availability refresh, so a session recorded on this provider can * fall back to the default model with * `Could not restore model /`. * * This runs from `session_start`, after registry composition, and repairs that * fallback by waiting (bounded) for the provider to become available again. */ export async function restoreSessionModel( options: RestoreSessionModelOptions, ): Promise { const recorded = recordedSessionModel(options.entries); if (!recorded || recorded.provider !== options.providerId) { return "skipped"; } if (options.activeModel?.provider === recorded.provider && options.activeModel.id === recorded.modelId) { return "skipped"; } const model = options.registry.find(recorded.provider, recorded.modelId); if (!model) { // The slug is gone from the catalog; let pi's fallback model stand. return "skipped"; } if (!options.registry.hasConfiguredAuth(model)) { // Provider-scoped, offline refresh: it rebuilds the configured-provider // snapshot without triggering a catalog fetch. await options.registry.refresh(options.providerId).catch(() => undefined); } const attempts = options.attempts ?? DEFAULT_RESTORE_ATTEMPTS; const retryDelayMs = options.retryDelayMs ?? DEFAULT_RESTORE_RETRY_DELAY_MS; const sleep = options.sleep ?? defaultSleep; for (let attempt = 0; attempt < attempts; attempt++) { if (options.registry.hasConfiguredAuth(model)) { options.setModel(model); return "restored"; } await sleep(retryDelayMs); } return "unavailable"; } function defaultSleep(ms: number): Promise { return new Promise((resolve) => setTimeout(resolve, ms)); }