Merge PR #6785: fix(coding-agent): preserve extension providers for plan subagents (@roboomp)

2026-07-27 15:57:45 +02:00
parent d3e55817fb 879707bd32
commit e4e8f8a40a
3 changed files with 122 additions and 4 deletions
@@ -20,6 +20,10 @@
 ### Fixed

 - Fixed `/usage` and the other large transcript command panels (`/session`, `/advisor status`, `/jobs`, `/changelog`, `/context`, `/memory view`) duplicating in native scrollback when invoked while an agent turn is streaming. These callsites mounted their finalized panel immediately via `present()` instead of deferring it until the turn ends via `presentCommandOutput()` (the path added in #5427 for `/tools`/`/mcp`), so the panel landed above a still-growing live block and was recommitted lower down ([#6767](https://github.com/can1357/oh-my-pi/issues/6767)).
+### Fixed
+
+- Fixed plan-mode task subagents unregistering extension-provided models, credentials, managers, and custom APIs from the shared parent `ModelRegistry` when restricted sessions intentionally skip extension loading ([#6783](https://github.com/can1357/oh-my-pi/issues/6783)).
+

 ## [17.1.5] - 2026-07-27

@@ -1987,10 +1987,12 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {}
 		// Process provider registrations queued during extension loading.
 		// This must happen before the runner is created so that models registered by
 		// extensions are available for model selection on session resume / fallback.
-		const activeExtensionSources = extensionsResult.extensions.map(extension => extension.path);
-		modelRegistry.syncExtensionSources(activeExtensionSources);
-		for (const sourceId of new Set(activeExtensionSources)) {
-			modelRegistry.clearSourceRegistrations(sourceId);
+		if (!restrictToolNames) {
+			const activeExtensionSources = extensionsResult.extensions.map(extension => extension.path);
+			modelRegistry.syncExtensionSources(activeExtensionSources);
+			for (const sourceId of new Set(activeExtensionSources)) {
+				modelRegistry.clearSourceRegistrations(sourceId);
+			}
 		}
 		if (extensionsResult.runtime.pendingProviderRegistrations.length > 0) {
 			for (const { name, config, sourceId } of extensionsResult.runtime.pendingProviderRegistrations) {
@@ -0,0 +1,112 @@
+import { afterEach, beforeEach, describe, expect, test } from "bun:test";
+import * as fs from "node:fs";
+import * as os from "node:os";
+import * as path from "node:path";
+import { createAssistantMessageEventStream, getCustomApi } from "@oh-my-pi/pi-ai";
+import { ModelRegistry } from "@oh-my-pi/pi-coding-agent/config/model-registry";
+import { Settings } from "@oh-my-pi/pi-coding-agent/config/settings";
+import {
+	type CreateAgentSessionOptions,
+	createAgentSession,
+	type ExtensionFactory,
+} from "@oh-my-pi/pi-coding-agent/sdk";
+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 { removeSyncWithRetries, Snowflake } from "@oh-my-pi/pi-utils";
+
+const providerName = "restricted-session-provider";
+const modelId = "restricted-session-model";
+const apiId = "restricted-session-api";
+const sourceId = "<inline-0>";
+
+describe("restricted sessions sharing extension providers", () => {
+	let tempDir: string;
+	let authStorage: AuthStorage;
+	let modelRegistry: ModelRegistry;
+	let settings: Settings;
+
+	beforeEach(async () => {
+		tempDir = path.join(os.tmpdir(), `pi-sdk-restricted-provider-${Snowflake.next()}`);
+		fs.mkdirSync(tempDir, { recursive: true });
+		authStorage = await AuthStorage.create(path.join(tempDir, "auth.db"));
+		modelRegistry = new ModelRegistry(authStorage, path.join(tempDir, "models.yml"));
+		settings = Settings.isolated();
+		settings.setModelRole("default", `${providerName}/${modelId}`);
+	});
+
+	afterEach(() => {
+		modelRegistry.clearSourceRegistrations(sourceId);
+		authStorage.close();
+		removeSyncWithRetries(tempDir);
+	});
+
+	const providerExtension: ExtensionFactory = pi => {
+		pi.registerProvider(providerName, {
+			baseUrl: "https://runtime.example.com/v1",
+			apiKey: "RUNTIME_KEY",
+			api: apiId,
+			streamSimple: () => createAssistantMessageEventStream(),
+			models: [
+				{
+					id: modelId,
+					name: "Restricted Session Model",
+					reasoning: false,
+					input: ["text"],
+					cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
+					contextWindow: 128000,
+					maxTokens: 8192,
+				},
+			],
+		});
+	};
+
+	function createOptions(): CreateAgentSessionOptions {
+		return {
+			cwd: tempDir,
+			agentDir: tempDir,
+			authStorage,
+			modelRegistry,
+			settings,
+			sessionManager: SessionManager.inMemory(),
+			disableExtensionDiscovery: true,
+			skills: [],
+			contextFiles: [],
+			promptTemplates: [],
+			slashCommands: [],
+			enableMCP: false,
+			enableLsp: false,
+			skipPythonPreflight: true,
+		};
+	}
+
+	test("does not unregister the parent's provider when extension loading is restricted", async () => {
+		const { session: parent } = await createAgentSession({
+			...createOptions(),
+			extensions: [providerExtension],
+		});
+
+		try {
+			expect(parent.model?.provider).toBe(providerName);
+			expect(modelRegistry.authStorage.hasAuth(providerName)).toBe(true);
+			expect(getCustomApi(apiId)).toBeDefined();
+
+			const { session: child } = await createAgentSession({
+				...createOptions(),
+				model: parent.model,
+				restrictToolNames: true,
+				toolNames: ["read"],
+			});
+
+			try {
+				expect(child.model?.provider).toBe(providerName);
+				expect(modelRegistry.find(providerName, modelId)).toBeDefined();
+				expect(modelRegistry.authStorage.hasAuth(providerName)).toBe(true);
+				expect(getCustomApi(apiId)).toBeDefined();
+			} finally {
+				await child.dispose();
+			}
+		} finally {
+			await parent.dispose();
+		}
+	});
+});