Add registry providers for APM data

x-pack/solutions/observability/plugins/apm/kibana.jsonc

@@ -59,6 +59,7 @@
       "licenseManagement",
       "profilingDataAccess",
       "observabilityAIAssistant",
+      "observabilityAgent",
       "share",
       "onechat"
     ],

x-pack/solutions/observability/plugins/apm/moon.yml

@@ -154,6 +154,7 @@ dependsOn:
   - '@kbn/onechat-server'
   - '@kbn/onechat-common'
   - '@kbn/onechat-plugin'
+  - '@kbn/observability-agent-plugin'
 tags:
   - plugin
   - prod

x-pack/solutions/observability/plugins/apm/server/agent_tools/utils/build_apm_tool_resources.ts

@@ -22,6 +22,7 @@ export interface ApmToolResources {
   randomSampler: Awaited<ReturnType<typeof getRandomSampler>>;
   mlClient: Awaited<ReturnType<typeof getMlClient>>;
   apmAlertsClient: ApmAlertsClient;
+  esClient: IScopedClusterClient;
 }
 
 export async function buildApmToolResources({
@@ -107,5 +108,5 @@ export async function buildApmToolResources({
     apmAlertsClientPromise,
   ]);
 
-  return { apmEventClient, randomSampler, mlClient, apmAlertsClient };
+  return { apmEventClient, randomSampler, mlClient, apmAlertsClient, esClient: esScoped };
 }

x-pack/solutions/observability/plugins/apm/server/observability_agent/register_data_providers.ts

@@ -0,0 +1,128 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import type { CoreSetup, Logger } from '@kbn/core/server';
+import { getRandomSampler } from '../lib/helpers/get_random_sampler';
+import { getApmServiceSummary } from '../routes/assistant_functions/get_apm_service_summary';
+import { getApmDownstreamDependencies } from '../routes/assistant_functions/get_apm_downstream_dependencies';
+import { getApmErrors } from '../routes/assistant_functions/get_observability_alert_details_context/get_apm_errors';
+import {
+  getExitSpanChangePoints,
+  getServiceChangePoints,
+} from '../routes/assistant_functions/get_changepoints';
+import { buildApmToolResources } from '../agent_tools/utils/build_apm_tool_resources';
+import type { APMPluginSetupDependencies, APMPluginStartDependencies } from '../types';
+
+export function registerDataProviders({
+  core,
+  plugins,
+  logger,
+}: {
+  core: CoreSetup<APMPluginStartDependencies>;
+  plugins: APMPluginSetupDependencies;
+  logger: Logger;
+}) {
+  const { observabilityAgent } = plugins;
+  if (!observabilityAgent) {
+    return;
+  }
+
+  observabilityAgent.registerDataProvider(
+    'apmServiceSummary',
+    async ({ request, serviceName, serviceEnvironment, start, end, transactionType }) => {
+      const { apmEventClient, apmAlertsClient, mlClient, esClient } = await buildApmToolResources({
+        core,
+        plugins,
+        request,
+        logger,
+      });
+
+      return getApmServiceSummary({
+        apmEventClient,
+        esClient: esClient.asCurrentUser,
+        apmAlertsClient,
+        mlClient,
+        logger,
+        arguments: {
+          'service.name': serviceName,
+          'service.environment': serviceEnvironment,
+          start,
+          end,
+          'transaction.type': transactionType,
+        },
+      });
+    }
+  );
+
+  observabilityAgent.registerDataProvider(
+    'apmDownstreamDependencies',
+    async ({ request, serviceName, serviceEnvironment, start, end }) => {
+      const { apmEventClient } = await buildApmToolResources({ core, plugins, request, logger });
+      const [coreStart] = await core.getStartServices();
+      const randomSampler = await getRandomSampler({ coreStart, probability: 1, request });
+
+      return getApmDownstreamDependencies({
+        apmEventClient,
+        randomSampler,
+        arguments: {
+          serviceName,
+          serviceEnvironment,
+          start,
+          end,
+        },
+      });
+    }
+  );
+
+  observabilityAgent.registerDataProvider(
+    'apmErrors',
+    async ({ request, serviceName, serviceEnvironment, start, end }) => {
+      const { apmEventClient } = await buildApmToolResources({ core, plugins, request, logger });
+      return getApmErrors({ apmEventClient, serviceName, serviceEnvironment, start, end });
+    }
+  );
+
+  observabilityAgent.registerDataProvider(
+    'apmExitSpanChangePoints',
+    async ({ request, serviceName, serviceEnvironment, start, end }) => {
+      const { apmEventClient } = await buildApmToolResources({ core, plugins, request, logger });
+
+      return getExitSpanChangePoints({
+        apmEventClient,
+        serviceName,
+        serviceEnvironment,
+        start,
+        end,
+      });
+    }
+  );
+
+  observabilityAgent.registerDataProvider(
+    'apmServiceChangePoints',
+    async ({
+      request,
+      serviceName,
+      serviceEnvironment,
+      transactionType,
+      transactionName,
+      start,
+      end,
+    }) => {
+      const { apmEventClient } = await buildApmToolResources({ core, plugins, request, logger });
+
+      return getServiceChangePoints({
+        apmEventClient,
+        serviceName,
+        serviceEnvironment,
+        transactionType,
+        transactionName,
+        start,
+        end,
+      });
+    }
+  );
+}

x-pack/solutions/observability/plugins/apm/server/plugin.ts

@@ -52,6 +52,7 @@ import type {
   APMPluginStartDependencies,
 } from './types';
 import { registerAgentTools } from './agent_tools';
+import { registerDataProviders } from './observability_agent/register_data_providers';
 
 export class APMPlugin
   implements Plugin<APMPluginSetup, void, APMPluginSetupDependencies, APMPluginStartDependencies>
@@ -260,6 +261,12 @@ export class APMPlugin
       this.logger?.debug('Successfully registered observability agent APM tools');
     }
 
+    registerDataProviders({
+      core,
+      plugins,
+      logger: this.logger!.get('observabilityAgent'),
+    });
+
     registerDeprecations({
       core,
       apmDeps: {

x-pack/solutions/observability/plugins/apm/server/types.ts

@@ -53,6 +53,7 @@ import type {
   ObservabilityAIAssistantServerStart,
 } from '@kbn/observability-ai-assistant-plugin/server';
 import type { OnechatPluginSetup, OnechatPluginStart } from '@kbn/onechat-plugin/server/types';
+import type { ObservabilityAgentPluginSetup } from '@kbn/observability-agent-plugin/server';
 import type {
   ProfilingDataAccessPluginSetup,
   ProfilingDataAccessPluginStart,
@@ -76,6 +77,7 @@ export interface APMPluginSetupDependencies {
   share: SharePluginSetup;
   logsDataAccess: LogsDataAccessPluginSetup;
   // optional dependencies
+  observabilityAgent?: ObservabilityAgentPluginSetup;
   observabilityAIAssistant?: ObservabilityAIAssistantServerSetup;
   onechat?: OnechatPluginSetup;
   actions?: ActionsPlugin['setup'];

x-pack/solutions/observability/plugins/apm/tsconfig.json

@@ -151,7 +151,8 @@
     "@kbn/zod",
     "@kbn/onechat-server",
     "@kbn/onechat-common",
-    "@kbn/onechat-plugin"
+    "@kbn/onechat-plugin",
+    "@kbn/observability-agent-plugin"
   ],
   "exclude": ["target/**/*"]
 }

x-pack/solutions/observability/plugins/observability_agent/server/data_registry/data_registry.ts

@@ -0,0 +1,42 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import type { Logger } from '@kbn/core/server';
+import type { ObservabilityAgentDataRegistryTypes } from './data_registry_types';
+
+export class ObservabilityAgentDataRegistry {
+  private readonly providers = new Map<string, (...args: any[]) => Promise<any>>();
+
+  constructor(private readonly logger: Logger) {}
+
+  public registerDataProvider<K extends keyof ObservabilityAgentDataRegistryTypes>(
+    id: K,
+    provider: ObservabilityAgentDataRegistryTypes[K]
+  ): void {
+    if (this.providers.has(id)) {
+      this.logger.warn(`Overwriting data provider for key: ${id}`);
+    } else {
+      this.logger.debug(`Registered data provider for key: ${id}`);
+    }
+
+    this.providers.set(id, provider);
+  }
+
+  public async getData<K extends keyof ObservabilityAgentDataRegistryTypes>(
+    id: K,
+    params: Parameters<ObservabilityAgentDataRegistryTypes[K]>[0]
+  ): Promise<ReturnType<ObservabilityAgentDataRegistryTypes[K]> | undefined> {
+    const provider = this.providers.get(id) as ObservabilityAgentDataRegistryTypes[K] | undefined;
+
+    if (!provider) {
+      this.logger.error(`No data provider registered for key: ${id}`);
+      return;
+    }
+
+    return provider(params) as Promise<ReturnType<ObservabilityAgentDataRegistryTypes[K]>>;
+  }
+}