🌊 [Feature identification] Run as background task

x-pack/platform/packages/shared/kbn-streams-schema/index.ts

@@ -20,6 +20,8 @@ export {
   routingDefinitionListSchema,
 } from './src/models/ingest/routing';
 
+export { getStreamTypeFromDefinition } from './src/helpers/get_stream_type_from_definition';
+export type { StreamType } from './src/helpers/get_stream_type_from_definition';
 export { isRootStreamDefinition } from './src/helpers/is_root';
 export { isOtelStream } from './src/helpers/is_otel_stream';
 export { getIndexPatternsForStream } from './src/helpers/hierarchy_helpers';

x-pack/platform/packages/shared/kbn-streams-schema/src/helpers/get_stream_type_from_definition.ts

@@ -0,0 +1,22 @@
+/*
+ * 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 { Streams } from '../models/streams';
+
+export type StreamType = 'wired' | 'classic' | 'unknown';
+
+export function getStreamTypeFromDefinition(definition: Streams.all.Definition): StreamType {
+  if (Streams.WiredStream.Definition.is(definition)) {
+    return 'wired';
+  }
+
+  if (Streams.ClassicStream.Definition.is(definition)) {
+    return 'classic';
+  }
+
+  return 'unknown';
+}

x-pack/platform/plugins/shared/streams/server/lib/streams/feature/feature_type_registry.ts

@@ -15,6 +15,11 @@ import type { StoredFeature } from './stored_feature';
 import { SystemFeatureHandler } from './handlers/system';
 import { FEATURE_TYPE } from './fields';
 
+export interface IdentifyFeaturesResult {
+  features: Feature[];
+  tokensUsed: ChatCompletionTokenCount;
+}
+
 export class FeatureTypeRegistry {
   private handlers = new Map<string, FeatureTypeHandler>();
 
@@ -58,7 +63,7 @@ export class FeatureTypeRegistry {
 
   async identifyFeatures(
     options: Omit<IdentifyFeaturesOptions, 'analysis'>
-  ): Promise<{ features: Feature[]; tokensUsed: ChatCompletionTokenCount }> {
+  ): Promise<IdentifyFeaturesResult> {
     options.logger.debug(`Identifying features for stream ${options.stream.name}`);
 
     options.logger.trace('Describing dataset for feature identification');

x-pack/platform/plugins/shared/streams/server/lib/tasks/acknowledging_incomplete_error.ts

@@ -0,0 +1,13 @@
+/*
+ * 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.
+ */
+
+export class AcknowledgingIncompleteError extends Error {
+  constructor(message: string) {
+    super(message);
+    this.name = 'AcknowledgingIncompleteError';
+  }
+}

x-pack/platform/plugins/shared/streams/server/lib/tasks/cancellable_task.ts

@@ -0,0 +1,85 @@
+/*
+ * 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 { RunContext } from '@kbn/task-manager-plugin/server';
+import type { RunFunction } from '@kbn/task-manager-plugin/server/task';
+import type { TaskContext } from './task_definitions';
+
+export function cancellableTask(
+  run: RunFunction,
+  runContext: RunContext,
+  taskContext: TaskContext
+) {
+  return async () => {
+    if (!runContext.fakeRequest) {
+      throw new Error('Request is required to run this task');
+    }
+
+    const { taskClient } = await taskContext.getScopedClients({
+      request: runContext.fakeRequest,
+    });
+
+    try {
+      let intervalId: NodeJS.Timeout;
+      const cancellationPromise = new Promise<'canceled'>((resolve) => {
+        taskContext.logger.debug('Starting cancellable task check loop');
+        intervalId = setInterval(async () => {
+          const task = await taskClient.get(runContext.taskInstance.id);
+          taskContext.logger.trace(
+            `Cancellable task check loop for task ${runContext.taskInstance.id}: status is ${task.status}`
+          );
+          if (task.status === 'being_canceled') {
+            runContext.abortController.abort();
+            await taskClient.update({
+              ...task,
+              status: 'canceled',
+            });
+            resolve('canceled' as const);
+          }
+        }, 5000);
+      });
+
+      taskContext.logger.debug(
+        `Running task ${runContext.taskInstance.id} with cancellation support (race)`
+      );
+      const result = await Promise.race([run(), cancellationPromise]).finally(() => {
+        clearInterval(intervalId);
+      });
+
+      if (result === 'canceled') {
+        taskContext.logger.debug(`Task ${runContext.taskInstance.id} canceled`);
+        return undefined;
+      }
+
+      taskContext.logger.debug(`Task ${runContext.taskInstance.id} completed`);
+      return result;
+    } catch (error) {
+      taskContext.logger.error(`Task ${runContext.taskInstance.id} failed unexpectedly`, { error });
+
+      try {
+        await taskClient.update({
+          id: runContext.taskInstance.id,
+          status: 'failed',
+          task: {
+            params: {},
+            error: error.message,
+          },
+          created_at: new Date().toISOString(),
+          space: '',
+          type: '',
+          stream: '',
+        });
+      } catch (updateError) {
+        taskContext.logger.error('Failed to update task status after error', {
+          error: updateError,
+        });
+      }
+
+      throw error;
+    }
+  };
+}

x-pack/platform/plugins/shared/streams/server/lib/tasks/cancellation_in_progress_error.ts

@@ -0,0 +1,13 @@
+/*
+ * 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.
+ */
+
+export class CancellationInProgressError extends Error {
+  constructor(message: string) {
+    super(message);
+    this.name = 'CancellationInProgressError';
+  }
+}

x-pack/platform/plugins/shared/streams/server/lib/tasks/is_stale.ts

@@ -0,0 +1,14 @@
+/*
+ * 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.
+ */
+
+const fiveMinutesInMs = 5 * 60 * 1000;
+
+export function isStale(taskCreatedAt: string) {
+  const createdAt = new Date(taskCreatedAt).getTime();
+  const now = Date.now();
+  return now - createdAt > fiveMinutesInMs;
+}

x-pack/platform/plugins/shared/streams/server/lib/tasks/task_client.ts

@@ -10,6 +10,8 @@ import { TaskPriority, type TaskManagerStartContract } from '@kbn/task-manager-p
 import { isNotFoundError, isResponseError } from '@kbn/es-errors';
 import type { TaskStorageClient } from './storage';
 import type { PersistedTask, TaskParams } from './types';
+import { CancellationInProgressError } from './cancellation_in_progress_error';
+import { AcknowledgingIncompleteError } from './acknowledging_incomplete_error';
 
 interface TaskRequest<TaskType, TParams extends {}> {
   task: Omit<PersistedTask & { type: TaskType }, 'status' | 'created_at' | 'task'>;
@@ -24,7 +26,9 @@ export class TaskClient<TaskType extends string> {
     private readonly logger: Logger
   ) {}
 
-  public async get<TPayload extends {} = {}>(id: string): Promise<PersistedTask<TPayload>> {
+  public async get<TParams extends {} = {}, TPayload extends {} = {}>(
+    id: string
+  ): Promise<PersistedTask<TParams, TPayload>> {
     try {
       this.logger.debug(`Getting task ${id}`);
 
@@ -37,7 +41,7 @@ export class TaskClient<TaskType extends string> {
         throw new Error(`Task ${id} has no source`);
       }
 
-      return response._source as PersistedTask<TPayload>;
+      return response._source as PersistedTask<TParams, TPayload>;
     } catch (error) {
       if (isNotFoundError(error)) {
         return {
@@ -47,6 +51,9 @@ export class TaskClient<TaskType extends string> {
           space: '',
           stream: '',
           type: '',
+          task: {
+            params: {} as TParams,
+          },
         };
       }
 
@@ -58,9 +65,17 @@ export class TaskClient<TaskType extends string> {
     task,
     params,
     request,
-  }: TaskRequest<TaskType, TParams>): Promise<PersistedTask> {
-    const taskDoc: PersistedTask = {
+  }: TaskRequest<TaskType, TParams>) {
+    const storedTask = await this.get(task.id);
+    if (storedTask.status === 'being_canceled') {
+      throw new CancellationInProgressError('Previous task run is still being canceled');
+    }
+
+    const taskDoc: PersistedTask<TParams> = {
       ...task,
+      task: {
+        params,
+      },
       status: 'in_progress',
       created_at: new Date().toISOString(),
     };
@@ -94,13 +109,44 @@ export class TaskClient<TaskType extends string> {
         throw error;
       }
     }
+  }
+
+  public async cancel(id: string) {
+    this.logger.debug(`Canceling task ${id}`);
+
+    const task = await this.get(id);
+    if (task.status !== 'in_progress') {
+      return;
+    }
+
+    await this.update({
+      ...task,
+      status: 'being_canceled',
+    });
+  }
+
+  public async acknowledge<TParams extends {} = {}, TPayload extends {} = {}>(id: string) {
+    const task = await this.get<TParams, TPayload>(id);
+
+    if (task.status !== 'completed') {
+      throw new AcknowledgingIncompleteError('Only completed tasks can be acknowledged');
+    }
+
+    this.logger.debug(`Acknowledging task ${id}`);
+
+    const taskDoc = {
+      ...task,
+      status: 'acknowledged' as const,
+    };
+
+    await this.update(taskDoc);
 
     return taskDoc;
   }
 
-  public async update<TPayload extends {} = {}>(
-    task: PersistedTask<TPayload>
-  ): Promise<PersistedTask<TPayload>> {
+  public async update<TParams extends {} = {}, TPayload extends {} = {}>(
+    task: PersistedTask<TParams, TPayload>
+  ) {
     this.logger.debug(`Updating task ${task.id}`);
 
     await this.storageClient.index({
@@ -109,7 +155,5 @@ export class TaskClient<TaskType extends string> {
       // This might cause issues if there are many updates in a short time from multiple tasks running concurrently
       refresh: true,
     });
-
-    return task;
   }
 }