✨ feat: support white list for discover assistant (#5216)

* ✨ feat: support white list for discover assistant

* improve error status

src/app/(backend)/webapi/assistant/store/route.ts

@@ -1,6 +1,5 @@
 import { NextResponse } from 'next/server';
 
-import { DEFAULT_LANG } from '@/const/locale';
 import { AssistantStore } from '@/server/modules/AssistantStore';
 
 export const runtime = 'edge';
@@ -11,18 +10,10 @@ export const GET = async (req: Request) => {
 
     const market = new AssistantStore();
 
-    let res: Response;
+    const data = await market.getAgentIndex(locale as any);
 
-    res = await fetch(market.getAgentIndexUrl(locale as any));
-
-    if (res.status === 404) {
-      res = await fetch(market.getAgentIndexUrl(DEFAULT_LANG));
-    }
-
-    const data = await res.json();
     return NextResponse.json(data);
-  } catch (e) {
-    console.error(e);
+  } catch {
     return new Response(`failed to fetch agent market index`, {
       headers: {
         'Access-Control-Allow-Origin': '*',

src/config/app.ts

@@ -47,6 +47,8 @@ export const getAppConfig = () => {
       PLUGIN_SETTINGS: z.string().optional(),
 
       APP_URL: z.string().optional(),
+      VERCEL_EDGE_CONFIG: z.string().optional(),
+
       CDN_USE_GLOBAL: z.boolean().optional(),
       CUSTOM_FONT_FAMILY: z.string().optional(),
       CUSTOM_FONT_URL: z.string().optional(),
@@ -75,6 +77,8 @@ export const getAppConfig = () => {
 
       PLUGIN_SETTINGS: process.env.PLUGIN_SETTINGS,
 
+      VERCEL_EDGE_CONFIG: process.env.VERCEL_EDGE_CONFIG,
+
       APP_URL,
       CUSTOM_FONT_FAMILY: process.env.CUSTOM_FONT_FAMILY,
       CUSTOM_FONT_URL: process.env.CUSTOM_FONT_URL,

src/server/modules/AssistantStore/index.test.ts

@@ -7,35 +7,35 @@ const baseURL = 'https://registry.npmmirror.com/@lobehub/agents-index/v1/files/p
 describe('AssistantStore', () => {
   it('should return the default index URL when no language is provided', () => {
     const agentMarket = new AssistantStore();
-    const url = agentMarket.getAgentIndexUrl();
+    const url = agentMarket['getAgentIndexUrl']();
     expect(url).toBe(`${baseURL}/index.en-US.json`);
   });
 
   it('should return the index URL for a not supported language', () => {
     const agentMarket = new AssistantStore();
-    const url = agentMarket.getAgentIndexUrl('xxx' as any);
+    const url = agentMarket['getAgentIndexUrl']('xxx' as any);
     expect(url).toBe('https://registry.npmmirror.com/@lobehub/agents-index/v1/files/public');
   });
 
   it('should return the zh-CN URL for zh locale', () => {
     const agentMarket = new AssistantStore();
-    const url = agentMarket.getAgentIndexUrl('zh' as any);
+    const url = agentMarket['getAgentIndexUrl']('zh' as any);
     expect(url).toBe(
       'https://registry.npmmirror.com/@lobehub/agents-index/v1/files/public/index.zh-CN.json',
     );
   });
 
   it('should return the default URL for en locale', () => {
     const agentMarket = new AssistantStore();
-    const url = agentMarket.getAgentIndexUrl('en' as any);
+    const url = agentMarket['getAgentIndexUrl']('en' as any);
     expect(url).toBe(
       'https://registry.npmmirror.com/@lobehub/agents-index/v1/files/public/index.en-US.json',
     );
   });
 
   it('should return the base URL if the provided language is not supported', () => {
     const agentMarket = new AssistantStore();
-    const url = agentMarket.getAgentIndexUrl('fr' as any);
+    const url = agentMarket['getAgentIndexUrl']('fr' as any);
     expect(url).toBe(baseURL);
   });
 

src/server/modules/AssistantStore/index.ts

@@ -3,6 +3,8 @@ import urlJoin from 'url-join';
 import { appEnv } from '@/config/app';
 import { DEFAULT_LANG, isLocaleNotSupport } from '@/const/locale';
 import { Locales, normalizeLocale } from '@/locales/resources';
+import { EdgeConfig } from '@/server/modules/EdgeConfig';
+import { AgentStoreIndex } from '@/types/discover';
 
 export class AssistantStore {
   private readonly baseUrl: string;
@@ -11,7 +13,7 @@ export class AssistantStore {
     this.baseUrl = baseUrl || appEnv.AGENTS_INDEX_URL;
   }
 
-  getAgentIndexUrl = (lang: Locales = DEFAULT_LANG) => {
+  private getAgentIndexUrl = (lang: Locales = DEFAULT_LANG) => {
     if (isLocaleNotSupport(lang)) return this.baseUrl;
 
     return urlJoin(this.baseUrl, `index.${normalizeLocale(lang)}.json`);
@@ -22,4 +24,40 @@ export class AssistantStore {
 
     return urlJoin(this.baseUrl, `${identifier}.${normalizeLocale(lang)}.json`);
   };
+
+  getAgentIndex = async (locale: Locales = DEFAULT_LANG, revalidate?: number) => {
+    try {
+      let res: Response;
+
+      res = await fetch(this.getAgentIndexUrl(locale as any), { next: { revalidate } });
+
+      if (res.status === 404) {
+        res = await fetch(this.getAgentIndexUrl(DEFAULT_LANG), { next: { revalidate } });
+      }
+
+      if (!res.ok) {
+        console.error('fetch agent index error:', await res.text());
+        return [];
+      }
+
+      const data: AgentStoreIndex = await res.json();
+
+      // Get the assistant whitelist from Edge Config
+      const edgeConfig = new EdgeConfig();
+
+      if (!!appEnv.VERCEL_EDGE_CONFIG) {
+        const assistantWhitelist = await edgeConfig.getAgentWhitelist();
+
+        if (assistantWhitelist && assistantWhitelist?.length > 0) {
+          data.agents = data.agents.filter((item) => assistantWhitelist.includes(item.identifier));
+        }
+      }
+
+      return data;
+    } catch (e) {
+      console.error('fetch agent index error:', e);
+
+      throw e;
+    }
+  };
 }

src/server/modules/EdgeConfig/index.ts

@@ -0,0 +1,23 @@
+import { EdgeConfigClient, createClient } from '@vercel/edge-config';
+
+import { appEnv } from '@/config/app';
+
+enum EdgeConfigKeys {
+  /**
+   * Assistant whitelist
+   */
+  AssistantWhitelist = 'assistant_whitelist',
+}
+
+export class EdgeConfig {
+  get client(): EdgeConfigClient {
+    if (!appEnv.VERCEL_EDGE_CONFIG) {
+      throw new Error('VERCEL_EDGE_CONFIG is not set');
+    }
+    return createClient(appEnv.VERCEL_EDGE_CONFIG);
+  }
+
+  getAgentWhitelist = async (): Promise<string[] | undefined> => {
+    return this.client.get<string[]>(EdgeConfigKeys.AssistantWhitelist);
+  };
+}

src/server/services/discover/index.ts

@@ -50,20 +50,9 @@ export class DiscoverService {
   };
 
   getAssistantList = async (locale: Locales): Promise<DiscoverAssistantItem[]> => {
-    let res = await fetch(this.assistantStore.getAgentIndexUrl(locale), {
-      next: { revalidate },
-    });
-
-    if (!res.ok) {
-      res = await fetch(this.assistantStore.getAgentIndexUrl(DEFAULT_LANG), {
-        next: { revalidate },
-      });
-    }
-
-    if (!res.ok) return [];
-
-    const json = await res.json();
+    const json = await this.assistantStore.getAgentIndex(locale, revalidate);
 
+    // @ts-expect-error 目前类型不一致，未来要统一
     return json.agents;
   };
 

src/types/discover.ts

@@ -154,3 +154,23 @@ export interface FilterBy {
   token?: number;
   vision?: boolean;
 }
+
+interface AgentIndexItem {
+  author: string;
+  createAt: string;
+  createdAt: string;
+  homepage: string;
+  identifier: string;
+  meta: {
+    avatar: string;
+    category: string;
+    description: string;
+    tags: string[];
+    title: string;
+  };
+}
+
+export interface AgentStoreIndex {
+  agents: AgentIndexItem[];
+  schemaVersion: number;
+}