feat(ai): complete models runtime migration

packages/ai/src/compat.ts

@@ -19,7 +19,6 @@ export * from "./api/mistral-conversations.lazy.ts";
 export * from "./api/openai-codex-responses.lazy.ts";
 export * from "./api/openai-completions.lazy.ts";
 export * from "./api/openai-responses.lazy.ts";
-export * from "./api-registry.ts";
 export * from "./env-api-keys.ts";
 export * from "./image-models.ts";
 export * from "./images.ts";
@@ -36,11 +35,12 @@ import { mistralConversationsApi } from "./api/mistral-conversations.lazy.ts";
 import { openAICodexResponsesApi } from "./api/openai-codex-responses.lazy.ts";
 import { openAICompletionsApi } from "./api/openai-completions.lazy.ts";
 import { openAIResponsesApi } from "./api/openai-responses.lazy.ts";
-import { clearApiProviders, getApiProvider, registerApiProvider } from "./api-registry.ts";
 import { getEnvApiKey } from "./env-api-keys.ts";
-import { getBuiltinModel, getBuiltinModels, getBuiltinProviders } from "./providers/all.ts";
+import { builtinModels, getBuiltinModel, getBuiltinModels, getBuiltinProviders } from "./providers/all.ts";
+import { createFauxCore, type FauxProviderRegistration, type RegisterFauxProviderOptions } from "./providers/faux.ts";
 import type {
 	Api,
+	ApiStreamOptions,
 	AssistantMessage,
 	AssistantMessageEventStream,
 	Context,
@@ -48,6 +48,7 @@ import type {
 	ProviderStreamOptions,
 	ProviderStreams,
 	SimpleStreamOptions,
+	StreamFunction,
 	StreamOptions,
 } from "./types.ts";
 
@@ -60,6 +61,113 @@ export const getModels = getBuiltinModels;
 /** @deprecated Static catalog read. Use `getBuiltinProviders` from "@earendil-works/pi-ai/providers/all" or `Models.getProviders()`. */
 export const getProviders = getBuiltinProviders;
 
+export type ApiStreamFunction = (
+	model: Model<Api>,
+	context: Context,
+	options?: StreamOptions,
+) => AssistantMessageEventStream;
+
+export type ApiStreamSimpleFunction = (
+	model: Model<Api>,
+	context: Context,
+	options?: SimpleStreamOptions,
+) => AssistantMessageEventStream;
+
+export interface ApiProvider<TApi extends Api = Api, TOptions extends StreamOptions = StreamOptions> {
+	api: TApi;
+	stream: StreamFunction<TApi, TOptions>;
+	streamSimple: StreamFunction<TApi, SimpleStreamOptions>;
+}
+
+interface ApiProviderInternal {
+	api: Api;
+	stream: ApiStreamFunction;
+	streamSimple: ApiStreamSimpleFunction;
+}
+
+type RegisteredApiProvider = {
+	provider: ApiProviderInternal;
+	sourceId?: string;
+};
+
+const apiProviderRegistry = new Map<string, RegisteredApiProvider>();
+
+function wrapStream<TApi extends Api, TOptions extends StreamOptions>(
+	api: TApi,
+	stream: StreamFunction<TApi, TOptions>,
+): ApiStreamFunction {
+	return (model, context, options) => {
+		if (model.api !== api) {
+			throw new Error(`Mismatched api: ${model.api} expected ${api}`);
+		}
+		return stream(model as Model<TApi>, context, options as TOptions);
+	};
+}
+
+function wrapStreamSimple<TApi extends Api>(
+	api: TApi,
+	streamSimple: StreamFunction<TApi, SimpleStreamOptions>,
+): ApiStreamSimpleFunction {
+	return (model, context, options) => {
+		if (model.api !== api) {
+			throw new Error(`Mismatched api: ${model.api} expected ${api}`);
+		}
+		return streamSimple(model as Model<TApi>, context, options);
+	};
+}
+
+export function registerApiProvider<TApi extends Api, TOptions extends StreamOptions>(
+	provider: ApiProvider<TApi, TOptions>,
+	sourceId?: string,
+): void {
+	apiProviderRegistry.set(provider.api, {
+		provider: {
+			api: provider.api,
+			stream: wrapStream(provider.api, provider.stream),
+			streamSimple: wrapStreamSimple(provider.api, provider.streamSimple),
+		},
+		sourceId,
+	});
+}
+
+export function getApiProvider(api: Api): ApiProviderInternal | undefined {
+	return apiProviderRegistry.get(api)?.provider;
+}
+
+export function getApiProviders(): ApiProviderInternal[] {
+	return Array.from(apiProviderRegistry.values(), (entry) => entry.provider);
+}
+
+export function unregisterApiProviders(sourceId: string): void {
+	for (const [api, entry] of apiProviderRegistry.entries()) {
+		if (entry.sourceId === sourceId) {
+			apiProviderRegistry.delete(api);
+		}
+	}
+}
+
+function clearApiProviders(): void {
+	apiProviderRegistry.clear();
+}
+
+export function registerFauxProvider(options: RegisterFauxProviderOptions = {}): FauxProviderRegistration {
+	const core = createFauxCore(options);
+	const sourceId = `faux-provider-${Math.random().toString(36).slice(2, 10)}`;
+	registerApiProvider({ api: core.api, stream: core.stream, streamSimple: core.streamSimple }, sourceId);
+	return {
+		api: core.api,
+		models: core.models,
+		getModel: core.getModel,
+		state: core.state,
+		setResponses: core.setResponses,
+		appendResponses: core.appendResponses,
+		getPendingResponseCount: core.getPendingResponseCount,
+		unregister() {
+			unregisterApiProviders(sourceId);
+		},
+	};
+}
+
 const BUILTIN_APIS: [Api, ProviderStreams][] = [
 	["anthropic-messages", anthropicMessagesApi()],
 	["openai-completions", openAICompletionsApi()],
@@ -72,6 +180,8 @@ const BUILTIN_APIS: [Api, ProviderStreams][] = [
 	["bedrock-converse-stream", bedrockConverseStreamApi()],
 ];
 
+const builtinApiProviderInstances = new Map<Api, ReturnType<typeof getApiProvider>>();
+
 /**
  * Registers the builtin API implementations into the api-registry without
  * clobbering existing entries: compat may load after a test or extension has
@@ -79,18 +189,23 @@ const BUILTIN_APIS: [Api, ProviderStreams][] = [
  */
 export function registerBuiltInApiProviders(): void {
 	for (const [api, streams] of BUILTIN_APIS) {
-		if (getApiProvider(api)) continue;
-		registerApiProvider({ api, stream: streams.stream, streamSimple: streams.streamSimple });
+		if (!getApiProvider(api)) {
+			registerApiProvider({ api, stream: streams.stream, streamSimple: streams.streamSimple });
+		}
+		builtinApiProviderInstances.set(api, getApiProvider(api));
 	}
 }
 
 export function resetApiProviders(): void {
 	clearApiProviders();
+	builtinApiProviderInstances.clear();
 	registerBuiltInApiProviders();
 }
 
 registerBuiltInApiProviders();
 
+const compatModels = builtinModels();
+
 function hasExplicitApiKey(apiKey: string | undefined): apiKey is string {
 	return typeof apiKey === "string" && apiKey.trim().length > 0;
 }
@@ -105,6 +220,11 @@ function withEnvApiKey<TOptions extends StreamOptions>(
 	return { ...options, apiKey } as TOptions;
 }
 
+function shouldUseBuiltinModels(model: Model<Api>): boolean {
+	const builtin = compatModels.getModel(model.provider, model.id);
+	return builtin?.api === model.api && getApiProvider(model.api) === builtinApiProviderInstances.get(model.api);
+}
+
 function resolveApiProvider(api: Api) {
 	const provider = getApiProvider(api);
 	if (!provider) {
@@ -118,6 +238,9 @@ export function stream<TApi extends Api>(
 	context: Context,
 	options?: ProviderStreamOptions,
 ): AssistantMessageEventStream {
+	if (shouldUseBuiltinModels(model)) {
+		return compatModels.stream(model, context, options as ApiStreamOptions<TApi> | undefined);
+	}
 	const provider = resolveApiProvider(model.api);
 	return provider.stream(model, context, withEnvApiKey(model, options) as StreamOptions);
 }
@@ -136,6 +259,9 @@ export function streamSimple<TApi extends Api>(
 	context: Context,
 	options?: SimpleStreamOptions,
 ): AssistantMessageEventStream {
+	if (shouldUseBuiltinModels(model)) {
+		return compatModels.streamSimple(model, context, options);
+	}
 	const provider = resolveApiProvider(model.api);
 	return provider.streamSimple(model, context, withEnvApiKey(model, options));
 }