From bd1c48e4d9b83751bd50f61b07a3e940d1a6bf47 Mon Sep 17 00:00:00 2001
From: Gustavo Madeira Santana <gumadeiras@gmail.com>
Date: Sat, 28 Mar 2026 19:09:27 -0400
Subject: [PATCH] Tests: lazy-load extension contract registries

---
 src/plugins/contracts/registry.ts             | 35 +++++++++++++++-
 src/plugins/contracts/suites.ts               |  7 ++--
 .../web-search-provider-contract.ts           | 40 ++++++++++++++-----
 3 files changed, 68 insertions(+), 14 deletions(-)

diff --git a/src/plugins/contracts/registry.ts b/src/plugins/contracts/registry.ts
index 3dde5fd4303..5f4cdbd8874 100644
--- a/src/plugins/contracts/registry.ts
+++ b/src/plugins/contracts/registry.ts
@@ -75,6 +75,7 @@ function uniqueStrings(values: readonly string[]): string[] {
 }
 
 let providerContractRegistryCache: ProviderContractEntry[] | null = null;
+let providerContractRegistryByPluginIdCache: Map<string, ProviderContractEntry[]> | null = null;
 let webSearchProviderContractRegistryCache: WebSearchProviderContractEntry[] | null = null;
 let speechProviderContractRegistryCache: SpeechProviderContractEntry[] | null = null;
 let mediaUnderstandingProviderContractRegistryCache:
@@ -89,8 +90,38 @@ export let providerContractLoadError: Error | undefined;
 function loadProviderContractEntriesForPluginIds(
   pluginIds: readonly string[],
 ): ProviderContractEntry[] {
-  const allowed = new Set(pluginIds);
-  return loadProviderContractRegistry().filter((entry) => allowed.has(entry.pluginId));
+  return pluginIds.flatMap((pluginId) => loadProviderContractEntriesForPluginId(pluginId));
+}
+
+function loadProviderContractEntriesForPluginId(pluginId: string): ProviderContractEntry[] {
+  if (providerContractRegistryCache) {
+    return providerContractRegistryCache.filter((entry) => entry.pluginId === pluginId);
+  }
+
+  const cache =
+    providerContractRegistryByPluginIdCache ?? new Map<string, ProviderContractEntry[]>();
+  providerContractRegistryByPluginIdCache = cache;
+  const cached = cache.get(pluginId);
+  if (cached) {
+    return cached;
+  }
+
+  try {
+    providerContractLoadError = undefined;
+    const entries = loadBundledCapabilityRuntimeRegistry({
+      pluginIds: [pluginId],
+      pluginSdkResolution: "dist",
+    }).providers.map((entry) => ({
+      pluginId: entry.pluginId,
+      provider: entry.provider,
+    }));
+    cache.set(pluginId, entries);
+    return entries;
+  } catch (error) {
+    providerContractLoadError = error instanceof Error ? error : new Error(String(error));
+    cache.set(pluginId, []);
+    return [];
+  }
 }
 
 function loadProviderContractRegistry(): ProviderContractEntry[] {
diff --git a/src/plugins/contracts/suites.ts b/src/plugins/contracts/suites.ts
index 4e0d93183d4..48b322af981 100644
--- a/src/plugins/contracts/suites.ts
+++ b/src/plugins/contracts/suites.ts
@@ -89,10 +89,11 @@ export function installProviderPluginContractSuite(params: { provider: Lazy<Prov
 
 export function installWebSearchProviderContractSuite(params: {
   provider: Lazy<WebSearchProviderPlugin>;
-  credentialValue: unknown;
+  credentialValue: Lazy<unknown>;
 }) {
   it("satisfies the base web search provider contract", () => {
     const provider = resolveLazy(params.provider);
+    const credentialValue = resolveLazy(params.credentialValue);
 
     expect(provider.id).toMatch(/^[a-z0-9][a-z0-9-]*$/);
     expect(provider.label.trim()).not.toBe("");
@@ -107,8 +108,8 @@ export function installWebSearchProviderContractSuite(params: {
     expect(provider.envVars.every((entry) => entry.trim().length > 0)).toBe(true);
 
     const searchConfigTarget: Record<string, unknown> = {};
-    provider.setCredentialValue(searchConfigTarget, params.credentialValue);
-    expect(provider.getCredentialValue(searchConfigTarget)).toEqual(params.credentialValue);
+    provider.setCredentialValue(searchConfigTarget, credentialValue);
+    expect(provider.getCredentialValue(searchConfigTarget)).toEqual(credentialValue);
 
     const config = {
       tools: {
diff --git a/test/helpers/extensions/web-search-provider-contract.ts b/test/helpers/extensions/web-search-provider-contract.ts
index 1e4b79ca9d7..1ed0db7e96f 100644
--- a/test/helpers/extensions/web-search-provider-contract.ts
+++ b/test/helpers/extensions/web-search-provider-contract.ts
@@ -1,23 +1,45 @@
 import { describe, expect, it } from "vitest";
-import { webSearchProviderContractRegistry } from "../../../src/plugins/contracts/registry.js";
+import {
+  pluginRegistrationContractRegistry,
+  webSearchProviderContractRegistry,
+} from "../../../src/plugins/contracts/registry.js";
 import { installWebSearchProviderContractSuite } from "../../../src/plugins/contracts/suites.js";
 
 export function describeWebSearchProviderContracts(pluginId: string) {
-  const providers = webSearchProviderContractRegistry.filter(
-    (entry) => entry.pluginId === pluginId,
-  );
+  const providerIds =
+    pluginRegistrationContractRegistry.find((entry) => entry.pluginId === pluginId)
+      ?.webSearchProviderIds ?? [];
+
+  const resolveProviders = () =>
+    webSearchProviderContractRegistry.filter((entry) => entry.pluginId === pluginId);
 
   describe(`${pluginId} web search provider contract registry load`, () => {
     it("loads bundled web search providers", () => {
-      expect(providers.length).toBeGreaterThan(0);
+      expect(resolveProviders().length).toBeGreaterThan(0);
     });
   });
 
-  for (const entry of providers) {
-    describe(`${pluginId}:${entry.provider.id} web search contract`, () => {
+  for (const providerId of providerIds) {
+    describe(`${pluginId}:${providerId} web search contract`, () => {
       installWebSearchProviderContractSuite({
-        provider: entry.provider,
-        credentialValue: entry.credentialValue,
+        provider: () => {
+          const entry = resolveProviders().find((provider) => provider.provider.id === providerId);
+          if (!entry) {
+            throw new Error(
+              `web search provider contract entry missing for ${pluginId}:${providerId}`,
+            );
+          }
+          return entry.provider;
+        },
+        credentialValue: () => {
+          const entry = resolveProviders().find((provider) => provider.provider.id === providerId);
+          if (!entry) {
+            throw new Error(
+              `web search provider contract entry missing for ${pluginId}:${providerId}`,
+            );
+          }
+          return entry.credentialValue;
+        },
       });
     });
   }