[ML] Inference API _services retrieves authorization information directly from EIS

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

@@ -23,6 +23,15 @@
 
 public class InferenceGetModelsWithElasticInferenceServiceIT extends BaseMockEISAuthServerTest {
 
+    /**
+     * This is done before the class because I've run into issues where another class that extends {@link BaseMockEISAuthServerTest}
+     * results in an authorization response not being queued up for the new Elasticsearch Node in time. When the node starts up, it
+     * retrieves authorization. If the request isn't queued up when that happens the tests will fail. From my testing locally it seems
+     * like the base class's static functionality to queue a response is only done once and not for each subclass.
+     *
+     * My understanding is that the @Before will be run after the node starts up and wouldn't be sufficient to handle
+     * this scenario. That is why this needs to be @BeforeClass.
+     */
     @BeforeClass
     public static void init() {
         // Ensure the mock EIS server has an authorized response ready

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

@@ -12,6 +12,7 @@
 import org.elasticsearch.client.Request;
 import org.elasticsearch.common.Strings;
 import org.elasticsearch.inference.TaskType;
+import org.junit.Before;
 import org.junit.BeforeClass;
 
 import java.io.IOException;
@@ -23,6 +24,23 @@
 
 public class InferenceGetServicesIT extends BaseMockEISAuthServerTest {
 
+    @Before
+    public void setUp() throws Exception {
+        super.setUp();
+        // Ensure the mock EIS server has an authorized response ready before each test because each test will
+        // use the services API which makes a call to EIS
+        mockEISServer.enqueueAuthorizeAllModelsResponse();
+    }
+
+    /**
+     * This is done before the class because I've run into issues where another class that extends {@link BaseMockEISAuthServerTest}
+     * results in an authorization response not being queued up for the new Elasticsearch Node in time. When the node starts up, it
+     * retrieves authorization. If the request isn't queued up when that happens the tests will fail. From my testing locally it seems
+     * like the base class's static functionality to queue a response is only done once and not for each subclass.
+     *
+     * My understanding is that the @Before will be run after the node starts up and wouldn't be sufficient to handle
+     * this scenario. That is why this needs to be @BeforeClass.
+     */
     @BeforeClass
     public static void init() {
         // Ensure the mock EIS server has an authorized response ready

x-pack/plugin/inference/qa/test-service-plugin/src/main/java/org/elasticsearch/xpack/inference/mock/TestCompletionServiceExtension.java

@@ -178,6 +178,7 @@ public static InferenceServiceConfiguration get() {
                     );
 
                     return new InferenceServiceConfiguration.Builder().setService(NAME)
+                        .setName(NAME)
                         .setTaskTypes(SUPPORTED_TASK_TYPES)
                         .setConfigurations(configurationMap)
                         .build();

x-pack/plugin/inference/qa/test-service-plugin/src/main/java/org/elasticsearch/xpack/inference/mock/TestDenseInferenceServiceExtension.java

@@ -283,6 +283,7 @@ public static InferenceServiceConfiguration get() {
                     );
 
                     return new InferenceServiceConfiguration.Builder().setService(NAME)
+                        .setName(NAME)
                         .setTaskTypes(supportedTaskTypes)
                         .setConfigurations(configurationMap)
                         .build();

x-pack/plugin/inference/qa/test-service-plugin/src/main/java/org/elasticsearch/xpack/inference/mock/TestRerankingServiceExtension.java

@@ -228,6 +228,7 @@ public static InferenceServiceConfiguration get() {
                     );
 
                     return new InferenceServiceConfiguration.Builder().setService(NAME)
+                        .setName(NAME)
                         .setTaskTypes(supportedTaskTypes)
                         .setConfigurations(configurationMap)
                         .build();

x-pack/plugin/inference/qa/test-service-plugin/src/main/java/org/elasticsearch/xpack/inference/mock/TestSparseInferenceServiceExtension.java

@@ -224,6 +224,7 @@ public static InferenceServiceConfiguration get() {
                     );
 
                     return new InferenceServiceConfiguration.Builder().setService(NAME)
+                        .setName(NAME)
                         .setTaskTypes(supportedTaskTypes)
                         .setConfigurations(configurationMap)
                         .build();

x-pack/plugin/inference/qa/test-service-plugin/src/main/java/org/elasticsearch/xpack/inference/mock/TestStreamingCompletionServiceExtension.java

@@ -326,6 +326,7 @@ public static InferenceServiceConfiguration get() {
                     );
 
                     return new InferenceServiceConfiguration.Builder().setService(NAME)
+                        .setName(NAME)
                         .setTaskTypes(supportedTaskTypes)
                         .setConfigurations(configurationMap)
                         .build();

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

@@ -176,7 +176,7 @@ public void testRemoves_DefaultChatCompletion_V1_WhenAuthorizationReturnsEmpty()
             try (var service = createElasticInferenceService()) {
                 ensureAuthorizationCallFinished(service);
 
-                assertThat(service.supportedStreamingTasks(), is(EnumSet.noneOf(TaskType.class)));
+                assertThat(service.supportedStreamingTasks(), is(EnumSet.of(TaskType.CHAT_COMPLETION)));
                 assertTrue(service.defaultConfigIds().isEmpty());
                 assertThat(service.supportedTaskTypes(), is(EnumSet.noneOf(TaskType.class)));
 
@@ -299,7 +299,7 @@ public void testRemoves_DefaultChatCompletion_V1_WhenAuthorizationDoesNotReturnA
             try (var service = createElasticInferenceService()) {
                 ensureAuthorizationCallFinished(service);
 
-                assertThat(service.supportedStreamingTasks(), is(EnumSet.noneOf(TaskType.class)));
+                assertThat(service.supportedStreamingTasks(), is(EnumSet.of(TaskType.CHAT_COMPLETION)));
                 assertThat(
                     service.defaultConfigIds(),
                     containsInAnyOrder(

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/InferencePlugin.java

@@ -88,6 +88,7 @@
 import org.elasticsearch.xpack.inference.external.http.retry.RetrySettings;
 import org.elasticsearch.xpack.inference.external.http.sender.HttpRequestSender;
 import org.elasticsearch.xpack.inference.external.http.sender.RequestExecutorServiceSettings;
+import org.elasticsearch.xpack.inference.external.http.sender.Sender;
 import org.elasticsearch.xpack.inference.highlight.SemanticTextHighlighter;
 import org.elasticsearch.xpack.inference.logging.ThrottlerManager;
 import org.elasticsearch.xpack.inference.mapper.OffsetSourceFieldMapper;
@@ -382,6 +383,8 @@ public Collection<?> createComponents(PluginServices services) {
             new TransportGetInferenceDiagnosticsAction.ClientManagers(httpClientManager, elasticInferenceServiceHttpClientManager)
         );
         components.add(inferenceStatsBinding);
+        components.add(authorizationHandler);
+        components.add(new PluginComponentBinding<>(Sender.class, elasicInferenceServiceFactory.get().createSender()));
         components.add(
             new InferenceEndpointRegistry(
                 services.clusterService(),

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/action/TransportGetInferenceServicesAction.java

@@ -7,36 +7,55 @@
 
 package org.elasticsearch.xpack.inference.action;
 
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
 import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.action.support.ActionFilters;
 import org.elasticsearch.action.support.HandledTransportAction;
+import org.elasticsearch.action.support.SubscribableListener;
 import org.elasticsearch.common.util.concurrent.EsExecutors;
+import org.elasticsearch.core.Nullable;
 import org.elasticsearch.inference.InferenceService;
 import org.elasticsearch.inference.InferenceServiceConfiguration;
 import org.elasticsearch.inference.InferenceServiceRegistry;
 import org.elasticsearch.inference.TaskType;
 import org.elasticsearch.injection.guice.Inject;
 import org.elasticsearch.tasks.Task;
+import org.elasticsearch.threadpool.ThreadPool;
 import org.elasticsearch.transport.TransportService;
 import org.elasticsearch.xpack.core.inference.action.GetInferenceServicesAction;
+import org.elasticsearch.xpack.inference.external.http.sender.Sender;
+import org.elasticsearch.xpack.inference.services.elastic.ElasticInferenceService;
+import org.elasticsearch.xpack.inference.services.elastic.authorization.ElasticInferenceServiceAuthorizationModel;
+import org.elasticsearch.xpack.inference.services.elastic.authorization.ElasticInferenceServiceAuthorizationRequestHandler;
 
 import java.util.ArrayList;
 import java.util.Comparator;
 import java.util.List;
 import java.util.Map;
 import java.util.stream.Collectors;
 
+import static org.elasticsearch.xpack.inference.InferencePlugin.UTILITY_THREAD_POOL_NAME;
+
 public class TransportGetInferenceServicesAction extends HandledTransportAction<
     GetInferenceServicesAction.Request,
     GetInferenceServicesAction.Response> {
 
+    private static final Logger logger = LogManager.getLogger(TransportGetInferenceServicesAction.class);
+
     private final InferenceServiceRegistry serviceRegistry;
+    private final ElasticInferenceServiceAuthorizationRequestHandler eisAuthorizationRequestHandler;
+    private final Sender eisSender;
+    private final ThreadPool threadPool;
 
     @Inject
     public TransportGetInferenceServicesAction(
         TransportService transportService,
         ActionFilters actionFilters,
-        InferenceServiceRegistry serviceRegistry
+        ThreadPool threadPool,
+        InferenceServiceRegistry serviceRegistry,
+        ElasticInferenceServiceAuthorizationRequestHandler eisAuthorizationRequestHandler,
+        Sender sender
     ) {
         super(
             GetInferenceServicesAction.NAME,
@@ -46,6 +65,9 @@ public TransportGetInferenceServicesAction(
             EsExecutors.DIRECT_EXECUTOR_SERVICE
         );
         this.serviceRegistry = serviceRegistry;
+        this.eisAuthorizationRequestHandler = eisAuthorizationRequestHandler;
+        this.eisSender = sender;
+        this.threadPool = threadPool;
     }
 
     @Override
@@ -69,41 +91,86 @@ private void getServiceConfigurationsForTaskType(
             .entrySet()
             .stream()
             .filter(
-                service -> service.getValue().hideFromConfigurationApi() == false
+                // Exclude EIS as the EIS specific configurations must be retrieved directly from EIS and merged in later
+                service -> service.getValue().name().equals(ElasticInferenceService.NAME) == false
+                    && service.getValue().hideFromConfigurationApi() == false
                     && service.getValue().supportedTaskTypes().contains(requestedTaskType)
             )
             .sorted(Comparator.comparing(service -> service.getValue().name()))
             .collect(Collectors.toCollection(ArrayList::new));
 
-        getServiceConfigurationsForServices(filteredServices, listener.delegateFailureAndWrap((delegate, configurations) -> {
-            delegate.onResponse(new GetInferenceServicesAction.Response(configurations));
-        }));
+        getServiceConfigurationsForServicesAndEis(listener, filteredServices, requestedTaskType);
     }
 
     private void getAllServiceConfigurations(ActionListener<GetInferenceServicesAction.Response> listener) {
         var availableServices = serviceRegistry.getServices()
             .entrySet()
             .stream()
-            .filter(service -> service.getValue().hideFromConfigurationApi() == false)
+            .filter(
+                // Exclude EIS as the EIS specific configurations must be retrieved directly from EIS and merged in later
+                service -> service.getValue().name().equals(ElasticInferenceService.NAME) == false
+                    && service.getValue().hideFromConfigurationApi() == false
+            )
             .sorted(Comparator.comparing(service -> service.getValue().name()))
             .collect(Collectors.toCollection(ArrayList::new));
-        getServiceConfigurationsForServices(availableServices, listener.delegateFailureAndWrap((delegate, configurations) -> {
-            delegate.onResponse(new GetInferenceServicesAction.Response(configurations));
-        }));
+
+        getServiceConfigurationsForServicesAndEis(listener, availableServices, null);
     }
 
-    private void getServiceConfigurationsForServices(
-        ArrayList<Map.Entry<String, InferenceService>> services,
-        ActionListener<List<InferenceServiceConfiguration>> listener
+    private void getServiceConfigurationsForServicesAndEis(
+        ActionListener<GetInferenceServicesAction.Response> listener,
+        ArrayList<Map.Entry<String, InferenceService>> availableServices,
+        @Nullable TaskType requestedTaskType
     ) {
-        try {
-            var serviceConfigurations = new ArrayList<InferenceServiceConfiguration>();
-            for (var service : services) {
-                serviceConfigurations.add(service.getValue().getConfiguration());
+        SubscribableListener.<ElasticInferenceServiceAuthorizationModel>newForked(authModelListener -> {
+            // Executing on a separate thread because there's a chance the authorization call needs to do some initialization for the Sender
+            threadPool.executor(UTILITY_THREAD_POOL_NAME).execute(() -> getEisAuthorization(authModelListener, eisSender));
+        }).<List<InferenceServiceConfiguration>>andThen((configurationListener, authorizationModel) -> {
+            var serviceConfigs = getServiceConfigurationsForServices(availableServices);
+
+            if (authorizationModel.isAuthorized() == false) {
+                configurationListener.onResponse(serviceConfigs);
+                return;
+            }
+
+            var config = ElasticInferenceService.createConfiguration(authorizationModel.getAuthorizedTaskTypes());
+            if (requestedTaskType != null && authorizationModel.getAuthorizedTaskTypes().contains(requestedTaskType) == false) {
+                configurationListener.onResponse(serviceConfigs);
+                return;
             }
-            listener.onResponse(serviceConfigurations.stream().toList());
-        } catch (Exception e) {
-            listener.onFailure(e);
+
+            serviceConfigs.add(config);
+            serviceConfigs.sort(Comparator.comparing(InferenceServiceConfiguration::getService));
+            configurationListener.onResponse(serviceConfigs);
+        })
+            .addListener(
+                listener.delegateFailureAndWrap(
+                    (delegate, configurations) -> delegate.onResponse(new GetInferenceServicesAction.Response(configurations))
+                )
+            );
+    }
+
+    private void getEisAuthorization(ActionListener<ElasticInferenceServiceAuthorizationModel> listener, Sender sender) {
+        var disabledServiceListener = listener.delegateResponse((delegate, e) -> {
+            logger.warn(
+                "Failed to retrieve authorization information from the "
+                    + "Elastic Inference Service while determining service configurations. Marking service as disabled.",
+                e
+            );
+            delegate.onResponse(ElasticInferenceServiceAuthorizationModel.newDisabledService());
+        });
+
+        eisAuthorizationRequestHandler.getAuthorization(disabledServiceListener, sender);
+    }
+
+    private List<InferenceServiceConfiguration> getServiceConfigurationsForServices(
+        ArrayList<Map.Entry<String, InferenceService>> services
+    ) {
+        var serviceConfigurations = new ArrayList<InferenceServiceConfiguration>();
+        for (var service : services) {
+            serviceConfigurations.add(service.getValue().getConfiguration());
         }
+
+        return serviceConfigurations;
     }
 }