[ML] Propagate auth token through EIS logic

muted-tests.yml

@@ -414,15 +414,6 @@ tests:
 - class: org.elasticsearch.indices.mapping.UpdateMappingIntegrationIT
   method: testUpdateMappingConcurrently
   issue: https://github.com/elastic/elasticsearch/issues/137758
-- class: org.elasticsearch.xpack.inference.integration.CCMPersistentStorageServiceIT
-  method: testDelete_RemovesCCMConfiguration
-  issue: https://github.com/elastic/elasticsearch/issues/137786
-- class: org.elasticsearch.xpack.inference.integration.CCMPersistentStorageServiceIT
-  method: testDelete_DoesNotThrow_WhenTheConfigurationDoesNotExist
-  issue: https://github.com/elastic/elasticsearch/issues/137797
-- class: org.elasticsearch.xpack.inference.integration.CCMServiceIT
-  method: testIsEnabled_ReturnsTrue_WhenCCMConfigurationIsPresent
-  issue: https://github.com/elastic/elasticsearch/issues/137798
 - class: org.elasticsearch.xpack.inference.external.http.sender.RequestExecutorServiceTests
   method: testChangingCapacity_DoesNotRejectsOverflowTasks_BecauseOfQueueFull
   issue: https://github.com/elastic/elasticsearch/issues/137823

x-pack/plugin/inference/qa/inference-service-tests/src/javaRestTest/java/org/elasticsearch/xpack/inference/BaseMockEISAuthServerTest.java

@@ -25,6 +25,7 @@
 import org.junit.rules.TestRule;
 
 import static org.elasticsearch.xpack.inference.InferenceBaseRestTest.getModel;
+import static org.elasticsearch.xpack.inference.services.elastic.ccm.CCMSettings.CCM_SUPPORTED_ENVIRONMENT;
 
 public class BaseMockEISAuthServerTest extends ESRestTestCase {
 
@@ -46,6 +47,9 @@ public class BaseMockEISAuthServerTest extends ESRestTestCase {
         // calls which would result in a test failure because the webserver is only expecting a single request
         // So to ensure we avoid that all together, this flag indicates that we'll only perform a single authorization request
         .setting("xpack.inference.elastic.periodic_authorization_enabled", "false")
+        // Setting to false so that the CCM logic will be skipped when running the tests, the authorization logic skip trying to determine
+        // if CCM is enabled
+        .setting(CCM_SUPPORTED_ENVIRONMENT.getKey(), "false")
         // This plugin is located in the inference/qa/test-service-plugin package, look for TestInferenceServicePlugin
         .plugin("inference-service-test")
         .user("x_pack_rest_user", "x-pack-test-password")

x-pack/plugin/inference/qa/inference-service-tests/src/javaRestTest/java/org/elasticsearch/xpack/inference/CCMCrudForbiddenIT.java

@@ -21,6 +21,7 @@
 import org.junit.ClassRule;
 
 import static org.elasticsearch.xpack.inference.services.elastic.ccm.CCMFeature.CCM_FORBIDDEN_EXCEPTION;
+import static org.elasticsearch.xpack.inference.services.elastic.ccm.CCMSettings.CCM_SUPPORTED_ENVIRONMENT;
 import static org.hamcrest.Matchers.containsString;
 import static org.hamcrest.Matchers.is;
 
@@ -31,7 +32,7 @@ public class CCMCrudForbiddenIT extends CCMRestBaseIT {
         .distribution(DistributionType.DEFAULT)
         .setting("xpack.license.self_generated.type", "basic")
         .setting("xpack.security.enabled", "true")
-        .setting("xpack.inference.elastic.allow_configuring_ccm", "false")
+        .setting(CCM_SUPPORTED_ENVIRONMENT.getKey(), "false")
         .user("x_pack_rest_user", "x-pack-test-password")
         .build();
 

x-pack/plugin/inference/qa/inference-service-tests/src/javaRestTest/java/org/elasticsearch/xpack/inference/CCMCrudIT.java

@@ -26,6 +26,7 @@
 import java.io.IOException;
 
 import static org.elasticsearch.xpack.inference.rest.Paths.INFERENCE_CCM_PATH;
+import static org.elasticsearch.xpack.inference.services.elastic.ccm.CCMSettings.CCM_SUPPORTED_ENVIRONMENT;
 import static org.hamcrest.Matchers.containsString;
 import static org.hamcrest.Matchers.is;
 
@@ -36,7 +37,7 @@ public class CCMCrudIT extends CCMRestBaseIT {
         .distribution(DistributionType.DEFAULT)
         .setting("xpack.license.self_generated.type", "basic")
         .setting("xpack.security.enabled", "true")
-        .setting("xpack.inference.elastic.allow_configuring_ccm", "true")
+        .setting(CCM_SUPPORTED_ENVIRONMENT.getKey(), "true")
         .user("x_pack_rest_user", "x-pack-test-password")
         .build();
 

x-pack/plugin/inference/src/internalClusterTest/java/org/elasticsearch/xpack/inference/integration/AuthorizationTaskExecutorIT.java

@@ -28,6 +28,7 @@
 import org.elasticsearch.xpack.inference.services.elastic.InternalPreconfiguredEndpoints;
 import org.elasticsearch.xpack.inference.services.elastic.authorization.AuthorizationPoller;
 import org.elasticsearch.xpack.inference.services.elastic.authorization.AuthorizationTaskExecutor;
+import org.elasticsearch.xpack.inference.services.elastic.ccm.CCMSettings;
 import org.junit.After;
 import org.junit.AfterClass;
 import org.junit.Before;
@@ -87,6 +88,10 @@ public void createComponents() {
 
     @After
     public void shutdown() {
+        removeEisPreconfiguredEndpoints(modelRegistry);
+    }
+
+    static void removeEisPreconfiguredEndpoints(ModelRegistry modelRegistry) {
         // Delete all the eis preconfigured endpoints
         var listener = new PlainActionFuture<Boolean>();
         modelRegistry.deleteModels(InternalPreconfiguredEndpoints.EIS_PRECONFIGURED_ENDPOINT_IDS, listener);
@@ -101,6 +106,8 @@ public static void cleanUpClass() {
     @Override
     protected Settings nodeSettings() {
         return Settings.builder()
+            // Disable CCM to ensure that only the authorization task executor is initialized in the inference plugin when it is created
+            .put(CCMSettings.CCM_SUPPORTED_ENVIRONMENT.getKey(), false)
             .put(ElasticInferenceServiceSettings.ELASTIC_INFERENCE_SERVICE_URL.getKey(), gatewayUrl)
             // Ensure that the polling logic only occurs once so we can deterministically control when an authorization response is
             // received
@@ -123,15 +130,23 @@ public void testCreatesEisChatCompletionEndpoint() throws Exception {
     }
 
     private void assertNoAuthorizedEisEndpoints() throws Exception {
-        waitForTask(AUTH_TASK_ACTION, admin());
+        assertNoAuthorizedEisEndpoints(admin(), authorizationTaskExecutor, modelRegistry);
+    }
+
+    static void assertNoAuthorizedEisEndpoints(
+        AdminClient adminClient,
+        AuthorizationTaskExecutor authorizationTaskExecutor,
+        ModelRegistry modelRegistry
+    ) throws Exception {
+        waitForTask(AUTH_TASK_ACTION, adminClient);
 
         assertBusy(() -> {
             var newPoller = authorizationTaskExecutor.getCurrentPollerTask();
             assertNotNull(newPoller);
             newPoller.waitForAuthorizationToComplete(TimeValue.THIRTY_SECONDS);
         });
 
-        var eisEndpoints = getEisEndpoints();
+        var eisEndpoints = getEisEndpoints(modelRegistry);
         assertThat(eisEndpoints, empty());
 
         for (String eisPreconfiguredEndpoints : InternalPreconfiguredEndpoints.EIS_PRECONFIGURED_ENDPOINT_IDS) {
@@ -153,7 +168,22 @@ public static TaskInfo waitForTask(String taskAction, AdminClient adminClient) t
         return taskRef.get();
     }
 
+    static void waitForNoTask(String taskAction, AdminClient adminClient) throws Exception {
+        var builder = new ListTasksRequestBuilder(adminClient.cluster());
+
+        assertBusy(() -> {
+            var response = builder.get();
+            var authPollerTask = response.getTasks().stream().filter(task -> task.action().equals(taskAction)).findFirst();
+            assertFalse(authPollerTask.isPresent());
+        });
+
+    }
+
     private List<UnparsedModel> getEisEndpoints() {
+        return getEisEndpoints(modelRegistry);
+    }
+
+    static List<UnparsedModel> getEisEndpoints(ModelRegistry modelRegistry) {
         var listener = new PlainActionFuture<List<UnparsedModel>>();
         modelRegistry.getAllModels(false, listener);
 
@@ -162,17 +192,26 @@ private List<UnparsedModel> getEisEndpoints() {
     }
 
     private void restartPollingTaskAndWaitForAuthResponse() throws Exception {
-        cancelAuthorizationTask(admin());
+        restartPollingTaskAndWaitForAuthResponse(admin(), authorizationTaskExecutor);
+    }
+
+    static void restartPollingTaskAndWaitForAuthResponse(AdminClient adminClient, AuthorizationTaskExecutor authTaskExecutor)
+        throws Exception {
+        cancelAuthorizationTask(adminClient);
 
         // wait for the new task to be recreated and an authorization response to be processed
+        waitForAuthorizationToComplete(authTaskExecutor);
+    }
+
+    static void waitForAuthorizationToComplete(AuthorizationTaskExecutor authTaskExecutor) throws Exception {
         assertBusy(() -> {
-            var newPoller = authorizationTaskExecutor.getCurrentPollerTask();
+            var newPoller = authTaskExecutor.getCurrentPollerTask();
             assertNotNull(newPoller);
             newPoller.waitForAuthorizationToComplete(TimeValue.THIRTY_SECONDS);
         });
     }
 
-    public static void cancelAuthorizationTask(AdminClient adminClient) throws Exception {
+    static void cancelAuthorizationTask(AdminClient adminClient) throws Exception {
         var pollerTask = waitForTask(AUTH_TASK_ACTION, adminClient);
         var builder = new CancelTasksRequestBuilder(adminClient.cluster());
 
@@ -202,7 +241,11 @@ public void testCreatesEisChatCompletion_DoesNotRemoveEndpointWhenNoLongerAuthor
     }
 
     private void assertChatCompletionEndpointExists() {
-        var eisEndpoints = getEisEndpoints();
+        assertChatCompletionEndpointExists(modelRegistry);
+    }
+
+    static void assertChatCompletionEndpointExists(ModelRegistry modelRegistry) {
+        var eisEndpoints = getEisEndpoints(modelRegistry);
         assertThat(eisEndpoints.size(), is(1));
 
         var rainbowSprinklesModel = eisEndpoints.get(0);
@@ -212,7 +255,7 @@ private void assertChatCompletionEndpointExists() {
         );
     }
 
-    private void assertChatCompletionUnparsedModel(UnparsedModel rainbowSprinklesModel) {
+    static void assertChatCompletionUnparsedModel(UnparsedModel rainbowSprinklesModel) {
         assertThat(rainbowSprinklesModel.taskType(), is(TaskType.CHAT_COMPLETION));
         assertThat(rainbowSprinklesModel.service(), is(ElasticInferenceService.NAME));
         assertThat(rainbowSprinklesModel.inferenceEntityId(), is(InternalPreconfiguredEndpoints.DEFAULT_CHAT_COMPLETION_ENDPOINT_ID_V1));

x-pack/plugin/inference/src/internalClusterTest/java/org/elasticsearch/xpack/inference/integration/AuthorizationTaskExecutorMultipleNodesIT.java

@@ -20,6 +20,7 @@
 import org.elasticsearch.xpack.inference.services.elastic.ElasticInferenceServiceSettings;
 import org.elasticsearch.xpack.inference.services.elastic.InternalPreconfiguredEndpoints;
 import org.elasticsearch.xpack.inference.services.elastic.authorization.AuthorizationPoller;
+import org.elasticsearch.xpack.inference.services.elastic.ccm.CCMSettings;
 import org.junit.AfterClass;
 import org.junit.Before;
 import org.junit.BeforeClass;
@@ -85,6 +86,8 @@ protected Collection<Class<? extends Plugin>> nodePlugins() {
     protected Settings nodeSettings(int nodeOrdinal, Settings otherSettings) {
         return Settings.builder()
             .put(super.nodeSettings(nodeOrdinal, otherSettings))
+            // Disable CCM to ensure that only the authorization task executor is initialized in the inference plugin when it is created
+            .put(CCMSettings.CCM_SUPPORTED_ENVIRONMENT.getKey(), false)
             .put(LicenseSettings.SELF_GENERATED_LICENSE_TYPE.getKey(), "trial")
             .put(ElasticInferenceServiceSettings.ELASTIC_INFERENCE_SERVICE_URL.getKey(), gatewayUrl)
             .put(ElasticInferenceServiceSettings.PERIODIC_AUTHORIZATION_ENABLED.getKey(), false)

x-pack/plugin/inference/src/internalClusterTest/java/org/elasticsearch/xpack/inference/integration/CCMServiceIT.java

@@ -9,16 +9,49 @@
 
 import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.action.support.PlainActionFuture;
+import org.elasticsearch.action.support.TestPlainActionFuture;
+import org.elasticsearch.common.settings.SecureString;
+import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.core.TimeValue;
+import org.elasticsearch.inference.TaskType;
+import org.elasticsearch.test.http.MockResponse;
+import org.elasticsearch.test.http.MockWebServer;
+import org.elasticsearch.xpack.inference.registry.ModelRegistry;
+import org.elasticsearch.xpack.inference.services.elastic.ElasticInferenceServiceSettings;
+import org.elasticsearch.xpack.inference.services.elastic.authorization.AuthorizationTaskExecutor;
 import org.elasticsearch.xpack.inference.services.elastic.ccm.CCMModel;
 import org.elasticsearch.xpack.inference.services.elastic.ccm.CCMService;
+import org.elasticsearch.xpack.inference.services.elastic.ccm.CCMSettings;
+import org.junit.After;
+import org.junit.AfterClass;
 import org.junit.Before;
+import org.junit.BeforeClass;
 
+import java.io.IOException;
 import java.util.concurrent.atomic.AtomicReference;
 
+import static org.elasticsearch.xpack.inference.external.http.Utils.getUrl;
+import static org.elasticsearch.xpack.inference.integration.AuthorizationTaskExecutorIT.AUTHORIZED_RAINBOW_SPRINKLES_RESPONSE;
+import static org.elasticsearch.xpack.inference.integration.AuthorizationTaskExecutorIT.AUTH_TASK_ACTION;
+import static org.elasticsearch.xpack.inference.integration.AuthorizationTaskExecutorIT.assertChatCompletionEndpointExists;
+import static org.elasticsearch.xpack.inference.integration.AuthorizationTaskExecutorIT.getEisEndpoints;
+import static org.elasticsearch.xpack.inference.integration.AuthorizationTaskExecutorIT.removeEisPreconfiguredEndpoints;
+import static org.elasticsearch.xpack.inference.integration.AuthorizationTaskExecutorIT.waitForAuthorizationToComplete;
+import static org.elasticsearch.xpack.inference.integration.AuthorizationTaskExecutorIT.waitForNoTask;
+import static org.elasticsearch.xpack.inference.integration.AuthorizationTaskExecutorIT.waitForTask;
+import static org.elasticsearch.xpack.inference.integration.ModelRegistryIT.buildElserModelConfig;
+import static org.elasticsearch.xpack.inference.registry.ModelRegistryTests.assertStoreModel;
+import static org.hamcrest.Matchers.empty;
+
 public class CCMServiceIT extends CCMSingleNodeIT {
     private static final AtomicReference<CCMService> ccmService = new AtomicReference<>();
 
+    private static final MockWebServer webServer = new MockWebServer();
+    private static String gatewayUrl;
+
+    private AuthorizationTaskExecutor authorizationTaskExecutor;
+    private ModelRegistry modelRegistry;
+
     public CCMServiceIT() {
         super(new Provider() {
             @Override
@@ -38,9 +71,47 @@ public void delete(ActionListener<Void> listener) {
         });
     }
 
+    @BeforeClass
+    public static void initClass() throws IOException {
+        webServer.start();
+        gatewayUrl = getUrl(webServer);
+    }
+
     @Before
     public void createComponents() {
         ccmService.set(node().injector().getInstance(CCMService.class));
+        modelRegistry = node().injector().getInstance(ModelRegistry.class);
+        authorizationTaskExecutor = node().injector().getInstance(AuthorizationTaskExecutor.class);
+    }
+
+    @After
+    public void shutdown() {
+        // disable CCM to clean up any stored configuration
+        disableCCM();
+
+        removeEisPreconfiguredEndpoints(modelRegistry);
+    }
+
+    private void disableCCM() {
+        var listener = new PlainActionFuture<Void>();
+        ccmService.get().disableCCM(listener);
+        listener.actionGet(TimeValue.THIRTY_SECONDS);
+    }
+
+    @AfterClass
+    public static void cleanUpClass() {
+        webServer.close();
+    }
+
+    @Override
+    protected Settings nodeSettings() {
+        return Settings.builder()
+            .put(CCMSettings.CCM_SUPPORTED_ENVIRONMENT.getKey(), true)
+            .put(ElasticInferenceServiceSettings.ELASTIC_INFERENCE_SERVICE_URL.getKey(), gatewayUrl)
+            // Ensure that the polling logic only occurs once so we can deterministically control when an authorization response is
+            // received
+            .put(ElasticInferenceServiceSettings.PERIODIC_AUTHORIZATION_ENABLED.getKey(), false)
+            .build();
     }
 
     public void testIsEnabled_ReturnsFalse_WhenNoCCMConfigurationStored() {
@@ -58,4 +129,48 @@ public void testIsEnabled_ReturnsTrue_WhenCCMConfigurationIsPresent() {
 
         assertTrue(listener.actionGet(TimeValue.THIRTY_SECONDS));
     }
+
+    public void testCreatesEisChatCompletionEndpoint() throws Exception {
+        disableCCM();
+        waitForNoTask(AUTH_TASK_ACTION, admin());
+
+        var eisEndpoints = getEisEndpoints(modelRegistry);
+        assertThat(eisEndpoints, empty());
+
+        webServer.enqueue(new MockResponse().setResponseCode(200).setBody(AUTHORIZED_RAINBOW_SPRINKLES_RESPONSE));
+        var listener = new TestPlainActionFuture<Void>();
+        ccmService.get().storeConfiguration(new CCMModel(new SecureString("secret".toCharArray())), listener);
+        listener.actionGet(TimeValue.THIRTY_SECONDS);
+
+        // Force a cluster state update to ensure the authorization task is created
+        forceClusterUpdate();
+
+        waitForTask(AUTH_TASK_ACTION, admin());
+        waitForAuthorizationToComplete(authorizationTaskExecutor);
+
+        assertChatCompletionEndpointExists(modelRegistry);
+    }
+
+    private void forceClusterUpdate() {
+        var model = buildElserModelConfig("test-store-model", TaskType.SPARSE_EMBEDDING);
+        assertStoreModel(modelRegistry, model);
+    }
+
+    public void testDisableCCM_RemovesAuthorizationTask() throws Exception {
+        disableCCM();
+        waitForNoTask(AUTH_TASK_ACTION, admin());
+
+        var listener = new TestPlainActionFuture<Void>();
+        ccmService.get().storeConfiguration(new CCMModel(new SecureString("secret".toCharArray())), listener);
+        listener.actionGet(TimeValue.THIRTY_SECONDS);
+
+        // Force a cluster state update to ensure the authorization task is created
+        forceClusterUpdate();
+
+        waitForTask(AUTH_TASK_ACTION, admin());
+        waitForAuthorizationToComplete(authorizationTaskExecutor);
+
+        disableCCM();
+        waitForNoTask(AUTH_TASK_ACTION, admin());
+    }
 }

x-pack/plugin/inference/src/internalClusterTest/java/org/elasticsearch/xpack/inference/integration/ModelRegistryIT.java

@@ -1041,7 +1041,7 @@ private void assertReturnModelIsModifiable(UnparsedModel unparsedModel) {
         }
     }
 
-    private Model buildElserModelConfig(String inferenceEntityId, TaskType taskType) {
+    static Model buildElserModelConfig(String inferenceEntityId, TaskType taskType) {
         return switch (taskType) {
             case SPARSE_EMBEDDING -> new org.elasticsearch.xpack.inference.services.elasticsearch.ElserInternalModel(
                 inferenceEntityId,

x-pack/plugin/inference/src/main/java/module-info.java

@@ -42,6 +42,7 @@
     exports org.elasticsearch.xpack.inference.rest;
     exports org.elasticsearch.xpack.inference.services;
     exports org.elasticsearch.xpack.inference.services.elastic.ccm;
+    exports org.elasticsearch.xpack.inference.services.elastic.authorization;
     exports org.elasticsearch.xpack.inference;
     exports org.elasticsearch.xpack.inference.action.task;
     exports org.elasticsearch.xpack.inference.telemetry;