[streams] stop llm reasoning when request is aborted (elastic#241550)

## Summary

We have several routes that triggers an llm workflow that can be long
running. If a request is closed early by the client, the process will
continue executing unnecessarily.

This change passes the request's abort signal to the relevant workflows
to prevent unneeded work in case the request is already terminated


### Testing

Trigger any ai workflow and abort the request (eg close the window). The
process should end with a `CancelledError`

---------

Co-authored-by: kibanamachine <[email protected]>

Author: klacabane

x-pack/platform/packages/shared/kbn-evals-suite-streams/evals/feature_identification.spec.ts

@@ -250,6 +250,7 @@ evaluate.describe('Streams feature identification', { tag: '@svlOblt' }, () => {
             stream,
             kql: '',
             dropUnmapped: true,
+            signal: new AbortController().signal,
           });
 
           const featuresWithAnalysis = await Promise.all(

x-pack/platform/packages/shared/kbn-streams-ai/src/description/generate_description.ts

@@ -21,13 +21,15 @@ export async function generateStreamDescription({
   end,
   esClient,
   inferenceClient,
+  signal,
 }: {
   stream: Streams.all.Definition;
   feature?: Feature;
   start: number;
   end: number;
   esClient: ElasticsearchClient;
   inferenceClient: BoundInferenceClient;
+  signal: AbortSignal;
 }): Promise<string> {
   const analysis = await describeDataset({
     start,
@@ -45,6 +47,7 @@ export async function generateStreamDescription({
       ),
     },
     prompt: GenerateStreamDescriptionPrompt,
+    abortSignal: signal,
   });
 
   return response.content;

x-pack/platform/packages/shared/kbn-streams-ai/src/features/identify_features.ts

@@ -30,6 +30,7 @@ export async function identifyFeatures({
   kql,
   inferenceClient,
   logger,
+  signal,
   dropUnmapped = false,
   maxSteps: initialMaxSteps,
 }: {
@@ -41,6 +42,7 @@ export async function identifyFeatures({
   kql?: string;
   inferenceClient: BoundInferenceClient;
   logger: Logger;
+  signal: AbortSignal;
   dropUnmapped?: boolean;
   maxSteps?: number;
 }): Promise<{ features: Omit<Feature, 'description'>[] }> {
@@ -121,6 +123,7 @@ export async function identifyFeatures({
         };
       },
     },
+    abortSignal: signal,
   });
 
   return {

x-pack/platform/packages/shared/kbn-streams-ai/src/significant_events/generate_significant_events.ts

@@ -33,14 +33,15 @@ export async function generateSignificantEvents({
   end,
   esClient,
   inferenceClient,
-  logger,
+  signal,
 }: {
   stream: Streams.all.Definition;
   feature?: Feature;
   start: number;
   end: number;
   esClient: ElasticsearchClient;
   inferenceClient: BoundInferenceClient;
+  signal: AbortSignal;
   logger: Logger;
 }): Promise<{
   queries: Query[];
@@ -88,6 +89,7 @@ export async function generateSignificantEvents({
         };
       },
     },
+    abortSignal: signal,
   });
 
   const queries = response.input.flatMap((message) => {

x-pack/platform/plugins/shared/streams/server/lib/significant_events/generate_significant_events.ts

@@ -23,14 +23,15 @@ interface Dependencies {
   inferenceClient: InferenceClient;
   esClient: ElasticsearchClient;
   logger: Logger;
+  signal: AbortSignal;
 }
 
 export async function generateSignificantEventDefinitions(
   params: Params,
   dependencies: Dependencies
 ): Promise<GeneratedSignificantEventQuery[]> {
   const { definition, connectorId, start, end, feature } = params;
-  const { inferenceClient, esClient, logger } = dependencies;
+  const { inferenceClient, esClient, logger, signal } = dependencies;
 
   const boundInferenceClient = inferenceClient.bindTo({
     connectorId,
@@ -44,6 +45,7 @@ export async function generateSignificantEventDefinitions(
     inferenceClient: boundInferenceClient,
     logger,
     feature,
+    signal,
   });
 
   return queries.map((query) => ({

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

@@ -18,6 +18,7 @@ export function runFeatureIdentification({
   logger,
   stream,
   features,
+  signal,
 }: {
   start: number;
   end: number;
@@ -26,6 +27,7 @@ export function runFeatureIdentification({
   logger: Logger;
   stream: Streams.all.Definition;
   features: Feature[];
+  signal: AbortSignal;
 }) {
   return identifyFeatures({
     start,
@@ -35,6 +37,7 @@ export function runFeatureIdentification({
     logger,
     stream,
     features,
+    signal,
     dropUnmapped: true,
   });
 }

x-pack/platform/plugins/shared/streams/server/routes/internal/streams/features/route.ts

@@ -23,6 +23,7 @@ import { STREAMS_API_PRIVILEGES } from '../../../../../common/constants';
 import { assertSignificantEventsAccess } from '../../../utils/assert_significant_events_access';
 import { runFeatureIdentification } from '../../../../lib/streams/feature/run_feature_identification';
 import type { IdentifiedFeaturesEvent, StreamDescriptionEvent } from './types';
+import { getRequestAbortSignal } from '../../../utils/get_request_abort_signal';
 
 const dateFromString = z.string().transform((input) => new Date(input));
 
@@ -318,6 +319,7 @@ export const identifyFeaturesRoute = createServerRoute({
     const esClient = scopedClusterClient.asCurrentUser;
 
     const boundInferenceClient = inferenceClient.bindTo({ connectorId });
+    const signal = getRequestAbortSignal(request);
 
     return from(
       runFeatureIdentification({
@@ -328,6 +330,7 @@ export const identifyFeaturesRoute = createServerRoute({
         logger,
         stream,
         features: hits,
+        signal,
       })
     ).pipe(
       switchMap(({ features }) => {
@@ -344,6 +347,7 @@ export const identifyFeaturesRoute = createServerRoute({
                   ...feature,
                   description: '',
                 },
+                signal,
               });
 
               return {
@@ -419,6 +423,7 @@ export const describeStreamRoute = createServerRoute({
         inferenceClient: inferenceClient.bindTo({ connectorId }),
         start: start.valueOf(),
         end: end.valueOf(),
+        signal: getRequestAbortSignal(request),
       })
     ).pipe(
       map((description) => {

x-pack/platform/plugins/shared/streams/server/routes/internal/streams/management/suggest_partitions_route.ts

@@ -15,6 +15,7 @@ import { STREAMS_TIERED_ML_FEATURE } from '../../../../../common';
 import { STREAMS_API_PRIVILEGES } from '../../../../../common/constants';
 import { SecurityError } from '../../../../lib/streams/errors/security_error';
 import { createServerRoute } from '../../../create_server_route';
+import { getRequestAbortSignal } from '../../../utils/get_request_abort_signal';
 
 export interface SuggestPartitionsParams {
   path: {
@@ -84,7 +85,7 @@ export const suggestPartitionsRoute = createServerRoute({
       start: params.body.start,
       end: params.body.end,
       maxSteps: 1, // Longer reasoning seems to add unnecessary conditions (and latency), instead of improving accuracy, so we limit the steps.
-      signal: new AbortController().signal,
+      signal: getRequestAbortSignal(request),
     });
 
     // Turn our promise into an Observable ServerSideEvent. The only reason we're streaming the

x-pack/platform/plugins/shared/streams/server/routes/internal/streams/processing/grok_suggestions_handler.ts

@@ -39,6 +39,7 @@ export interface ProcessingGrokSuggestionsHandlerDeps {
   scopedClusterClient: IScopedClusterClient;
   streamsClient: StreamsClient;
   fieldsMetadataClient: IFieldsMetadataClient;
+  signal: AbortSignal;
 }
 
 export const processingGrokSuggestionsSchema = z.object({
@@ -65,6 +66,7 @@ export const handleProcessingGrokSuggestions = async ({
   inferenceClient,
   streamsClient,
   fieldsMetadataClient,
+  signal,
 }: ProcessingGrokSuggestionsHandlerDeps) => {
   const stream = await streamsClient.getStream(params.path.name);
   const isWiredStream = Streams.WiredStream.Definition.is(stream);
@@ -76,6 +78,7 @@ export const handleProcessingGrokSuggestions = async ({
       sample_messages: params.body.sample_messages,
       review_fields: JSON.stringify(params.body.review_fields),
     },
+    abortSignal: signal,
   });
   const reviewResult = response.toolCalls[0].function.arguments;
 

x-pack/platform/plugins/shared/streams/server/routes/internal/streams/processing/route.ts

@@ -26,6 +26,7 @@ import {
   handleProcessingGrokSuggestions,
   processingGrokSuggestionsSchema,
 } from './grok_suggestions_handler';
+import { getRequestAbortSignal } from '../../../utils/get_request_abort_signal';
 
 const paramsSchema = z.object({
   path: z.object({ name: z.string() }),
@@ -110,6 +111,7 @@ export const processingGrokSuggestionRoute = createServerRoute({
         streamsClient,
         scopedClusterClient,
         fieldsMetadataClient,
+        signal: getRequestAbortSignal(request),
       })
     ).pipe(
       map((grokProcessor) => ({