From 55edc9c3e85827cfdd71fbd7e114b2f37d1179d9 Mon Sep 17 00:00:00 2001
From: Samiul Monir <samiul.monir@elastic.co>
Date: Tue, 5 Mar 2024 11:04:39 -0500
Subject: [PATCH] Adding restriciton on AI Playground

---
 .../enterprise_search_content/index.tsx       | 11 +++++---
 .../public/applications/shared/layout/nav.tsx | 26 +++++++++++--------
 .../routes/enterprise_search/ai_playground.ts |  7 +++--
 3 files changed, 28 insertions(+), 16 deletions(-)

diff --git a/x-pack/plugins/enterprise_search/public/applications/enterprise_search_content/index.tsx b/x-pack/plugins/enterprise_search/public/applications/enterprise_search_content/index.tsx
index d41cbb2ac5a10..150842b3ae5b5 100644
--- a/x-pack/plugins/enterprise_search/public/applications/enterprise_search_content/index.tsx
+++ b/x-pack/plugins/enterprise_search/public/applications/enterprise_search_content/index.tsx
@@ -75,6 +75,9 @@ export const EnterpriseSearchContent: React.FC<InitialAppData> = (props) => {
 };
 
 export const EnterpriseSearchContentConfigured: React.FC<Required<InitialAppData>> = () => {
+  const {
+    productFeatures: { showAIPlayground },
+  } = useValues(KibanaLogic);
   return (
     <Routes>
       <Redirect exact from={ROOT_PATH} to={SEARCH_INDICES_PATH} />
@@ -93,9 +96,11 @@ export const EnterpriseSearchContentConfigured: React.FC<Required<InitialAppData
       <Route path={SETTINGS_PATH}>
         <Settings />
       </Route>
-      <Route path={AI_PLAYGROUND_PATH}>
-        <AIPlayground />
-      </Route>
+      {showAIPlayground && (
+        <Route path={AI_PLAYGROUND_PATH}>
+          <AIPlayground />
+        </Route>
+      )}
       <Route>
         <NotFound />
       </Route>
diff --git a/x-pack/plugins/enterprise_search/public/applications/shared/layout/nav.tsx b/x-pack/plugins/enterprise_search/public/applications/shared/layout/nav.tsx
index ea73265bc3b1a..259f317b9b447 100644
--- a/x-pack/plugins/enterprise_search/public/applications/shared/layout/nav.tsx
+++ b/x-pack/plugins/enterprise_search/public/applications/shared/layout/nav.tsx
@@ -109,17 +109,21 @@ export const useEnterpriseSearchNav = () => {
               },
             ]
           : []),
-        {
-          id: 'ai-playground',
-          name: i18n.translate('xpack.enterpriseSearch.nav.aiPlaygroundTitle', {
-            defaultMessage: 'AI Playground',
-          }),
-          ...generateNavLink({
-            shouldNotCreateHref: true,
-            shouldShowActiveForSubroutes: true,
-            to: ENTERPRISE_SEARCH_CONTENT_PLUGIN.URL + AI_PLAYGROUND_PATH,
-          }),
-        },
+        ...(productFeatures.showAIPlayground
+          ? [
+              {
+                id: 'ai-playground',
+                name: i18n.translate('xpack.enterpriseSearch.nav.aiPlaygroundTitle', {
+                  defaultMessage: 'AI Playground',
+                }),
+                ...generateNavLink({
+                  shouldNotCreateHref: true,
+                  shouldShowActiveForSubroutes: true,
+                  to: ENTERPRISE_SEARCH_CONTENT_PLUGIN.URL + AI_PLAYGROUND_PATH,
+                }),
+              },
+            ]
+          : []),
       ],
       name: i18n.translate('xpack.enterpriseSearch.nav.contentTitle', {
         defaultMessage: 'Content',
diff --git a/x-pack/plugins/enterprise_search/server/routes/enterprise_search/ai_playground.ts b/x-pack/plugins/enterprise_search/server/routes/enterprise_search/ai_playground.ts
index cdd9c62962cae..00be72c265dfe 100644
--- a/x-pack/plugins/enterprise_search/server/routes/enterprise_search/ai_playground.ts
+++ b/x-pack/plugins/enterprise_search/server/routes/enterprise_search/ai_playground.ts
@@ -15,9 +15,12 @@ import { streamFactory } from '@kbn/ml-response-stream/server';
 
 import { RouteDependencies } from '../../plugin';
 import { elasticsearchErrorHandler } from '../../utils/elasticsearch_error_handler';
-import { createApiKey } from '@kbn/enterprise-search-plugin/server/lib/analytics/create_api_key';
 
-export function registerAIPlaygroundRoutes({ log, router }: RouteDependencies) {
+export function registerAIPlaygroundRoutes({ log, router, config }: RouteDependencies) {
+  if (!config.showAIPlayground) {
+    return;
+  }
+
   router.post(
     {
       path: '/internal/enterprise_search/ai_playground/query_source_fields',