@gajae-code/coding-agent 0.15.5 → 0.16.0

package/src/sdk/host/model-pin.ts

@@ -17,9 +17,19 @@ export interface SdkHostModelResolveContext {
 }
 
 /** Registry loader owned by the SDK host, not by machine-facing adapters. */
-export type SdkHostModelRegistryLoader = (
+export type SdkHostModelRegistryLoader = ((
 	context?: SdkHostModelResolveContext,
-) => ModelRegistry | Promise<ModelRegistry>;
+) => ModelRegistry | Promise<ModelRegistry>) & {
+	acquire?: (context?: SdkHostModelResolveContext) => () => void;
+	dispose?: () => Promise<void>;
+};
+
+type OwnedModelRegistry = {
+	registry: ModelRegistry;
+	storage: AuthStorage;
+};
+
+const MAX_CACHED_MODEL_REGISTRIES = 32;
 
 /**
  * Builds an offline model resolver for a single SDK host process.
@@ -32,16 +42,112 @@ export function createSdkHostModelRegistryLoader(
 	modelsPath?: string,
 	loadSettings?: (context?: SdkHostModelResolveContext) => Promise<Pick<Settings, "get" | "getGlobal">>,
 ): SdkHostModelRegistryLoader {
-	const cachedRegistries = new Map<string, Promise<ModelRegistry>>();
-	return async (context?: SdkHostModelResolveContext) => {
-		const scopeKey = context?.cwd ?? "";
+	const cachedRegistries = new Map<string, Promise<OwnedModelRegistry>>();
+	const retiredRegistries = new Map<string, Promise<OwnedModelRegistry>[]>();
+	const activeScopes = new Map<string, number>();
+	const disposedEntries = new WeakSet<Promise<OwnedModelRegistry>>();
+	const pendingDisposals = new Set<Promise<void>>();
+	const disposalBarrier = Promise.withResolvers<void>();
+	const registryAgentDir = path.resolve(path.dirname(modelsPath ?? "."));
+	let disposed = false;
+	const scopeKeyFor = (context?: SdkHostModelResolveContext): string =>
+		context?.cwd === undefined ? "" : path.resolve(context.cwd);
+	const disposeEntry = async (entry: Promise<OwnedModelRegistry>): Promise<void> => {
+		let owned: OwnedModelRegistry;
+		try {
+			owned = await entry;
+		} catch {
+			// Initialization failures close their storage before rejecting. Disposal is
+			// best-effort and must not block broker shutdown on one broken scope.
+			return;
+		}
+		try {
+			owned.registry.dispose();
+		} finally {
+			owned.storage.close();
+		}
+	};
+	const maybeResolveDisposal = (): void => {
+		if (
+			disposed &&
+			cachedRegistries.size === 0 &&
+			retiredRegistries.size === 0 &&
+			activeScopes.size === 0 &&
+			pendingDisposals.size === 0
+		)
+			disposalBarrier.resolve();
+	};
+	const disposeEntries = (entries: Promise<OwnedModelRegistry>[]): void => {
+		const seen = new Set<Promise<OwnedModelRegistry>>();
+		const unique = entries.filter(entry => {
+			if (seen.has(entry) || disposedEntries.has(entry)) return false;
+			seen.add(entry);
+			return true;
+		});
+		if (unique.length === 0) {
+			maybeResolveDisposal();
+			return;
+		}
+		for (const entry of unique) disposedEntries.add(entry);
+		const pending = Promise.all(unique.map(entry => disposeEntry(entry))).then(
+			() => undefined,
+			() => undefined,
+		);
+		pendingDisposals.add(pending);
+		void pending.finally(() => {
+			pendingDisposals.delete(pending);
+			maybeResolveDisposal();
+		});
+	};
+	const retireEntry = (scopeKey: string, entry: Promise<OwnedModelRegistry>): void => {
+		if ((activeScopes.get(scopeKey) ?? 0) > 0) {
+			const retired = retiredRegistries.get(scopeKey) ?? [];
+			if (!retired.includes(entry)) retired.push(entry);
+			retiredRegistries.set(scopeKey, retired);
+			return;
+		}
+		disposeEntries([entry]);
+	};
+	const disposeRetired = (scopeKey: string): void => {
+		const retired = retiredRegistries.get(scopeKey);
+		if (!retired || (activeScopes.get(scopeKey) ?? 0) > 0) return;
+		retiredRegistries.delete(scopeKey);
+		disposeEntries(retired);
+	};
+	const evictIdleEntries = (): Promise<OwnedModelRegistry>[] => {
+		const evicted: Promise<OwnedModelRegistry>[] = [];
+		while (cachedRegistries.size > MAX_CACHED_MODEL_REGISTRIES) {
+			const oldest = [...cachedRegistries.keys()].find(key => (activeScopes.get(key) ?? 0) === 0);
+			if (oldest === undefined) break;
+			const entry = cachedRegistries.get(oldest);
+			cachedRegistries.delete(oldest);
+			if (entry) evicted.push(entry);
+		}
+		return evicted;
+	};
+	const loadRegistry = async (context?: SdkHostModelResolveContext): Promise<ModelRegistry> => {
+		if (disposed) throw new Error("SDK host model registry loader is disposed.");
+		const scopeKey = scopeKeyFor(context);
 		let cachedRegistry = cachedRegistries.get(scopeKey);
 		if (cachedRegistry === undefined) {
-			const initializing = Promise.all([discoverStorage(), loadSettings?.(context)]).then(
-				([storage, registrySettings]) => new ModelRegistry(storage, modelsPath, registrySettings),
-			);
+			const initializing = discoverStorage().then(async storage => {
+				try {
+					const registrySettings = await loadSettings?.(context);
+					return {
+						storage,
+						registry: new ModelRegistry(storage, modelsPath, registrySettings, {
+							agentDir: registryAgentDir,
+							automaticRefresh: false,
+						}),
+					};
+				} catch (error) {
+					storage.close();
+					throw error;
+				}
+			});
 			cachedRegistries.set(scopeKey, initializing);
 			cachedRegistry = initializing;
+			disposeEntries(evictIdleEntries());
 			try {
 				await initializing;
 			} catch (error) {
@@ -49,21 +155,67 @@ export function createSdkHostModelRegistryLoader(
 				throw error;
 			}
 		}
-		const registry = await cachedRegistry;
-		if (loadSettings) registry.setScopedSettings(await loadSettings(context));
+		const owned = await cachedRegistry;
+		if (disposed || disposedEntries.has(cachedRegistry) || cachedRegistries.get(scopeKey) !== cachedRegistry) {
+			retireEntry(scopeKey, cachedRegistry);
+			throw new Error("SDK host model registry loader is disposed.");
+		}
+		cachedRegistries.delete(scopeKey);
+		cachedRegistries.set(scopeKey, cachedRegistry);
+		const registry = owned.registry;
+		if (loadSettings) {
+			const registrySettings = await loadSettings(context);
+			if (disposed || disposedEntries.has(cachedRegistry)) {
+				retireEntry(scopeKey, cachedRegistry);
+				throw new Error("SDK host model registry loader is disposed.");
+			}
+			registry.setScopedSettings(registrySettings);
+		}
 		await registry.refresh("offline");
+		if (disposed || disposedEntries.has(cachedRegistry)) {
+			retireEntry(scopeKey, cachedRegistry);
+			throw new Error("SDK host model registry loader is disposed.");
+		}
 		return registry;
 	};
+	return Object.assign(loadRegistry, {
+		dispose: async (): Promise<void> => {
+			if (disposed) return await disposalBarrier.promise;
+			disposed = true;
+			for (const [scopeKey, entry] of cachedRegistries) retireEntry(scopeKey, entry);
+			cachedRegistries.clear();
+			for (const scopeKey of retiredRegistries.keys()) disposeRetired(scopeKey);
+			maybeResolveDisposal();
+			await disposalBarrier.promise;
+		},
+		acquire: (context?: SdkHostModelResolveContext): (() => void) => {
+			const scopeKey = scopeKeyFor(context);
+			activeScopes.set(scopeKey, (activeScopes.get(scopeKey) ?? 0) + 1);
+			let released = false;
+			return () => {
+				if (released) return;
+				released = true;
+				const next = (activeScopes.get(scopeKey) ?? 1) - 1;
+				if (next <= 0) activeScopes.delete(scopeKey);
+				else activeScopes.set(scopeKey, next);
+				disposeRetired(scopeKey);
+				disposeEntries(evictIdleEntries());
+				maybeResolveDisposal();
+			};
+		},
+	});
 }
 
 export type SdkHostModelResolution =
 	| { ok: true; model: string | null }
 	| { ok: false; reason: "unknown_model"; model: string; error: string };
 
-export type SdkHostModelResolver = (
+export type SdkHostModelResolver = ((
 	raw: unknown,
 	context?: SdkHostModelResolveContext,
-) => Promise<SdkHostModelResolution>;
+) => Promise<SdkHostModelResolution>) & {
+	dispose?: () => Promise<void>;
+};
 
 /** Resolve the explicit model pin at the SDK host boundary. */
 export async function resolveSdkHostModel(
@@ -74,19 +226,24 @@ export async function resolveSdkHostModel(
 	if (raw === undefined || raw === null) return { ok: true, model: null };
 	const requested = typeof raw === "string" ? raw : "";
 	const echoed = requested.trim().slice(0, MAX_ECHOED_MODEL_LENGTH);
-	const registry = await loadRegistry(context);
-	const resolved: ResolveCliModelResult = resolveCliModel({ cliModel: requested, modelRegistry: registry });
-	if (!resolved.model)
+	const release = loadRegistry.acquire?.(context) ?? (() => {});
+	try {
+		const registry = await loadRegistry(context);
+		const resolved: ResolveCliModelResult = resolveCliModel({ cliModel: requested, modelRegistry: registry });
+		if (!resolved.model)
+			return {
+				ok: false,
+				reason: "unknown_model",
+				model: echoed,
+				error: resolved.error ?? "No models available. Check your installation or add models to models.json.",
+			};
 		return {
-			ok: false,
-			reason: "unknown_model",
-			model: echoed,
-			error: resolved.error ?? "No models available. Check your installation or add models to models.json.",
+			ok: true,
+			model: formatModelSelectorValue(formatModelString(resolved.model), resolved.thinkingLevel),
 		};
-	return {
-		ok: true,
-		model: formatModelSelectorValue(formatModelString(resolved.model), resolved.thinkingLevel),
-	};
+	} finally {
+		release();
+	}
 }
 
 /** Default resolver used by the SDK broker host. */
@@ -96,6 +253,9 @@ export function createDefaultSdkHostModelResolver(agentDir: string): SdkHostMode
 		path.join(agentDir, "models.yml"),
 		context => Settings.loadReadonly({ agentDir, cwd: context?.cwd }),
 	);
-	return (raw: unknown, context?: SdkHostModelResolveContext): Promise<SdkHostModelResolution> =>
-		resolveSdkHostModel(raw, loadRegistry, context);
+	return Object.assign(
+		(raw: unknown, context?: SdkHostModelResolveContext): Promise<SdkHostModelResolution> =>
+			resolveSdkHostModel(raw, loadRegistry, context),
+		{ dispose: loadRegistry.dispose },
+	);
 }

package/src/sdk/host/sdk-run-capability.ts

@@ -0,0 +1 @@
+export { createSdkRunCapability } from "../../session/sdk-run-capability-internal";