diff --git a/packages/ai/src/registry/oauth/index.ts b/packages/ai/src/registry/oauth/index.ts index 4b47d01de..bd5961458 100644 --- a/packages/ai/src/registry/oauth/index.ts +++ b/packages/ai/src/registry/oauth/index.ts @@ -34,6 +34,13 @@ export function registerOAuthProvider(provider: OAuthProviderInterface): void { customOAuthProviders.set(provider.id, provider); } +/** + * Remove a custom OAuth provider by ID. + */ +export function unregisterOAuthProvider(id: string): void { + customOAuthProviders.delete(id); +} + /** * Get a custom OAuth provider by ID. */ diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index 68fd251e3..aa016442f 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -18,6 +18,7 @@ ### Fixed - Fixed `/reload-plugins` and the Agent Control Center leaving the running session's `task` tool on its create-time agent definitions; refreshed `.omp/agents/*.md` names and descriptions now reach existing tools without restarting ([#7940](https://github.com/can1357/oh-my-pi/issues/7940)). +- Fixed legacy Pi extensions failing to load when they call `pi.unregisterProvider()`, and made provider replacement calls take effect immediately after extension runtime initialization ([#7914](https://github.com/can1357/oh-my-pi/issues/7914)). - Fixed proxy discovery preferring the bundled catalog name over the proxy-reported name, so `omp models refresh` now updates stale display names (e.g. a proxy serving `longcat-2.0` as `"LongCat"` no longer shows the raw id). - Fixed the compiled binary build on Windows: `Bun.Glob.scan` yields backslash-separated paths, which the legacy Pi virtual module used verbatim for export keys and generated identifiers, producing invalid JavaScript. - Fixed Ctrl+O (`app.tools.expand`) not expanding truncated tool output while a tool-approval prompt or other selection dialog held keyboard focus, by promoting the shortcut to a global input listener that fires regardless of focus (it still defers to fullscreen overlays and the tree selector's own Ctrl+O filter cycle) ([#7837](https://github.com/can1357/oh-my-pi/issues/7837)). diff --git a/packages/coding-agent/src/config/model-registry.ts b/packages/coding-agent/src/config/model-registry.ts index 45a87627b..5e23a21ca 100644 --- a/packages/coding-agent/src/config/model-registry.ts +++ b/packages/coding-agent/src/config/model-registry.ts @@ -65,7 +65,7 @@ const BUILT_IN_DISCOVERY_CACHE_TTL_MS = 2 * 60 * 60 * 1000; const BUILT_IN_DISCOVERY_NON_AUTHORITATIVE_RETRY_MS = 5 * 60 * 1000; import type { ApiKeyResolver, FetchImpl } from "@oh-my-pi/pi-ai"; -import { registerOAuthProvider, unregisterOAuthProviders } from "@oh-my-pi/pi-ai/oauth"; +import { registerOAuthProvider, unregisterOAuthProvider, unregisterOAuthProviders } from "@oh-my-pi/pi-ai/oauth"; import type { OAuthCredentials, OAuthLoginCallbacks } from "@oh-my-pi/pi-ai/oauth/types"; import { setCodexAttestationProvider } from "@oh-my-pi/pi-ai/providers/openai-codex-responses"; import { getProviderDefinition } from "@oh-my-pi/pi-ai/registry"; @@ -2514,6 +2514,26 @@ export class ModelRegistry { this.#reloadStaticModels(); } + /** + * Remove one extension-registered provider and restore its static models. + */ + unregisterProvider(providerName: string): void { + const sourceId = this.#runtimeProviderSourceByName.get(providerName); + if (sourceId) { + const sourceProviders = this.#runtimeProvidersBySource.get(sourceId); + sourceProviders?.delete(providerName); + if (sourceProviders?.size === 0) { + this.#runtimeProvidersBySource.delete(sourceId); + } + this.#runtimeProviderSourceByName.delete(providerName); + } + unregisterOAuthProvider(providerName); + this.#ensureFullSnapshot(); + this.#clearRuntimeProviderState(providerName); + this.#lastStaticLoadMtime = null; + this.#reloadStaticModels(); + } + /** * Remove registrations for extension sources that are no longer active. */ diff --git a/packages/coding-agent/src/extensibility/extensions/loader.ts b/packages/coding-agent/src/extensibility/extensions/loader.ts index 237cbf7a9..4a641f91c 100644 --- a/packages/coding-agent/src/extensibility/extensions/loader.ts +++ b/packages/coding-agent/src/extensibility/extensions/loader.ts @@ -73,6 +73,15 @@ export class ExtensionRuntime implements IExtensionRuntime { flagValues = new Map(); pendingProviderRegistrations: Array<{ name: string; config: ProviderConfig; sourceId: string }> = []; + registerProvider(name: string, config: ProviderConfig, sourceId: string): void { + this.pendingProviderRegistrations.push({ name, config, sourceId }); + } + + unregisterProvider(name: string): void { + const remaining = this.pendingProviderRegistrations.filter(registration => registration.name !== name); + this.pendingProviderRegistrations.splice(0, this.pendingProviderRegistrations.length, ...remaining); + } + sendMessage(): void { throw new ExtensionRuntimeNotInitializedError(); } @@ -290,7 +299,11 @@ class ConcreteExtensionAPI implements ExtensionAPI, IExtensionRuntime { } registerProvider(name: string, config: ProviderConfig): void { - this.runtime.pendingProviderRegistrations.push({ name, config, sourceId: this.extension.path }); + this.runtime.registerProvider(name, config, this.extension.path); + } + + unregisterProvider(name: string): void { + this.runtime.unregisterProvider(name, this.extension.path); } } @@ -313,20 +326,24 @@ function createExtension(extensionPath: string, resolvedPath: string): Extension /** * Runs an extension factory with provider registration rollback on failure. - * Records the number of pending provider registrations before the factory runs, - * and restores that checkpoint if the factory throws. + * Restores the complete registration queue when the factory throws because an + * extension may unregister entries queued by an earlier extension. */ async function runExtensionFactory( factory: ExtensionFactory, api: ExtensionAPI, runtime: IExtensionRuntime, ): Promise { - const providerRegistrationCheckpoint = runtime.pendingProviderRegistrations.length; + const providerRegistrationCheckpoint = [...runtime.pendingProviderRegistrations]; try { await factory(api); } catch (error) { - runtime.pendingProviderRegistrations.length = providerRegistrationCheckpoint; + runtime.pendingProviderRegistrations.splice( + 0, + runtime.pendingProviderRegistrations.length, + ...providerRegistrationCheckpoint, + ); throw error; } } diff --git a/packages/coding-agent/src/extensibility/extensions/runner.ts b/packages/coding-agent/src/extensibility/extensions/runner.ts index af1f931c7..500f4bd50 100644 --- a/packages/coding-agent/src/extensibility/extensions/runner.ts +++ b/packages/coding-agent/src/extensibility/extensions/runner.ts @@ -530,6 +530,12 @@ export class ExtensionRunner { this.runtime.setServiceTier = actions.setServiceTier ?? throwUnsupportedServiceTierAction; this.runtime.getSessionName = actions.getSessionName; this.runtime.setSessionName = actions.setSessionName; + this.runtime.registerProvider = (name, config, sourceId) => { + this.modelRegistry.registerProvider(name, config, sourceId); + }; + this.runtime.unregisterProvider = name => { + this.modelRegistry.unregisterProvider(name); + }; // Context actions (required) this.#getModel = contextActions.getModel; diff --git a/packages/coding-agent/src/extensibility/extensions/types.ts b/packages/coding-agent/src/extensibility/extensions/types.ts index 7163dee83..6d46fedfb 100644 --- a/packages/coding-agent/src/extensibility/extensions/types.ts +++ b/packages/coding-agent/src/extensibility/extensions/types.ts @@ -1373,6 +1373,14 @@ export interface ExtensionAPI { */ registerProvider(name: string, config: ProviderConfig): void; + /** + * Unregister a provider previously registered by an extension. + * + * Removes extension-provided models and restores overridden built-in models. + * Has no effect when the provider is not registered. + */ + unregisterProvider(name: string): void; + /** Shared event bus for extension communication. */ events: EventBus; } @@ -1516,6 +1524,10 @@ export interface ExtensionRuntimeState { flagValues: Map; /** Provider registrations queued during extension loading, processed during session initialization */ pendingProviderRegistrations: Array<{ name: string; config: ProviderConfig; sourceId: string }>; + /** Queue a provider registration until initialization, then apply it immediately. */ + registerProvider(name: string, config: ProviderConfig, sourceId: string): void; + /** Remove a queued or initialized provider registration. */ + unregisterProvider(name: string, sourceId: string): void; } /** Action implementations for ExtensionAPI methods. */ diff --git a/packages/coding-agent/test/extension-provider-registration-rollback.test.ts b/packages/coding-agent/test/extension-provider-registration-rollback.test.ts index 56e1a5bfc..fc40c5a44 100644 --- a/packages/coding-agent/test/extension-provider-registration-rollback.test.ts +++ b/packages/coding-agent/test/extension-provider-registration-rollback.test.ts @@ -1,7 +1,13 @@ import { describe, expect, test } from "bun:test"; +import { unregisterOAuthProvider } from "@oh-my-pi/pi-ai/oauth"; +import { ModelRegistry } from "@oh-my-pi/pi-coding-agent/config/model-registry"; import { ExtensionRuntime, loadExtensionFromFactory } from "@oh-my-pi/pi-coding-agent/extensibility/extensions/loader"; +import { ExtensionRunner } from "@oh-my-pi/pi-coding-agent/extensibility/extensions/runner"; import type { ProviderConfig } from "@oh-my-pi/pi-coding-agent/extensibility/extensions/types"; +import { AuthStorage } from "@oh-my-pi/pi-coding-agent/session/auth-storage"; +import { SessionManager } from "@oh-my-pi/pi-coding-agent/session/session-manager"; import { EventBus } from "@oh-my-pi/pi-coding-agent/utils/event-bus"; +import { TempDir } from "@oh-my-pi/pi-utils"; const testProviderConfig: ProviderConfig = { baseUrl: "https://example.invalid/v1", @@ -41,6 +47,33 @@ describe("extension provider registration rollback", () => { expect(runtime.pendingProviderRegistrations).toEqual([]); }); + test("replaces a queued provider after unregistering it", async () => { + const runtime = new ExtensionRuntime(); + const events = new EventBus(); + + await loadExtensionFromFactory( + pi => { + pi.registerProvider("cliproxyapi", testProviderConfig); + pi.unregisterProvider("cliproxyapi"); + pi.registerProvider("cliproxyapi", { + baseUrl: "https://replacement.example.invalid/v1", + }); + }, + process.cwd(), + events, + runtime, + "pi-cliproxyapi-provider@1.4.13", + ); + + expect(runtime.pendingProviderRegistrations).toEqual([ + { + name: "cliproxyapi", + config: { baseUrl: "https://replacement.example.invalid/v1" }, + sourceId: "pi-cliproxyapi-provider@1.4.13", + }, + ]); + }); + test("preserves provider registrations from earlier successful extensions", async () => { const runtime = new ExtensionRuntime(); const events = new EventBus(); @@ -71,6 +104,36 @@ describe("extension provider registration rollback", () => { expect(runtime.pendingProviderRegistrations.map(r => r.name)).toEqual(["working-provider"]); }); + test("restores an earlier registration when unregistering extension fails", async () => { + const runtime = new ExtensionRuntime(); + const events = new EventBus(); + + await loadExtensionFromFactory( + pi => { + pi.registerProvider("working-provider", testProviderConfig); + }, + process.cwd(), + events, + runtime, + "working-extension", + ); + + await expect( + loadExtensionFromFactory( + pi => { + pi.unregisterProvider("working-provider"); + throw new Error("failed after unregistering"); + }, + process.cwd(), + events, + runtime, + "broken-extension", + ), + ).rejects.toThrow("failed after unregistering"); + + expect(runtime.pendingProviderRegistrations.map(registration => registration.name)).toEqual(["working-provider"]); + }); + test("keeps provider registrations when extension initialization succeeds", async () => { const runtime = new ExtensionRuntime(); const events = new EventBus(); @@ -93,6 +156,85 @@ describe("extension provider registration rollback", () => { expect(runtime.pendingProviderRegistrations.map(r => r.name)).toEqual(["provider-one", "provider-two"]); }); + test("applies provider replacement after runtime initialization", async () => { + const tempDir = TempDir.createSync("@provider-replacement-"); + const authStorage = await AuthStorage.create(tempDir.join("auth.db")); + try { + const modelRegistry = new ModelRegistry(authStorage, tempDir.join("models.json")); + modelRegistry.registerProvider("cliproxyapi", testProviderConfig, "pi-cliproxyapi-provider"); + + const runtime = new ExtensionRuntime(); + const events = new EventBus(); + let replaceProvider: (() => void) | undefined; + const extension = await loadExtensionFromFactory( + pi => { + replaceProvider = () => { + pi.unregisterProvider("cliproxyapi"); + pi.registerProvider("cliproxyapi", { + baseUrl: "https://replacement.example.invalid/v1", + api: "openai-completions", + models: testProviderConfig.models, + oauth: { + name: "CLIProxyAPI", + login: async () => "test-token", + }, + }); + }; + }, + process.cwd(), + events, + runtime, + "pi-cliproxyapi-provider", + ); + const runner = new ExtensionRunner( + [extension], + runtime, + process.cwd(), + SessionManager.inMemory(), + modelRegistry, + ); + runner.initialize( + { + sendMessage: () => {}, + sendUserMessage: () => {}, + appendEntry: () => {}, + setLabel: () => {}, + getActiveTools: () => [], + getAllTools: () => [], + setActiveTools: async () => {}, + getCommands: () => [], + setModel: async () => false, + getThinkingLevel: () => undefined, + setThinkingLevel: () => {}, + getSessionName: () => undefined, + setSessionName: async () => {}, + }, + { + getModel: () => undefined, + isIdle: () => true, + abort: () => {}, + hasPendingMessages: () => false, + shutdown: () => {}, + getContextUsage: () => undefined, + compact: async () => {}, + getSystemPrompt: () => [], + }, + ); + + if (!replaceProvider) throw new Error("Extension did not expose its provider replacement action"); + replaceProvider(); + + expect(modelRegistry.authStorage.hasAuth("cliproxyapi")).toBe(false); + expect(modelRegistry.find("cliproxyapi", "test-model")?.baseUrl).toBe( + "https://replacement.example.invalid/v1", + ); + } finally { + unregisterOAuthProvider("cliproxyapi"); + authStorage.close(); + tempDir.removeSync(); + } + }); + test("rolls back every provider added by the failed extension", async () => { const runtime = new ExtensionRuntime(); const events = new EventBus(); diff --git a/packages/coding-agent/test/model-registry-runtime-cleanup.test.ts b/packages/coding-agent/test/model-registry-runtime-cleanup.test.ts index 26d0c2757..e7b06a069 100644 --- a/packages/coding-agent/test/model-registry-runtime-cleanup.test.ts +++ b/packages/coding-agent/test/model-registry-runtime-cleanup.test.ts @@ -3,6 +3,7 @@ import * as fs from "node:fs"; import * as os from "node:os"; import * as path from "node:path"; import { type AssistantMessageEventStream, clearCustomApis, getCustomApi } from "@oh-my-pi/pi-ai"; +import { getOAuthProvider } from "@oh-my-pi/pi-ai/oauth"; import { ModelRegistry, type ProviderConfigInput } from "@oh-my-pi/pi-coding-agent/config/model-registry"; import { AuthStorage } from "@oh-my-pi/pi-coding-agent/session/auth-storage"; import { removeSyncWithRetries, Snowflake } from "@oh-my-pi/pi-utils"; @@ -63,4 +64,41 @@ describe("ModelRegistry runtime source cleanup", () => { expect(registry.authStorage.hasAuth("runtime-provider")).toBe(false); expect(getCustomApi("custom-runtime-cleanup-api")).toBeUndefined(); }); + + test("unregisterProvider removes only the named provider and its login entry", () => { + const registry = new ModelRegistry(authStorage, modelsJsonPath); + registry.registerProvider( + "runtime-provider", + { + baseUrl: "https://runtime.example.com/v1", + apiKey: "RUNTIME_KEY", + api: "custom-runtime-cleanup-api", + streamSimple, + models: [baseModel], + oauth: { + name: "Runtime Provider", + login: async () => "runtime-token", + }, + }, + sourceId, + ); + registry.registerProvider( + "peer-provider", + { + baseUrl: "https://peer.example.com/v1", + apiKey: "PEER_KEY", + api: "openai-completions", + models: [{ ...baseModel, id: "peer-model" }], + }, + sourceId, + ); + + expect(getOAuthProvider("runtime-provider")).toBeDefined(); + registry.unregisterProvider("runtime-provider"); + + expect(registry.find("runtime-provider", "runtime-model")).toBeUndefined(); + expect(registry.authStorage.hasAuth("runtime-provider")).toBe(false); + expect(getOAuthProvider("runtime-provider")).toBeUndefined(); + expect(registry.find("peer-provider", "peer-model")).toBeDefined(); + }); });