diff --git a/x-pack/platform/plugins/shared/observability_ai_assistant/server/routes/knowledge_base/route.ts b/x-pack/platform/plugins/shared/observability_ai_assistant/server/routes/knowledge_base/route.ts
index 6647e88d48d1f..9a6b3f89f7f69 100644
--- a/x-pack/platform/plugins/shared/observability_ai_assistant/server/routes/knowledge_base/route.ts
+++ b/x-pack/platform/plugins/shared/observability_ai_assistant/server/routes/knowledge_base/route.ts
@@ -93,11 +93,6 @@ const warmupModelKnowledgeBase = createObservabilityAIAssistantServerRoute({
 
 const reIndexKnowledgeBase = createObservabilityAIAssistantServerRoute({
   endpoint: 'POST /internal/observability_ai_assistant/kb/reindex',
-  params: t.type({
-    query: t.type({
-      inference_id: t.string,
-    }),
-  }),
   security: {
     authz: {
       requiredPrivileges: ['ai_assistant'],
@@ -105,8 +100,7 @@ const reIndexKnowledgeBase = createObservabilityAIAssistantServerRoute({
   },
   handler: async (resources): Promise<{ success: boolean }> => {
     const client = await resources.service.getClient({ request: resources.request });
-    const { inference_id: inferenceId } = resources.params.query;
-    await client.reIndexKnowledgeBaseWithLock(inferenceId);
+    await client.reIndexKnowledgeBaseWithLock();
     return { success: true };
   },
 });
diff --git a/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/client/index.ts b/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/client/index.ts
index 7c3412cc73968..7d65a2c1857f6 100644
--- a/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/client/index.ts
+++ b/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/client/index.ts
@@ -709,7 +709,6 @@ export class ObservabilityAIAssistantClient {
           core,
           logger,
           esClient,
-          inferenceId: nextInferenceId,
         });
         await populateMissingSemanticTextFieldWithLock({
           core,
@@ -742,12 +741,11 @@ export class ObservabilityAIAssistantClient {
     });
   };
 
-  reIndexKnowledgeBaseWithLock = (inferenceId: string) => {
+  reIndexKnowledgeBaseWithLock = () => {
     return reIndexKnowledgeBaseWithLock({
       core: this.dependencies.core,
       esClient: this.dependencies.esClient,
       logger: this.dependencies.logger,
-      inferenceId,
     });
   };
 
diff --git a/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/knowledge_base_service/create_knowledge_base_index.ts b/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/knowledge_base_service/create_knowledge_base_index.ts
index 796389bf20dfa..38c86d35c789c 100644
--- a/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/knowledge_base_service/create_knowledge_base_index.ts
+++ b/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/knowledge_base_service/create_knowledge_base_index.ts
@@ -12,28 +12,16 @@ import { Logger } from '@kbn/logging';
 export async function createKnowledgeBaseIndex({
   esClient,
   logger,
-  inferenceId,
   indexName,
 }: {
   esClient: { asInternalUser: ElasticsearchClient };
   logger: Logger;
-  inferenceId: string;
   indexName: string;
 }) {
   logger.debug(`Creating knowledge base write index "${indexName}"`);
 
   try {
-    await esClient.asInternalUser.indices.create({
-      index: indexName,
-      mappings: {
-        properties: {
-          semantic_text: {
-            type: 'semantic_text',
-            inference_id: inferenceId,
-          },
-        },
-      },
-    });
+    await esClient.asInternalUser.indices.create({ index: indexName });
   } catch (error) {
     if (
       error instanceof errors.ResponseError &&
@@ -43,6 +31,7 @@ export async function createKnowledgeBaseIndex({
         `Write index "${indexName}" already exists. Please delete it before creating a new index.`
       );
     }
+    logger.error(`Failed to create write index "${indexName}": ${error.message}`);
     throw error;
   }
 }
diff --git a/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/knowledge_base_service/index.ts b/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/knowledge_base_service/index.ts
index 857cbd09ec241..0552f925fb949 100644
--- a/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/knowledge_base_service/index.ts
+++ b/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/knowledge_base_service/index.ts
@@ -30,7 +30,6 @@ import { recallFromSearchConnectors } from './recall_from_search_connectors';
 import { ObservabilityAIAssistantPluginStartDependencies } from '../../types';
 import { ObservabilityAIAssistantConfig } from '../../config';
 import { hasKbWriteIndex } from './has_kb_index';
-import { getInferenceIdFromWriteIndex } from './get_inference_id_from_write_index';
 import { reIndexKnowledgeBaseWithLock } from './reindex_knowledge_base';
 import { isSemanticTextUnsupportedError } from '../startup_migrations/run_startup_migrations';
 
@@ -440,13 +439,10 @@ export class KnowledgeBaseService {
       }
 
       if (isSemanticTextUnsupportedError(error)) {
-        const inferenceId = await getInferenceIdFromWriteIndex(this.dependencies.esClient);
-
         reIndexKnowledgeBaseWithLock({
           core: this.dependencies.core,
           logger: this.dependencies.logger,
           esClient: this.dependencies.esClient,
-          inferenceId,
         }).catch((e) => {
           if (isLockAcquisitionError(e)) {
             this.dependencies.logger.info(`Re-indexing operation is already in progress`);
diff --git a/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/knowledge_base_service/reindex_knowledge_base.ts b/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/knowledge_base_service/reindex_knowledge_base.ts
index 931c45833cd8b..d9fca4130d735 100644
--- a/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/knowledge_base_service/reindex_knowledge_base.ts
+++ b/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/knowledge_base_service/reindex_knowledge_base.ts
@@ -21,29 +21,25 @@ export async function reIndexKnowledgeBaseWithLock({
   core,
   logger,
   esClient,
-  inferenceId,
 }: {
   core: CoreSetup<ObservabilityAIAssistantPluginStartDependencies>;
   logger: Logger;
   esClient: {
     asInternalUser: ElasticsearchClient;
   };
-  inferenceId: string;
 }): Promise<void> {
   const lmService = new LockManagerService(core, logger);
   return lmService.withLock(KB_REINDEXING_LOCK_ID, () =>
-    reIndexKnowledgeBase({ logger, esClient, inferenceId })
+    reIndexKnowledgeBase({ logger, esClient })
   );
 }
 
 async function reIndexKnowledgeBase({
   logger,
   esClient,
-  inferenceId,
 }: {
   logger: Logger;
   esClient: { asInternalUser: ElasticsearchClient };
-  inferenceId: string;
 }): Promise<void> {
   const activeReindexingTask = await getActiveReindexingTaskId(esClient);
   if (activeReindexingTask) {
@@ -57,7 +53,7 @@ async function reIndexKnowledgeBase({
     logger,
   });
 
-  await createKnowledgeBaseIndex({ esClient, logger, inferenceId, indexName: nextWriteIndexName });
+  await createKnowledgeBaseIndex({ esClient, logger, indexName: nextWriteIndexName });
 
   logger.info(
     `Re-indexing knowledge base from "${currentWriteIndexName}" to index "${nextWriteIndexName}"...`
diff --git a/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/knowledge_base_service/update_knowledge_base_index_alias.ts b/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/knowledge_base_service/update_knowledge_base_index_alias.ts
index cc9c665430932..3d893d6c2a6f3 100644
--- a/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/knowledge_base_service/update_knowledge_base_index_alias.ts
+++ b/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/knowledge_base_service/update_knowledge_base_index_alias.ts
@@ -24,10 +24,15 @@ export async function updateKnowledgeBaseWriteIndexAlias({
     `Updating write index alias from "${currentWriteIndexName}" to "${nextWriteIndexName}"`
   );
   const alias = resourceNames.writeIndexAlias.kb;
-  await esClient.asInternalUser.indices.updateAliases({
-    actions: [
-      { remove: { index: currentWriteIndexName, alias } },
-      { add: { index: nextWriteIndexName, alias, is_write_index: true } },
-    ],
-  });
+  try {
+    await esClient.asInternalUser.indices.updateAliases({
+      actions: [
+        { remove: { index: currentWriteIndexName, alias } },
+        { add: { index: nextWriteIndexName, alias, is_write_index: true } },
+      ],
+    });
+  } catch (error) {
+    logger.error(`Failed to update write index alias: ${error.message}`);
+    throw error;
+  }
 }
diff --git a/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/startup_migrations/run_startup_migrations.ts b/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/startup_migrations/run_startup_migrations.ts
index 34ccae06ca0a8..201ae87e4722d 100644
--- a/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/startup_migrations/run_startup_migrations.ts
+++ b/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/startup_migrations/run_startup_migrations.ts
@@ -16,7 +16,6 @@ import { ObservabilityAIAssistantConfig } from '../../config';
 import { reIndexKnowledgeBaseWithLock } from '../knowledge_base_service/reindex_knowledge_base';
 import { populateMissingSemanticTextFieldWithLock } from './populate_missing_semantic_text_fields';
 import { hasKbWriteIndex } from '../knowledge_base_service/has_kb_index';
-import { getInferenceIdFromWriteIndex } from '../knowledge_base_service/get_inference_id_from_write_index';
 import { updateExistingIndexAssets } from '../index_assets/update_existing_index_assets';
 
 const PLUGIN_STARTUP_LOCK_ID = 'observability_ai_assistant:startup_migrations';
@@ -56,8 +55,7 @@ export async function runStartupMigrations({
       });
 
       if (!isKbSemanticTextCompatible) {
-        const inferenceId = await getInferenceIdFromWriteIndex(esClient);
-        await reIndexKnowledgeBaseWithLock({ core, logger, esClient, inferenceId });
+        await reIndexKnowledgeBaseWithLock({ core, logger, esClient });
       }
 
       await pRetry(
diff --git a/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/complete/functions/alerts.spec.ts b/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/complete/functions/alerts.spec.ts
index 1f9a19b3a6137..043297bb6a150 100644
--- a/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/complete/functions/alerts.spec.ts
+++ b/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/complete/functions/alerts.spec.ts
@@ -47,7 +47,7 @@ export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderCon
   const observabilityAIAssistantAPIClient = getService('observabilityAIAssistantApi');
 
   describe('alerts', function () {
-    // Fails on MKI: https://github.com/elastic/kibana/issues/205581
+    // LLM Proxy is not yet support in MKI: https://github.com/elastic/obs-ai-assistant-team/issues/199
     this.tags(['skipCloud']);
     let proxy: LlmProxy;
     let connectorId: string;
diff --git a/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/complete/functions/context.spec.ts b/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/complete/functions/context.spec.ts
index 438a7dfe2c9d5..471c5455fe00d 100644
--- a/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/complete/functions/context.spec.ts
+++ b/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/complete/functions/context.spec.ts
@@ -73,7 +73,8 @@ export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderCon
   const log = getService('log');
 
   describe('context', function () {
-    this.tags(['failsOnMKI']);
+    // LLM Proxy is not yet support in MKI: https://github.com/elastic/obs-ai-assistant-team/issues/199
+    this.tags(['skipCloud']);
     let llmProxy: LlmProxy;
     let connectorId: string;
     let messageAddedEvents: MessageAddEvent[];
@@ -84,7 +85,7 @@ export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderCon
       connectorId = await observabilityAIAssistantAPIClient.createProxyActionConnector({
         port: llmProxy.getPort(),
       });
-      await restoreIndexAssets(observabilityAIAssistantAPIClient, es);
+      await restoreIndexAssets(getService);
       await deployTinyElserAndSetupKb(getService);
       await addSampleDocsToInternalKb(getService, sampleDocsForInternalKb);
 
diff --git a/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/complete/functions/elasticsearch.spec.ts b/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/complete/functions/elasticsearch.spec.ts
index 9298debf545fa..2a9cdde47e29c 100644
--- a/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/complete/functions/elasticsearch.spec.ts
+++ b/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/complete/functions/elasticsearch.spec.ts
@@ -26,7 +26,7 @@ export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderCon
   const observabilityAIAssistantAPIClient = getService('observabilityAIAssistantApi');
 
   describe('elasticsearch', function () {
-    // Fails on MKI: https://github.com/elastic/kibana/issues/205581
+    // LLM Proxy is not yet support in MKI: https://github.com/elastic/obs-ai-assistant-team/issues/199
     this.tags(['skipCloud']);
     let proxy: LlmProxy;
     let connectorId: string;
diff --git a/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/complete/functions/summarize.spec.ts b/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/complete/functions/summarize.spec.ts
index ccfb5a7749fc9..b5f20f060c2cd 100644
--- a/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/complete/functions/summarize.spec.ts
+++ b/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/complete/functions/summarize.spec.ts
@@ -25,7 +25,7 @@ export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderCon
   const observabilityAIAssistantAPIClient = getService('observabilityAIAssistantApi');
 
   describe('summarize', function () {
-    // Fails on MKI: https://github.com/elastic/kibana/issues/205581
+    // LLM Proxy is not yet support in MKI: https://github.com/elastic/obs-ai-assistant-team/issues/199
     this.tags(['skipCloud']);
     let proxy: LlmProxy;
     let connectorId: string;
diff --git a/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/distributed_lock_manager/distributed_lock_manager.spec.ts b/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/distributed_lock_manager/distributed_lock_manager.spec.ts
index ef6a9c26cb29b..fdf211f5d4efb 100644
--- a/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/distributed_lock_manager/distributed_lock_manager.spec.ts
+++ b/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/distributed_lock_manager/distributed_lock_manager.spec.ts
@@ -37,8 +37,8 @@ export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderCon
   const logger = getLoggerMock(log);
 
   describe('LockManager', function () {
-    // see details: https://github.com/elastic/kibana/issues/219091
-    this.tags(['failsOnMKI']);
+    // These tests should be moved to Jest Integration tests: https://github.com/elastic/kibana/issues/216690
+    this.tags(['skipCloud']);
     before(async () => {
       // delete existing index mappings to ensure we start from a clean state
       await deleteLockIndexAssets(es, log);
@@ -52,7 +52,6 @@ export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderCon
     });
 
     describe('Manual locking API', function () {
-      this.tags(['failsOnMKI']);
       before(async () => {
         await clearAllLocks(es, log);
       });
@@ -454,7 +453,6 @@ export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderCon
     });
 
     describe('withLock API', function () {
-      this.tags(['failsOnMKI']);
       before(async () => {
         await clearAllLocks(es, log);
       });
diff --git a/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/index_assets/index_assets.spec.ts b/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/index_assets/index_assets.spec.ts
index 1f3c6a3b0d955..09fc4761817bd 100644
--- a/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/index_assets/index_assets.spec.ts
+++ b/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/index_assets/index_assets.spec.ts
@@ -16,7 +16,7 @@ export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderCon
 
   describe('index assets: creating mappings, templates, aliases and write indices', () => {
     before(async () => {
-      await restoreIndexAssets(observabilityAIAssistantAPIClient, es);
+      await restoreIndexAssets(getService);
     });
 
     for (const componentTemplateName of Object.values(resourceNames.componentTemplate)) {
diff --git a/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/knowledge_base/knowledge_base_8.10_upgrade_test.spec.ts b/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/knowledge_base/knowledge_base_8.10_upgrade_test.spec.ts
index c833f84d302da..248f926cb8a3d 100644
--- a/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/knowledge_base/knowledge_base_8.10_upgrade_test.spec.ts
+++ b/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/knowledge_base/knowledge_base_8.10_upgrade_test.spec.ts
@@ -48,7 +48,7 @@ export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderCon
 
     after(async () => {
       log.info('Restoring index assets');
-      await restoreIndexAssets(observabilityAIAssistantAPIClient, es);
+      await restoreIndexAssets(getService);
 
       log.info('Tearing down tiny ELSER model and inference endpoint');
       await teardownTinyElserModelAndInferenceEndpoint(getService);
@@ -57,7 +57,7 @@ export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderCon
     describe('before running migrations', () => {
       before(async () => {
         log.info('Delete index assets');
-        await deleteIndexAssets(es);
+        await deleteIndexAssets(getService);
 
         log.info('Restoring snapshot');
         await restoreKbSnapshot({
@@ -99,7 +99,7 @@ export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderCon
 
     describe('after running migrations', () => {
       beforeEach(async () => {
-        await deleteIndexAssets(es);
+        await deleteIndexAssets(getService);
         await restoreKbSnapshot({
           log,
           es,
diff --git a/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/knowledge_base/knowledge_base_8.16_upgrade_test.spec.ts b/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/knowledge_base/knowledge_base_8.16_upgrade_test.spec.ts
index 55c9cc88406cd..ccc827383f420 100644
--- a/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/knowledge_base/knowledge_base_8.16_upgrade_test.spec.ts
+++ b/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/knowledge_base/knowledge_base_8.16_upgrade_test.spec.ts
@@ -41,7 +41,7 @@ export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderCon
 
     before(async () => {
       await teardownTinyElserModelAndInferenceEndpoint(getService);
-      await deleteIndexAssets(es);
+      await deleteIndexAssets(getService);
       await restoreKbSnapshot({
         log,
         es,
@@ -56,7 +56,7 @@ export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderCon
 
     after(async () => {
       await teardownTinyElserModelAndInferenceEndpoint(getService);
-      await restoreIndexAssets(observabilityAIAssistantAPIClient, es);
+      await restoreIndexAssets(getService);
     });
 
     describe('before migrating', () => {
diff --git a/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/knowledge_base/knowledge_base_8.18_upgrade_test.spec.ts b/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/knowledge_base/knowledge_base_8.18_upgrade_test.spec.ts
index 65ef4a2bb0bcf..c37baa03c7c85 100644
--- a/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/knowledge_base/knowledge_base_8.18_upgrade_test.spec.ts
+++ b/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/knowledge_base/knowledge_base_8.18_upgrade_test.spec.ts
@@ -29,7 +29,6 @@ export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderCon
   const es = getService('es');
   const retry = getService('retry');
   const log = getService('log');
-  const ml = getService('ml');
 
   // In 8.18 inference happens via the custom inference endpoint "obs_ai_assistant_kb_inference"
   // In 8.19 / 9.1 the custom inference endpoint ("obs_ai_assistant_kb_inference") is replaced with the preconfigured endpoint ".elser-2-elasticsearch"
@@ -39,12 +38,12 @@ export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderCon
     this.tags(['skipServerless']);
 
     before(async () => {
-      await importModel(ml, { modelId: TINY_ELSER_MODEL_ID });
+      await importModel(getService, { modelId: TINY_ELSER_MODEL_ID });
       await createTinyElserInferenceEndpoint(getService, {
         inferenceId: LEGACY_CUSTOM_INFERENCE_ID,
       });
 
-      await deleteIndexAssets(es);
+      await deleteIndexAssets(getService);
       await restoreKbSnapshot({
         log,
         es,
@@ -57,7 +56,7 @@ export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderCon
     });
 
     after(async () => {
-      await restoreIndexAssets(observabilityAIAssistantAPIClient, es);
+      await restoreIndexAssets(getService);
       await deleteModel(getService, { modelId: TINY_ELSER_MODEL_ID });
       await deleteInferenceEndpoint(getService, { inferenceId: LEGACY_CUSTOM_INFERENCE_ID });
     });
diff --git a/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/knowledge_base/knowledge_base_change_model_from_elser_to_e5.spec.ts b/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/knowledge_base/knowledge_base_change_model_from_elser_to_e5.spec.ts
index 74a4d94f25f69..48d141b46e8a5 100644
--- a/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/knowledge_base/knowledge_base_change_model_from_elser_to_e5.spec.ts
+++ b/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/knowledge_base/knowledge_base_change_model_from_elser_to_e5.spec.ts
@@ -54,7 +54,7 @@ export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderCon
     let e5WriteIndex: string;
 
     before(async () => {
-      await importModel(ml, { modelId: TINY_ELSER_MODEL_ID });
+      await importModel(getService, { modelId: TINY_ELSER_MODEL_ID });
       await createTinyElserInferenceEndpoint(getService, { inferenceId: TINY_ELSER_INFERENCE_ID });
       await setupKnowledgeBase(observabilityAIAssistantAPIClient, TINY_ELSER_INFERENCE_ID);
       await waitForKnowledgeBaseReady(getService);
@@ -71,7 +71,7 @@ export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderCon
       elserWriteIndex = await getConcreteWriteIndexFromAlias(es);
 
       // setup KB with E5-like model
-      await importModel(ml, { modelId: TINY_TEXT_EMBEDDING_MODEL_ID });
+      await importModel(getService, { modelId: TINY_TEXT_EMBEDDING_MODEL_ID });
       await ml.api.startTrainedModelDeploymentES(TINY_TEXT_EMBEDDING_MODEL_ID);
       await createTinyTextEmbeddingInferenceEndpoint(getService, {
         inferenceId: TINY_TEXT_EMBEDDING_INFERENCE_ID,
@@ -99,7 +99,7 @@ export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderCon
       await deleteModel(getService, { modelId: TINY_TEXT_EMBEDDING_MODEL_ID });
       await deleteInferenceEndpoint(getService, { inferenceId: TINY_TEXT_EMBEDDING_INFERENCE_ID });
 
-      await restoreIndexAssets(observabilityAIAssistantAPIClient, es);
+      await restoreIndexAssets(getService);
     });
 
     describe('when model is ELSER', () => {
diff --git a/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/knowledge_base/knowledge_base_setup.spec.ts b/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/knowledge_base/knowledge_base_setup.spec.ts
index 3500826d9ead0..7e0a5c9466386 100644
--- a/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/knowledge_base/knowledge_base_setup.spec.ts
+++ b/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/knowledge_base/knowledge_base_setup.spec.ts
@@ -29,18 +29,19 @@ import {
 export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderContext) {
   const es = getService('es');
   const retry = getService('retry');
-  const ml = getService('ml');
   const observabilityAIAssistantAPIClient = getService('observabilityAIAssistantApi');
 
   describe('Knowledge base: POST /internal/observability_ai_assistant/kb/setup', function () {
     before(async () => {
-      await teardownTinyElserModelAndInferenceEndpoint(getService);
-      await restoreIndexAssets(observabilityAIAssistantAPIClient, es);
+      await restoreIndexAssets(getService);
     });
 
     afterEach(async () => {
+      await restoreIndexAssets(getService);
+    });
+
+    after(async () => {
       await teardownTinyElserModelAndInferenceEndpoint(getService);
-      await restoreIndexAssets(observabilityAIAssistantAPIClient, es);
     });
 
     it('returns 200 when model is deployed', async () => {
@@ -65,6 +66,7 @@ export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderCon
 
         before(async () => {
           // setup KB initially
+
           await deployTinyElserAndSetupKb(getService);
 
           // setup KB with custom inference endpoint
@@ -113,8 +115,8 @@ export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderCon
       const customInferenceId = 'my_custom_inference_id';
 
       before(async () => {
-        await restoreIndexAssets(observabilityAIAssistantAPIClient, es);
-        await importModel(ml, { modelId: TINY_ELSER_MODEL_ID });
+        await restoreIndexAssets(getService);
+        await importModel(getService, { modelId: TINY_ELSER_MODEL_ID });
         await createTinyElserInferenceEndpoint(getService, {
           inferenceId: customInferenceId,
         });
diff --git a/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/knowledge_base/knowledge_base_user_instructions.spec.ts b/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/knowledge_base/knowledge_base_user_instructions.spec.ts
index badc0919922cd..9bb3e56f6a719 100644
--- a/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/knowledge_base/knowledge_base_user_instructions.spec.ts
+++ b/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/knowledge_base/knowledge_base_user_instructions.spec.ts
@@ -250,7 +250,7 @@ export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderCon
 
     describe('when a user instruction exist and a conversation is created', () => {
       // Fails on MKI because the LLM Proxy does not yet work there: https://github.com/elastic/obs-ai-assistant-team/issues/199
-      this.tags(['failsOnMKI']);
+      this.tags(['skipCloud']);
 
       let proxy: LlmProxy;
       let connectorId: string;
@@ -395,7 +395,7 @@ export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderCon
 
     describe('Forwarding User Instructions via System Message to the LLM', () => {
       // Fails on MKI because the LLM Proxy does not yet work there: https://github.com/elastic/obs-ai-assistant-team/issues/199
-      this.tags(['failsOnMKI']);
+      this.tags(['skipCloud']);
 
       let proxy: LlmProxy;
       let connectorId: string;
diff --git a/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/utils/index_assets.ts b/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/utils/index_assets.ts
index fed69ef98ab3a..4cd1d9e4b0f3f 100644
--- a/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/utils/index_assets.ts
+++ b/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/utils/index_assets.ts
@@ -11,6 +11,7 @@ import {
   getResourceName,
   resourceNames,
 } from '@kbn/observability-ai-assistant-plugin/server/service';
+import { DeploymentAgnosticFtrProviderContext } from '../../../../ftr_provider_context';
 import type { ObservabilityAIAssistantApiClient } from '../../../../services/observability_ai_assistant_api';
 import { TINY_ELSER_INFERENCE_ID } from './model_and_inference';
 
@@ -37,14 +38,25 @@ export async function createOrUpdateIndexAssets(
   expect(status).to.be(200);
 }
 
-export async function deleteIndexAssets(es: Client) {
+export async function deleteIndexAssets(
+  getService: DeploymentAgnosticFtrProviderContext['getService']
+) {
+  const es = getService('es');
+  const log = getService('log');
+
   // delete write indices
   const response = await es.indices.get({ index: getResourceName('*') });
   const indicesToDelete = Object.keys(response);
+
   if (indicesToDelete.length > 0) {
-    await es.indices.delete({ index: indicesToDelete, ignore_unavailable: true }).catch((err) => {
-      // ignore `IndexNotFoundException` error thrown by ES serverless: https://github.com/elastic/elasticsearch/blob/f1f745966f9c6b9d9fcad5242efb9a494d11e526/server/src/main/java/org/elasticsearch/cluster/metadata/Metadata.java#L2120-L2124
-    });
+    log.debug(`Deleting indices: ${indicesToDelete.join(', ')}`);
+    try {
+      await Promise.all(
+        indicesToDelete.map(async (index) => es.indices.delete({ index, ignore_unavailable: true }))
+      );
+    } catch (err) {
+      log.error(`Failed to delete indices: ${err}`);
+    }
   }
 
   await es.indices.deleteIndexTemplate({ name: getResourceName('*') }, { ignore: [404] });
@@ -52,10 +64,11 @@ export async function deleteIndexAssets(es: Client) {
 }
 
 export async function restoreIndexAssets(
-  observabilityAIAssistantAPIClient: ObservabilityAIAssistantApiClient,
-  es: Client
+  getService: DeploymentAgnosticFtrProviderContext['getService']
 ) {
-  await deleteIndexAssets(es);
+  const observabilityAIAssistantAPIClient = getService('observabilityAIAssistantApi');
+
+  await deleteIndexAssets(getService);
   await createOrUpdateIndexAssets(observabilityAIAssistantAPIClient);
 }
 
diff --git a/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/utils/knowledge_base.ts b/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/utils/knowledge_base.ts
index 5cf451471ca39..3e47f56e35464 100644
--- a/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/utils/knowledge_base.ts
+++ b/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/utils/knowledge_base.ts
@@ -181,11 +181,6 @@ export async function reIndexKnowledgeBase(
 ) {
   return observabilityAIAssistantAPIClient.admin({
     endpoint: 'POST /internal/observability_ai_assistant/kb/reindex',
-    params: {
-      query: {
-        inference_id: TINY_ELSER_INFERENCE_ID,
-      },
-    },
   });
 }
 
diff --git a/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/utils/model_and_inference.ts b/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/utils/model_and_inference.ts
index 9dcbc1998c6af..2a610e0e45296 100644
--- a/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/utils/model_and_inference.ts
+++ b/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/utils/model_and_inference.ts
@@ -5,11 +5,10 @@
  * 2.0.
  */
 
-import { Client } from '@elastic/elasticsearch';
+import { Client, errors } from '@elastic/elasticsearch';
 import { ToolingLog } from '@kbn/tooling-log';
 import { InferenceTaskType } from '@elastic/elasticsearch/lib/api/types';
 import { DeploymentAgnosticFtrProviderContext } from '../../../../ftr_provider_context';
-import { MachineLearningProvider } from '../../../../../services/ml';
 import { SUPPORTED_TRAINED_MODELS } from '../../../../../../functional/services/ml/api';
 import { setupKnowledgeBase, waitForKnowledgeBaseReady } from './knowledge_base';
 
@@ -24,24 +23,36 @@ export const TINY_ELSER_INFERENCE_ID = 'pt_tiny_elser_inference_id';
 export const TINY_TEXT_EMBEDDING_INFERENCE_ID = 'pt_tiny_text_embedding_inference_id';
 
 export async function importModel(
-  ml: ReturnType<typeof MachineLearningProvider>,
+  getService: DeploymentAgnosticFtrProviderContext['getService'],
   {
     modelId,
   }: {
     modelId: typeof TINY_ELSER_MODEL_ID | typeof TINY_TEXT_EMBEDDING_MODEL_ID;
   }
 ) {
+  const ml = getService('ml');
+  const log = getService('log');
+
   const config = ml.api.getTrainedModelConfig(modelId);
   await ml.api.assureMlStatsIndexExists();
-  await ml.api.importTrainedModel(modelId, modelId, config);
+
+  try {
+    await ml.api.importTrainedModel(modelId, modelId, config);
+  } catch (error) {
+    if (error.message.includes('resource_already_exists_exception')) {
+      log.info(`Model "${modelId}" is already imported. Skipping import.`);
+      return;
+    }
+
+    log.error(`Could not import model "${modelId}": ${error}`);
+    throw error;
+  }
 }
 
 export async function setupTinyElserModelAndInferenceEndpoint(
   getService: DeploymentAgnosticFtrProviderContext['getService']
 ) {
-  const ml = getService('ml');
-
-  await importModel(ml, { modelId: TINY_ELSER_MODEL_ID });
+  await importModel(getService, { modelId: TINY_ELSER_MODEL_ID });
   await createTinyElserInferenceEndpoint(getService, { inferenceId: TINY_ELSER_INFERENCE_ID });
 }
 
@@ -88,8 +99,11 @@ export async function deployTinyElserAndSetupKb(
   getService: DeploymentAgnosticFtrProviderContext['getService']
 ) {
   const observabilityAIAssistantAPIClient = getService('observabilityAIAssistantApi');
+  const log = getService('log');
 
   await setupTinyElserModelAndInferenceEndpoint(getService);
+
+  log.debug(`Setting up knowledge base with inference endpoint "${TINY_ELSER_INFERENCE_ID}"`);
   const { status, body } = await setupKnowledgeBase(
     observabilityAIAssistantAPIClient,
     TINY_ELSER_INFERENCE_ID
@@ -115,10 +129,12 @@ export async function deleteInferenceEndpoint(
     log.info(`Inference endpoint "${inferenceId}" deleted.`);
   } catch (e) {
     if (e.message.includes('resource_not_found_exception')) {
-      log.debug(`Inference endpoint "${inferenceId}" was already deleted.`);
-    } else {
-      log.error(`Could not delete inference endpoint "${inferenceId}": ${e}`);
+      log.debug(`Inference endpoint "${inferenceId}" was already deleted. Skipping deletion.`);
+      return;
     }
+
+    log.error(`Could not delete inference endpoint "${inferenceId}": ${e}`);
+    throw e;
   }
 }
 
@@ -152,9 +168,18 @@ export async function createInferenceEndpoint({
 
     log.info(`Inference endpoint ${inferenceId} created.`);
     return res;
-  } catch (e) {
-    log.error(`Error creating inference endpoint "${inferenceId}": ${e}`);
-    throw e;
+  } catch (error) {
+    if (
+      error instanceof errors.ResponseError &&
+      (error.body?.error?.type === 'resource_not_found_exception' ||
+        error.body?.error?.type === 'status_exception')
+    ) {
+      log.debug(`Inference endpoint "${inferenceId}" already exists. Skipping creation.`);
+      return;
+    }
+
+    log.error(`Error creating inference endpoint "${inferenceId}": ${error}`);
+    throw error;
   }
 }