smashimo/oh-my-pi
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(extensions): restored provider unregistration

Browse Source 
Added the upstream unregisterProvider lifecycle to queued and initialized extension runtimes. Provider removal now clears runtime model/auth state before replacement, while failed factories restore the prior registration queue.

Fixes #7914
This commit is contained in:
roboomp
2026-08-07 15:13:19 +00:00
committed by can1357
parent 336975c46c
commit a09dfd0ba8
8 changed files with 249 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/ai/src/registry/oauth/index.ts
+7
View File
@@ -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.
*/
packages/coding-agent/CHANGELOG.md
+1
View File
@@ -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)).
packages/coding-agent/src/config/model-registry.ts
+21 -1
View File
@@ -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.
*/
packages/coding-agent/src/extensibility/extensions/loader.ts
+22 -5
View File
@@ -73,6 +73,15 @@ export class ExtensionRuntime implements IExtensionRuntime {
flagValues = new Map<string, boolean | string>();
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<void> {
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;
}
}
packages/coding-agent/src/extensibility/extensions/runner.ts
+6
View File
@@ -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;
packages/coding-agent/src/extensibility/extensions/types.ts
+12
View File
@@ -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<string, boolean | string>;
/** 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. */
packages/coding-agent/test/extension-provider-registration-rollback.test.ts
+142
View File
@@ -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();
packages/coding-agent/test/model-registry-runtime-cleanup.test.ts
+38
View File
@@ -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();
});
});