mirror of
https://github.com/badlogic/pi-mono.git
synced 2026-11-05 23:37:25 +00:00
fddc968b96
registerNativeProvider() did not update the auth snapshot, so the provider stayed unconfigured until an async refresh saved a snapshot that included it. At startup that refresh can resolve before its result is saved, so initial model selection saw the provider as unconfigured and fell back to another provider or warned that no models are available. Share registerProvider's provisional entry through markProvisionallyConfigured() so native providers with a stored credential or a models.json apiKey are configured immediately. fixes #9962
376 lines
11 KiB
TypeScript
376 lines
11 KiB
TypeScript
import { mkdtempSync, rmSync, writeFileSync } from "node:fs";
|
|
import { tmpdir } from "node:os";
|
|
import { join } from "node:path";
|
|
import {
|
|
createAssistantMessageEventStream,
|
|
type DeferredCancelOptions,
|
|
type DeferredFetchOptions,
|
|
InMemoryModelsStore,
|
|
type Model,
|
|
type Provider,
|
|
} from "@earendil-works/pi-ai";
|
|
import { describe, expect, it } from "vitest";
|
|
import { AuthStorage } from "../src/core/auth-storage.ts";
|
|
import { ModelRegistry } from "../src/core/model-registry.ts";
|
|
import { ModelRuntime } from "../src/core/model-runtime.ts";
|
|
|
|
function model(id: string): Model<"openai-completions"> {
|
|
return {
|
|
id,
|
|
name: id,
|
|
api: "openai-completions",
|
|
provider: "extension-oauth",
|
|
baseUrl: "https://example.test/v1",
|
|
reasoning: false,
|
|
input: ["text"],
|
|
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
|
|
contextWindow: 1000,
|
|
maxTokens: 100,
|
|
};
|
|
}
|
|
|
|
describe("extension provider model lifecycle", () => {
|
|
it("registers native pi-ai providers with their auth implementation", async () => {
|
|
const runtime = await ModelRuntime.create({
|
|
credentials: AuthStorage.inMemory(),
|
|
modelsStore: new InMemoryModelsStore(),
|
|
modelsPath: null,
|
|
allowModelNetwork: false,
|
|
});
|
|
const nativeModel = {
|
|
...model("native"),
|
|
provider: "extension-native",
|
|
baseUrl: "https://fallback.test/v1",
|
|
};
|
|
const provider: Provider = {
|
|
id: "extension-native",
|
|
name: "Extension Native",
|
|
auth: {
|
|
apiKey: {
|
|
name: "Native setup",
|
|
login: async (interaction) => ({
|
|
type: "api_key",
|
|
key: await interaction.prompt({ type: "secret", message: "API key" }),
|
|
}),
|
|
check: async ({ credential }) =>
|
|
credential?.key ? { type: "api_key", source: "stored native key" } : undefined,
|
|
resolve: async ({ credential }) =>
|
|
credential?.key
|
|
? {
|
|
auth: { apiKey: credential.key, baseUrl: "https://resolved.test/v1" },
|
|
source: "stored native key",
|
|
}
|
|
: undefined,
|
|
},
|
|
},
|
|
getModels: () => [nativeModel],
|
|
stream: () => {
|
|
throw new Error("unused");
|
|
},
|
|
streamSimple: () => {
|
|
throw new Error("unused");
|
|
},
|
|
};
|
|
|
|
runtime.registerNativeProvider(provider);
|
|
const registry = new ModelRegistry(runtime);
|
|
expect(registry.getProvider("extension-native")).toBe(provider);
|
|
expect(registry.getRegisteredNativeProvider("extension-native")).toBe(provider);
|
|
expect(registry.getRegisteredProviderIds()).toContain("extension-native");
|
|
expect(registry.find("extension-native", "native")).toBeDefined();
|
|
|
|
await runtime.login("extension-native", "api_key", {
|
|
prompt: async () => "secret",
|
|
notify: () => {},
|
|
});
|
|
expect(await registry.getProviderAuth("extension-native")).toMatchObject({
|
|
auth: { apiKey: "secret", baseUrl: "https://resolved.test/v1" },
|
|
});
|
|
|
|
registry.unregisterProvider("extension-native");
|
|
expect(registry.getProvider("extension-native")).toBeUndefined();
|
|
});
|
|
|
|
// Regression for #9962: initial model selection reads the snapshot before the async refresh finishes.
|
|
it("marks a native provider with a stored credential as configured when it registers", async () => {
|
|
const runtime = await ModelRuntime.create({
|
|
credentials: AuthStorage.inMemory({
|
|
"extension-native": {
|
|
type: "oauth",
|
|
access: "access",
|
|
refresh: "refresh",
|
|
expires: Date.now() + 3_600_000,
|
|
},
|
|
}),
|
|
modelsStore: new InMemoryModelsStore(),
|
|
modelsPath: null,
|
|
allowModelNetwork: false,
|
|
});
|
|
const nativeModel = { ...model("native"), provider: "extension-native" };
|
|
const unused = () => {
|
|
throw new Error("unused");
|
|
};
|
|
const provider: Provider = {
|
|
id: "extension-native",
|
|
name: "Extension Native",
|
|
auth: {
|
|
oauth: {
|
|
name: "Native OAuth",
|
|
login: unused,
|
|
refresh: async (credential) => credential,
|
|
toAuth: async (credential) => ({ apiKey: credential.access }),
|
|
},
|
|
},
|
|
getModels: () => [nativeModel],
|
|
stream: unused,
|
|
streamSimple: unused,
|
|
};
|
|
|
|
runtime.registerNativeProvider(provider);
|
|
|
|
expect(runtime.hasConfiguredAuth("extension-native")).toBe(true);
|
|
expect(runtime.isUsingOAuth("extension-native")).toBe(true);
|
|
expect(runtime.getAvailableSnapshot().map((m) => `${m.provider}/${m.id}`)).toContain("extension-native/native");
|
|
await runtime.refresh({ allowNetwork: false });
|
|
expect(runtime.hasConfiguredAuth("extension-native")).toBe(true);
|
|
});
|
|
|
|
it("preserves native deferred methods through provider overlays", async () => {
|
|
const tempDir = mkdtempSync(join(tmpdir(), "pi-native-provider-deferred-"));
|
|
const modelsPath = join(tempDir, "models.json");
|
|
writeFileSync(
|
|
modelsPath,
|
|
JSON.stringify({
|
|
providers: {
|
|
"extension-native-deferred": { baseUrl: "https://overlay.test/v1" },
|
|
},
|
|
}),
|
|
);
|
|
try {
|
|
const runtime = await ModelRuntime.create({
|
|
credentials: AuthStorage.inMemory(),
|
|
modelsStore: new InMemoryModelsStore(),
|
|
modelsPath,
|
|
allowModelNetwork: false,
|
|
});
|
|
const nativeModel = {
|
|
...model("native-deferred"),
|
|
provider: "extension-native-deferred",
|
|
baseUrl: "https://native.test/v1",
|
|
};
|
|
let fetchedBaseUrl: string | undefined;
|
|
let fetchedOptions: DeferredFetchOptions | undefined;
|
|
let cancelledId: string | undefined;
|
|
let cancelledOptions: DeferredCancelOptions | undefined;
|
|
const provider: Provider = {
|
|
id: "extension-native-deferred",
|
|
name: "Extension Native Deferred",
|
|
auth: {
|
|
apiKey: {
|
|
name: "Native key",
|
|
resolve: async () => ({ auth: { apiKey: "key" }, source: "native" }),
|
|
},
|
|
},
|
|
getModels: () => [nativeModel],
|
|
stream: () => {
|
|
throw new Error("unused");
|
|
},
|
|
streamSimple: () => {
|
|
throw new Error("unused");
|
|
},
|
|
fetchDeferred: (requestModel, _handle, options) => {
|
|
fetchedBaseUrl = requestModel.baseUrl;
|
|
fetchedOptions = options;
|
|
const message = {
|
|
role: "assistant" as const,
|
|
content: [],
|
|
api: requestModel.api,
|
|
provider: requestModel.provider,
|
|
model: requestModel.id,
|
|
usage: {
|
|
input: 0,
|
|
output: 0,
|
|
cacheRead: 0,
|
|
cacheWrite: 0,
|
|
totalTokens: 0,
|
|
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
|
|
},
|
|
stopReason: "stop" as const,
|
|
timestamp: 0,
|
|
};
|
|
const stream = createAssistantMessageEventStream();
|
|
stream.push({ type: "start", partial: message });
|
|
stream.push({ type: "done", reason: "stop", message });
|
|
stream.end(message);
|
|
return stream;
|
|
},
|
|
cancelDeferred: async (_requestModel, handle, options) => {
|
|
cancelledId = handle.id;
|
|
cancelledOptions = options;
|
|
},
|
|
};
|
|
|
|
runtime.registerNativeProvider(provider);
|
|
const composedModel = runtime.getModel(provider.id, nativeModel.id);
|
|
expect(composedModel).toBeDefined();
|
|
|
|
await runtime.fetchDeferred(
|
|
composedModel!,
|
|
{
|
|
provider: provider.id,
|
|
modelId: nativeModel.id,
|
|
api: nativeModel.api,
|
|
id: "fetch-id",
|
|
},
|
|
{
|
|
wait: 25,
|
|
headers: { "X-Fetch": "fetch" },
|
|
transformHeaders: (headers) => ({ ...headers, "X-Transformed": "fetch" }),
|
|
},
|
|
);
|
|
await runtime.cancelDeferred(
|
|
composedModel!,
|
|
{
|
|
provider: provider.id,
|
|
modelId: nativeModel.id,
|
|
api: nativeModel.api,
|
|
id: "cancel-id",
|
|
},
|
|
{
|
|
timeoutMs: 100,
|
|
transformHeaders: (headers) => ({ ...headers, "X-Transformed": "cancel" }),
|
|
},
|
|
);
|
|
|
|
expect(fetchedBaseUrl).toBe("https://overlay.test/v1");
|
|
expect(fetchedOptions).toMatchObject({
|
|
apiKey: "key",
|
|
wait: 25,
|
|
headers: { "X-Fetch": "fetch", "X-Transformed": "fetch" },
|
|
});
|
|
expect(cancelledId).toBe("cancel-id");
|
|
expect(cancelledOptions).toMatchObject({
|
|
apiKey: "key",
|
|
timeoutMs: 100,
|
|
headers: { "X-Transformed": "cancel" },
|
|
});
|
|
} finally {
|
|
rmSync(tempDir, { recursive: true, force: true });
|
|
}
|
|
});
|
|
|
|
it("applies models.json overrides above native providers", async () => {
|
|
const tempDir = mkdtempSync(join(tmpdir(), "pi-native-provider-"));
|
|
const modelsPath = join(tempDir, "models.json");
|
|
writeFileSync(
|
|
modelsPath,
|
|
JSON.stringify({
|
|
providers: {
|
|
"extension-native": {
|
|
modelOverrides: {
|
|
native: { contextWindow: 4242 },
|
|
},
|
|
},
|
|
},
|
|
}),
|
|
);
|
|
try {
|
|
const runtime = await ModelRuntime.create({
|
|
credentials: AuthStorage.inMemory(),
|
|
modelsStore: new InMemoryModelsStore(),
|
|
modelsPath,
|
|
allowModelNetwork: false,
|
|
});
|
|
const nativeModel = {
|
|
...model("native"),
|
|
provider: "extension-native",
|
|
baseUrl: "https://native.test/v1",
|
|
};
|
|
runtime.registerNativeProvider({
|
|
id: "extension-native",
|
|
name: "Extension Native",
|
|
auth: {
|
|
apiKey: {
|
|
name: "Native key",
|
|
resolve: async () => ({ auth: { apiKey: "key" }, source: "native" }),
|
|
},
|
|
},
|
|
getModels: () => [nativeModel],
|
|
stream: () => {
|
|
throw new Error("unused");
|
|
},
|
|
streamSimple: () => {
|
|
throw new Error("unused");
|
|
},
|
|
});
|
|
|
|
expect(runtime.getModel("extension-native", "native")?.contextWindow).toBe(4242);
|
|
} finally {
|
|
rmSync(tempDir, { recursive: true, force: true });
|
|
}
|
|
});
|
|
|
|
it("publishes refreshModels results without forcing ModelsStore persistence", async () => {
|
|
const modelsStore = new InMemoryModelsStore();
|
|
const runtime = await ModelRuntime.create({
|
|
credentials: AuthStorage.inMemory(),
|
|
modelsStore,
|
|
modelsPath: null,
|
|
allowModelNetwork: false,
|
|
});
|
|
runtime.registerProvider("extension-dynamic", {
|
|
baseUrl: "http://localhost:8080/v1",
|
|
apiKey: "local",
|
|
api: "openai-completions",
|
|
refreshModels: async () => [
|
|
{
|
|
...model("live"),
|
|
provider: "extension-dynamic",
|
|
baseUrl: "http://localhost:8080/v1",
|
|
},
|
|
],
|
|
});
|
|
|
|
await runtime.refresh({ allowNetwork: false });
|
|
expect(runtime.getModel("extension-dynamic", "live")).toBeDefined();
|
|
expect(await modelsStore.read("extension-dynamic")).toBeUndefined();
|
|
});
|
|
|
|
it("applies legacy OAuth modifyModels after async credential initialization", async () => {
|
|
const runtime = await ModelRuntime.create({
|
|
credentials: AuthStorage.inMemory({
|
|
"extension-oauth": {
|
|
type: "oauth",
|
|
access: "access",
|
|
refresh: "refresh",
|
|
expires: Date.now() + 60_000,
|
|
},
|
|
}),
|
|
modelsStore: new InMemoryModelsStore(),
|
|
modelsPath: null,
|
|
allowModelNetwork: false,
|
|
});
|
|
runtime.registerProvider("extension-oauth", {
|
|
baseUrl: "https://example.test/v1",
|
|
api: "openai-completions",
|
|
models: [model("base")],
|
|
oauth: {
|
|
name: "Extension OAuth",
|
|
login: async () => {
|
|
throw new Error("not used");
|
|
},
|
|
refreshToken: async (credential) => credential,
|
|
getApiKey: (credential) => credential.access,
|
|
modifyModels: (models, credential) =>
|
|
credential.access === "access" ? [...models, model("credential-model")] : models,
|
|
},
|
|
});
|
|
|
|
await runtime.refresh({ allowNetwork: false });
|
|
expect(runtime.getModel("extension-oauth", "base")).toBeDefined();
|
|
expect(runtime.getModel("extension-oauth", "credential-model")).toBeDefined();
|
|
|
|
await runtime.logout("extension-oauth");
|
|
expect(runtime.getModel("extension-oauth", "credential-model")).toBeUndefined();
|
|
});
|
|
});
|