feat(KB-196): load extension providers for dashboard model selector

- Load pi extensions on dashboard startup to register custom model providers
- Include extension-provided models in the dashboard model selector
- Add comprehensive tests for extension model loading and error handling
- Create changeset for patch release fixing model selector extensions
This commit is contained in:
gsxdsm
2026-03-31 03:21:16 -07:00
parent 2554904421
commit d34951a913
6 changed files with 159 additions and 12 deletions

View File

@@ -101,13 +101,24 @@ vi.mock("@kb/engine", async (importOriginal) => {
// ── Mock @mariozechner/pi-coding-agent ──────────────────────────────
const mockAuthStorage = { getAuth: vi.fn(), setAuth: vi.fn() };
const mockModelRegistry = { getModels: vi.fn().mockResolvedValue([]) };
const mockModelRegistry = {
getModels: vi.fn().mockResolvedValue([]),
registerProvider: vi.fn(),
refresh: vi.fn(),
};
const mockDiscoverAndLoadExtensions = vi.fn().mockResolvedValue({
runtime: { pendingProviderRegistrations: [] },
errors: [],
});
const mockCreateExtensionRuntime = vi.fn();
vi.mock("@mariozechner/pi-coding-agent", () => ({
AuthStorage: {
create: vi.fn(() => mockAuthStorage),
},
ModelRegistry: vi.fn().mockImplementation(() => mockModelRegistry),
discoverAndLoadExtensions: mockDiscoverAndLoadExtensions,
createExtensionRuntime: mockCreateExtensionRuntime,
}));
// ── Import module under test (after mocks) ──────────────────────────
@@ -119,6 +130,10 @@ const { runDashboard } = await import("../dashboard.js");
describe("runDashboard — AuthStorage & ModelRegistry wiring", () => {
beforeEach(async () => {
vi.clearAllMocks();
mockDiscoverAndLoadExtensions.mockResolvedValue({
runtime: { pendingProviderRegistrations: [] },
errors: [],
});
const { TaskStore } = await import("@kb/core");
(TaskStore as ReturnType<typeof vi.fn>).mockImplementation(() => makeMockStore());
});
@@ -150,4 +165,81 @@ describe("runDashboard — AuthStorage & ModelRegistry wiring", () => {
expect(ModelRegistry).toHaveBeenCalledTimes(1);
expect(ModelRegistry).toHaveBeenCalledWith(mockAuthStorage);
});
it("discovers extensions and registers extension providers", async () => {
mockDiscoverAndLoadExtensions.mockResolvedValueOnce({
runtime: {
pendingProviderRegistrations: [
{
name: "custom-anthropic",
config: { models: [{ id: "claude-sonnet-4-5" }] },
extensionPath: "/extensions/custom-anthropic",
},
],
},
errors: [],
});
await runDashboard(0, {});
expect(mockDiscoverAndLoadExtensions).toHaveBeenCalledWith([], expect.any(String), undefined);
expect(mockModelRegistry.registerProvider).toHaveBeenCalledWith(
"custom-anthropic",
expect.objectContaining({ models: [{ id: "claude-sonnet-4-5" }] }),
);
expect(mockModelRegistry.refresh).toHaveBeenCalled();
});
it("logs extension load errors without aborting startup", async () => {
const consoleSpy = vi.spyOn(console, "log").mockImplementation(() => {});
mockDiscoverAndLoadExtensions.mockResolvedValueOnce({
runtime: { pendingProviderRegistrations: [] },
errors: [{ path: "/extensions/bad", error: "Invalid manifest" }],
});
await runDashboard(0, {});
expect(consoleSpy).toHaveBeenCalledWith("[extensions] Failed to load /extensions/bad: Invalid manifest");
expect(consoleSpy).toHaveBeenCalledWith(expect.stringContaining("kb board"));
consoleSpy.mockRestore();
});
it("falls back gracefully when extension discovery throws", async () => {
const consoleSpy = vi.spyOn(console, "log").mockImplementation(() => {});
mockDiscoverAndLoadExtensions.mockRejectedValueOnce(new Error("boom"));
await runDashboard(0, {});
expect(mockCreateExtensionRuntime).toHaveBeenCalledTimes(1);
expect(mockModelRegistry.refresh).toHaveBeenCalled();
expect(consoleSpy).toHaveBeenCalledWith("[extensions] Failed to discover extensions: boom");
consoleSpy.mockRestore();
});
it("logs provider registration errors without aborting startup", async () => {
const consoleSpy = vi.spyOn(console, "log").mockImplementation(() => {});
mockDiscoverAndLoadExtensions.mockResolvedValueOnce({
runtime: {
pendingProviderRegistrations: [
{
name: "duplicate-provider",
config: { models: [{ id: "model-a" }] },
extensionPath: "/extensions/duplicate-provider",
},
],
},
errors: [],
});
mockModelRegistry.registerProvider.mockImplementationOnce(() => {
throw new Error("duplicate provider");
});
await runDashboard(0, {});
expect(consoleSpy).toHaveBeenCalledWith(
"[extensions] Failed to register provider from /extensions/duplicate-provider: duplicate provider",
);
expect(mockModelRegistry.refresh).toHaveBeenCalled();
consoleSpy.mockRestore();
});
});

View File

@@ -164,6 +164,28 @@ vi.mock("@kb/engine", async (importOriginal) => {
};
});
// ── Mock @mariozechner/pi-coding-agent ──────────────────────────────
const mockAuthStorage = { getAuth: vi.fn(), setAuth: vi.fn() };
const mockModelRegistry = {
registerProvider: vi.fn(),
refresh: vi.fn(),
};
const mockDiscoverAndLoadExtensions = vi.fn().mockResolvedValue({
runtime: { pendingProviderRegistrations: [] },
errors: [],
});
const mockCreateExtensionRuntime = vi.fn();
vi.mock("@mariozechner/pi-coding-agent", () => ({
AuthStorage: {
create: vi.fn(() => mockAuthStorage),
},
ModelRegistry: vi.fn().mockImplementation(() => mockModelRegistry),
discoverAndLoadExtensions: mockDiscoverAndLoadExtensions,
createExtensionRuntime: mockCreateExtensionRuntime,
}));
// ── Import module under test (after mocks) ──────────────────────────
const { runDashboard, processPullRequestMergeTask, getMergeStrategy, getTaskBranchName } = await import("./dashboard.js");

View File

@@ -5,7 +5,7 @@ import { TaskStore, AutomationStore } from "@kb/core";
import type { Settings, TaskDetail, PrInfo } from "@kb/core";
import { createServer, GitHubClient } from "@kb/dashboard";
import { TriageProcessor, TaskExecutor, Scheduler, AgentSemaphore, WorktreePool, aiMergeTask, UsageLimitPauser, PRIORITY_MERGE, scanIdleWorktrees, cleanupOrphanedWorktrees, NtfyNotifier, PrMonitor, PrCommentHandler, CronRunner, StuckTaskDetector } from "@kb/engine";
import { AuthStorage, ModelRegistry } from "@mariozechner/pi-coding-agent";
import { AuthStorage, ModelRegistry, discoverAndLoadExtensions, createExtensionRuntime } from "@mariozechner/pi-coding-agent";
/**
* Prompt the user for a port number interactively.
@@ -453,6 +453,31 @@ export async function runDashboard(port: number, opts: { paused?: boolean; dev?:
const authStorage = AuthStorage.create();
const modelRegistry = new ModelRegistry(authStorage);
try {
const extensionsResult = await discoverAndLoadExtensions([], cwd, undefined);
for (const { path, error } of extensionsResult.errors) {
console.log(`[extensions] Failed to load ${path}: ${error}`);
}
for (const { name, config, extensionPath } of extensionsResult.runtime.pendingProviderRegistrations) {
try {
modelRegistry.registerProvider(name, config);
} catch (error) {
const message = error instanceof Error ? error.message : String(error);
console.log(`[extensions] Failed to register provider from ${extensionPath}: ${message}`);
}
}
extensionsResult.runtime.pendingProviderRegistrations = [];
modelRegistry.refresh();
} catch (error) {
const message = error instanceof Error ? error.message : String(error);
console.log(`[extensions] Failed to discover extensions: ${message}`);
createExtensionRuntime();
modelRegistry.refresh();
}
// Start the web server with AI merge, auth, and model registry wired in
const app = createServer(store, { onMerge, authStorage, modelRegistry, automationStore });