feat: Implement workflow cancellation (elastic#233884)

## Summary

Closes: elastic/security-team#13502



https://github.com/user-attachments/assets/b8d4b5b0-71df-4395-9ae2-6fd9d52dc73c


### Checklist

Check the PR satisfies following conditions. 

Reviewers should verify this PR satisfies this list as well.

- [ ] Any text added follows [EUI's writing
guidelines](https://elastic.github.io/eui/#/guidelines/writing), uses
sentence case text and includes [i18n
support](https://github.com/elastic/kibana/blob/main/src/platform/packages/shared/kbn-i18n/README.md)
- [ ]
[Documentation](https://www.elastic.co/guide/en/kibana/master/development-documentation.html)
was added for features that require explanation or tutorials
- [ ] [Unit or functional
tests](https://www.elastic.co/guide/en/kibana/master/development-tests.html)
were updated or added to match the most common scenarios
- [ ] If a plugin configuration key changed, check if it needs to be
allowlisted in the cloud and added to the [docker
list](https://github.com/elastic/kibana/blob/main/src/dev/build/tasks/os_packages/docker_generator/resources/base/bin/kibana-docker)
- [ ] This was checked for breaking HTTP API changes, and any breaking
changes have been approved by the breaking-change committee. The
`release_note:breaking` label should be applied in these situations.
- [ ] [Flaky Test
Runner](https://ci-stats.kibana.dev/trigger_flaky_test_runner/1) was
used on any tests changed
- [ ] The PR description includes the appropriate Release Notes section,
and the correct `release_note:*` label is applied per the
[guidelines](https://www.elastic.co/guide/en/kibana/master/contributing.html#kibana-release-notes-process)
- [ ] Review the [backport
guidelines](https://docs.google.com/document/d/1VyN5k91e5OVumlc0Gb9RPa3h1ewuPE705nRtioPiTvY/edit?usp=sharing)
and apply applicable `backport:*` labels.

### Identify risks

Does this PR introduce any risks? For example, consider risks like hard
to test bugs, performance regression, potential of data loss.

Describe the risk, its severity, and mitigation for each identified
risk. Invite stakeholders and evaluate how to proceed before merging.

- [ ] [See some risk
examples](https://github.com/elastic/kibana/blob/main/RISK_MATRIX.mdx)
- [ ] ...

Author: skynetigor

src/platform/packages/shared/kbn-workflows/common/errors.ts

@@ -0,0 +1,15 @@
+/*
+ * 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", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+export class WorkflowExecutionNotFoundError extends Error {
+  constructor(executionId: string) {
+    super(`Workflow execution with id "${executionId}" not found.`);
+    this.name = 'WorkflowExecutionNotFoundError';
+  }
+}

src/platform/packages/shared/kbn-workflows/types/v1.ts

@@ -41,6 +41,9 @@ export interface EsWorkflowExecution {
   createdBy: string;
   startedAt: string;
   finishedAt: string;
+  cancelRequested: boolean;
+  cancelledAt?: string;
+  cancelledBy?: string;
   duration: number;
   triggeredBy?: string; // 'manual' or 'scheduled'
   traceId?: string; // APM trace ID for observability

src/platform/plugins/shared/workflows_execution_engine/server/plugin.ts

@@ -17,6 +17,7 @@ import type {
 import type { EsWorkflowExecution, WorkflowExecutionEngineModel } from '@kbn/workflows';
 import { ExecutionStatus } from '@kbn/workflows';
 import { convertToWorkflowGraph } from '@kbn/workflows/graph';
+import { WorkflowExecutionNotFoundError } from '@kbn/workflows/common/errors';
 
 import type { Client } from '@elastic/elasticsearch';
 import type { PluginStartContract as ActionsPluginStartContract } from '@kbn/actions-plugin/server';
@@ -84,19 +85,25 @@ export class WorkflowsExecutionEnginePlugin
               const esClient = coreStart.elasticsearch.client.asInternalUser as Client;
               const workflowExecutionRepository = new WorkflowExecutionRepository(esClient);
 
-              const { workflowRuntime, workflowLogger, nodesFactory } = await createContainer(
-                workflowRunId,
-                spaceId,
-                actions,
-                taskManager,
-                esClient,
-                logger,
-                config,
-                workflowExecutionRepository
-              );
-              await workflowRuntime.start();
+              const { workflowRuntime, workflowExecutionState, workflowLogger, nodesFactory } =
+                await createContainer(
+                  workflowRunId,
+                  spaceId,
+                  actions,
+                  taskManager,
+                  esClient,
+                  logger,
+                  config,
+                  workflowExecutionRepository
+                );
 
-              await workflowExecutionLoop(workflowRuntime, workflowLogger, nodesFactory);
+              await workflowRuntime.start();
+              await workflowExecutionLoop(
+                workflowRuntime,
+                workflowExecutionState,
+                workflowLogger,
+                nodesFactory
+              );
             },
             async cancel() {
               // Cancel function for the task
@@ -124,19 +131,25 @@ export class WorkflowsExecutionEnginePlugin
               const esClient = coreStart.elasticsearch.client.asInternalUser as Client;
               const workflowExecutionRepository = new WorkflowExecutionRepository(esClient);
 
-              const { workflowRuntime, workflowLogger, nodesFactory } = await createContainer(
-                workflowRunId,
-                spaceId,
-                actions,
-                taskManager,
-                esClient,
-                logger,
-                config,
-                workflowExecutionRepository
-              );
+              const { workflowRuntime, workflowExecutionState, workflowLogger, nodesFactory } =
+                await createContainer(
+                  workflowRunId,
+                  spaceId,
+                  actions,
+                  taskManager,
+                  esClient,
+                  logger,
+                  config,
+                  workflowExecutionRepository
+                );
               await workflowRuntime.resume();
 
-              await workflowExecutionLoop(workflowRuntime, workflowLogger, nodesFactory);
+              await workflowExecutionLoop(
+                workflowRuntime,
+                workflowExecutionState,
+                workflowLogger,
+                nodesFactory
+              );
             },
             async cancel() {},
           };
@@ -199,8 +212,48 @@ export class WorkflowsExecutionEnginePlugin
       };
     };
 
+    const cancelWorkflowExecution = async (workflowExecutionId: string, spaceId: string) => {
+      const esClient = core.elasticsearch.client.asInternalUser as Client;
+      const workflowExecutionRepository = new WorkflowExecutionRepository(esClient);
+      const workflowExecution = await workflowExecutionRepository.getWorkflowExecutionById(
+        workflowExecutionId,
+        spaceId
+      );
+
+      if (!workflowExecution) {
+        throw new WorkflowExecutionNotFoundError(workflowExecutionId);
+      }
+
+      if (
+        [ExecutionStatus.CANCELLED, ExecutionStatus.COMPLETED, ExecutionStatus.FAILED].includes(
+          workflowExecution.status
+        )
+      ) {
+        // Already in a terminal state or being canceled
+        return;
+      }
+
+      // Request cancellation
+      await workflowExecutionRepository.updateWorkflowExecution({
+        id: workflowExecution.id,
+        cancelRequested: true,
+        cancelledAt: new Date().toISOString(),
+        cancelledBy: 'system', // TODO: set user if available
+      });
+
+      if (
+        [ExecutionStatus.WAITING, ExecutionStatus.WAITING_FOR_INPUT].includes(
+          workflowExecution.status
+        )
+      ) {
+        // TODO: handle WAITING states
+        // It should clean up resume tasks, etc
+      }
+    };
+
     return {
       executeWorkflow,
+      cancelWorkflowExecution,
     };
   }
 
@@ -281,6 +334,7 @@ async function createContainer(
 
   return {
     workflowRuntime,
+    workflowExecutionState,
     contextManager,
     connectorExecutor,
     workflowLogger,

src/platform/plugins/shared/workflows_execution_engine/server/types.ts

@@ -25,6 +25,7 @@ export interface WorkflowsExecutionEnginePluginStart {
     workflow: WorkflowExecutionEngineModel,
     context: Record<string, any>
   ): Promise<ExecuteWorkflowResponse>;
+  cancelWorkflowExecution(workflowExecutionId: string, spaceId: string): Promise<void>;
 }
 
 export interface WorkflowsExecutionEnginePluginSetupDeps {

src/platform/plugins/shared/workflows_execution_engine/server/workflow_context_manager/tests/workflow_execution_state.test.ts

@@ -154,6 +154,12 @@ describe('WorkflowExecutionState', () => {
   });
 
   describe('flush', () => {
+    beforeEach(() => {
+      workflowExecutionRepository.getWorkflowExecutionById = jest
+        .fn()
+        .mockResolvedValue({} as EsWorkflowExecution);
+    });
+
     it('should flush workflow execution changes', async () => {
       const updatedWorkflowExecution = {
         id: 'test-workflow-execution-id',
@@ -172,6 +178,20 @@ describe('WorkflowExecutionState', () => {
       );
     });
 
+    it('should flush workflow execution changes with execution id even if execution id is not in change', async () => {
+      const updatedWorkflowExecution = {} as EsWorkflowExecution;
+
+      underTest.updateWorkflowExecution(updatedWorkflowExecution);
+
+      await underTest.flush();
+
+      expect(workflowExecutionRepository.updateWorkflowExecution).toHaveBeenCalledWith(
+        expect.objectContaining({
+          id: 'test-workflow-execution-id',
+        })
+      );
+    });
+
     it('should flush new step executions', async () => {
       const mockUuid = 'fake-uuid';
       jest.requireMock('uuid').v4.mockImplementation(() => mockUuid);
@@ -188,12 +208,10 @@ describe('WorkflowExecutionState', () => {
       expect(stepExecutionRepository.createStepExecution).toHaveBeenCalledWith(
         expect.objectContaining({
           id: mockUuid,
-          workflowRunId: 'test-workflow-execution-id',
-          workflowId: 'test-workflow-id',
           stepId: 'test-step-execution-id',
           status: ExecutionStatus.RUNNING,
           startedAt: '2025-08-05T20:00:00.000Z',
-        })
+        } as EsWorkflowStepExecution)
       );
     });
 
@@ -221,19 +239,17 @@ describe('WorkflowExecutionState', () => {
       await underTest.flush();
 
       expect(stepExecutionRepository.updateStepExecutions).toHaveBeenCalledWith([
-        expect.objectContaining({
+        {
           id: fakeUuid,
-          workflowRunId: 'test-workflow-execution-id',
-          workflowId: 'test-workflow-id',
           stepId: 'test-step-execution-id',
           status: ExecutionStatus.COMPLETED,
           completedAt: '2025-08-05T20:01:00.000Z',
           executionTimeMs: 60000,
-        }),
+        } as EsWorkflowStepExecution,
       ]);
     });
 
-    it('should be able to create step executions that changed multiple times', async () => {
+    it('should be able to create step executions that changed multiple times by merging changes', async () => {
       // create initial step execution
       const fakeUuid = 'fake-uuid-1';
       jest.requireMock('uuid').v4.mockImplementationOnce(() => fakeUuid);
@@ -249,21 +265,20 @@ describe('WorkflowExecutionState', () => {
         stepId: 'test-step-execution-id',
         status: ExecutionStatus.COMPLETED,
         completedAt: '2025-08-05T20:01:00.000Z',
-        executionTimeMs: 60000,
+        executionTimeMs: 2000,
       } as EsWorkflowStepExecution);
 
       await underTest.flush();
 
       expect(stepExecutionRepository.createStepExecution).toHaveBeenCalledWith(
         expect.objectContaining({
           id: fakeUuid,
-          workflowRunId: 'test-workflow-execution-id',
-          workflowId: 'test-workflow-id',
           stepId: 'test-step-execution-id',
           status: ExecutionStatus.COMPLETED,
           completedAt: '2025-08-05T20:01:00.000Z',
-          executionTimeMs: 60000,
-        })
+          startedAt: '2025-08-05T20:00:00.000Z',
+          executionTimeMs: 2000,
+        } as EsWorkflowStepExecution)
       );
     });
 
@@ -312,6 +327,31 @@ describe('WorkflowExecutionState', () => {
       expect(stepExecutionRepository.createStepExecution).toHaveBeenCalledTimes(2); // create the first step execution and the second one
       expect(stepExecutionRepository.updateStepExecutions).toHaveBeenCalledTimes(1);
     });
+
+    it('should sync workflow execution with latest from repository', async () => {
+      workflowExecutionRepository.getWorkflowExecutionById = jest.fn().mockResolvedValue({
+        id: 'test-workflow-execution-id',
+        status: ExecutionStatus.CANCELLED,
+        cancelledAt: '2025-08-05T20:02:00.000Z',
+        cancelledBy: 'user-123',
+      } as EsWorkflowExecution);
+      underTest.updateWorkflowExecution({
+        status: ExecutionStatus.SKIPPED,
+      });
+
+      await underTest.flush();
+
+      expect(workflowExecutionRepository.getWorkflowExecutionById).toHaveBeenCalledWith(
+        'test-workflow-execution-id',
+        undefined
+      );
+      expect(underTest.getWorkflowExecution()).toEqual({
+        id: 'test-workflow-execution-id',
+        status: ExecutionStatus.CANCELLED,
+        cancelledAt: '2025-08-05T20:02:00.000Z',
+        cancelledBy: 'user-123',
+      } as EsWorkflowExecution);
+    });
   });
 
   describe('getStepExecutionsByStepId', () => {