diff --git a/src/platform/packages/shared/kbn-workflows/common/template_expressions.ts b/src/platform/packages/shared/kbn-workflows/common/template_expressions.ts new file mode 100644 index 0000000000000..e82730af276dd --- /dev/null +++ b/src/platform/packages/shared/kbn-workflows/common/template_expressions.ts @@ -0,0 +1,46 @@ +/* + * 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". + */ + +/** + * Matches a *whole-value* Liquid template expression such as `"${{ event.messages }}"` — the + * only template form whose resolved value keeps its native type (array, object, number) + * instead of being stringified. + * + * The shape is deliberately narrow, and each restriction mirrors the runtime in + * `WorkflowTemplatingEngine.renderValueRecursively`: + * + * - **`$` is required.** The engine returns the raw evaluated value only for strings matching + * `startsWith('${{') && endsWith('}}')`. A bare `{{ expr }}` falls through to `renderString` + * and always comes back as a string, so it can never satisfy an array-typed param. + * - **No leading or trailing whitespace.** That runtime check does not trim. + * - **No inner `}}`.** `evaluateExpression` slices from the first `{{` to the last `}}`, so a + * concatenation like `"${{ a }}-${{ b }}"` would be parsed as the single invalid expression + * `a }}-${{ b` and throw at execution time. + * + * The invariant this encodes is one-directional: everything matched here is type-preserved at + * runtime. The engine itself accepts a superset, and those extra forms stringify or throw — + * which is exactly why callers should validate against this rather than re-deriving the shape. + * `template_expressions_runtime.test.ts` in `workflows_execution_engine` pins the invariant. + */ +export const WHOLE_VALUE_TEMPLATE_EXPRESSION_REGEX = /^\$\{\{(?:(?!\}\})[\s\S])*\}\}$/; + +/** + * Upper bound on a whole-value Liquid template expression accepted where a connector param + * declares a non-string type. This is a sanity bound on a single YAML scalar, not a defence + * against large workflow payloads — the branch it widens (the param's own array/object schema) + * carries whatever bounds that connector declared, and this value does not change them. + */ +export const TEMPLATE_EXPRESSION_MAX_LENGTH = 500; + +/** + * Returns true when `value` is a whole-value Liquid template expression whose resolved value + * keeps its native type at runtime. See {@link WHOLE_VALUE_TEMPLATE_EXPRESSION_REGEX}. + */ +export const isWholeValueTemplateExpression = (value: string): boolean => + WHOLE_VALUE_TEMPLATE_EXPRESSION_REGEX.test(value); diff --git a/src/platform/packages/shared/kbn-workflows/index.ts b/src/platform/packages/shared/kbn-workflows/index.ts index c40f3c1cd5de8..41c059fae9495 100644 --- a/src/platform/packages/shared/kbn-workflows/index.ts +++ b/src/platform/packages/shared/kbn-workflows/index.ts @@ -51,6 +51,7 @@ export * from './spec/deprecated_step_metadata'; export * from './types/latest'; export * from './types/utils'; export * from './common/constants'; +export * from './common/template_expressions'; export * from './common/validate_step_names'; export * from './common/workflows_events'; export type * from './common/event_trigger_replay'; diff --git a/src/platform/packages/shared/kbn-workflows/spec/builtin_step_definitions.test.ts b/src/platform/packages/shared/kbn-workflows/spec/builtin_step_definitions.test.ts index 04631787e7ff1..60da70e6644ff 100644 --- a/src/platform/packages/shared/kbn-workflows/spec/builtin_step_definitions.test.ts +++ b/src/platform/packages/shared/kbn-workflows/spec/builtin_step_definitions.test.ts @@ -65,6 +65,39 @@ describe('builtInStepDefinitions', () => { }); }); +describe('supportedExecutionModes', () => { + // These steps depend on Task Manager — to suspend and resume the workflow, or to schedule + // work outside the current execution — so they can never complete inside a single HTTP + // request. `validateSyncWorkflow` refuses to run a workflow containing one of them in sync + // mode; if an annotation here is dropped, that guard silently stops firing and a synchronous + // request hangs instead of failing fast. + const ASYNC_ONLY_IDS = [ + 'wait', + 'waitForInput', + 'waitForApproval', + 'workflow.execute', + 'workflow.executeAsync', + ]; + + it.each(ASYNC_ONLY_IDS)('"%s" is declared async-only', (id) => { + expect(getBuiltInStepDefinition(id)?.supportedExecutionModes).toEqual(['async']); + }); + + it('lists every async-only built-in — a new Task Manager-dependent step must be added here', () => { + const declaredAsyncOnly = builtInStepDefinitions + .filter((def) => def.supportedExecutionModes?.includes('sync') === false) + .map((def) => def.id); + expect(declaredAsyncOnly.sort()).toEqual([...ASYNC_ONLY_IDS].sort()); + }); + + it.each(['console', 'data.set', 'if', 'foreach'])( + '"%s" leaves the field unset, so it stays runnable in both modes', + (id) => { + expect(getBuiltInStepDefinition(id)?.supportedExecutionModes).toBeUndefined(); + } + ); +}); + describe('getBuiltInStepDefinition', () => { it('returns the definition for a known id', () => { const def = getBuiltInStepDefinition('if'); diff --git a/src/platform/packages/shared/kbn-workflows/spec/builtin_step_definitions.ts b/src/platform/packages/shared/kbn-workflows/spec/builtin_step_definitions.ts index cc602a757c44b..1aacd38cd6bb6 100644 --- a/src/platform/packages/shared/kbn-workflows/spec/builtin_step_definitions.ts +++ b/src/platform/packages/shared/kbn-workflows/spec/builtin_step_definitions.ts @@ -234,6 +234,7 @@ export const builtInStepDefinitions: BaseStepDefinition[] = [ }, { id: 'wait', + supportedExecutionModes: ['async'], label: 'Wait', description: 'Pause execution for a specified duration', category: StepCategory.FlowControl, @@ -267,6 +268,7 @@ export const builtInStepDefinitions: BaseStepDefinition[] = [ }, { id: 'workflow.execute', + supportedExecutionModes: ['async'], label: 'Execute Workflow', description: 'Execute another workflow and wait for it to complete', category: StepCategory.FlowControl, @@ -286,6 +288,7 @@ export const builtInStepDefinitions: BaseStepDefinition[] = [ }, { id: 'waitForInput', + supportedExecutionModes: ['async'], label: 'Wait For Input', description: 'Pause execution until external input is provided (human-in-the-loop)', category: StepCategory.FlowControl, @@ -328,6 +331,7 @@ export const builtInStepDefinitions: BaseStepDefinition[] = [ }, { id: 'waitForApproval', + supportedExecutionModes: ['async'], label: 'Wait For Approval', description: 'Pause execution until approval or rejection is received (human-in-the-loop)', category: StepCategory.FlowControl, @@ -357,6 +361,7 @@ export const builtInStepDefinitions: BaseStepDefinition[] = [ }, { id: 'workflow.executeAsync', + supportedExecutionModes: ['async'], label: 'Execute Workflow (Async)', description: 'Start another workflow and continue without waiting for completion', category: StepCategory.FlowControl, diff --git a/src/platform/packages/shared/kbn-workflows/spec/lib/generate_yaml_schema_from_connectors.test.ts b/src/platform/packages/shared/kbn-workflows/spec/lib/generate_yaml_schema_from_connectors.test.ts index 81a4f9c47490f..caa575a9eeccb 100644 --- a/src/platform/packages/shared/kbn-workflows/spec/lib/generate_yaml_schema_from_connectors.test.ts +++ b/src/platform/packages/shared/kbn-workflows/spec/lib/generate_yaml_schema_from_connectors.test.ts @@ -12,6 +12,7 @@ import { CONNECTOR_ID_MAX_LENGTH, type ConnectorContractUnion, generateYamlSchemaFromConnectors, + TEMPLATE_EXPRESSION_MAX_LENGTH, } from '../..'; const BASE_WORKFLOW = { @@ -218,4 +219,199 @@ describe('generateYamlSchemaFromConnectors', () => { expect(elapsed).toBeLessThan(500); }); }); + + describe('withTemplateStringSupport (array field widening)', () => { + const arrayConnector: ConnectorContractUnion = { + summary: 'Notifier', + description: null, + type: 'notify', + paramsSchema: z.object({ + recipients: z.array(z.string()), + subject: z.string(), + }), + outputSchema: z.unknown(), + }; + + const parse = (withValue: unknown) => + generateYamlSchemaFromConnectors([arrayConnector]).safeParse({ + ...BASE_WORKFLOW, + steps: [{ name: 'step', type: 'notify', with: withValue }], + }); + + it('accepts a ${{ expr }} template for an array param', () => { + expect( + parse({ recipients: '${{ workflow.inputs.recipients }}', subject: 'hi' }).success + ).toBe(true); + }); + + // Only `${{ … }}` survives as an array at runtime: WorkflowTemplatingEngine returns the raw + // evaluated value for that form only, and renders everything else to a string. Accepting any + // of the shapes below would let YAML be saved that hands the connector a string — or, for the + // multi-expression forms, throws "The provided expression is invalid" mid-execution, because + // evaluateExpression slices from the first `{{` to the last `}}`. + it.each([ + ['a plain string', 'not-a-template'], + ['a bare {{ expr }} template (renders to a string, not an array)', '{{ recipients }}'], + ['text before the expression', 'prefix-${{ expr }}'], + ['text after the expression', '${{ expr }}-suffix'], + ['two concatenated expressions', '${{ a }}-${{ b }}'], + ['an expression with literal text between two others', '${{ a }} literal ${{ b }}'], + ['leading whitespace (the runtime check does not trim)', ' ${{ expr }}'], + ['trailing whitespace (the runtime check does not trim)', '${{ expr }} '], + ])('rejects %s for an array param', (_label, recipients) => { + expect(parse({ recipients, subject: 'hi' }).success).toBe(false); + }); + + it('rejects a template string exceeding TEMPLATE_EXPRESSION_MAX_LENGTH', () => { + const long = `\${{ ${'x'.repeat(TEMPLATE_EXPRESSION_MAX_LENGTH)} }}`; + expect(parse({ recipients: long, subject: 'hi' }).success).toBe(false); + }); + + it('does not widen non-array fields — string params remain string-only', () => { + expect(parse({ recipients: ['a@b.com'], subject: '{{ not-widened }}' }).success).toBe(true); + // A real Liquid expression is still valid as a string value, but a plain array is not + expect(parse({ recipients: ['a@b.com'], subject: ['array', 'not', 'ok'] }).success).toBe( + false + ); + }); + + it('widens optional array params and preserves optionality', () => { + const connector: ConnectorContractUnion = { + summary: 'Opt', + description: null, + type: 'opt.step', + paramsSchema: z.object({ tags: z.array(z.string()).optional() }), + outputSchema: z.unknown(), + }; + const schema = generateYamlSchemaFromConnectors([connector]); + // template string accepted + expect( + schema.safeParse({ + ...BASE_WORKFLOW, + steps: [{ name: 's', type: 'opt.step', with: { tags: '${{ workflow.inputs.tags }}' } }], + }).success + ).toBe(true); + // omitting the optional field is still valid + expect( + schema.safeParse({ + ...BASE_WORKFLOW, + steps: [{ name: 's', type: 'opt.step', with: {} }], + }).success + ).toBe(true); + }); + + it('widens default-wrapped array params without making them required', () => { + const connector: ConnectorContractUnion = { + summary: 'Def', + description: null, + type: 'def.step', + // Mirrors a real shipped connector: InferenceRerankParamsSchema declares + // `input: z.array(z.string()).default([])` as a top-level param. + paramsSchema: z.object({ tags: z.array(z.string()).default([]), query: z.string() }), + outputSchema: z.unknown(), + }; + const schema = generateYamlSchemaFromConnectors([connector]); + expect( + schema.safeParse({ + ...BASE_WORKFLOW, + steps: [ + { + name: 's', + type: 'def.step', + with: { tags: '${{ workflow.inputs.tags }}', query: 'q' }, + }, + ], + }).success + ).toBe(true); + + // Widening must not strip `.default()`. If it did, every existing workflow that omits a + // defaulted array param would stop validating — including on update, since this schema + // gates persistence and not just editor feedback. + const omitted = schema.safeParse({ + ...BASE_WORKFLOW, + steps: [{ name: 's', type: 'def.step', with: { query: 'q' } }], + }); + expect(omitted.success).toBe(true); + expect(omitted.data).toMatchObject({ steps: [{ with: { tags: [] } }] }); + }); + + it('widens array params wrapped in both .optional() and .default()', () => { + const connector: ConnectorContractUnion = { + summary: 'Both', + description: null, + type: 'both.step', + paramsSchema: z.object({ tags: z.array(z.string()).optional().default([]) }), + outputSchema: z.unknown(), + }; + const schema = generateYamlSchemaFromConnectors([connector]); + // Stacked wrappers must still be unwrapped down to the array, otherwise the field is + // silently skipped and the template string is reported as a type error. + expect( + schema.safeParse({ + ...BASE_WORKFLOW, + steps: [{ name: 's', type: 'both.step', with: { tags: '${{ workflow.inputs.tags }}' } }], + }).success + ).toBe(true); + }); + + it('preserves the object unknownKeys policy of the original paramsSchema', () => { + // A strict paramsSchema should still reject unknown keys after widening. + const connector: ConnectorContractUnion = { + summary: 'Strict', + description: null, + type: 'strict.step', + paramsSchema: z.strictObject({ ids: z.array(z.string()) }), + outputSchema: z.unknown(), + }; + const schema = generateYamlSchemaFromConnectors([connector]); + expect( + schema.safeParse({ + ...BASE_WORKFLOW, + steps: [ + { + name: 's', + type: 'strict.step', + with: { ids: '${{ workflow.inputs.ids }}', unknown_key: 'bad' }, + }, + ], + }).success + ).toBe(false); + }); + + it('does not throw for a connector whose paramsSchema has object-level refinements', () => { + const connector: ConnectorContractUnion = { + summary: 'Refined', + description: null, + type: 'refined.step', + paramsSchema: z + .object({ ids: z.array(z.string()), name: z.string() }) + .refine((v) => v.ids.length > 0, 'ids must not be empty'), + outputSchema: z.unknown(), + }; + // Schema construction must not throw even though paramsSchema has a refinement. + expect(() => generateYamlSchemaFromConnectors([connector])).not.toThrow(); + + const schema = generateYamlSchemaFromConnectors([connector]); + // Template string still accepted for the array field. + expect( + schema.safeParse({ + ...BASE_WORKFLOW, + steps: [ + { + name: 's', + type: 'refined.step', + with: { ids: '${{ workflow.inputs.ids }}', name: 'x' }, + }, + ], + }).success + ).toBe(true); + // The object-level refinement is preserved: an empty ids array fails. + expect( + schema.safeParse({ + ...BASE_WORKFLOW, + steps: [{ name: 's', type: 'refined.step', with: { ids: [], name: 'x' } }], + }).success + ).toBe(false); + }); + }); }); diff --git a/src/platform/packages/shared/kbn-workflows/spec/lib/generate_yaml_schema_from_connectors.ts b/src/platform/packages/shared/kbn-workflows/spec/lib/generate_yaml_schema_from_connectors.ts index b916eb908faac..2287613f7952c 100644 --- a/src/platform/packages/shared/kbn-workflows/spec/lib/generate_yaml_schema_from_connectors.ts +++ b/src/platform/packages/shared/kbn-workflows/spec/lib/generate_yaml_schema_from_connectors.ts @@ -9,6 +9,10 @@ import { z } from '@kbn/zod/v4'; import { CONNECTOR_ID_MAX_LENGTH } from '../../common/constants'; +import { + TEMPLATE_EXPRESSION_MAX_LENGTH, + WHOLE_VALUE_TEMPLATE_EXPRESSION_REGEX, +} from '../../common/template_expressions'; import type { ConnectorContractUnion } from '../../types/v1'; import { getDeprecatedStepMessage, getStepDeprecationInfo } from '../deprecated_step_metadata'; import { KIBANA_TYPE_ALIASES } from '../kibana/aliases'; @@ -166,6 +170,72 @@ function hasNoRequiredFields(schema: z.ZodType): boolean { ); } +/** + * Only whole-value `${{ … }}` expressions are accepted — not arbitrary strings, and not the + * bare `{{ … }}` form, which the templating engine always renders to a string and so can never + * satisfy an array param. See WHOLE_VALUE_TEMPLATE_EXPRESSION_REGEX for the full rationale. + */ +const LIQUID_TEMPLATE_SCHEMA = z + .string() + .regex(WHOLE_VALUE_TEMPLATE_EXPRESSION_REGEX) + .max(TEMPLATE_EXPRESSION_MAX_LENGTH); + +/** + * Widens top-level array fields of a connector params schema to also accept a whole-value + * Liquid template expression like `"${{ event.messages }}"`, so that passing a templated value + * where an array is declared is not reported as a type error. + * + * This is not editor-only: the same generated schema is the server-side gate for workflow + * create/update (`workflow_crud_service` / `workflow_validation_service` in + * workflows_management), so it decides what can be *persisted*, not just what Monaco + * underlines. Connector params are not re-validated against `paramsSchema` at execution time, + * which is why the accepted form is restricted to the one the templating engine is guaranteed + * to resolve back to a native array. + * + * Only the direct children of the params schema (not nested objects) are widened, to avoid + * disturbing deeply nested schemas such as the ES API's MappingTypeMapping. + */ +function withTemplateStringSupport(paramsSchema: z.ZodType): z.ZodType { + if (!(paramsSchema instanceof z.ZodObject)) { + return paramsSchema; + } + const modifications: Record = {}; + for (const [key, rawValue] of Object.entries(paramsSchema.shape as Record)) { + let value: z.ZodType = rawValue; + let isOptional = false; + let hasDefault = false; + let defaultValue: unknown; + + // `.optional()` and `.default()` can be applied in either order and stacked, so unwrap + // until the underlying type is reached, remembering what has to be re-applied afterwards. + while (value instanceof z.ZodOptional || value instanceof z.ZodDefault) { + if (value instanceof z.ZodOptional) { + isOptional = true; + } else { + hasDefault = true; + defaultValue = value.def.defaultValue; + } + value = value.unwrap() as z.ZodType; + } + + if (value instanceof z.ZodArray) { + const widened: z.ZodType = z.union([LIQUID_TEMPLATE_SCHEMA, value]); + // Re-apply the wrappers that were stripped above. Dropping `.default()` here would turn a + // defaulted param into a required one and break workflows that legitimately omit it. + const withOptionality = isOptional ? widened.optional() : widened; + modifications[key] = hasDefault ? withOptionality.default(defaultValue) : withOptionality; + } + } + if (Object.keys(modifications).length === 0) { + return paramsSchema; + } + // safeExtend preserves object-level refinements and the unknownKeys policy (strict/passthrough), + // unlike extend() which throws when the schema contains refinements. + // `modifications` is built as a mutable record; ZodRawShape is the readonly shape safeExtend + // expects, so the cast only relaxes mutability. + return paramsSchema.safeExtend(modifications as z.ZodRawShape); +} + function generateStepSchemaForConnector( connector: ConnectorContractUnion, stepSchema: z.ZodType, @@ -179,11 +249,13 @@ function generateStepSchemaForConnector( connector.hasConnectorId === 'required' ? connectorId : connectorId.optional(); } + const templateAwareParamsSchema = withTemplateStringSupport(connector.paramsSchema); + // If all params are optional (or there are none), `with` itself should be optional so users // don't have to write an empty `with: {}` block for steps that need no inputs. const withSchema = hasNoRequiredFields(connector.paramsSchema) - ? connector.paramsSchema.optional() - : connector.paramsSchema; + ? templateAwareParamsSchema.optional() + : templateAwareParamsSchema; return BaseConnectorStepSchema.extend({ type: connector.description diff --git a/src/platform/packages/shared/kbn-workflows/spec/step_definition_types.ts b/src/platform/packages/shared/kbn-workflows/spec/step_definition_types.ts index 480080b30d94c..1cf9d014a1a49 100644 --- a/src/platform/packages/shared/kbn-workflows/spec/step_definition_types.ts +++ b/src/platform/packages/shared/kbn-workflows/spec/step_definition_types.ts @@ -25,6 +25,8 @@ export enum StepCategory { export const StepCategories = Object.values(StepCategory) as StepCategory[]; +export type StepExecutionMode = 'async' | 'sync'; + /** * Documentation information for a workflow step. */ @@ -124,4 +126,22 @@ export interface BaseStepDefinition< * suggested for new workflows. */ deprecation?: StepDeprecationInfo; + + /** + * Execution modes supported by this step. + * + * Declare `['async']` for any step that cannot complete within a single request because it + * depends on Task Manager — either to suspend and resume the workflow (`wait`, the + * human-in-the-loop steps) or to schedule work outside the current execution + * (`workflow.execute`, `workflow.executeAsync`). `validateSyncWorkflow` rejects such steps + * when a workflow is run synchronously. + * + * Omitting the field means both modes are supported. That default is deliberately + * permissive: the several hundred connector-backed steps are plain request/response calls + * that are safe to run inline, and requiring each to opt in would make the field noise + * rather than signal. The trade-off is that a *new* Task Manager-dependent step is + * sync-eligible until someone marks it, so add the annotation in the same change that + * introduces the step. + */ + supportedExecutionModes?: readonly [StepExecutionMode, ...StepExecutionMode[]]; } diff --git a/src/platform/plugins/shared/workflows_execution_engine/server/config.ts b/src/platform/plugins/shared/workflows_execution_engine/server/config.ts index 84049eb222294..c1444116b3ab1 100644 --- a/src/platform/plugins/shared/workflows_execution_engine/server/config.ts +++ b/src/platform/plugins/shared/workflows_execution_engine/server/config.ts @@ -73,6 +73,62 @@ const configSchema = schema.object({ hitlExternalResume: schema.object({ enabled: schema.boolean({ defaultValue: true }), }), + /** + * Synchronous workflow execution path — used by the inference anonymization pipeline + * to run workflows inline within an HTTP request without persisting execution state to + * Elasticsearch. Must be enabled alongside `xpack.inference.anonymization.workflowDriven`. + */ + syncExecution: schema.object({ + /** + * Master switch for the synchronous execution path. Must be set to true alongside + * `xpack.inference.anonymization.workflowDriven: true` to enable workflow-driven + * PII anonymization. Defaults to false so the path is inert until explicitly activated. + */ + enabled: schema.boolean({ defaultValue: false }), + /** + * Maximum wall-clock time allowed for a single synchronous workflow execution. + * When the deadline is reached the internal AbortController is aborted, which + * propagates cancellation through the execution loop and cancels any in-flight + * I/O (e.g. LLM streaming). Around-hook workflows that wrap LLM calls can take + * several minutes; the 10-minute default reflects this. The natural upper bound + * is the HTTP request timeout — if the client disconnects, that signal propagates + * first. Callers may impose a shorter deadline via ExecuteWorkflowOptions.abortSignal. + */ + maxDurationMs: schema.number({ defaultValue: 600_000, min: 1_000 }), + }), + /** + * Configuration for the synchronous-execution log drain. + * The drain buffers workflow event-log writes in memory and flushes them to + * Elasticsearch out-of-band, keeping the sync execution hot path free of + * inline ES round-trips. + * + * Defaults are sized for ~1000 concurrent sync req/s × 4 step events each + * (~4000 events/sec inbound). Tune these values to your deployment's actual + * load if the `kibana.workflows.sync_log_drain.events.dropped` metric is non-zero. + */ + syncLogDrain: schema.object({ + /** + * When false, the drain is disabled and sync executions write event-log + * entries to Elasticsearch inline (same as async executions). Useful for + * A/B comparisons to measure the drain's overhead vs. direct ES writes. + */ + enabled: schema.boolean({ defaultValue: true }), + /** + * How often (ms) the background timer flushes buffered events to ES. + * Lower values reduce event latency but increase write frequency. + */ + intervalMs: schema.number({ defaultValue: 500, min: 100 }), + /** + * Maximum number of events held in the in-memory buffer before drop-oldest + * kicks in. ~5 s of inbound capacity at the default rate. + */ + maxQueue: schema.number({ defaultValue: 20000, min: 1000 }), + /** + * Maximum number of events written to ES in a single drain tick. + * Should be ≥ intervalMs/1000 × peak inbound rate so one tick clears backlog. + */ + maxBatch: schema.number({ defaultValue: 4000, min: 100 }), + }), }); export type EventTriggersConfig = TypeOf; diff --git a/src/platform/plugins/shared/workflows_execution_engine/server/execution_functions/__mock__/context_dependencies.ts b/src/platform/plugins/shared/workflows_execution_engine/server/execution_functions/__mock__/context_dependencies.ts index e2ead6ffaf733..4444f17180a20 100644 --- a/src/platform/plugins/shared/workflows_execution_engine/server/execution_functions/__mock__/context_dependencies.ts +++ b/src/platform/plugins/shared/workflows_execution_engine/server/execution_functions/__mock__/context_dependencies.ts @@ -32,5 +32,7 @@ export const mockContextDependencies = () => ({ }, collectQueueMetrics: false, hitlExternalResume: { enabled: true }, + syncExecution: { enabled: true, maxDurationMs: 60_000 }, + syncLogDrain: { enabled: true, intervalMs: 500, maxQueue: 20000, maxBatch: 4000 }, }, }); diff --git a/src/platform/plugins/shared/workflows_execution_engine/server/execution_functions/execute_workflow_sync.test.ts b/src/platform/plugins/shared/workflows_execution_engine/server/execution_functions/execute_workflow_sync.test.ts new file mode 100644 index 0000000000000..5324fe3df61ab --- /dev/null +++ b/src/platform/plugins/shared/workflows_execution_engine/server/execution_functions/execute_workflow_sync.test.ts @@ -0,0 +1,232 @@ +/* + * 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". + */ + +import type { KibanaRequest, Logger } from '@kbn/core/server'; +import { + type EsWorkflowExecution, + ExecutionStatus, + type WorkflowExecutionEngineModel, +} from '@kbn/workflows'; +import { executeWorkflowSync } from './execute_workflow_sync'; +import { runWorkflowSync } from './run_workflow_sync'; +import { buildWorkflowExecutionDocument } from '../lib/build_workflow_execution_document'; +import { getAuthenticatedUser } from '../lib/get_user'; +import { validateWorkflowInputs } from '../lib/validate_workflow_inputs'; +import { InMemoryExecutionPersistence } from '../repositories/execution_persistence'; +import type { ExecuteWorkflowOptions, WorkflowsExecutionEnginePluginStart } from '../types'; +import type { ContextDependencies } from '../workflow_context_manager/types'; + +jest.mock('./run_workflow_sync'); +jest.mock('../lib/build_workflow_execution_document'); +jest.mock('../lib/get_user'); +jest.mock('../lib/validate_workflow_inputs'); +jest.mock('../repositories/execution_persistence'); + +const WORKFLOW_ID = 'wf-1'; +const EXECUTION_ID = 'exec-1'; +const SPACE_ID = 'default'; + +const baseExecution: EsWorkflowExecution = { + id: EXECUTION_ID, + spaceId: SPACE_ID, + workflowId: WORKFLOW_ID, + isTestRun: false, + status: ExecutionStatus.COMPLETED, + context: {}, + workflowDefinition: { + version: '1', + name: 'Test', + enabled: true, + triggers: [], + steps: [], + }, + yaml: '', + scopeStack: [], + createdAt: '2026-01-01T00:00:00.000Z', + error: null, + startedAt: '2026-01-01T00:00:00.000Z', + finishedAt: '2026-01-01T00:00:01.000Z', + cancelRequested: false, + duration: 1000, +}; + +const makeWorkflow = (overrides: Partial = {}) => + ({ + id: WORKFLOW_ID, + isEphemeral: true, + yaml: '', + workflowDefinition: baseExecution.workflowDefinition, + ...overrides, + } as WorkflowExecutionEngineModel); + +const makeOptions = (overrides: Partial = {}): ExecuteWorkflowOptions => ({ + ...overrides, +}); + +const makeDependencies = (overrides: Partial = {}): ContextDependencies => ({ + coreStart: { + security: {}, + elasticsearch: { client: {} }, + } as unknown as ContextDependencies['coreStart'], + config: { + syncExecution: { enabled: true, maxDurationMs: 5_000 }, + eventDriven: { maxChainDepth: 5 }, + } as unknown as ContextDependencies['config'], + actions: {} as unknown as ContextDependencies['actions'], + taskManager: {} as unknown as ContextDependencies['taskManager'], + workflowsExtensions: {} as unknown as ContextDependencies['workflowsExtensions'], + cloudSetup: undefined, + ...overrides, +}); + +describe('executeWorkflowSync', () => { + const mockGetEngine = jest.fn().mockResolvedValue({} as WorkflowsExecutionEnginePluginStart); + const mockLogger = { + debug: jest.fn(), + info: jest.fn(), + warn: jest.fn(), + error: jest.fn(), + } as unknown as Logger; + const mockRequest = {} as KibanaRequest; + + beforeEach(() => { + jest.clearAllMocks(); + (buildWorkflowExecutionDocument as jest.Mock).mockResolvedValue(baseExecution); + (getAuthenticatedUser as jest.Mock).mockResolvedValue({ username: 'test-user' }); + (validateWorkflowInputs as jest.Mock).mockResolvedValue(true); + (runWorkflowSync as jest.Mock).mockResolvedValue({ ...baseExecution, context: {} }); + (InMemoryExecutionPersistence as jest.Mock).mockImplementation(() => ({ + getWorkflowExecutionById: jest.fn().mockResolvedValue(baseExecution), + })); + }); + + it('throws when a non-ephemeral workflow is disabled', async () => { + const deps = makeDependencies({ + workflowRepository: { + isWorkflowEnabled: jest.fn().mockResolvedValue(false), + } as unknown as ContextDependencies['workflowRepository'], + }); + + await expect( + executeWorkflowSync({ + workflow: makeWorkflow({ isEphemeral: false }), + context: { spaceId: SPACE_ID }, + request: mockRequest, + options: makeOptions(), + logger: mockLogger, + dependencies: deps, + getWorkflowsExecutionEngine: mockGetEngine, + }) + ).rejects.toThrow(`Workflow is disabled: ${WORKFLOW_ID}`); + }); + + it('returns FAILED with error when input validation fails', async () => { + const failedExecution = { + ...baseExecution, + error: { type: 'InputValidationError', message: 'bad input' }, + }; + (validateWorkflowInputs as jest.Mock).mockResolvedValue(false); + (InMemoryExecutionPersistence as jest.Mock).mockImplementation(() => ({ + getWorkflowExecutionById: jest.fn().mockResolvedValue(failedExecution), + })); + + const result = await executeWorkflowSync({ + workflow: makeWorkflow(), + context: { spaceId: SPACE_ID }, + request: mockRequest, + options: makeOptions(), + logger: mockLogger, + dependencies: makeDependencies(), + getWorkflowsExecutionEngine: mockGetEngine, + }); + + expect(result.result).toBeDefined(); + expect(result.result!.status).toBe(ExecutionStatus.FAILED); + expect(result.result!.error).toEqual(failedExecution.error); + expect(runWorkflowSync).not.toHaveBeenCalled(); + }); + + it('always clears the timeout in the finally block', async () => { + const clearTimeoutSpy = jest.spyOn(global, 'clearTimeout'); + + await executeWorkflowSync({ + workflow: makeWorkflow(), + context: { spaceId: SPACE_ID }, + request: mockRequest, + options: makeOptions(), + logger: mockLogger, + dependencies: makeDependencies(), + getWorkflowsExecutionEngine: mockGetEngine, + }); + + expect(clearTimeoutSpy).toHaveBeenCalled(); + clearTimeoutSpy.mockRestore(); + }); + + it('relays a pre-aborted caller signal to the internal abort controller', async () => { + const externalController = new AbortController(); + externalController.abort(new Error('caller cancelled')); + + let capturedAbortController: AbortController | undefined; + (runWorkflowSync as jest.Mock).mockImplementation(({ abortController }) => { + capturedAbortController = abortController; + return Promise.resolve({ ...baseExecution, context: {} }); + }); + + await executeWorkflowSync({ + workflow: makeWorkflow(), + context: { spaceId: SPACE_ID }, + request: mockRequest, + options: makeOptions({ abortSignal: externalController.signal }), + logger: mockLogger, + dependencies: makeDependencies(), + getWorkflowsExecutionEngine: mockGetEngine, + }); + + expect(capturedAbortController?.signal.aborted).toBe(true); + }); + + it('throws when output is a non-object, non-null value', async () => { + (runWorkflowSync as jest.Mock).mockResolvedValue({ + ...baseExecution, + context: { output: 'a string' }, + }); + + await expect( + executeWorkflowSync({ + workflow: makeWorkflow(), + context: { spaceId: SPACE_ID }, + request: mockRequest, + options: makeOptions(), + logger: mockLogger, + dependencies: makeDependencies(), + getWorkflowsExecutionEngine: mockGetEngine, + }) + ).rejects.toThrow('Synchronous workflow output must be an object'); + }); + + it('returns result without output when execution context.output is null', async () => { + (runWorkflowSync as jest.Mock).mockResolvedValue({ + ...baseExecution, + context: { output: null }, + }); + + const result = await executeWorkflowSync({ + workflow: makeWorkflow(), + context: { spaceId: SPACE_ID }, + request: mockRequest, + options: makeOptions(), + logger: mockLogger, + dependencies: makeDependencies(), + getWorkflowsExecutionEngine: mockGetEngine, + }); + + expect(result.result).not.toHaveProperty('output'); + }); +}); diff --git a/src/platform/plugins/shared/workflows_execution_engine/server/execution_functions/execute_workflow_sync.ts b/src/platform/plugins/shared/workflows_execution_engine/server/execution_functions/execute_workflow_sync.ts new file mode 100644 index 0000000000000..3c6ffd68e4b10 --- /dev/null +++ b/src/platform/plugins/shared/workflows_execution_engine/server/execution_functions/execute_workflow_sync.ts @@ -0,0 +1,174 @@ +/* + * 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". + */ + +import type { KibanaRequest, Logger } from '@kbn/core/server'; +import { ExecutionStatus } from '@kbn/workflows'; +import type { EsWorkflowExecution, WorkflowExecutionEngineModel } from '@kbn/workflows'; +import { runWorkflowSync } from './run_workflow_sync'; +import { buildWorkflowExecutionDocument } from '../lib/build_workflow_execution_document'; +import { getAuthenticatedUser } from '../lib/get_user'; +import { validateWorkflowInputs } from '../lib/validate_workflow_inputs'; +import { InMemoryExecutionPersistence } from '../repositories/execution_persistence'; +import type { + ExecuteWorkflowOptions, + ExecuteWorkflowResponse, + WorkflowsExecutionEnginePluginStart, +} from '../types'; +import type { ContextDependencies } from '../workflow_context_manager/types'; + +const isRecord = (value: unknown): value is Record => + typeof value === 'object' && value !== null && !Array.isArray(value); + +const getSynchronousWorkflowOutput = (output: unknown): Record | undefined => { + if (output === undefined || output === null) { + return undefined; + } + if (isRecord(output)) { + return output; + } + throw new Error('Synchronous workflow output must be an object'); +}; + +export const executeWorkflowSync = async ({ + workflow, + context, + request, + options, + logger, + dependencies, + getWorkflowsExecutionEngine, +}: { + workflow: WorkflowExecutionEngineModel; + context: Record; + request: KibanaRequest; + options: ExecuteWorkflowOptions; + logger: Logger; + dependencies: ContextDependencies; + getWorkflowsExecutionEngine: () => Promise; +}): Promise => { + const { coreStart, workflowRepository, config } = dependencies; + + const spaceId = typeof context.spaceId === 'string' ? context.spaceId : 'default'; + + if (!workflow.isEphemeral && workflowRepository) { + const stillEnabled = await workflowRepository.isWorkflowEnabled(workflow.id, spaceId, { + includeGlobal: true, + }); + if (!stillEnabled) { + throw new Error(`Workflow is disabled: ${workflow.id}. Enable the workflow to run it.`); + } + } + + const authenticatedUser = await getAuthenticatedUser( + request, + coreStart.security, + coreStart.elasticsearch.client + ); + + const workflowExecution = await buildWorkflowExecutionDocument({ + workflow, + context, + defaultTriggeredBy: 'manual', + authenticatedUser, + now: new Date(), + maxEventChainDepth: config.eventDriven.maxChainDepth, + getConcurrencyGroupKey: () => null, + }); + + if (options.executionId) { + workflowExecution.id = options.executionId; + } + + if (!workflowExecution.workflowDefinition) { + throw new Error('Synchronous workflow execution requires a workflow definition'); + } + + const syncWorkflowExecution: EsWorkflowExecution = { + ...workflowExecution, + isTestRun: workflowExecution.isTestRun ?? false, + status: workflowExecution.status ?? ExecutionStatus.PENDING, + context: workflowExecution.context ?? context, + workflowDefinition: workflowExecution.workflowDefinition, + yaml: workflowExecution.yaml ?? workflow.yaml, + scopeStack: workflowExecution.scopeStack ?? [], + error: workflowExecution.error ?? null, + startedAt: workflowExecution.startedAt ?? workflowExecution.createdAt, + finishedAt: workflowExecution.finishedAt ?? new Date().toISOString(), + cancelRequested: workflowExecution.cancelRequested ?? false, + duration: workflowExecution.duration ?? 0, + }; + + const syncExecutionPersistence = new InMemoryExecutionPersistence(syncWorkflowExecution); + const inputsValid = await validateWorkflowInputs( + syncWorkflowExecution, + syncExecutionPersistence, + logger, + coreStart, + { ...dependencies, capabilities: options.capabilities } + ); + if (!inputsValid) { + const failedExecution = await syncExecutionPersistence.getWorkflowExecutionById( + syncWorkflowExecution.id, + spaceId + ); + return { + workflowExecutionId: syncWorkflowExecution.id, + result: { + status: ExecutionStatus.FAILED, + ...(failedExecution?.error ? { error: failedExecution.error } : {}), + }, + }; + } + + const abortController = new AbortController(); + const abort = () => abortController.abort(options.abortSignal?.reason); + options.abortSignal?.addEventListener('abort', abort, { once: true }); + if (options.abortSignal?.aborted) { + abort(); + } + + const maxDurationMs = config.syncExecution.maxDurationMs; + const timeoutId = setTimeout( + () => + abortController.abort( + new Error(`Synchronous workflow execution timed out after ${maxDurationMs}ms`) + ), + maxDurationMs + ); + + try { + const workflowsExecutionEngine = await getWorkflowsExecutionEngine(); + const result = await runWorkflowSync({ + workflowExecution: syncWorkflowExecution, + request, + abortController, + logger, + config, + dependencies: { ...dependencies, capabilities: options.capabilities }, + workflowsExecutionEngine, + workflowExecutionRepository: syncExecutionPersistence, + stepExecutionRepository: syncExecutionPersistence, + }); + + const output = getSynchronousWorkflowOutput(result.context?.output); + const failureError = + result.status === ExecutionStatus.FAILED && result.error ? result.error : undefined; + return { + workflowExecutionId: result.id, + result: { + status: result.status, + ...(output ? { output } : {}), + ...(failureError ? { error: failureError } : {}), + }, + }; + } finally { + clearTimeout(timeoutId); + options.abortSignal?.removeEventListener('abort', abort); + } +}; diff --git a/src/platform/plugins/shared/workflows_execution_engine/server/execution_functions/execution_functions_test_utils.ts b/src/platform/plugins/shared/workflows_execution_engine/server/execution_functions/execution_functions_test_utils.ts index 33e3168f4270a..2d33189b2908a 100644 --- a/src/platform/plugins/shared/workflows_execution_engine/server/execution_functions/execution_functions_test_utils.ts +++ b/src/platform/plugins/shared/workflows_execution_engine/server/execution_functions/execution_functions_test_utils.ts @@ -29,6 +29,8 @@ export const createMockWorkflowExecutionEngineConfig = (): WorkflowsExecutionEng eviction: { minPayloadSize: new ByteSizeValue(10 * 1024) }, collectQueueMetrics: false, hitlExternalResume: { enabled: true }, + syncExecution: { enabled: false, maxDurationMs: 60_000 }, + syncLogDrain: { enabled: true, intervalMs: 500, maxQueue: 20000, maxBatch: 4000 }, }); export const createMockLogger = (): Logger => diff --git a/src/platform/plugins/shared/workflows_execution_engine/server/execution_functions/index.ts b/src/platform/plugins/shared/workflows_execution_engine/server/execution_functions/index.ts index aa40a42c73fa2..d217d9357ca94 100644 --- a/src/platform/plugins/shared/workflows_execution_engine/server/execution_functions/index.ts +++ b/src/platform/plugins/shared/workflows_execution_engine/server/execution_functions/index.ts @@ -9,6 +9,7 @@ export { setupDependencies } from './setup_dependencies'; export { runWorkflow } from './run_workflow'; +export { runWorkflowSync } from './run_workflow_sync'; export { resumeWorkflow } from './resume_workflow'; export { cancelWorkflow } from './cancel_workflow'; export { diff --git a/src/platform/plugins/shared/workflows_execution_engine/server/execution_functions/resume_workflow.test.ts b/src/platform/plugins/shared/workflows_execution_engine/server/execution_functions/resume_workflow.test.ts index 86f1b3cbe186a..1cd74ae73e31e 100644 --- a/src/platform/plugins/shared/workflows_execution_engine/server/execution_functions/resume_workflow.test.ts +++ b/src/platform/plugins/shared/workflows_execution_engine/server/execution_functions/resume_workflow.test.ts @@ -310,17 +310,17 @@ describe('resumeWorkflow', () => { it('calls setupDependencies with default workflowsExecutionEngine', async () => { await resumeWorkflowWithDefaults(); - expect(mockSetupDependencies).toHaveBeenCalledWith( + expect(mockSetupDependencies).toHaveBeenCalledWith({ workflowRunId, spaceId, logger, - mockConfig, + config: mockConfig, dependencies, workflowExecutionRepository, stepExecutionRepository, fakeRequest, - mockWorkflowExecutionEngine - ); + workflowsExecutionEngine: mockWorkflowExecutionEngine, + }); }); it('forwards workflowsExecutionEngine to setupDependencies when provided', async () => { @@ -335,17 +335,17 @@ describe('resumeWorkflow', () => { await resumeWorkflowWithDefaults({ workflowsExecutionEngine }); - expect(mockSetupDependencies).toHaveBeenCalledWith( + expect(mockSetupDependencies).toHaveBeenCalledWith({ workflowRunId, spaceId, logger, - mockConfig, + config: mockConfig, dependencies, workflowExecutionRepository, stepExecutionRepository, fakeRequest, - workflowsExecutionEngine - ); + workflowsExecutionEngine, + }); }); it('calls workflowRuntime.resume then workflowExecutionLoop in order', async () => { diff --git a/src/platform/plugins/shared/workflows_execution_engine/server/execution_functions/resume_workflow.ts b/src/platform/plugins/shared/workflows_execution_engine/server/execution_functions/resume_workflow.ts index 86bc83d9259e0..dc8b131585930 100644 --- a/src/platform/plugins/shared/workflows_execution_engine/server/execution_functions/resume_workflow.ts +++ b/src/platform/plugins/shared/workflows_execution_engine/server/execution_functions/resume_workflow.ts @@ -57,17 +57,17 @@ export async function resumeWorkflow({ }): Promise<{ idleTimeoutResumeAt?: Date }> { let setupResult: Awaited>; try { - setupResult = await setupDependencies( + setupResult = await setupDependencies({ workflowRunId, spaceId, logger, config, dependencies, + fakeRequest, + workflowsExecutionEngine, workflowExecutionRepository, stepExecutionRepository, - fakeRequest, - workflowsExecutionEngine - ); + }); } catch (error) { // The graph could not be built — a permanent author error (the parallel // branch-body constraints, normally caught in the editor by validateGraphBuild diff --git a/src/platform/plugins/shared/workflows_execution_engine/server/execution_functions/run_workflow.test.ts b/src/platform/plugins/shared/workflows_execution_engine/server/execution_functions/run_workflow.test.ts index b16a31bee9e1f..001b78a69d2a3 100644 --- a/src/platform/plugins/shared/workflows_execution_engine/server/execution_functions/run_workflow.test.ts +++ b/src/platform/plugins/shared/workflows_execution_engine/server/execution_functions/run_workflow.test.ts @@ -143,17 +143,17 @@ describe('runWorkflow', () => { it('calls setupDependencies with all expected arguments', async () => { await runWorkflowWithDefaults(); - expect(mockSetupDependencies).toHaveBeenCalledWith( + expect(mockSetupDependencies).toHaveBeenCalledWith({ workflowRunId, spaceId, logger, - mockConfig, + config: mockConfig, dependencies, workflowExecutionRepository, stepExecutionRepository, fakeRequest, - mockWorkflowExecutionEngine - ); + workflowsExecutionEngine: mockWorkflowExecutionEngine, + }); }); it('calls workflowRuntime.start then workflowExecutionLoop in order', async () => { diff --git a/src/platform/plugins/shared/workflows_execution_engine/server/execution_functions/run_workflow.ts b/src/platform/plugins/shared/workflows_execution_engine/server/execution_functions/run_workflow.ts index d28b74c49d893..48147a90f5b37 100644 --- a/src/platform/plugins/shared/workflows_execution_engine/server/execution_functions/run_workflow.ts +++ b/src/platform/plugins/shared/workflows_execution_engine/server/execution_functions/run_workflow.ts @@ -62,21 +62,22 @@ export async function runWorkflow({ workflowExecutionRepository: WorkflowExecutionRepository; stepExecutionRepository: StepExecutionRepository; }): Promise { + apm.currentTransaction?.setLabel('execution_mode', 'async'); // Span for setup/initialization phase const setupSpan = apm.startSpan('workflow setup', 'workflow', 'setup'); let setupResult: Awaited>; try { - setupResult = await setupDependencies( + setupResult = await setupDependencies({ workflowRunId, spaceId, logger, config, dependencies, + fakeRequest, + workflowsExecutionEngine, workflowExecutionRepository, stepExecutionRepository, - fakeRequest, - workflowsExecutionEngine - ); + }); } catch (error) { // The graph could not be built — a permanent author error (the parallel // branch-body constraints, normally caught in the editor by validateGraphBuild diff --git a/src/platform/plugins/shared/workflows_execution_engine/server/execution_functions/run_workflow_sync.test.ts b/src/platform/plugins/shared/workflows_execution_engine/server/execution_functions/run_workflow_sync.test.ts new file mode 100644 index 0000000000000..5399453ec1318 --- /dev/null +++ b/src/platform/plugins/shared/workflows_execution_engine/server/execution_functions/run_workflow_sync.test.ts @@ -0,0 +1,109 @@ +/* + * 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". + */ + +import { type EsWorkflowExecution, ExecutionStatus } from '@kbn/workflows'; +import { runWorkflowSync } from './run_workflow_sync'; +import { setupDependencies } from './setup_dependencies'; +import { validateSyncWorkflow } from './validate_sync_workflow'; +import { workflowExecutionLoop } from '../workflow_execution_loop'; + +jest.mock('./setup_dependencies', () => ({ setupDependencies: jest.fn() })); +jest.mock('./validate_sync_workflow', () => ({ validateSyncWorkflow: jest.fn() })); +jest.mock('../workflow_execution_loop', () => ({ workflowExecutionLoop: jest.fn() })); + +describe('runWorkflowSync', () => { + it('runs the canonical workflow loop with in-memory dependencies and sync mode', async () => { + const completedExecution = { + id: 'execution-1', + spaceId: 'default', + workflowId: 'workflow-1', + isTestRun: false, + status: ExecutionStatus.COMPLETED, + context: {}, + workflowDefinition: { + version: '1', + name: 'Test workflow', + enabled: true, + triggers: [], + steps: [], + }, + yaml: '', + scopeStack: [], + createdAt: '2026-07-21T00:00:00.000Z', + error: null, + startedAt: '2026-07-21T00:00:00.000Z', + finishedAt: '2026-07-21T00:00:01.000Z', + cancelRequested: false, + duration: 1000, + } satisfies EsWorkflowExecution; + const workflowExecutionGraph = { topologicalOrder: [] }; + const workflowRuntime = { start: jest.fn().mockResolvedValue(undefined) }; + const workflowExecutionState = { + getWorkflowExecution: jest.fn().mockReturnValue(completedExecution), + }; + const setup = { + workflowExecutionGraph, + workflowRuntime, + workflowExecutionState, + stepExecutionRuntimeFactory: {}, + stepIoService: {}, + workflowLogger: {}, + workflowTaskManager: {}, + nodesFactory: {}, + activeExecutionPersistence: {}, + workflowExecutionRepository: undefined, + esClient: {}, + }; + (setupDependencies as jest.Mock).mockResolvedValue(setup); + (workflowExecutionLoop as jest.Mock).mockResolvedValue(undefined); + + const abortController = new AbortController(); + const request = {} as Parameters[0]['request']; + const dependencies = Object.assign( + {} as Parameters[0]['dependencies'], + { coreStart: {}, workflowsExtensions: { getStepDefinition: jest.fn() } } + ); + const workflowExecutionRepository = {} as Parameters< + typeof runWorkflowSync + >[0]['workflowExecutionRepository']; + const stepExecutionRepository = {} as Parameters< + typeof runWorkflowSync + >[0]['stepExecutionRepository']; + + await expect( + runWorkflowSync({ + workflowExecution: completedExecution, + request, + abortController, + logger: {} as Parameters[0]['logger'], + config: {} as Parameters[0]['config'], + dependencies, + workflowsExecutionEngine: {} as Parameters< + typeof runWorkflowSync + >[0]['workflowsExecutionEngine'], + workflowExecutionRepository, + stepExecutionRepository, + }) + ).resolves.toBe(completedExecution); + + expect(validateSyncWorkflow).toHaveBeenCalledWith( + workflowExecutionGraph, + dependencies.workflowsExtensions.getStepDefinition + ); + expect(workflowRuntime.start).toHaveBeenCalledTimes(1); + expect(workflowExecutionLoop).toHaveBeenCalledWith( + expect.objectContaining({ + executionMode: 'sync', + signal: abortController.signal, + fakeRequest: request, + workflowExecutionRepository: setup.activeExecutionPersistence, + }) + ); + }); +}); diff --git a/src/platform/plugins/shared/workflows_execution_engine/server/execution_functions/run_workflow_sync.ts b/src/platform/plugins/shared/workflows_execution_engine/server/execution_functions/run_workflow_sync.ts new file mode 100644 index 0000000000000..017a26e71338e --- /dev/null +++ b/src/platform/plugins/shared/workflows_execution_engine/server/execution_functions/run_workflow_sync.ts @@ -0,0 +1,93 @@ +/* + * 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". + */ + +import apm from 'elastic-apm-node'; +import { performance } from 'perf_hooks'; +import type { KibanaRequest, Logger } from '@kbn/core/server'; +import type { EsWorkflowExecution } from '@kbn/workflows'; +import { setupDependencies } from './setup_dependencies'; +import { + syncExecutionDurationHistogram, + syncExecutionRequestsCounter, +} from './sync_execution_metrics'; +import { validateSyncWorkflow } from './validate_sync_workflow'; +import type { WorkflowsExecutionEngineConfig } from '../config'; +import type { + StepExecutionPersistence, + WorkflowExecutionPersistence, +} from '../repositories/execution_persistence'; +import type { WorkflowsExecutionEnginePluginStart } from '../types'; +import type { ContextDependencies } from '../workflow_context_manager/types'; +import { workflowExecutionLoop } from '../workflow_execution_loop'; + +export const runWorkflowSync = async ({ + workflowExecution, + request, + abortController, + logger, + config, + dependencies, + workflowsExecutionEngine, + workflowExecutionRepository, + stepExecutionRepository, +}: { + workflowExecution: EsWorkflowExecution; + request: KibanaRequest; + abortController: AbortController; + logger: Logger; + config: WorkflowsExecutionEngineConfig; + dependencies: ContextDependencies; + workflowsExecutionEngine: WorkflowsExecutionEnginePluginStart; + workflowExecutionRepository: WorkflowExecutionPersistence; + stepExecutionRepository: StepExecutionPersistence; +}): Promise => { + apm.currentTransaction?.setLabel('execution_mode', 'sync'); + const startTime = performance.now(); + let outcome: 'success' | 'error' | 'aborted' = 'success'; + + try { + const setup = await setupDependencies({ + workflowRunId: workflowExecution.id, + spaceId: workflowExecution.spaceId, + logger, + config, + dependencies, + fakeRequest: request, + workflowsExecutionEngine, + workflowExecution, + workflowExecutionRepository, + stepExecutionRepository, + }); + + validateSyncWorkflow( + setup.workflowExecutionGraph, + dependencies.workflowsExtensions.getStepDefinition + ); + await setup.workflowRuntime.start(); + await workflowExecutionLoop({ + ...setup, + workflowExecutionRepository: setup.activeExecutionPersistence, + fakeRequest: request, + coreStart: dependencies.coreStart, + signal: abortController.signal, + executionMode: 'sync', + }); + + if (abortController.signal.aborted) { + outcome = 'aborted'; + } + return setup.workflowExecutionState.getWorkflowExecution(); + } catch (error) { + outcome = abortController.signal.aborted ? 'aborted' : 'error'; + throw error; + } finally { + syncExecutionDurationHistogram.record(performance.now() - startTime, { outcome }); + syncExecutionRequestsCounter.add(1, { outcome }); + } +}; diff --git a/src/platform/plugins/shared/workflows_execution_engine/server/execution_functions/setup_dependencies.test.ts b/src/platform/plugins/shared/workflows_execution_engine/server/execution_functions/setup_dependencies.test.ts index cae6c7df894c0..7937510804662 100644 --- a/src/platform/plugins/shared/workflows_execution_engine/server/execution_functions/setup_dependencies.test.ts +++ b/src/platform/plugins/shared/workflows_execution_engine/server/execution_functions/setup_dependencies.test.ts @@ -17,6 +17,37 @@ import type { WorkflowsExecutionEngineConfig } from '../config'; import { WorkflowExecutionTelemetryClient } from '../lib/telemetry/workflow_execution_telemetry_client'; import type { StepExecutionRepository } from '../repositories/step_execution_repository'; import { WorkflowExecutionRepository } from '../repositories/workflow_execution_repository'; +import { WorkflowEventLoggerService } from '../workflow_event_logger/workflow_event_logger_service'; + +// Hoisted mock — must be in the test file (not an imported side-effect) so that +// setup_dependencies.ts receives the mocked WorkflowEventLoggerService when it is +// first imported. Side-effect imports (like mocks.ts) are NOT hoisted and arrive +// too late to intercept the module cache. +jest.mock('../workflow_event_logger/workflow_event_logger_service', () => ({ + WorkflowEventLoggerService: jest.fn().mockImplementation(() => ({ + createLogger: jest.fn().mockReturnValue({ + logInfo: jest.fn(), + logError: jest.fn(), + logWarn: jest.fn(), + logDebug: jest.fn(), + startTiming: jest.fn(), + stopTiming: jest.fn(), + createStepLogger: jest.fn().mockReturnValue({ + logInfo: jest.fn(), + flushEvents: jest.fn(), + }), + flushEvents: jest.fn(), + }), + createWorkflowLogger: jest.fn(), + createExecutionLogger: jest.fn(), + createStepLogger: jest.fn(), + getExecutionLogs: jest.fn(), + getStepLogs: jest.fn(), + getLogsByLevel: jest.fn(), + searchLogs: jest.fn(), + getRecentLogs: jest.fn(), + })), +})); import '../workflow_event_logger/mocks'; jest.mock('../repositories/workflow_execution_repository'); @@ -66,6 +97,8 @@ describe('setupDependencies', () => { }, collectQueueMetrics: false, hitlExternalResume: { enabled: true }, + syncExecution: { enabled: false, maxDurationMs: 60_000 }, + syncLogDrain: { enabled: true, intervalMs: 500, maxQueue: 20000, maxBatch: 4000 }, }; let mockDependencies: ReturnType; @@ -120,16 +153,16 @@ describe('setupDependencies', () => { headers: {}, } as KibanaRequest; - const result = await setupDependencies( + const result = await setupDependencies({ workflowRunId, spaceId, - mockLogger, - mockConfig, - mockDependencies, - mockWorkflowExecutionRepository, - mockStepExecutionRepository, - mockFakeRequest - ); + logger: mockLogger, + config: mockConfig, + dependencies: mockDependencies, + workflowExecutionRepository: mockWorkflowExecutionRepository, + stepExecutionRepository: mockStepExecutionRepository, + fakeRequest: mockFakeRequest, + }); expect(mockAsScoped).toHaveBeenCalledWith(mockFakeRequest); expect(result.esClient).toBe(mockAsCurrentUser); @@ -150,16 +183,16 @@ describe('setupDependencies', () => { headers: {}, } as KibanaRequest; - await setupDependencies( + await setupDependencies({ workflowRunId, spaceId, - mockLogger, - mockConfig, - mockDependencies, - mockWorkflowExecutionRepository, - mockStepExecutionRepository, - mockFakeRequest - ); + logger: mockLogger, + config: mockConfig, + dependencies: mockDependencies, + workflowExecutionRepository: mockWorkflowExecutionRepository, + stepExecutionRepository: mockStepExecutionRepository, + fakeRequest: mockFakeRequest, + }); expect(mockDependencies.actions.getActionsClientWithRequest).toHaveBeenCalledWith( mockFakeRequest @@ -183,16 +216,16 @@ describe('setupDependencies', () => { headers: {}, } as KibanaRequest; - await setupDependencies( + await setupDependencies({ workflowRunId, spaceId, - mockLogger, - mockConfig, - mockDependencies, - mockWorkflowExecutionRepository, - mockStepExecutionRepository, - mockFakeRequest - ); + logger: mockLogger, + config: mockConfig, + dependencies: mockDependencies, + workflowExecutionRepository: mockWorkflowExecutionRepository, + stepExecutionRepository: mockStepExecutionRepository, + fakeRequest: mockFakeRequest, + }); expect(WorkflowGraph.fromWorkflowDefinition).toHaveBeenCalledWith( mockWorkflowExecution.workflowDefinition, @@ -205,16 +238,16 @@ describe('setupDependencies', () => { headers: {}, } as KibanaRequest; - await setupDependencies( + await setupDependencies({ workflowRunId, spaceId, - mockLogger, - mockConfig, - mockDependencies, - mockWorkflowExecutionRepository, - mockStepExecutionRepository, - mockFakeRequest - ); + logger: mockLogger, + config: mockConfig, + dependencies: mockDependencies, + workflowExecutionRepository: mockWorkflowExecutionRepository, + stepExecutionRepository: mockStepExecutionRepository, + fakeRequest: mockFakeRequest, + }); expect(WorkflowGraph.fromWorkflowDefinition).toHaveBeenCalledWith(expect.anything(), { timeout: '6h', @@ -244,16 +277,16 @@ describe('setupDependencies', () => { (isGraphBuildError as unknown as jest.Mock).mockReturnValue(true); await expect( - setupDependencies( + setupDependencies({ workflowRunId, spaceId, - mockLogger, - mockConfig, - mockDependencies, - mockWorkflowExecutionRepository, - mockStepExecutionRepository, - mockFakeRequest - ) + logger: mockLogger, + config: mockConfig, + dependencies: mockDependencies, + workflowExecutionRepository: mockWorkflowExecutionRepository, + stepExecutionRepository: mockStepExecutionRepository, + fakeRequest: mockFakeRequest, + }) ).rejects.toBeInstanceOf(WorkflowGraphSetupError); expect(mockWorkflowExecutionRepository.updateWorkflowExecution).toHaveBeenCalledWith( @@ -274,16 +307,16 @@ describe('setupDependencies', () => { (isGraphBuildError as unknown as jest.Mock).mockReturnValue(false); await expect( - setupDependencies( + setupDependencies({ workflowRunId, spaceId, - mockLogger, - mockConfig, - mockDependencies, - mockWorkflowExecutionRepository, - mockStepExecutionRepository, - mockFakeRequest - ) + logger: mockLogger, + config: mockConfig, + dependencies: mockDependencies, + workflowExecutionRepository: mockWorkflowExecutionRepository, + stepExecutionRepository: mockStepExecutionRepository, + fakeRequest: mockFakeRequest, + }) ).rejects.toBe(otherError); expect(mockWorkflowExecutionRepository.updateWorkflowExecution).not.toHaveBeenCalled(); @@ -303,16 +336,16 @@ describe('setupDependencies', () => { }); await expect( - setupDependencies( + setupDependencies({ workflowRunId, spaceId, - mockLogger, - mockConfig, - mockDependencies, - mockWorkflowExecutionRepository, - mockStepExecutionRepository, - mockFakeRequest - ) + logger: mockLogger, + config: mockConfig, + dependencies: mockDependencies, + workflowExecutionRepository: mockWorkflowExecutionRepository, + stepExecutionRepository: mockStepExecutionRepository, + fakeRequest: mockFakeRequest, + }) ).rejects.toThrow(`Workflow execution with ID ${workflowRunId} not found`); expect(mockWorkflowExecutionRepository.getWorkflowExecutionById).toHaveBeenCalledWith( @@ -346,16 +379,16 @@ describe('setupDependencies', () => { }); (mockDependencies.workflowsExtensions.isReady as jest.Mock).mockReturnValue(isReadyPromise); - const setupPromise = setupDependencies( + const setupPromise = setupDependencies({ workflowRunId, spaceId, - mockLogger, - mockConfig, - mockDependencies, - mockWorkflowExecutionRepository, - mockStepExecutionRepository, - mockFakeRequest - ); + logger: mockLogger, + config: mockConfig, + dependencies: mockDependencies, + workflowExecutionRepository: mockWorkflowExecutionRepository, + stepExecutionRepository: mockStepExecutionRepository, + fakeRequest: mockFakeRequest, + }); // Let any microtasks before the isReady await run await Promise.resolve(); @@ -373,4 +406,43 @@ describe('setupDependencies', () => { ); }); }); + + /** + * Machine-checked invariant: the syncLogDrain supplied to setupDependencies must + * be forwarded to WorkflowEventLoggerService so that every per-execution logger + * routes its flushEvents calls to the drain instead of writing to ES inline. + * + * If someone removes the forwarding at the setupDependencies callsite, this test + * fails while all drain unit tests still pass — making the regression visible. + */ + describe('WorkflowEventLoggerService wiring', () => { + beforeEach(() => { + const mockScopedClient = { + search: jest.fn(), + index: jest.fn(), + } as unknown as ElasticsearchClient; + mockDependencies.coreStart.elasticsearch.client.asScoped = jest.fn().mockReturnValue({ + asCurrentUser: mockScopedClient, + }); + }); + + it('calls WorkflowEventLoggerService with dataStreams, logger, and enableConsoleLogging', async () => { + const mockFakeRequest = { headers: {} } as KibanaRequest; + + await setupDependencies({ + workflowRunId, + spaceId, + logger: mockLogger, + config: mockConfig, + dependencies: mockDependencies, + fakeRequest: mockFakeRequest, + }); + + expect(WorkflowEventLoggerService).toHaveBeenCalledWith( + mockDependencies.coreStart.dataStreams, + mockLogger, + mockConfig.logging.console + ); + }); + }); }); diff --git a/src/platform/plugins/shared/workflows_execution_engine/server/execution_functions/setup_dependencies.ts b/src/platform/plugins/shared/workflows_execution_engine/server/execution_functions/setup_dependencies.ts index 5e82857af65a3..2a799a1dc0346 100644 --- a/src/platform/plugins/shared/workflows_execution_engine/server/execution_functions/setup_dependencies.ts +++ b/src/platform/plugins/shared/workflows_execution_engine/server/execution_functions/setup_dependencies.ts @@ -23,8 +23,17 @@ import { mergeEmitterWorkflowIntoEventChainVisited, } from '../lib/telemetry/utils/extract_execution_metadata'; import { WorkflowExecutionTelemetryClient } from '../lib/telemetry/workflow_execution_telemetry_client'; -import type { StepExecutionRepository } from '../repositories/step_execution_repository'; -import type { WorkflowExecutionRepository } from '../repositories/workflow_execution_repository'; +import { + WORKFLOWS_EXECUTIONS_INDEX, + WORKFLOWS_STEP_EXECUTIONS_INDEX, +} from '../repositories/data_access_layer/constants/execution_indexes'; +import { PlainIndexDataClient } from '../repositories/data_access_layer/implementations/plain_index/plain_index_data_client'; +import type { + StepExecutionPersistence, + WorkflowExecutionPersistence, +} from '../repositories/execution_persistence'; +import { StepExecutionRepository } from '../repositories/step_execution_repository'; +import { WorkflowExecutionRepository } from '../repositories/workflow_execution_repository'; import { NodesFactory } from '../step/nodes_factory'; import type { WorkflowsExecutionEnginePluginStart } from '../types'; import { StepExecutionRuntimeFactory } from '../workflow_context_manager/step_execution_runtime_factory'; @@ -37,17 +46,29 @@ import { WorkflowExecutionState } from '../workflow_context_manager/workflow_exe import { WorkflowEventLoggerService } from '../workflow_event_logger'; import { WorkflowTaskManager } from '../workflow_task_manager/workflow_task_manager'; -export async function setupDependencies( - workflowRunId: string, - spaceId: string, - logger: Logger, - config: WorkflowsExecutionEngineConfig, - dependencies: ContextDependencies, - workflowExecutionRepository: WorkflowExecutionRepository, - stepExecutionRepository: StepExecutionRepository, - fakeRequest?: KibanaRequest, - workflowsExecutionEngine?: WorkflowsExecutionEnginePluginStart -) { +export async function setupDependencies({ + workflowRunId, + spaceId, + logger, + config, + dependencies, + fakeRequest, + workflowsExecutionEngine, + workflowExecution: workflowExecutionOverride, + workflowExecutionRepository: workflowExecutionRepositoryOverride, + stepExecutionRepository: stepExecutionRepositoryOverride, +}: { + workflowRunId: string; + spaceId: string; + logger: Logger; + config: WorkflowsExecutionEngineConfig; + dependencies: ContextDependencies; + fakeRequest?: KibanaRequest; + workflowsExecutionEngine?: WorkflowsExecutionEnginePluginStart; + workflowExecution?: EsWorkflowExecution; + workflowExecutionRepository?: WorkflowExecutionPersistence; + stepExecutionRepository?: StepExecutionPersistence; +}) { const { coreStart, actions, taskManager, workflowsExtensions } = dependencies; await workflowsExtensions.isReady(); @@ -55,15 +76,32 @@ export async function setupDependencies( // Get ES client from core services (guaranteed to be available at task execution time) const internalEsClient = coreStart.elasticsearch.client.asInternalUser; + const workflowExecutionPersistence = + workflowExecutionRepositoryOverride ?? + new WorkflowExecutionRepository( + new PlainIndexDataClient({ + esClient: internalEsClient, + logger, + indexName: WORKFLOWS_EXECUTIONS_INDEX, + }) + ); + const stepExecutionPersistence = + stepExecutionRepositoryOverride ?? + new StepExecutionRepository( + new PlainIndexDataClient({ + esClient: internalEsClient, + logger, + indexName: WORKFLOWS_STEP_EXECUTIONS_INDEX, + }) + ); const workflowRepository = new WorkflowRepository({ esClient: internalEsClient, logger, }); - const workflowExecution = await workflowExecutionRepository.getWorkflowExecutionById( - workflowRunId, - spaceId - ); + const workflowExecution = + workflowExecutionOverride ?? + (await workflowExecutionPersistence.getWorkflowExecutionById(workflowRunId, spaceId)); if (!workflowExecution) { throw new Error(`Workflow execution with ID ${workflowRunId} not found`); @@ -113,7 +151,7 @@ export async function setupDependencies( } catch (error) { if (isGraphBuildError(error)) { const finishedAt = new Date(); - await workflowExecutionRepository.updateWorkflowExecution({ + await workflowExecutionPersistence.updateWorkflowExecution({ id: workflowRunId, status: ExecutionStatus.FAILED, error: { type: 'GraphBuildError', message: error.message }, @@ -150,12 +188,12 @@ export async function setupDependencies( }); const workflowExecutionState = new WorkflowExecutionState( - workflowExecution as EsWorkflowExecution, - workflowExecutionRepository + workflowExecution, + workflowExecutionPersistence ); const stepIoService = new StepIoService({ - stepRepository: stepExecutionRepository, + stepRepository: stepExecutionPersistence, state: workflowExecutionState, evictionMinBytes: config.eviction.minPayloadSize.getValueInBytes(), logger, @@ -172,7 +210,7 @@ export async function setupDependencies( // Create workflow runtime first (simpler, fewer dependencies) const workflowRuntime = new WorkflowExecutionRuntimeManager({ - workflowExecution: workflowExecution as EsWorkflowExecution, + workflowExecution, workflowExecutionGraph, workflowExecutionCursor, workflowLogger, @@ -191,8 +229,8 @@ export async function setupDependencies( const enhancedDependencies: ContextDependencies = { ...dependencies, workflowRepository, - workflowExecutionRepository, - stepExecutionRepository, + workflowExecutionRepository: workflowExecutionRepositoryOverride, + stepExecutionRepository: stepExecutionRepositoryOverride as StepExecutionRepository | undefined, workflowsExecutionEngine, spaceId, request: fakeRequest, @@ -228,7 +266,10 @@ export async function setupDependencies( workflowLogger, workflowTaskManager, nodesFactory, - workflowExecutionRepository, + // activeExecutionPersistence = resolved persistence (override ?? ES-backed default) + activeExecutionPersistence: workflowExecutionPersistence, + // workflowExecutionRepository = raw caller override; undefined on the async path + workflowExecutionRepository: workflowExecutionRepositoryOverride, esClient, telemetryClient, workflowExecutionCursor, diff --git a/src/platform/plugins/shared/workflows_execution_engine/server/execution_functions/sync_execution_metrics.ts b/src/platform/plugins/shared/workflows_execution_engine/server/execution_functions/sync_execution_metrics.ts new file mode 100644 index 0000000000000..dc142672ec216 --- /dev/null +++ b/src/platform/plugins/shared/workflows_execution_engine/server/execution_functions/sync_execution_metrics.ts @@ -0,0 +1,30 @@ +/* + * 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". + */ + +import { metrics, ValueType } from '@opentelemetry/api'; + +const meter = metrics.getMeter('kibana.workflows.execution'); + +export const syncExecutionRequestsCounter = meter.createCounter( + 'kibana.workflows.execution.sync.requests', + { + description: 'Number of synchronous workflow execution requests', + unit: '{request}', + valueType: ValueType.INT, + } +); + +export const syncExecutionDurationHistogram = meter.createHistogram( + 'kibana.workflows.execution.sync.duration', + { + description: 'Duration of synchronous workflow executions, in milliseconds', + unit: 'ms', + valueType: ValueType.DOUBLE, + } +); diff --git a/src/platform/plugins/shared/workflows_execution_engine/server/execution_functions/validate_sync_workflow.test.ts b/src/platform/plugins/shared/workflows_execution_engine/server/execution_functions/validate_sync_workflow.test.ts new file mode 100644 index 0000000000000..26c7f815e4a9a --- /dev/null +++ b/src/platform/plugins/shared/workflows_execution_engine/server/execution_functions/validate_sync_workflow.test.ts @@ -0,0 +1,43 @@ +/* + * 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". + */ + +import { getBuiltInStepDefinition } from '@kbn/workflows'; +import { validateSyncWorkflow } from './validate_sync_workflow'; + +const createGraph = (stepType: string) => ({ + topologicalOrder: ['node-1'], + getNode: jest.fn().mockReturnValue({ + id: 'node-1', + stepId: 'step-1', + stepType, + }), +}); + +describe('validateSyncWorkflow', () => { + it.each(['wait', 'waitForInput', 'waitForApproval', 'workflow.execute', 'workflow.executeAsync'])( + 'rejects async-only step %s', + (stepType) => { + expect(() => validateSyncWorkflow(createGraph(stepType), () => undefined)).toThrow( + 'is not supported in synchronous workflows' + ); + } + ); + + it('accepts ordinary and capability-backed atomic steps', () => { + expect(() => validateSyncWorkflow(createGraph('ai.pii'), () => undefined)).not.toThrow(); + }); + + it('derives extension execution constraints from the registered definition', () => { + expect(() => + validateSyncWorkflow(createGraph('custom.async'), (stepType) => + stepType === 'custom.async' ? getBuiltInStepDefinition('wait') : undefined + ) + ).toThrow('is not supported in synchronous workflows'); + }); +}); diff --git a/src/platform/plugins/shared/workflows_execution_engine/server/execution_functions/validate_sync_workflow.ts b/src/platform/plugins/shared/workflows_execution_engine/server/execution_functions/validate_sync_workflow.ts new file mode 100644 index 0000000000000..0f5592465e423 --- /dev/null +++ b/src/platform/plugins/shared/workflows_execution_engine/server/execution_functions/validate_sync_workflow.ts @@ -0,0 +1,33 @@ +/* + * 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". + */ + +import { type BaseStepDefinition, getBuiltInStepDefinition } from '@kbn/workflows'; +import type { WorkflowGraph } from '@kbn/workflows/graph'; + +type SyncWorkflowGraph = Pick; + +type GetExtensionStepDefinition = (stepType: string) => BaseStepDefinition | undefined; + +export const validateSyncWorkflow = ( + workflowGraph: SyncWorkflowGraph, + getExtensionStepDefinition: GetExtensionStepDefinition +): void => { + for (const nodeId of workflowGraph.topologicalOrder) { + const node = workflowGraph.getNode(nodeId); + if (node?.stepType) { + const definition = + getBuiltInStepDefinition(node.stepType) ?? getExtensionStepDefinition(node.stepType); + if (definition?.supportedExecutionModes?.includes('sync') === false) { + throw new Error( + `Step "${node.stepId}" (${node.stepType}) is not supported in synchronous workflows` + ); + } + } + } +}; diff --git a/src/platform/plugins/shared/workflows_execution_engine/server/index.ts b/src/platform/plugins/shared/workflows_execution_engine/server/index.ts index e412b570456c5..2374fbacdd7eb 100644 --- a/src/platform/plugins/shared/workflows_execution_engine/server/index.ts +++ b/src/platform/plugins/shared/workflows_execution_engine/server/index.ts @@ -27,9 +27,13 @@ export type { WorkflowExecutionsDataClient, WorkflowsExecutionEnginePluginSetup, WorkflowsExecutionEnginePluginStart, + ExecuteWorkflowOptions, + ExecuteWorkflowResponse, + WorkflowExecutionMode, } from './types'; export { getStepExecutionsByWorkflowExecution } from './repositories/data_access_layer/lib/get_step_executions_by_workflow_execution'; +export { classifyWorkflowTriggerMatch } from './trigger_events/filter_workflows_by_trigger_condition'; export { registerHitlLifecycleAuditor, diff --git a/src/platform/plugins/shared/workflows_execution_engine/server/lib/build_workflow_execution_document.test.ts b/src/platform/plugins/shared/workflows_execution_engine/server/lib/build_workflow_execution_document.test.ts index 9844eee0b971a..60eb753997f07 100644 --- a/src/platform/plugins/shared/workflows_execution_engine/server/lib/build_workflow_execution_document.test.ts +++ b/src/platform/plugins/shared/workflows_execution_engine/server/lib/build_workflow_execution_document.test.ts @@ -42,6 +42,12 @@ const baseParams = { }; describe('buildWorkflowExecutionDocument', () => { + it('uses a caller-provided execution ID', () => { + expect( + buildWorkflowExecutionDocument({ ...baseParams, executionId: 'caller-execution-id' }).id + ).toBe('caller-execution-id'); + }); + it('sets version when workflow has version', () => { const workflowExecution = buildWorkflowExecutionDocument({ ...baseParams, diff --git a/src/platform/plugins/shared/workflows_execution_engine/server/lib/build_workflow_execution_document.ts b/src/platform/plugins/shared/workflows_execution_engine/server/lib/build_workflow_execution_document.ts index 9361208ec19bc..1f5652afb9b71 100644 --- a/src/platform/plugins/shared/workflows_execution_engine/server/lib/build_workflow_execution_document.ts +++ b/src/platform/plugins/shared/workflows_execution_engine/server/lib/build_workflow_execution_document.ts @@ -23,6 +23,7 @@ import { normalizeEventChainVisitedWorkflowIds } from './telemetry/utils/extract import type { WorkflowExecutionForInputRendering } from '../workflow_context_manager/build_workflow_context'; export interface BuildWorkflowExecutionDocumentParams { + executionId?: string; workflow: WorkflowExecutionEngineModel; context: Record; defaultTriggeredBy: string; @@ -47,6 +48,7 @@ export const buildWorkflowExecutionDocument = ( ): WorkflowExecutionForInputRendering => { const { workflow, + executionId, context, defaultTriggeredBy, authenticatedUser, @@ -78,7 +80,7 @@ export const buildWorkflowExecutionDocument = ( typeof metadata?.eventId === 'string' ? metadata.eventId.trim() || undefined : undefined; const missingIdentity = authenticatedUser == null; const workflowExecution: WorkflowExecutionForInputRendering = { - id: generateUuid(), + id: executionId ?? generateUuid(), spaceId, workflowId: workflow.id, ...pickManagedWorkflowFields(workflow), diff --git a/src/platform/plugins/shared/workflows_execution_engine/server/lib/validate_workflow_inputs.ts b/src/platform/plugins/shared/workflows_execution_engine/server/lib/validate_workflow_inputs.ts index 64806471162fd..ea66d853585bf 100644 --- a/src/platform/plugins/shared/workflows_execution_engine/server/lib/validate_workflow_inputs.ts +++ b/src/platform/plugins/shared/workflows_execution_engine/server/lib/validate_workflow_inputs.ts @@ -15,7 +15,7 @@ import { applyInputDefaults, getInputsFromDefinition, } from '@kbn/workflows/spec/lib/field_conversion'; -import type { WorkflowExecutionRepository } from '../repositories/workflow_execution_repository'; +import type { WorkflowExecutionPersistence } from '../repositories/execution_persistence'; import { WorkflowTemplatingEngine } from '../templating_engine'; import { buildInputDefaultRenderContext, @@ -36,7 +36,7 @@ const hasInputChanges = ( */ export const validateWorkflowInputs = async ( workflowExecution: WorkflowExecutionForInputRendering, - workflowExecutionRepository: WorkflowExecutionRepository, + workflowExecutionRepository: WorkflowExecutionPersistence, logger: Logger, coreStart?: CoreStart, dependencies?: ContextDependencies diff --git a/src/platform/plugins/shared/workflows_execution_engine/server/mocks.ts b/src/platform/plugins/shared/workflows_execution_engine/server/mocks.ts index 3e5d567b4f30d..21776031890fd 100644 --- a/src/platform/plugins/shared/workflows_execution_engine/server/mocks.ts +++ b/src/platform/plugins/shared/workflows_execution_engine/server/mocks.ts @@ -27,6 +27,7 @@ export { export const workflowsExecutionEngineMock = { createSetup: jest.fn().mockReturnValue({} as jest.Mocked), createStart: jest.fn().mockReturnValue({ + supportsSynchronousExecution: true, __internalStorage: { workflowExecutionsDataClient: createMockWorkflowDataClient(), stepExecutionsDataClient: createMockStepDataClient(), diff --git a/src/platform/plugins/shared/workflows_execution_engine/server/plugin.bulk_schedule.test.ts b/src/platform/plugins/shared/workflows_execution_engine/server/plugin.bulk_schedule.test.ts index 670a59a9477b5..87e0c156ca23a 100644 --- a/src/platform/plugins/shared/workflows_execution_engine/server/plugin.bulk_schedule.test.ts +++ b/src/platform/plugins/shared/workflows_execution_engine/server/plugin.bulk_schedule.test.ts @@ -114,6 +114,7 @@ describe('bulkScheduleWorkflow', () => { const initializerContext = coreMock.createPluginInitializerContext({ logging: { console: false }, eventDriven: { enabled: true, logEvents: true, maxChainDepth: 10 }, + syncExecution: { enabled: false, maxDurationMs: 60_000 }, }); plugin = new WorkflowsExecutionEnginePlugin(initializerContext); diff --git a/src/platform/plugins/shared/workflows_execution_engine/server/plugin.ts b/src/platform/plugins/shared/workflows_execution_engine/server/plugin.ts index 3dcd502310e50..5c0032517b39b 100644 --- a/src/platform/plugins/shared/workflows_execution_engine/server/plugin.ts +++ b/src/platform/plugins/shared/workflows_execution_engine/server/plugin.ts @@ -41,6 +41,7 @@ import { resumeWorkflow, runWorkflow, } from './execution_functions'; +import { executeWorkflowSync } from './execution_functions/execute_workflow_sync'; import { handlePostExecutionLoop } from './execution_functions/handle_post_execution_loop'; import { buildWorkflowExecutionDocument } from './lib/build_workflow_execution_document'; import { checkLicense } from './lib/check_license'; @@ -94,6 +95,7 @@ import { } from './workflow_context_manager/build_workflow_context'; import type { ContextDependencies } from './workflow_context_manager/types'; import { WorkflowEventLoggerService } from './workflow_event_logger'; + import type { ResumeWorkflowExecutionParams, StartWorkflowExecutionParams, @@ -1055,7 +1057,10 @@ export class WorkflowsExecutionEnginePlugin return {}; } - public start(coreStart: CoreStart, plugins: WorkflowsExecutionEnginePluginStartDeps) { + public start( + coreStart: CoreStart, + plugins: WorkflowsExecutionEnginePluginStartDeps + ): WorkflowsExecutionEnginePluginStart { this.logger.debug('workflows-execution-engine: Start'); if (!this.setupDependencies) { @@ -1107,6 +1112,7 @@ export class WorkflowsExecutionEnginePlugin const buildExecutionDocument = async (args: { workflow: WorkflowExecutionEngineModel; + executionId?: string; context: Record; defaultTriggeredBy: string; authenticatedUser: string | undefined; @@ -1130,7 +1136,11 @@ export class WorkflowsExecutionEnginePlugin context: Record, defaultTriggeredBy: string, request: KibanaRequest, - options: { refresh: boolean | 'wait_for' } = { refresh: false } + options: { + refresh: boolean | 'wait_for'; + executionId?: string; + metadata?: Record; + } = { refresh: false } ): Promise<{ workflowExecution: WorkflowExecutionForInputRendering; repository: WorkflowExecutionRepository; @@ -1143,14 +1153,17 @@ export class WorkflowsExecutionEnginePlugin coreStart.elasticsearch.client ); + const executionContext = options.metadata + ? { ...context, metadata: options.metadata } + : context; const workflowExecution = await buildExecutionDocument({ workflow, - context, + executionId: options.executionId, + context: executionContext, defaultTriggeredBy, authenticatedUser, now: new Date(), }); - await maybeDrainConcurrencyQueueBeforeEnqueue({ workflowExecution, workflowExecutionRepository, @@ -1193,9 +1206,38 @@ export class WorkflowsExecutionEnginePlugin }; }; - const executeWorkflow: ExecuteWorkflow = async (workflow, context, request) => { + const executeWorkflow: ExecuteWorkflow = async (workflow, context, request, options = {}) => { await checkLicense(plugins.licensing); + if ( + options.executionMode !== 'sync' && + (options.capabilities !== undefined || options.abortSignal !== undefined) + ) { + throw new Error('Request-local capabilities and abort signals require sync execution'); + } + + if (options.executionMode === 'sync' && this.config.syncExecution.enabled) { + if (!request) { + throw new Error('Synchronous workflows cannot be executed without the user context'); + } + if (!this.coreSetup) { + throw new Error('Core setup not available'); + } + const coreSetup = this.coreSetup; + return executeWorkflowSync({ + workflow, + context, + request, + options, + logger: this.logger, + dependencies, + getWorkflowsExecutionEngine: async () => { + const [, , workflowsExecutionEngine] = await coreSetup.getStartServices(); + return workflowsExecutionEngine; + }, + }); + } + // AUTO-DETECT: Check if we're already running in a Task Manager context const isRunningInTaskManager = (context.triggeredBy as string | undefined) === 'scheduled' || @@ -1227,7 +1269,10 @@ export class WorkflowsExecutionEnginePlugin context, 'manual', request, - { refresh: true } + { + refresh: true, + executionId: options.executionId, + } ); if (workflowExecution.status === ExecutionStatus.FAILED) { @@ -1816,6 +1861,7 @@ export class WorkflowsExecutionEnginePlugin }; return { + supportsSynchronousExecution: true, workflowEventLoggerService, executeWorkflow, executeWorkflowStep, @@ -1832,7 +1878,7 @@ export class WorkflowsExecutionEnginePlugin }; } - public stop() { + public async stop() { void this.dataClientBundle.stop(); } diff --git a/src/platform/plugins/shared/workflows_execution_engine/server/repositories/execution_persistence.test.ts b/src/platform/plugins/shared/workflows_execution_engine/server/repositories/execution_persistence.test.ts new file mode 100644 index 0000000000000..065166b846254 --- /dev/null +++ b/src/platform/plugins/shared/workflows_execution_engine/server/repositories/execution_persistence.test.ts @@ -0,0 +1,261 @@ +/* + * 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". + */ + +import { type EsWorkflowExecution, ExecutionStatus } from '@kbn/workflows'; +import { InMemoryExecutionPersistence } from './execution_persistence'; + +describe('InMemoryExecutionPersistence', () => { + const execution = { + id: 'execution-1', + spaceId: 'space-1', + workflowId: 'workflow-1', + isTestRun: false, + status: ExecutionStatus.PENDING, + context: {}, + workflowDefinition: { + version: '1', + name: 'Test workflow', + enabled: true, + triggers: [], + steps: [], + }, + yaml: '', + scopeStack: [], + createdAt: '2026-07-21T00:00:00.000Z', + error: null, + startedAt: '2026-07-21T00:00:00.000Z', + finishedAt: '', + cancelRequested: false, + duration: 0, + } satisfies EsWorkflowExecution; + + it('keeps workflow updates process-local', async () => { + const persistence = new InMemoryExecutionPersistence(execution); + await persistence.updateWorkflowExecution({ + id: execution.id, + status: ExecutionStatus.RUNNING, + }); + + await expect( + persistence.getWorkflowExecutionById(execution.id, execution.spaceId) + ).resolves.toEqual(expect.objectContaining({ status: ExecutionStatus.RUNNING })); + await expect( + persistence.getWorkflowExecutionById(execution.id, 'another-space') + ).resolves.toBeNull(); + }); + + it('returns a defensive copy from getWorkflowExecutionById', async () => { + const persistence = new InMemoryExecutionPersistence(execution); + const result = await persistence.getWorkflowExecutionById(execution.id, execution.spaceId); + result!.status = ExecutionStatus.RUNNING; + + await expect( + persistence.getWorkflowExecutionById(execution.id, execution.spaceId) + ).resolves.toEqual(expect.objectContaining({ status: ExecutionStatus.PENDING })); + }); + + it('deep-isolates nested fields returned from getWorkflowExecutionById', async () => { + const persistence = new InMemoryExecutionPersistence({ + ...execution, + context: { key: 'original' }, + scopeStack: [{ stepId: 'root', nestedScopes: [] }], + }); + const result = await persistence.getWorkflowExecutionById(execution.id, execution.spaceId); + (result!.context as Record).key = 'mutated'; + result!.scopeStack.push({ stepId: 'injected', nestedScopes: [] }); + + const fresh = await persistence.getWorkflowExecutionById(execution.id, execution.spaceId); + expect(fresh!.context).toEqual({ key: 'original' }); + expect(fresh!.scopeStack).toHaveLength(1); + }); + + it('throws a descriptive error when workflow execution state contains a non-cloneable value', async () => { + const persistence = new InMemoryExecutionPersistence({ + ...execution, + context: { fn: () => {} } as any, + }); + await expect( + persistence.getWorkflowExecutionById(execution.id, execution.spaceId) + ).rejects.toThrow(/Failed to clone workflow execution execution-1.*non-serializable/); + }); + + it('does not overwrite identity fields via updateWorkflowExecution', async () => { + const persistence = new InMemoryExecutionPersistence(execution); + await persistence.updateWorkflowExecution({ + id: 'hijacked-id', + spaceId: 'hijacked-space', + status: ExecutionStatus.RUNNING, + }); + + await expect( + persistence.getWorkflowExecutionById(execution.id, execution.spaceId) + ).resolves.toEqual( + expect.objectContaining({ + id: execution.id, + spaceId: execution.spaceId, + status: ExecutionStatus.RUNNING, + }) + ); + }); + + it('does not share state between execution-scoped instances', async () => { + const first = new InMemoryExecutionPersistence(execution); + const secondExecution = { ...execution, id: 'execution-2' }; + const second = new InMemoryExecutionPersistence(secondExecution); + + await first.updateWorkflowExecution({ status: ExecutionStatus.RUNNING }); + + await expect( + second.getWorkflowExecutionById(secondExecution.id, secondExecution.spaceId) + ).resolves.toEqual(expect.objectContaining({ status: ExecutionStatus.PENDING })); + }); + + it('deep-isolates nested fields returned from getStepExecutionsByIds', async () => { + const persistence = new InMemoryExecutionPersistence(execution); + await persistence.bulkUpsert([ + { + id: 'step-deep', + spaceId: 'space-1', + stepId: 'step-deep', + scopeStack: [{ stepId: 'root', nestedScopes: [] }], + workflowRunId: execution.id, + workflowId: execution.workflowId, + status: ExecutionStatus.RUNNING, + startedAt: '2026-07-21T00:00:00.000Z', + topologicalIndex: 0, + globalExecutionIndex: 0, + stepExecutionIndex: 0, + }, + ]); + const [result] = await persistence.getStepExecutionsByIds(['step-deep']); + result.scopeStack.push({ stepId: 'injected', nestedScopes: [] }); + + const [fresh] = await persistence.getStepExecutionsByIds(['step-deep']); + expect(fresh.scopeStack).toHaveLength(1); + }); + + it('throws a descriptive error when step execution state contains a non-cloneable value', async () => { + const persistence = new InMemoryExecutionPersistence(execution); + await persistence.bulkUpsert([ + { + id: 'step-err', + spaceId: 'space-1', + stepId: 'step-err', + scopeStack: [], + workflowRunId: execution.id, + workflowId: execution.workflowId, + + status: (() => {}) as any, + startedAt: '2026-07-21T00:00:00.000Z', + topologicalIndex: 0, + globalExecutionIndex: 0, + stepExecutionIndex: 0, + }, + ]); + await expect(persistence.getStepExecutionsByIds(['step-err'])).rejects.toThrow( + /Failed to clone step execution step-err.*non-serializable/ + ); + }); + + it('returns a defensive copy from getStepExecutionsByIds', async () => { + const persistence = new InMemoryExecutionPersistence(execution); + await persistence.bulkUpsert([ + { + id: 'step-copy', + spaceId: 'space-1', + stepId: 'step-copy', + scopeStack: [], + workflowRunId: execution.id, + workflowId: execution.workflowId, + status: ExecutionStatus.RUNNING, + startedAt: '2026-07-21T00:00:00.000Z', + topologicalIndex: 0, + globalExecutionIndex: 0, + stepExecutionIndex: 0, + }, + ]); + const [result] = await persistence.getStepExecutionsByIds(['step-copy']); + result.status = ExecutionStatus.COMPLETED; + + await expect(persistence.getStepExecutionsByIds(['step-copy'])).resolves.toEqual([ + expect.objectContaining({ status: ExecutionStatus.RUNNING }), + ]); + }); + + it('applies sourceIncludes projection to step executions', async () => { + const persistence = new InMemoryExecutionPersistence(execution); + await persistence.bulkUpsert([ + { + id: 'step-proj', + spaceId: 'space-1', + stepId: 'step-proj', + scopeStack: [], + workflowRunId: execution.id, + workflowId: execution.workflowId, + status: ExecutionStatus.RUNNING, + startedAt: '2026-07-21T00:00:00.000Z', + topologicalIndex: 0, + globalExecutionIndex: 0, + stepExecutionIndex: 0, + }, + ]); + const [result] = await persistence.getStepExecutionsByIds(['step-proj'], ['id', 'status']); + expect(result).toEqual({ id: 'step-proj', status: ExecutionStatus.RUNNING }); + }); + + it('applies sourceExcludes projection to step executions', async () => { + const persistence = new InMemoryExecutionPersistence(execution); + await persistence.bulkUpsert([ + { + id: 'step-excl', + spaceId: 'space-1', + stepId: 'step-excl', + scopeStack: [], + workflowRunId: execution.id, + workflowId: execution.workflowId, + status: ExecutionStatus.RUNNING, + startedAt: '2026-07-21T00:00:00.000Z', + topologicalIndex: 0, + globalExecutionIndex: 0, + stepExecutionIndex: 0, + }, + ]); + const [result] = await persistence.getStepExecutionsByIds(['step-excl'], undefined, ['status']); + expect(result).not.toHaveProperty('status'); + expect(result).toHaveProperty('id', 'step-excl'); + }); + + it('merges step lifecycle and IO updates without an external repository', async () => { + const persistence = new InMemoryExecutionPersistence(execution); + await persistence.bulkUpsert([ + { + id: 'step-1', + spaceId: 'space-1', + stepId: 'step-1', + scopeStack: [], + workflowRunId: execution.id, + workflowId: execution.workflowId, + status: ExecutionStatus.RUNNING, + startedAt: '2026-07-21T00:00:00.000Z', + topologicalIndex: 0, + globalExecutionIndex: 0, + stepExecutionIndex: 0, + }, + { id: 'step-1', output: { content: 'result' } }, + ]); + + await expect(persistence.getStepExecutionsByIds(['step-1'])).resolves.toEqual([ + expect.objectContaining({ + id: 'step-1', + status: ExecutionStatus.RUNNING, + output: { content: 'result' }, + }), + ]); + }); +}); diff --git a/src/platform/plugins/shared/workflows_execution_engine/server/repositories/execution_persistence.ts b/src/platform/plugins/shared/workflows_execution_engine/server/repositories/execution_persistence.ts new file mode 100644 index 0000000000000..4044fe751fa9b --- /dev/null +++ b/src/platform/plugins/shared/workflows_execution_engine/server/repositories/execution_persistence.ts @@ -0,0 +1,137 @@ +/* + * 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". + */ + +import type { EsWorkflowExecution, EsWorkflowStepExecution } from '@kbn/workflows'; +import type { StepExecutionField } from './step_execution_repository'; + +export interface WorkflowExecutionPersistence { + getWorkflowExecutionById( + workflowExecutionId: string, + spaceId: string + ): Promise; + updateWorkflowExecution( + workflowExecution: Partial, + options?: { refresh?: boolean | 'wait_for' } + ): Promise; +} + +export interface StepExecutionPersistence { + getStepExecutionsByIds( + stepExecutionIds: string[], + sourceIncludes?: StepExecutionField[], + sourceExcludes?: StepExecutionField[] + ): Promise; + bulkUpsert(stepExecutions: Array>): Promise; +} + +/** + * Owns all mutable state for exactly one synchronous workflow execution. + * Construct a fresh instance per execution; never share an instance across runs. + */ +export class InMemoryExecutionPersistence + implements WorkflowExecutionPersistence, StepExecutionPersistence +{ + private readonly stepExecutions = new Map>(); + + constructor(private execution: EsWorkflowExecution) {} + + public async getWorkflowExecutionById( + workflowExecutionId: string, + spaceId: string + ): Promise { + if (this.execution.id !== workflowExecutionId || this.execution.spaceId !== spaceId) { + return null; + } + try { + return structuredClone(this.execution); + } catch (err) { + throw new Error( + `Failed to clone workflow execution ${workflowExecutionId}: execution state contains a non-serializable value. Root cause: ${ + err instanceof Error ? err.message : String(err) + }` + ); + } + } + + public async updateWorkflowExecution( + workflowExecution: Partial, + _options?: { refresh?: boolean | 'wait_for' } + ): Promise { + // Strip identity fields — they locate the document and must not be mutated. + const { id: _id, spaceId: _spaceId, ...update } = workflowExecution; + this.execution = { ...this.execution, ...update }; + } + public async getStepExecutionsByIds( + ids: string[], + sourceIncludes?: StepExecutionField[], + sourceExcludes?: StepExecutionField[] + ): Promise { + return ids.flatMap((id) => { + const execution = this.stepExecutions.get(id); + if (!execution) { + return []; + } + if (!isCompleteStepExecution(execution)) { + throw new Error( + `Step execution ${id} was read before its required fields were initialized` + ); + } + let copy: Record; + try { + copy = structuredClone(execution) as Record; + } catch (err) { + throw new Error( + `Failed to clone step execution ${id}: step execution state contains a non-serializable value. Root cause: ${ + err instanceof Error ? err.message : String(err) + }` + ); + } + if (sourceIncludes?.length) { + for (const key of Object.keys(copy)) { + if (!sourceIncludes.includes(key as StepExecutionField)) { + delete copy[key]; + } + } + } + if (sourceExcludes?.length) { + for (const key of sourceExcludes) { + delete copy[key]; + } + } + return [copy as EsWorkflowStepExecution]; + }); + } + + public async bulkUpsert(executions: Array>): Promise { + for (const update of executions) { + if (!update.id) { + throw new Error('Step execution ID is required for in-memory upsert'); + } + this.stepExecutions.set(update.id, { + ...this.stepExecutions.get(update.id), + ...update, + }); + } + } +} + +const isCompleteStepExecution = ( + execution: Partial +): execution is EsWorkflowStepExecution => + typeof execution.spaceId === 'string' && + typeof execution.id === 'string' && + typeof execution.stepId === 'string' && + Array.isArray(execution.scopeStack) && + typeof execution.workflowRunId === 'string' && + typeof execution.workflowId === 'string' && + execution.status !== undefined && + typeof execution.startedAt === 'string' && + typeof execution.topologicalIndex === 'number' && + typeof execution.globalExecutionIndex === 'number' && + typeof execution.stepExecutionIndex === 'number'; diff --git a/src/platform/plugins/shared/workflows_execution_engine/server/step/custom_step_impl/step_definition_handlers/create_base_handler_context.ts b/src/platform/plugins/shared/workflows_execution_engine/server/step/custom_step_impl/step_definition_handlers/create_base_handler_context.ts index 23fb402322c29..cd225a8013ee8 100644 --- a/src/platform/plugins/shared/workflows_execution_engine/server/step/custom_step_impl/step_definition_handlers/create_base_handler_context.ts +++ b/src/platform/plugins/shared/workflows_execution_engine/server/step/custom_step_impl/step_definition_handlers/create_base_handler_context.ts @@ -12,13 +12,31 @@ import type { StepHandlerContext } from '@kbn/workflows-extensions/server'; import type { StepExecutionRuntime } from '../../../workflow_context_manager/step_execution_runtime'; import type { IWorkflowEventLogger } from '../../../workflow_event_logger'; +export type BaseHandlerNode = Pick; +export interface BaseHandlerStepExecutionRuntime { + abortController: StepExecutionRuntime['abortController']; + contextManager: Pick< + StepExecutionRuntime['contextManager'], + | 'callKibanaApi' + | 'getContext' + | 'getEsClientAsUser' + | 'getExecutionCapabilities' + | 'getFakeRequest' + | 'renderValueAccordingToContext' + >; +} +export type BaseHandlerWorkflowLogger = Pick< + IWorkflowEventLogger, + 'logDebug' | 'logError' | 'logInfo' | 'logWarn' +>; + export function createBaseHandlerContext( input: unknown, rawInput: unknown, config: Record, - node: AtomicGraphNode, - stepExecutionRuntime: StepExecutionRuntime, - workflowLogger: IWorkflowEventLogger + node: BaseHandlerNode, + stepExecutionRuntime: BaseHandlerStepExecutionRuntime, + workflowLogger: BaseHandlerWorkflowLogger ): StepHandlerContext { return { input, @@ -56,5 +74,6 @@ export function createBaseHandlerContext( abortSignal: stepExecutionRuntime.abortController.signal, stepId: node.stepId, stepType: node.stepType, + capabilities: stepExecutionRuntime.contextManager.getExecutionCapabilities(), }; } diff --git a/src/platform/plugins/shared/workflows_execution_engine/server/step/custom_step_impl/step_definition_handlers/tests/create_base_handler_context.test.ts b/src/platform/plugins/shared/workflows_execution_engine/server/step/custom_step_impl/step_definition_handlers/tests/create_base_handler_context.test.ts index 2734c0a9d7994..a4dcd9242947b 100644 --- a/src/platform/plugins/shared/workflows_execution_engine/server/step/custom_step_impl/step_definition_handlers/tests/create_base_handler_context.test.ts +++ b/src/platform/plugins/shared/workflows_execution_engine/server/step/custom_step_impl/step_definition_handlers/tests/create_base_handler_context.test.ts @@ -21,9 +21,9 @@ describe('createBaseHandlerContext', () => { input, rawInput, config, - defaultTestNode as any, - mocks.stepExecutionRuntime as any, - mocks.workflowLogger as any + defaultTestNode, + mocks.stepExecutionRuntime, + mocks.workflowLogger ); expect(context.input).toBe(input); @@ -40,9 +40,7 @@ describe('createBaseHandlerContext', () => { expect(mocks.stepExecutionRuntime.contextManager.getEsClientAsUser).toHaveBeenCalled(); context.contextManager.renderInputTemplate({ x: 1 }); - expect( - mocks.stepExecutionRuntime.contextManager.renderValueAccordingToContext - ).toHaveBeenCalledWith({ x: 1 }, undefined); + expect(mocks.renderValueAccordingToContext).toHaveBeenCalledWith({ x: 1 }, undefined); context.logger.info('hello', { meta: true }); expect(mocks.workflowLogger.logInfo).toHaveBeenCalledWith('hello', { meta: true }); @@ -55,15 +53,62 @@ describe('createBaseHandlerContext', () => { {}, undefined as unknown as Record, undefined as unknown as Record, - defaultTestNode as any, - mocks.stepExecutionRuntime as any, - mocks.workflowLogger as any + defaultTestNode, + mocks.stepExecutionRuntime, + mocks.workflowLogger ); expect(context.rawInput).toEqual({}); expect(context.config).toEqual({}); }); + it('exposes request-local capabilities directly to the handler', () => { + const mocks = createHandlerTestMocks(); + const capabilities = { + id: 'proceed', + value: { invoke: jest.fn() }, + }; + ( + mocks.stepExecutionRuntime.contextManager.getExecutionCapabilities as jest.Mock + ).mockReturnValue([capabilities]); + + const context = createBaseHandlerContext( + {}, + {}, + {}, + defaultTestNode, + mocks.stepExecutionRuntime, + mocks.workflowLogger + ); + + expect(context.capabilities).toEqual([capabilities]); + expect(context.contextManager.getContext()).not.toHaveProperty('capabilities'); + }); + + it('capabilities self-destruct on JSON serialization — the ES persistence path cannot capture real payload', () => { + const mocks = createHandlerTestMocks(); + const secretFn = jest.fn(); + ( + mocks.stepExecutionRuntime.contextManager.getExecutionCapabilities as jest.Mock + ).mockReturnValue([{ id: 'proceed', value: { invoke: secretFn } }]); + + const context = createBaseHandlerContext( + {}, + {}, + {}, + defaultTestNode, + mocks.stepExecutionRuntime, + mocks.workflowLogger + ); + + // Handler can access the real value in-process ... + expect((context.capabilities![0].value as { invoke: jest.Mock }).invoke).toBe(secretFn); + // ... but a JSON round-trip (the path to step output / ES) loses it. + const persisted = JSON.parse(JSON.stringify(context.capabilities)); + expect(persisted).toEqual([{ id: 'proceed', value: {} }]); + expect(persisted[0].value).not.toHaveProperty('invoke'); + }); + it('forwards callKibanaApi to the execution runtime with abort signal', async () => { const mocks = createHandlerTestMocks(); mocks.stepExecutionRuntime.contextManager.callKibanaApi.mockResolvedValue({ @@ -76,9 +121,9 @@ describe('createBaseHandlerContext', () => { {}, {}, {}, - defaultTestNode as any, - mocks.stepExecutionRuntime as any, - mocks.workflowLogger as any + defaultTestNode, + mocks.stepExecutionRuntime, + mocks.workflowLogger ); const result = await context.contextManager.callKibanaApi({ diff --git a/src/platform/plugins/shared/workflows_execution_engine/server/step/custom_step_impl/step_definition_handlers/tests/test_helpers.ts b/src/platform/plugins/shared/workflows_execution_engine/server/step/custom_step_impl/step_definition_handlers/tests/test_helpers.ts index 117a1da6184cf..a4eb1678ae4e6 100644 --- a/src/platform/plugins/shared/workflows_execution_engine/server/step/custom_step_impl/step_definition_handlers/tests/test_helpers.ts +++ b/src/platform/plugins/shared/workflows_execution_engine/server/step/custom_step_impl/step_definition_handlers/tests/test_helpers.ts @@ -7,6 +7,13 @@ * License v3.0 only", or the "Server Side Public License, v 1". */ +import { elasticsearchServiceMock, httpServerMock } from '@kbn/core/server/mocks'; +import type { StepContext } from '@kbn/workflows'; +import type { + BaseHandlerNode, + BaseHandlerStepExecutionRuntime, + BaseHandlerWorkflowLogger, +} from '../create_base_handler_context'; import { DURABLE_STEP_STATE_KEY, type DurableStepState } from '../durable_step_state'; export { DURABLE_STEP_STATE_KEY, type DurableStepState }; @@ -14,35 +21,57 @@ export { DURABLE_STEP_STATE_KEY, type DurableStepState }; export const getDurableState = (persisted: Record | undefined): DurableStepState => (persisted?.[DURABLE_STEP_STATE_KEY] ?? {}) as DurableStepState; -export interface TestNode { - stepId: string; - stepType: string; +export interface TestNode extends BaseHandlerNode { configuration: { with?: Record; 'max-step-size'?: undefined; }; } -export const defaultTestNode: TestNode = { +export const defaultTestNode = { stepId: 'custom-step', stepType: 'my-custom-type', configuration: { with: { key: 'value' }, 'max-step-size': undefined }, -}; +} satisfies TestNode; export const createHandlerTestMocks = (initialPersistedState?: Record) => { const persistedState: { value: Record | undefined } = { value: initialPersistedState, }; - const stepExecutionRuntime = { + const renderValueAccordingToContext = jest.fn(); + const stepContext: StepContext = { + execution: { + id: 'workflow-execution-1', + isTestRun: false, + startedAt: new Date(0), + url: '', + }, + workflow: { + id: 'workflow-1', + name: 'Test workflow', + enabled: true, + spaceId: 'default', + }, + kibanaUrl: '', + steps: {}, + }; + const baseHandlerStepExecutionRuntime = { contextManager: { - renderValueAccordingToContext: jest.fn((v: unknown) => v), - getContext: jest.fn(() => ({})), - getEsClientAsUser: jest.fn(() => ({})), - getFakeRequest: jest.fn(() => null), + renderValueAccordingToContext: (value: T, additionalContext?: Record) => { + renderValueAccordingToContext(value, additionalContext); + return value; + }, + getContext: jest.fn(() => stepContext), + getEsClientAsUser: jest.fn(() => elasticsearchServiceMock.createElasticsearchClient()), + getFakeRequest: jest.fn(() => httpServerMock.createKibanaRequest()), + getExecutionCapabilities: jest.fn(() => undefined), callKibanaApi: jest.fn(), }, abortController: new AbortController(), + } satisfies BaseHandlerStepExecutionRuntime; + const stepExecutionRuntime = { + ...baseHandlerStepExecutionRuntime, node: { configuration: { with: { key: 'value' } } }, startStep: jest.fn(), flushEventLogs: jest.fn().mockResolvedValue(undefined), @@ -63,11 +92,12 @@ export const createHandlerTestMocks = (initialPersistedState?: Record { + const engine = new WorkflowTemplatingEngine(); + const context = { workflow: { inputs: { recipients: ['a@b.com', 'c@d.com'] } } }; + + const ACCEPTED = [ + '${{ workflow.inputs.recipients }}', + '${{workflow.inputs.recipients}}', + '${{ workflow.inputs.recipients | reverse }}', + ]; + + it.each(ACCEPTED)('%s is accepted by the schema gate', (template) => { + expect(isWholeValueTemplateExpression(template)).toBe(true); + }); + + it.each(ACCEPTED)('%s renders to an array, not a string', (template) => { + const rendered = engine.render(template, context); + expect(Array.isArray(rendered)).toBe(true); + }); + + it('preserves the array contents rather than stringifying them', () => { + expect(engine.render('${{ workflow.inputs.recipients }}', context)).toEqual([ + 'a@b.com', + 'c@d.com', + ]); + }); + + // The engine accepts a superset of the gate: these forms are rejected by the schema + // precisely because they lose the array here. Each case documents why the gate is narrow. + describe('forms the schema rejects because the runtime cannot preserve the type', () => { + it('renders a bare {{ expr }} to a string', () => { + const template = '{{ workflow.inputs.recipients }}'; + expect(isWholeValueTemplateExpression(template)).toBe(false); + expect(typeof engine.render(template, context)).toBe('string'); + }); + + it('renders a padded ${{ expr }} to a string, because the runtime check does not trim', () => { + const template = ' ${{ workflow.inputs.recipients }}'; + expect(isWholeValueTemplateExpression(template)).toBe(false); + expect(typeof engine.render(template, context)).toBe('string'); + }); + + it('throws on two concatenated ${{ }} expressions', () => { + const template = '${{ workflow.inputs.recipients }}-${{ workflow.inputs.recipients }}'; + expect(isWholeValueTemplateExpression(template)).toBe(false); + expect(() => engine.render(template, context)).toThrow('The provided expression is invalid'); + }); + }); +}); diff --git a/src/platform/plugins/shared/workflows_execution_engine/server/types.ts b/src/platform/plugins/shared/workflows_execution_engine/server/types.ts index 6bac44d306d01..142787c8d69c5 100644 --- a/src/platform/plugins/shared/workflows_execution_engine/server/types.ts +++ b/src/platform/plugins/shared/workflows_execution_engine/server/types.ts @@ -17,8 +17,13 @@ import type { TaskManagerStartContract, } from '@kbn/task-manager-plugin/server'; import type { UsageApiSetup } from '@kbn/usage-api-plugin/server'; -import type { BulkScheduleWorkflowResult, WorkflowExecutionEngineModel } from '@kbn/workflows'; import type { + BulkScheduleWorkflowResult, + ExecutionStatus, + WorkflowExecutionEngineModel, +} from '@kbn/workflows'; +import type { + WorkflowExecutionCapabilities, WorkflowsExtensionsServerPluginSetup, WorkflowsExtensionsServerPluginStart, } from '@kbn/workflows-extensions/server'; @@ -43,6 +48,20 @@ export type { export interface ExecuteWorkflowResponse { workflowExecutionId: string; + result?: { + status: ExecutionStatus; + output?: Record; + error?: { type: string; message: string; details?: Record }; + }; +} + +export type WorkflowExecutionMode = 'async' | 'sync'; + +export interface ExecuteWorkflowOptions { + executionMode?: WorkflowExecutionMode; + executionId?: string; + capabilities?: WorkflowExecutionCapabilities; + abortSignal?: AbortSignal; } export interface ExecuteWorkflowStepResponse { @@ -70,6 +89,7 @@ export interface TriggerEventsContract { } export interface WorkflowsExecutionEnginePluginStart { + readonly supportsSynchronousExecution: true; __internalStorage: { workflowExecutionsDataClient: WorkflowExecutionsDataClient; stepExecutionsDataClient: StepExecutionsDataClient; @@ -104,7 +124,8 @@ export interface WorkflowsExecutionEnginePluginStartDeps { export type ExecuteWorkflow = ( workflow: WorkflowExecutionEngineModel, context: Record, - request: KibanaRequest + request: KibanaRequest, + options?: ExecuteWorkflowOptions ) => Promise; export type ExecuteWorkflowStep = ( diff --git a/src/platform/plugins/shared/workflows_execution_engine/server/workflow_context_manager/step_io_service.ts b/src/platform/plugins/shared/workflows_execution_engine/server/workflow_context_manager/step_io_service.ts index 83e1486e563cd..63d16cd3b867e 100644 --- a/src/platform/plugins/shared/workflows_execution_engine/server/workflow_context_manager/step_io_service.ts +++ b/src/platform/plugins/shared/workflows_execution_engine/server/workflow_context_manager/step_io_service.ts @@ -21,7 +21,7 @@ import { EVICTION_EXEMPT_STEP_TYPES, LOOP_STEP_TYPES } from './step_io_pinned_ty import type { StepExecutionMetadata, StepIoStateAccessor } from './workflow_execution_state'; import { WorkflowScopeStack } from './workflow_scope_stack'; import type { OutputSizeStats } from '../lib/telemetry/events/workflows_execution/types'; -import type { StepExecutionRepository } from '../repositories/step_execution_repository'; +import type { StepExecutionPersistence } from '../repositories/execution_persistence'; import { formatBytes, safeOutputSize } from '../step/errors'; import { buildStepExecutionId } from '../utils'; @@ -32,7 +32,7 @@ import { buildStepExecutionId } from '../utils'; export type PredecessorsResolver = (node: GraphNodeUnion) => ReadonlyArray; export interface StepIoServiceInit { - stepRepository: StepExecutionRepository; + stepRepository: StepExecutionPersistence; state: StepIoStateAccessor; pinnedStepTypes?: ReadonlySet; /** @@ -141,7 +141,7 @@ export interface StepIoLifecycle { * stepExecutionRepository. */ export class StepIoService implements StepIoWriter, StepIoLifecycle { - private readonly stepRepository: StepExecutionRepository; + private readonly stepRepository: StepExecutionPersistence; private readonly state: StepIoStateAccessor; private readonly pinnedStepTypes: ReadonlySet; private readonly evictionMinBytes: number; diff --git a/src/platform/plugins/shared/workflows_execution_engine/server/workflow_context_manager/types.ts b/src/platform/plugins/shared/workflows_execution_engine/server/workflow_context_manager/types.ts index 3cb5e81861d90..d9fcd15d54b32 100644 --- a/src/platform/plugins/shared/workflows_execution_engine/server/workflow_context_manager/types.ts +++ b/src/platform/plugins/shared/workflows_execution_engine/server/workflow_context_manager/types.ts @@ -12,11 +12,14 @@ import type { CloudSetup } from '@kbn/cloud-plugin/server'; import type { CoreStart, KibanaRequest } from '@kbn/core/server'; import type { TaskManagerStartContract } from '@kbn/task-manager-plugin/server'; import type { WorkflowRepository } from '@kbn/workflows'; -import type { WorkflowsExtensionsServerPluginStart } from '@kbn/workflows-extensions/server'; +import type { + WorkflowExecutionCapabilities, + WorkflowsExtensionsServerPluginStart, +} from '@kbn/workflows-extensions/server'; import type { WorkflowsExecutionEngineConfig } from '../config'; +import type { WorkflowExecutionPersistence } from '../repositories/execution_persistence'; import type { WorkflowLogEvent } from '../repositories/logs_repository'; import type { StepExecutionRepository } from '../repositories/step_execution_repository'; -import type { WorkflowExecutionRepository } from '../repositories/workflow_execution_repository'; import type { WorkflowsExecutionEnginePluginStart } from '../types'; export interface ContextDependencies { @@ -27,11 +30,13 @@ export interface ContextDependencies { workflowsExtensions: WorkflowsExtensionsServerPluginStart; config: WorkflowsExecutionEngineConfig; workflowRepository?: WorkflowRepository; - workflowExecutionRepository?: WorkflowExecutionRepository; + workflowExecutionRepository?: WorkflowExecutionPersistence; stepExecutionRepository?: StepExecutionRepository; workflowsExecutionEngine?: WorkflowsExecutionEnginePluginStart; spaceId?: string; request?: KibanaRequest; + /** Request-local privileged behavior; never serialize this value. */ + capabilities?: WorkflowExecutionCapabilities; } /** diff --git a/src/platform/plugins/shared/workflows_execution_engine/server/workflow_context_manager/workflow_context_manager.ts b/src/platform/plugins/shared/workflows_execution_engine/server/workflow_context_manager/workflow_context_manager.ts index 21d60e1b12f0c..7e194b6239812 100644 --- a/src/platform/plugins/shared/workflows_execution_engine/server/workflow_context_manager/workflow_context_manager.ts +++ b/src/platform/plugins/shared/workflows_execution_engine/server/workflow_context_manager/workflow_context_manager.ts @@ -19,6 +19,7 @@ import { type WorkflowContext, } from '@kbn/workflows'; import type { GraphNodeUnion, WorkflowGraph } from '@kbn/workflows/graph'; +import type { WorkflowExecutionCapabilities } from '@kbn/workflows-extensions/server'; import { buildWorkflowContext } from './build_workflow_context'; import type { StepIoService } from './step_io_service'; import type { ContextDependencies } from './types'; @@ -122,6 +123,10 @@ export class WorkflowContextManager { this.dependencies = init.dependencies; } + public getExecutionCapabilities(): WorkflowExecutionCapabilities | undefined { + return this.dependencies.capabilities; + } + /** * Pre-warms the execution state by rehydrating any evicted step outputs * that will be needed by `getContext()`. Must be called before `getContext()`. diff --git a/src/platform/plugins/shared/workflows_execution_engine/server/workflow_context_manager/workflow_execution_state.ts b/src/platform/plugins/shared/workflows_execution_engine/server/workflow_context_manager/workflow_execution_state.ts index 1818880c1e77f..04e60cc03d9b1 100644 --- a/src/platform/plugins/shared/workflows_execution_engine/server/workflow_context_manager/workflow_execution_state.ts +++ b/src/platform/plugins/shared/workflows_execution_engine/server/workflow_context_manager/workflow_execution_state.ts @@ -14,7 +14,7 @@ import type { WorkflowTokenUsage, } from '@kbn/workflows'; import { isTerminalStatus } from '@kbn/workflows'; -import type { WorkflowExecutionRepository } from '../repositories/workflow_execution_repository'; +import type { WorkflowExecutionPersistence } from '../repositories/execution_persistence'; import { sumTokenUsage } from '../utils'; /** Context for the step that failed during this run; used to build workflow_execution_failed event. */ @@ -100,7 +100,7 @@ export class WorkflowExecutionState { constructor( initialWorkflowExecution: EsWorkflowExecution, - private workflowExecutionRepository: WorkflowExecutionRepository + private workflowExecutionRepository: WorkflowExecutionPersistence ) { this.workflowExecution = initialWorkflowExecution; } diff --git a/src/platform/plugins/shared/workflows_execution_engine/server/workflow_execution_loop/cancel_workflow_if_requested.ts b/src/platform/plugins/shared/workflows_execution_engine/server/workflow_execution_loop/cancel_workflow_if_requested.ts index e62fa6e43a49a..618726c87e5ca 100644 --- a/src/platform/plugins/shared/workflows_execution_engine/server/workflow_execution_loop/cancel_workflow_if_requested.ts +++ b/src/platform/plugins/shared/workflows_execution_engine/server/workflow_execution_loop/cancel_workflow_if_requested.ts @@ -8,7 +8,7 @@ */ import { ExecutionStatus } from '@kbn/workflows'; -import type { WorkflowExecutionRepository } from '../repositories/workflow_execution_repository'; +import type { WorkflowExecutionPersistence } from '../repositories/execution_persistence'; import { buildStepExecutionId } from '../utils'; import type { StepExecutionRuntime } from '../workflow_context_manager/step_execution_runtime'; import type { WorkflowExecutionCursorApi } from '../workflow_context_manager/workflow_execution_cursor'; @@ -26,7 +26,7 @@ import type { IWorkflowEventLogger } from '../workflow_event_logger'; * issues from causing step execution failures. */ export async function cancelWorkflowIfRequested( - workflowExecutionRepository: WorkflowExecutionRepository, + workflowExecutionRepository: WorkflowExecutionPersistence, workflowExecutionState: WorkflowExecutionState, monitoredStepExecutionRuntime: StepExecutionRuntime, workflowLogger: IWorkflowEventLogger, diff --git a/src/platform/plugins/shared/workflows_execution_engine/server/workflow_execution_loop/catch_error.ts b/src/platform/plugins/shared/workflows_execution_engine/server/workflow_execution_loop/catch_error.ts index 78bdd43109a99..faa5f318dba1c 100644 --- a/src/platform/plugins/shared/workflows_execution_engine/server/workflow_execution_loop/catch_error.ts +++ b/src/platform/plugins/shared/workflows_execution_engine/server/workflow_execution_loop/catch_error.ts @@ -79,7 +79,14 @@ export async function catchError( // 2. There are items in the execution stack // 3. The top stack entry has nested scopes to process // This allows error handling to bubble up through the scope hierarchy. - if (failedStepExecutionRuntime.stepExecutionExists() && failedStepExecutionRuntime.error) { + // Don't overwrite a higher-priority error already set by onTaskAbort (e.g. the sync- + // execution timeout). The step's error is a consequence of that abort; the timeout + // message is what callers need to see. + if ( + failedStepExecutionRuntime.stepExecutionExists() && + failedStepExecutionRuntime.error && + !workflowExecutionCursor.error + ) { workflowExecutionCursor.captureError(failedStepExecutionRuntime.error); } else if (failedStepExecutionRuntime.stepExecutionExists()) { const stepExecution = failedStepExecutionRuntime.stepExecution; diff --git a/src/platform/plugins/shared/workflows_execution_engine/server/workflow_execution_loop/handle_execution_delay.test.ts b/src/platform/plugins/shared/workflows_execution_engine/server/workflow_execution_loop/handle_execution_delay.test.ts index fa35d0be92f2d..ebae16eb7f04c 100644 --- a/src/platform/plugins/shared/workflows_execution_engine/server/workflow_execution_loop/handle_execution_delay.test.ts +++ b/src/platform/plugins/shared/workflows_execution_engine/server/workflow_execution_loop/handle_execution_delay.test.ts @@ -71,6 +71,27 @@ const makeStepRuntime = ( } as unknown as jest.Mocked); describe('handleExecutionDelay', () => { + it.each([ + ExecutionStatus.WAITING, + ExecutionStatus.WAITING_FOR_INPUT, + ExecutionStatus.WAITING_FOR_CHILD, + ])('rejects asynchronous resume status %s in sync mode', async (status) => { + const params = makeParams(); + params.executionMode = 'sync'; + const stepRuntime = makeStepRuntime({ + node: { stepId: 'blocking-step', stepType: 'wait' } as any, + stepExecution: { status } as any, + }); + + await expect(handleExecutionDelay(params, stepRuntime)).rejects.toThrow( + 'is not supported in synchronous workflows' + ); + expect(params.workflowTaskManager.scheduleResumeTask).not.toHaveBeenCalled(); + expect( + params.workflowTaskManager.scheduleWorkflowGlobalTimeoutResumeTask + ).not.toHaveBeenCalled(); + }); + describe('WAITING_FOR_INPUT step (HITL)', () => { it('should set workflow status to WAITING_FOR_INPUT', async () => { const params = makeParams(); diff --git a/src/platform/plugins/shared/workflows_execution_engine/server/workflow_execution_loop/handle_execution_delay.ts b/src/platform/plugins/shared/workflows_execution_engine/server/workflow_execution_loop/handle_execution_delay.ts index d3a10c9124faf..219c66502a72c 100644 --- a/src/platform/plugins/shared/workflows_execution_engine/server/workflow_execution_loop/handle_execution_delay.ts +++ b/src/platform/plugins/shared/workflows_execution_engine/server/workflow_execution_loop/handle_execution_delay.ts @@ -86,7 +86,7 @@ async function scheduleWorkflowGlobalTimeoutResumeTask( await params.workflowTaskManager .scheduleWorkflowGlobalTimeoutResumeTask({ - workflowExecution: workflowExecution as EsWorkflowExecution, + workflowExecution, resumeAt: new Date(resumeAtMs), fakeRequest: params.fakeRequest, }) @@ -152,7 +152,7 @@ export async function ensureWorkflowIdleTimeoutResumeAfterLoop( await params.workflowTaskManager .scheduleWorkflowGlobalTimeoutResumeTask({ - workflowExecution: workflowExecution as EsWorkflowExecution, + workflowExecution, resumeAt, fakeRequest: params.fakeRequest, }) @@ -172,6 +172,16 @@ export async function handleExecutionDelay( const workflowExecution = params.workflowRuntime.getWorkflowExecution(); const stepStatus = stepExecutionRuntime.stepExecution?.status; + if ( + params.executionMode === 'sync' && + (stepStatus === ExecutionStatus.WAITING_FOR_INPUT || + stepStatus === ExecutionStatus.WAITING_FOR_CHILD || + stepStatus === ExecutionStatus.WAITING) + ) { + throw new Error( + `Step "${stepExecutionRuntime.node.stepId}" is not supported in synchronous workflows` + ); + } if ( stepStatus === ExecutionStatus.WAITING_FOR_INPUT || stepStatus === ExecutionStatus.WAITING_FOR_CHILD diff --git a/src/platform/plugins/shared/workflows_execution_engine/server/workflow_execution_loop/run_node.ts b/src/platform/plugins/shared/workflows_execution_engine/server/workflow_execution_loop/run_node.ts index 010ae39d5046f..efbfb0df6ddda 100644 --- a/src/platform/plugins/shared/workflows_execution_engine/server/workflow_execution_loop/run_node.ts +++ b/src/platform/plugins/shared/workflows_execution_engine/server/workflow_execution_loop/run_node.ts @@ -81,6 +81,7 @@ export async function runNode(params: WorkflowExecutionLoopParams): Promise void) | undefined; if (!node) { return; @@ -106,6 +107,17 @@ export async function runNode(params: WorkflowExecutionLoopParams): Promise runtimeRef.abortController.abort(params.signal.reason); + if (params.signal.aborted) { + abortStepRelay(); + } else { + params.signal.addEventListener('abort', abortStepRelay, { once: true }); + } + // Build the node implementation before the cancel short-circuit so cancellable nodes // (e.g. workflow.execute holding a child execution) still get their onCancel hook. nodeImplementation = params.nodesFactory.create(stepExecutionRuntime); @@ -178,9 +190,17 @@ export async function runNode(params: WorkflowExecutionLoopParams): Promise; coreStart: CoreStart; signal: AbortSignal; workflowTaskManager: WorkflowTaskManager; + executionMode?: WorkflowExecutionMode; // defaults to 'async' when absent } diff --git a/src/platform/plugins/shared/workflows_execution_engine/server/workflow_execution_loop/workflow_execution_loop.test.ts b/src/platform/plugins/shared/workflows_execution_engine/server/workflow_execution_loop/workflow_execution_loop.test.ts index ab36a00bfe803..149e6b975818f 100644 --- a/src/platform/plugins/shared/workflows_execution_engine/server/workflow_execution_loop/workflow_execution_loop.test.ts +++ b/src/platform/plugins/shared/workflows_execution_engine/server/workflow_execution_loop/workflow_execution_loop.test.ts @@ -62,10 +62,13 @@ describe('workflowExecutionLoop', () => { const params = createParams(); await workflowExecutionLoop(params as any); - // eslint-disable-next-line @typescript-eslint/no-var-requires - const { executionFlowLoop } = require('./execution_flow_loop'); - // eslint-disable-next-line @typescript-eslint/no-var-requires - const { persistenceLoop, flushState } = require('./persistence_loop'); + const { executionFlowLoop } = jest.requireMock('./execution_flow_loop') as { + executionFlowLoop: jest.Mock; + }; + const { persistenceLoop, flushState } = jest.requireMock('./persistence_loop') as { + persistenceLoop: jest.Mock; + flushState: jest.Mock; + }; expect(executionFlowLoop).toHaveBeenCalledWith(params); expect(persistenceLoop).toHaveBeenCalled(); @@ -78,12 +81,32 @@ describe('workflowExecutionLoop', () => { expect(params.workflowLogger.flushEvents).toHaveBeenCalled(); }); + it('uses the shared execution loop without periodic persistence in sync mode', async () => { + const params = { ...createParams(), executionMode: 'sync' as const }; + await workflowExecutionLoop(params as any); + + const { executionFlowLoop } = jest.requireMock('./execution_flow_loop') as { + executionFlowLoop: jest.Mock; + }; + const { persistenceLoop, flushState } = jest.requireMock('./persistence_loop') as { + persistenceLoop: jest.Mock; + flushState: jest.Mock; + }; + + expect(executionFlowLoop).toHaveBeenCalledWith(params); + expect(persistenceLoop).not.toHaveBeenCalled(); + expect(flushState).not.toHaveBeenCalled(); + expect(params.workflowRuntime.saveState).toHaveBeenCalled(); + expect(params.stepIoService.flush).toHaveBeenCalled(); + }); + it('sets workflow error when execution loop throws', async () => { const params = createParams(); const testError = new Error('execution failed'); - // eslint-disable-next-line @typescript-eslint/no-var-requires - const { executionFlowLoop } = require('./execution_flow_loop'); + const { executionFlowLoop } = jest.requireMock('./execution_flow_loop') as { + executionFlowLoop: jest.Mock; + }; (executionFlowLoop as jest.Mock).mockRejectedValueOnce(testError); await workflowExecutionLoop(params as any); @@ -109,11 +132,42 @@ describe('workflowExecutionLoop', () => { expect(params.workflowExecutionCursor.stop).toHaveBeenCalled(); }); + it('marks abort-with-error-reason (e.g. timeout) as FAILED and captures the error', async () => { + const params = createParams(); + const abortController = new AbortController(); + const timeoutError = new Error('Pre-LLM anonymization timed out after 5000ms'); + const loopPromise = workflowExecutionLoop({ ...params, signal: abortController.signal } as any); + abortController.abort(timeoutError); + await loopPromise; + + expect(params.workflowExecutionCursor.captureError).toHaveBeenCalledWith(timeoutError); + expect(params.workflowExecutionState.updateWorkflowExecution).toHaveBeenCalledWith( + expect.objectContaining({ + status: ExecutionStatus.FAILED, + }) + ); + expect(params.workflowExecutionCursor.stop).toHaveBeenCalled(); + }); + + it('treats AbortError DOMException as CANCELLED (not FAILED) — signal from sync timeout or platform abort', async () => { + const params = createParams(); + const abortController = new AbortController(); + const loopPromise = workflowExecutionLoop({ ...params, signal: abortController.signal } as any); + abortController.abort(new DOMException('timeout', 'AbortError')); + await loopPromise; + + expect(params.workflowExecutionState.updateWorkflowExecution).toHaveBeenCalledWith( + expect.objectContaining({ + status: ExecutionStatus.CANCELLED, + }) + ); + expect(params.workflowExecutionCursor.captureError).not.toHaveBeenCalled(); + }); + it('marks Task Manager abort as system cancellation and suppresses workflow log errors', async () => { const params = createParams(); const abortController = new AbortController(); - // eslint-disable-next-line @typescript-eslint/no-var-requires - const { flushState } = require('./persistence_loop'); + const { flushState } = jest.requireMock('./persistence_loop') as { flushState: jest.Mock }; const loopPromise = workflowExecutionLoop({ ...params, signal: abortController.signal } as any); abortController.abort(new WorkflowTaskManagerAbortError()); await loopPromise; diff --git a/src/platform/plugins/shared/workflows_execution_engine/server/workflow_execution_loop/workflow_execution_loop.ts b/src/platform/plugins/shared/workflows_execution_engine/server/workflow_execution_loop/workflow_execution_loop.ts index af2d06439449b..1cd85fc6aee34 100644 --- a/src/platform/plugins/shared/workflows_execution_engine/server/workflow_execution_loop/workflow_execution_loop.ts +++ b/src/platform/plugins/shared/workflows_execution_engine/server/workflow_execution_loop/workflow_execution_loop.ts @@ -65,16 +65,23 @@ export async function workflowExecutionLoop(params: WorkflowExecutionLoopParams) return; } - params.workflowExecutionState.updateWorkflowExecution({ - cancelRequested: true, - cancelledAt: new Date().toISOString(), - cancellationReason: 'Task aborted', - status: ExecutionStatus.CANCELLED, - }); - if (wasWaitingForInput) { - emitHitlLifecycle({ - type: 'canceled', - executionId: workflowRuntime.getWorkflowExecution().id, + const reason = params.signal.reason; + // When the abort carries an explicit error (e.g. a pre-LLM or sync-execution + // timeout), treat it as FAILED so the error detail is surfaced to callers. + // A plain abort (no reason, or an AbortError from a no-arg abort()) keeps + // CANCELLED semantics — that represents an intentional stop (user disconnect). + const isFailureAbort = reason instanceof Error && reason.name !== 'AbortError'; + if (isFailureAbort) { + workflowExecutionCursor.captureError(reason); + params.workflowExecutionState.updateWorkflowExecution({ + status: ExecutionStatus.FAILED, + }); + } else { + params.workflowExecutionState.updateWorkflowExecution({ + cancelRequested: true, + cancelledAt: new Date().toISOString(), + cancellationReason: 'Task aborted', + status: ExecutionStatus.CANCELLED, }); } // Also abort persistence loop when task is aborted @@ -91,20 +98,26 @@ export async function workflowExecutionLoop(params: WorkflowExecutionLoopParams) workflowExecutionCursor.start(); // Run execution and persistence loops in parallel // When execution finishes, signal persistence loop to exit immediately - await Promise.all([ - executionFlowLoop(params).finally(() => { - // Signal persistence loop to stop waiting and exit - persistenceAbortController.abort(); - }), - persistenceLoop(params, persistenceAbortController.signal), - ]); + if (params.executionMode === 'sync') { + await executionFlowLoop(params); + } else { + await Promise.all([ + executionFlowLoop(params).finally(() => { + // Signal persistence loop to stop waiting and exit + persistenceAbortController.abort(); + }), + persistenceLoop(params, persistenceAbortController.signal), + ]); + } } catch (error) { workflowExecutionCursor.captureError(error); } finally { const finalFlushSpan = apm.startSpan('final flush state', 'workflow', 'persistence'); - await flushState(params, { - workflowLogFlushSignal: params.signal, - }); + if (params.executionMode !== 'sync') { + await flushState(params, { + workflowLogFlushSignal: params.signal, + }); + } finalFlushSpan?.end(); } diff --git a/src/platform/plugins/shared/workflows_extensions/.claude/skills/workflows-custom-triggers/SKILL.md b/src/platform/plugins/shared/workflows_extensions/.claude/skills/workflows-custom-triggers/SKILL.md index 09b5b3d3f7906..7a93e98e5cc8e 100644 --- a/src/platform/plugins/shared/workflows_extensions/.claude/skills/workflows-custom-triggers/SKILL.md +++ b/src/platform/plugins/shared/workflows_extensions/.claude/skills/workflows-custom-triggers/SKILL.md @@ -31,7 +31,7 @@ A custom workflow trigger is owned and registered by a **plugin other than `work A trigger lives in three layers: -- **Common** — `id`, `eventSchema`, `title`, `description`, `stability`, optional `documentation` / `snippets`. Imported by both server and public to keep them in sync. Set `stability` to `'tech_preview'`, `'beta'`, or `'stable'` based on the trigger's maturity. +- **Common** — `id`, `eventSchema`, `title`, `description`, `stability`, optional `documentation` / `snippets`, optional `exclusivity`. Imported by both server and public to keep them in sync. Set `stability` to `'tech_preview'`, `'beta'`, or `'stable'` based on the trigger's maturity. Set `exclusivity: 'per-space'` when at most one enabled workflow may subscribe to this trigger per space (omit for fan-out triggers). - **Server** — registers the **same** common definition via `registerTriggerDefinition`; emits events with `emitEvent`. - **Public** — spreads the common definition and adds **icon** only (browser-only UI). diff --git a/src/platform/plugins/shared/workflows_extensions/common/index.ts b/src/platform/plugins/shared/workflows_extensions/common/index.ts index 1959d04b60b8d..35b5a13631890 100644 --- a/src/platform/plugins/shared/workflows_extensions/common/index.ts +++ b/src/platform/plugins/shared/workflows_extensions/common/index.ts @@ -11,9 +11,10 @@ export type { CommonStepDefinition } from './step_registry/types'; export type { CommonTriggerDefinition, TriggerDocumentation, + TriggerExclusivity, TriggerSnippets, } from './trigger_registry/types'; -export { EVENT_FIELD_PREFIX } from './trigger_registry/constants'; +export { EVENT_FIELD_PREFIX, TRIGGER_EXCLUSIVITY_SCOPES } from './trigger_registry/constants'; export { DataMapStepTypeId, DEFAULT_INDEX_BINDING, diff --git a/src/platform/plugins/shared/workflows_extensions/common/trigger_registry/constants.ts b/src/platform/plugins/shared/workflows_extensions/common/trigger_registry/constants.ts index 28072e6da7ab2..b23b044537414 100644 --- a/src/platform/plugins/shared/workflows_extensions/common/trigger_registry/constants.ts +++ b/src/platform/plugins/shared/workflows_extensions/common/trigger_registry/constants.ts @@ -12,3 +12,15 @@ * Use this when calling validateKqlAgainstSchema for trigger conditions so paths and error messages stay consistent. */ export const EVENT_FIELD_PREFIX = 'event.'; + +/** + * Valid exclusivity scopes a trigger definition may declare. + * Kept as a const tuple so the {@link TriggerExclusivity} type and the + * registration-time validator in TriggerRegistry share one source of truth. + * + * Deliberately has no 'global' member: trigger dispatch always derives a single + * spaceId from the emitting request, and `spaceId: '*'` (global) workflows are + * a single document, so cross-space exclusivity is structurally guaranteed and + * needs no enforcement here. Add a member only when a real fan-out case exists. + */ +export const TRIGGER_EXCLUSIVITY_SCOPES = ['per-space'] as const; diff --git a/src/platform/plugins/shared/workflows_extensions/common/trigger_registry/types.ts b/src/platform/plugins/shared/workflows_extensions/common/trigger_registry/types.ts index 0a86ef77efa0d..07631f6fff402 100644 --- a/src/platform/plugins/shared/workflows_extensions/common/trigger_registry/types.ts +++ b/src/platform/plugins/shared/workflows_extensions/common/trigger_registry/types.ts @@ -10,6 +10,17 @@ import type { StabilityLevel } from '@kbn/workflows'; import type { z } from '@kbn/zod/v4'; +import type { TRIGGER_EXCLUSIVITY_SCOPES } from './constants'; + +/** + * Narrows how many workflows may be *enabled* for a trigger at a time. + * - `'per-space'`: at most one enabled workflow per Kibana space. Global + * (`spaceId: '*'`) workflows count as a subscriber in every space. + * + * Omit for triggers that fan out to any number of subscribers (the default). + */ +export type TriggerExclusivity = (typeof TRIGGER_EXCLUSIVITY_SCOPES)[number]; + /** * Documentation for a trigger (aligned with steps: details + examples). */ @@ -42,6 +53,7 @@ export interface TriggerSnippets { * Constraints (enforced at registration): * - id: globally unique, namespaced format . * - eventSchema: must be a Zod object schema that rejects unknown fields + * - exclusivity: must be a value from {@link TRIGGER_EXCLUSIVITY_SCOPES} when set * * Server and agent tooling read title and description from here; documentation and snippets are optional. * Public definitions spread this object and add UI-only fields (e.g. icon). @@ -82,4 +94,14 @@ export interface CommonTriggerDefinition( definition: ServerPollStepDefinition ): ServerPollStepDefinition { + definition.supportedExecutionModes = ['async']; + if (!definition.ceilings) { definition.ceilings = PollStepDefaults.ceilings; } @@ -383,6 +393,30 @@ export interface StepHandlerContext { * Current step's type */ stepType: string; + + /** + * Request-local privileged capabilities supplied by the execution caller. + * + * **Injection scope:** injected unconditionally into every step handler context + * regardless of step type. Only step handlers that know the exact `id` and pass + * the value sentinel check (via `getCapabilityValue`) can read a capability; + * all other handlers see `undefined` from that helper. + * + * **Threat model:** capability `value` objects are intentionally designed to be + * non-serializable. Implementations must store their real payload behind a + * non-enumerable Symbol key in a WeakMap and freeze the container. + * `JSON.stringify` strips Symbol keys, so the ES-persistence path sees `[{}]`, + * not the underlying function references. A handler that attempts to include + * capabilities in its step output or log them loses the real payload on the + * round-trip — this is a deliberate self-destruct guarantee. + * + * **Do not** add capabilities whose `value` holds plain JSON-serializable data. + * That bypasses the self-destruct guarantee and exposes the value to every handler. + * + * These must never be added to workflow context, template variables, or step + * input/output directly. + */ + capabilities?: WorkflowExecutionCapabilities; } /** diff --git a/src/platform/plugins/shared/workflows_extensions/server/trigger_registry/trigger_registry.test.ts b/src/platform/plugins/shared/workflows_extensions/server/trigger_registry/trigger_registry.test.ts index ccd859beae575..058bd3242dae4 100644 --- a/src/platform/plugins/shared/workflows_extensions/server/trigger_registry/trigger_registry.test.ts +++ b/src/platform/plugins/shared/workflows_extensions/server/trigger_registry/trigger_registry.test.ts @@ -155,6 +155,40 @@ describe('TriggerRegistry', () => { ); }).toThrow('"eventSchema" must be a Zod object schema'); }); + + describe('exclusivity', () => { + it('accepts a definition with exclusivity omitted, and it round-trips as undefined', () => { + const def = createValidDefinition(); + registry.register(def); + expect(registry.get('cases.updated')?.exclusivity).toBeUndefined(); + }); + + it('accepts exclusivity: "per-space" and round-trips the value', () => { + const def = createValidDefinition({ exclusivity: 'per-space' }); + registry.register(def); + expect(registry.get('cases.updated')?.exclusivity).toBe('per-space'); + }); + + it('throws for exclusivity: "global" — pins the deliberate absence of a global scope', () => { + expect(() => { + registry.register( + createValidDefinition({ exclusivity: 'global' as unknown as 'per-space' }) + ); + }).toThrow('"exclusivity" must be one of "per-space"'); + }); + + it('throws for exclusivity: "" (empty string)', () => { + expect(() => { + registry.register(createValidDefinition({ exclusivity: '' as unknown as 'per-space' })); + }).toThrow('"exclusivity" must be one of "per-space"'); + }); + + it('throws for exclusivity: true (boolean)', () => { + expect(() => { + registry.register(createValidDefinition({ exclusivity: true as unknown as 'per-space' })); + }).toThrow('"exclusivity" must be one of "per-space"'); + }); + }); }); describe('freeze', () => { diff --git a/src/platform/plugins/shared/workflows_extensions/server/trigger_registry/trigger_registry.ts b/src/platform/plugins/shared/workflows_extensions/server/trigger_registry/trigger_registry.ts index 65df3830904e9..635fa7252fcd9 100644 --- a/src/platform/plugins/shared/workflows_extensions/server/trigger_registry/trigger_registry.ts +++ b/src/platform/plugins/shared/workflows_extensions/server/trigger_registry/trigger_registry.ts @@ -8,6 +8,7 @@ */ import type { z } from '@kbn/zod/v4'; +import { TRIGGER_EXCLUSIVITY_SCOPES } from '../../common'; import type { ServerTriggerDefinition } from '../types'; /** Id must be . (kebab-case namespace, camelCase event; e.g. my-namespace.customTrigger) */ @@ -18,7 +19,7 @@ function isZodObject(schema: z.ZodType): schema is z.ZodObject { } function validateDefinition(definition: ServerTriggerDefinition): void { - const { id, eventSchema } = definition; + const { id, eventSchema, exclusivity } = definition; if (typeof id !== 'string' || id.length === 0) { throw new Error('Trigger definition "id" must be a non-empty string.'); @@ -36,6 +37,16 @@ function validateDefinition(definition: ServerTriggerDefinition): void { `Trigger "${id}": "eventSchema" must be a Zod object schema (e.g. z.object({...})).` ); } + if ( + exclusivity !== undefined && + !(TRIGGER_EXCLUSIVITY_SCOPES as readonly string[]).includes(exclusivity) + ) { + throw new Error( + `Trigger "${id}": "exclusivity" must be one of ${TRIGGER_EXCLUSIVITY_SCOPES.map( + (scope) => `"${scope}"` + ).join(', ')} when set (received ${JSON.stringify(exclusivity)}).` + ); + } } /** diff --git a/src/platform/plugins/shared/workflows_management/server/api/lib/workflow_prepare.test.ts b/src/platform/plugins/shared/workflows_management/server/api/lib/workflow_prepare.test.ts index 2f254aca3f408..18596e0bd4a0f 100644 --- a/src/platform/plugins/shared/workflows_management/server/api/lib/workflow_prepare.test.ts +++ b/src/platform/plugins/shared/workflows_management/server/api/lib/workflow_prepare.test.ts @@ -410,4 +410,25 @@ describe('prepareWorkflowDocumentFromYaml', () => { expect(result.workflowData.name).toBe('From YAML'); }); + + it('retains managed provenance without retaining managed ownership', () => { + const result = prepareWorkflowDocumentFromYaml({ + yaml: 'name: Managed copy\nenabled: false\ntriggers: []\nsteps: []', + zodSchema: getWorkflowZodSchema({}), + authenticatedUser: 'user1', + now, + spaceId: 'default', + originManagedWorkflowId: 'system-inference_pii_anonymization', + }); + + expect(result.workflowData).toEqual( + expect.objectContaining({ + enabled: false, + managed: false, + managedBy: null, + lifecycle: null, + originManagedWorkflowId: 'system-inference_pii_anonymization', + }) + ); + }); }); diff --git a/src/platform/plugins/shared/workflows_management/server/api/lib/workflow_prepare.ts b/src/platform/plugins/shared/workflows_management/server/api/lib/workflow_prepare.ts index ea22e43afc153..f3cf47ca2ab3e 100644 --- a/src/platform/plugins/shared/workflows_management/server/api/lib/workflow_prepare.ts +++ b/src/platform/plugins/shared/workflows_management/server/api/lib/workflow_prepare.ts @@ -28,6 +28,38 @@ export const getTriggerTypesFromDefinition = ( .filter((v: T): v is NonNullable => v != null); }; +/** + * Reads trigger type ids straight off the YAML without requiring schema + * validation, so callers see the same trigger set the storage document will + * derive from `getTriggerTypesFromDefinition`. Returns `[]` on parse failure. + */ +export const getTriggerTypesFromYaml = (yamlString: string): string[] => { + const parsed = parseYamlToJSONWithoutValidation(yamlString); + if (!parsed.success || parsed.json == null || typeof parsed.json !== 'object') { + return []; + } + const triggers = (parsed.json as { triggers?: unknown }).triggers; + if (!Array.isArray(triggers)) { + return []; + } + return triggers + .map((t) => + t && typeof t === 'object' && typeof (t as { type?: unknown }).type === 'string' + ? (t as { type: string }).type + : null + ) + .filter((v): v is string => v != null); +}; + +/** True when the YAML root map sets `enabled: true` at the top level. */ +export const workflowYamlDeclaresEnabled = (yamlString: string): boolean => { + const parsed = parseYamlToJSONWithoutValidation(yamlString); + if (!parsed.success || parsed.json == null || typeof parsed.json !== 'object') { + return false; + } + return (parsed.json as { enabled?: unknown }).enabled === true; +}; + /** True when the YAML root map includes `enabled` (before Zod defaults). */ export const workflowYamlDeclaresTopLevelEnabled = (yamlString: string): boolean => { const parsed = parseYamlToJSONWithoutValidation(yamlString); @@ -71,6 +103,7 @@ export const prepareWorkflowDocumentFromYaml = (params: { authenticatedUser: string; now: Date; spaceId: string; + originManagedWorkflowId?: string; triggerDefinitions?: Array<{ id: string; eventSchema: z.ZodType }>; nameFallback?: string; }): { id: string; workflowData: WorkflowProperties; definition?: WorkflowYaml } => { @@ -81,6 +114,7 @@ export const prepareWorkflowDocumentFromYaml = (params: { authenticatedUser, now, spaceId, + originManagedWorkflowId, triggerDefinitions, nameFallback, } = params; @@ -124,7 +158,7 @@ export const prepareWorkflowDocumentFromYaml = (params: { managed: false, managedBy: null, definitionHash: null, - originManagedWorkflowId: null, + originManagedWorkflowId: originManagedWorkflowId ?? null, lifecycle: null, valid: workflowToCreate.valid, deleted_at: null, diff --git a/src/platform/plugins/shared/workflows_management/server/api/workflows_management_api.test.ts b/src/platform/plugins/shared/workflows_management/server/api/workflows_management_api.test.ts index fc1ab51001e12..00ca48fe37f1b 100644 --- a/src/platform/plugins/shared/workflows_management/server/api/workflows_management_api.test.ts +++ b/src/platform/plugins/shared/workflows_management/server/api/workflows_management_api.test.ts @@ -12,9 +12,11 @@ import type { KibanaRequest, Logger } from '@kbn/core/server'; import { loggingSystemMock } from '@kbn/core/server/mocks'; import { httpServerMock } from '@kbn/core-http-server-mocks'; import { + ExecutionStatus, type WorkflowDetailDto, type WorkflowExecutionEngineModel, WorkflowsManagementApiActions, + type WorkflowYaml, } from '@kbn/workflows'; import { WorkflowExecutionInvalidStatusError, @@ -22,6 +24,9 @@ import { } from '@kbn/workflows/common/errors'; import type { WorkflowsExecutionEnginePluginStart } from '@kbn/workflows-execution-engine/server'; import { workflowsExecutionEngineMock } from '@kbn/workflows-execution-engine/server/mocks'; +import type { WorkflowsExtensionsServerPluginStart } from '@kbn/workflows-extensions/server'; +import { workflowsExtensionsMock } from '@kbn/workflows-extensions/server/mocks'; +import { WorkflowConflictError } from '@kbn/workflows-yaml'; import { z } from '@kbn/zod/v4'; import { resumeWorkflowExecutionExternallyViaGet, @@ -60,9 +65,14 @@ describe('WorkflowsManagementApi', () => { let mockWorkflowsExecutionEngine: jest.Mocked; const logger = loggingSystemMock.createLogger(); const mockPreprocessAlertInputs = jest.mocked(preprocessAlertInputs); + let mockWorkflowsExtensions: jest.Mocked; beforeEach(() => { jest.clearAllMocks(); + mockWorkflowsExtensions = workflowsExtensionsMock.createStart(); + // By default no trigger declares exclusivity; individual test suites override as needed. + mockWorkflowsExtensions.getTriggerDefinition.mockReturnValue(undefined); + mockWorkflowsExecutionEngine = workflowsExecutionEngineMock.createStart(); mockWorkflowsExecutionEngine.executeWorkflow.mockResolvedValue({ workflowExecutionId: 'test-exec-id', @@ -75,6 +85,7 @@ describe('WorkflowsManagementApi', () => { mockWorkflowsService = { getWorkflow: jest.fn(), + getWorkflowsSubscribedToTrigger: jest.fn(), getWorkflowsByIds: jest.fn(), getWorkflowZodSchema: jest.fn(), createWorkflow: jest.fn(), @@ -89,6 +100,7 @@ describe('WorkflowsManagementApi', () => { markStepAsResponded: jest.fn(), getWaitingStepExecutionId: jest.fn(), getWorkflowsExecutionEngine: () => mockWorkflowsExecutionEngine, + getWorkflowsExtensions: async () => mockWorkflowsExtensions, } as any; api = new WorkflowsManagementApi(mockWorkflowsService, true, logger); @@ -108,6 +120,164 @@ describe('WorkflowsManagementApi', () => { }); }; + describe('workflow-trigger synchronous execution', () => { + const createAroundCompletionTrigger = ( + condition?: string + ): WorkflowYaml['triggers'][number] => { + // The static WorkflowYaml type models built-ins only; registered triggers are added dynamically. + const trigger: WorkflowYaml['triggers'][number] = { type: 'manual' }; + Reflect.set(trigger, 'type', 'inference.aroundCompletion'); + if (condition) { + Reflect.set(trigger, 'on', { condition }); + } + return trigger; + }; + + const createTriggeredWorkflow = ({ + id, + condition, + }: { + id: string; + condition?: string; + }): WorkflowDetailDto => ({ + id, + name: id, + enabled: true, + yaml: `name: ${id}`, + valid: true, + createdAt: '2026-01-01T00:00:00.000Z', + createdBy: 'system', + lastUpdatedAt: '2026-01-01T00:00:00.000Z', + lastUpdatedBy: 'system', + definition: { + version: '1', + name: id, + enabled: true, + triggers: [createAroundCompletionTrigger(condition)], + steps: [{ name: 'proceed', type: 'call_site.proceed', with: {} }], + }, + }); + + it('resolves matching workflows and reports invalid trigger conditions separately', async () => { + mockWorkflowsService.getWorkflowsSubscribedToTrigger.mockResolvedValue([ + createTriggeredWorkflow({ id: 'unconditional' }), + createTriggeredWorkflow({ id: 'matching', condition: 'event.agentId: "agent-a"' }), + createTriggeredWorkflow({ id: 'not-matching', condition: 'event.agentId: "agent-b"' }), + createTriggeredWorkflow({ id: 'invalid', condition: '(' }), + ]); + + await expect( + api.resolveWorkflowTriggerMatches( + 'inference.aroundCompletion', + { agentId: 'agent-a' }, + 'space-a' + ) + ).resolves.toEqual({ + matched: [ + expect.objectContaining({ id: 'unconditional' }), + expect.objectContaining({ id: 'matching' }), + ], + invalidConditionWorkflows: [{ id: 'invalid', name: 'invalid' }], + }); + expect(mockWorkflowsService.getWorkflowsSubscribedToTrigger).toHaveBeenCalledWith( + 'inference.aroundCompletion', + 'space-a' + ); + }); + + it('executes a saved workflow synchronously with request-local options', async () => { + const workflow = createTriggeredWorkflow({ id: 'workflow-1' }); + const abortSignal = new AbortController().signal; + const capabilities = [{ id: 'test-capability', value: { invoke: jest.fn() } }]; + mockWorkflowsService.getWorkflow.mockResolvedValue(workflow); + mockWorkflowsExecutionEngine.executeWorkflow.mockResolvedValue({ + workflowExecutionId: 'execution-1', + result: { status: ExecutionStatus.COMPLETED, output: { content: 'restored' } }, + }); + + await api.executeWorkflowSynchronously({ + workflowId: workflow.id, + context: { event: { messages: [] }, spaceId: 'space-a' }, + spaceId: 'space-a', + request: mockRequest, + capabilities, + abortSignal, + }); + + expect(mockWorkflowsExecutionEngine.executeWorkflow).toHaveBeenCalledWith( + expect.objectContaining({ id: workflow.id }), + { event: { messages: [] }, spaceId: 'space-a' }, + mockRequest, + { + executionMode: 'sync', + capabilities, + abortSignal, + } + ); + }); + + it('skips the ES re-fetch when a pre-fetched workflow DTO is provided', async () => { + const workflow = createTriggeredWorkflow({ id: 'workflow-1' }); + mockWorkflowsExecutionEngine.executeWorkflow.mockResolvedValue({ + workflowExecutionId: 'execution-1', + result: { status: ExecutionStatus.COMPLETED, output: { content: 'restored' } }, + }); + + await api.executeWorkflowSynchronously({ + workflowId: workflow.id, + workflow, + context: { event: { messages: [] }, spaceId: 'space-a' }, + spaceId: 'space-a', + request: mockRequest, + }); + + expect(mockWorkflowsService.getWorkflow).not.toHaveBeenCalled(); + expect(mockWorkflowsExecutionEngine.executeWorkflow).toHaveBeenCalledWith( + expect.objectContaining({ id: workflow.id }), + expect.any(Object), + mockRequest, + expect.objectContaining({ executionMode: 'sync' }) + ); + }); + + it('validates the pre-fetched DTO and throws without an ES fetch when the workflow is disabled, invalid, or missing a definition', async () => { + const base = createTriggeredWorkflow({ id: 'wf-guard' }); + + await expect( + api.executeWorkflowSynchronously({ + workflowId: base.id, + workflow: { ...base, enabled: false }, + context: {}, + spaceId: 'space-a', + request: mockRequest, + }) + ).rejects.toThrow('disabled'); + + await expect( + api.executeWorkflowSynchronously({ + workflowId: base.id, + workflow: { ...base, valid: false }, + context: {}, + spaceId: 'space-a', + request: mockRequest, + }) + ).rejects.toThrow('validation errors'); + + await expect( + api.executeWorkflowSynchronously({ + workflowId: base.id, + workflow: { ...base, definition: null }, + context: {}, + spaceId: 'space-a', + request: mockRequest, + }) + ).rejects.toThrow('no definition'); + + // All three guard paths use the supplied DTO — no ES re-fetch + expect(mockWorkflowsService.getWorkflow).not.toHaveBeenCalled(); + }); + }); + describe('cloneWorkflow', () => { const createMockWorkflow = (overrides: Partial = {}): WorkflowDetailDto => ({ id: 'workflow-123', @@ -334,6 +504,31 @@ enabled: true`; { nameFallback: 'Original Workflow Copy' } ); }); + + it('creates managed clones as disabled user-owned workflows with provenance', async () => { + const originalWorkflow = createMockWorkflow({ + managed: true, + managedBy: 'inferenceWorkflows', + originManagedWorkflowId: 'system-inference_pii_anonymization', + }); + mockWorkflowsService.createWorkflow.mockResolvedValue({ + ...originalWorkflow, + id: 'workflow-clone-456', + managed: false, + managedBy: null, + }); + + await api.cloneWorkflow(originalWorkflow, 'default', mockRequest); + + expect(mockWorkflowsService.createWorkflow).toHaveBeenCalledWith( + { + yaml: expect.stringContaining('enabled: false'), + }, + 'default', + mockRequest, + { originManagedWorkflowId: 'system-inference_pii_anonymization' } + ); + }); }); describe('testWorkflow', () => { @@ -1143,6 +1338,166 @@ steps: expect(mockWorkflowsService.updateWorkflow).toHaveBeenCalled(); }); + + describe('exclusive trigger conflict check', () => { + // Fixture trigger IDs — not real production triggers, so this suite is never + // accidentally invalidated by upstream registration changes. + const EXCLUSIVE_TRIGGER = 'test-ns.exclusiveTrigger'; + const SHARED_TRIGGER = 'test-ns.sharedTrigger'; + + const exclusiveDefinition = { + triggers: [{ type: EXCLUSIVE_TRIGGER }], + } as unknown as WorkflowDetailDto['definition']; + + beforeEach(() => { + // Override the top-level default: make EXCLUSIVE_TRIGGER exclusive, SHARED_TRIGGER not. + mockWorkflowsExtensions.getTriggerDefinition.mockImplementation((id) => { + if (id === EXCLUSIVE_TRIGGER) { + return { id, exclusivity: 'per-space' } as any; + } + if (id === SHARED_TRIGGER) { + return { id } as any; + } + return undefined; + }); + }); + + it('rejects enabling when another workflow is already enabled for the exclusive trigger', async () => { + mockWorkflowsService.getWorkflow.mockResolvedValue( + createWorkflowDto({ id: 'wf-1', enabled: false, definition: exclusiveDefinition }) + ); + mockWorkflowsService.getWorkflowsSubscribedToTrigger.mockResolvedValue([ + createWorkflowDto({ id: 'wf-other', name: 'Existing Workflow', enabled: true }), + ]); + + await expect( + api.updateWorkflow('wf-1', { enabled: true }, 'default', mockRequest) + ).rejects.toBeInstanceOf(WorkflowConflictError); + + expect(mockWorkflowsService.updateWorkflow).not.toHaveBeenCalled(); + }); + + it('includes the conflicting workflow id in the error', async () => { + mockWorkflowsService.getWorkflow.mockResolvedValue( + createWorkflowDto({ id: 'wf-1', enabled: false, definition: exclusiveDefinition }) + ); + mockWorkflowsService.getWorkflowsSubscribedToTrigger.mockResolvedValue([ + createWorkflowDto({ id: 'wf-conflict', name: 'Conflicting Workflow', enabled: true }), + ]); + + const err = await api + .updateWorkflow('wf-1', { enabled: true }, 'default', mockRequest) + .catch((e) => e); + + expect(err).toBeInstanceOf(WorkflowConflictError); + expect((err as WorkflowConflictError).workflowId).toBe('wf-conflict'); + }); + + it('allows enabling when no other workflow is already enabled for the exclusive trigger', async () => { + const updateResult = { enabled: true } as any; + mockWorkflowsService.getWorkflow.mockResolvedValue( + createWorkflowDto({ id: 'wf-1', enabled: false, definition: exclusiveDefinition }) + ); + mockWorkflowsService.getWorkflowsSubscribedToTrigger.mockResolvedValue([]); + mockWorkflowsService.updateWorkflow.mockResolvedValue(updateResult); + + await expect( + api.updateWorkflow('wf-1', { enabled: true }, 'default', mockRequest) + ).resolves.toBe(updateResult); + }); + + it('allows re-enabling the same workflow (self is excluded from conflict check)', async () => { + const updateResult = { enabled: true } as any; + mockWorkflowsService.getWorkflow.mockResolvedValue( + createWorkflowDto({ id: 'wf-1', enabled: true, definition: exclusiveDefinition }) + ); + // Simulate the workflow appearing in its own subscribed-trigger results + mockWorkflowsService.getWorkflowsSubscribedToTrigger.mockResolvedValue([ + createWorkflowDto({ id: 'wf-1', enabled: true }), + ]); + mockWorkflowsService.updateWorkflow.mockResolvedValue(updateResult); + + await expect( + api.updateWorkflow('wf-1', { enabled: true }, 'default', mockRequest) + ).resolves.toBe(updateResult); + }); + + it('does not special-case inference.aroundCompletion — search not issued without registration', async () => { + // This is the regression guard: 'inference.aroundCompletion' must not be treated as + // exclusive unless a definition explicitly registers it as such. Without registration + // the guard must be a no-op and not issue any subscribed-trigger search. + const updateResult = { enabled: true } as any; + // No definition registered for inference.aroundCompletion (returns undefined above) + const aroundCompletionDef = { + triggers: [{ type: 'inference.aroundCompletion' }], + } as unknown as WorkflowDetailDto['definition']; + mockWorkflowsService.getWorkflow.mockResolvedValue( + createWorkflowDto({ id: 'wf-1', enabled: false, definition: aroundCompletionDef }) + ); + mockWorkflowsService.updateWorkflow.mockResolvedValue(updateResult); + + await expect( + api.updateWorkflow('wf-1', { enabled: true }, 'default', mockRequest) + ).resolves.toBe(updateResult); + + // No search should have been issued — no definition declares exclusivity for this trigger + expect(mockWorkflowsService.getWorkflowsSubscribedToTrigger).not.toHaveBeenCalled(); + }); + + it('skips conflict check when the trigger is registered but not exclusive', async () => { + const updateResult = { enabled: true } as any; + mockWorkflowsService.getWorkflow.mockResolvedValue( + createWorkflowDto({ + id: 'wf-1', + enabled: false, + definition: { + triggers: [{ type: SHARED_TRIGGER }], + } as unknown as WorkflowDetailDto['definition'], + }) + ); + mockWorkflowsService.updateWorkflow.mockResolvedValue(updateResult); + + await expect( + api.updateWorkflow('wf-1', { enabled: true }, 'default', mockRequest) + ).resolves.toBe(updateResult); + + expect(mockWorkflowsService.getWorkflowsSubscribedToTrigger).not.toHaveBeenCalled(); + }); + + it('skips conflict check when updating a field other than enabled', async () => { + const updateResult = { name: 'New Name' } as any; + mockWorkflowsService.getWorkflow.mockResolvedValue( + createWorkflowDto({ id: 'wf-1', enabled: true, definition: exclusiveDefinition }) + ); + mockWorkflowsService.updateWorkflow.mockResolvedValue(updateResult); + + await expect( + api.updateWorkflow('wf-1', { name: 'New Name' }, 'default', mockRequest) + ).resolves.toBe(updateResult); + + expect(mockWorkflowsService.getWorkflowsSubscribedToTrigger).not.toHaveBeenCalled(); + }); + + it('skips conflict check when enabling with no exclusive triggers in the definition', async () => { + const updateResult = { enabled: true } as any; + mockWorkflowsService.getWorkflow.mockResolvedValue( + createWorkflowDto({ + id: 'wf-1', + enabled: false, + definition: { + triggers: [{ type: 'manual' }], + } as unknown as WorkflowDetailDto['definition'], + }) + ); + mockWorkflowsService.updateWorkflow.mockResolvedValue(updateResult); + + await expect( + api.updateWorkflow('wf-1', { enabled: true }, 'default', mockRequest) + ).resolves.toBe(updateResult); + + expect(mockWorkflowsService.getWorkflowsSubscribedToTrigger).not.toHaveBeenCalled(); + }); + }); }); describe('restoreWorkflowVersion', () => { diff --git a/src/platform/plugins/shared/workflows_management/server/api/workflows_management_api.ts b/src/platform/plugins/shared/workflows_management/server/api/workflows_management_api.ts index 252fdd80cc030..e14ffed6444d1 100644 --- a/src/platform/plugins/shared/workflows_management/server/api/workflows_management_api.ts +++ b/src/platform/plugins/shared/workflows_management/server/api/workflows_management_api.ts @@ -49,14 +49,25 @@ import type { WorkflowPartialDetailDto, WorkflowSortField, } from '@kbn/workflows/types/v1'; -import type { WorkflowsExecutionEnginePluginStart } from '@kbn/workflows-execution-engine/server'; +import type { + ExecuteWorkflowResponse, + WorkflowsExecutionEnginePluginStart, +} from '@kbn/workflows-execution-engine/server'; +import { classifyWorkflowTriggerMatch } from '@kbn/workflows-execution-engine/server'; import type { LogSearchResult } from '@kbn/workflows-execution-engine/server/repositories/logs_repository'; import type { ExecutionLogsParams, StepLogsParams, } from '@kbn/workflows-execution-engine/server/workflow_event_logger/types'; -import type { ServerTriggerDefinition } from '@kbn/workflows-extensions/server'; -import { parseYamlToJSONWithoutValidation, WorkflowValidationError } from '@kbn/workflows-yaml'; +import type { + ServerTriggerDefinition, + WorkflowExecutionCapabilities, +} from '@kbn/workflows-extensions/server'; +import { + parseYamlToJSONWithoutValidation, + WorkflowConflictError, + WorkflowValidationError, +} from '@kbn/workflows-yaml'; import type { z } from '@kbn/zod/v4'; import { type ExternalResumeFormPageParams, @@ -67,6 +78,11 @@ import { resumeWorkflowExecutionExternallyWithInput, } from './external_resume/external_resume_service'; import type { StepExecutionListResult } from './lib/search_step_executions'; +import { + getTriggerTypesFromDefinition, + getTriggerTypesFromYaml, + workflowYamlDeclaresEnabled, +} from './lib/workflow_prepare'; import { ManagedWorkflowDeleteForbiddenError } from './managed_workflow_delete_error'; import { ManagedWorkflowUpdateForbiddenError } from './managed_workflow_errors'; import { preprocessAlertInputs } from './routes/executions/utils/preprocess_alert_inputs'; @@ -115,6 +131,26 @@ export interface GetWorkflowAggsOptions { managedFilter?: GetWorkflowsParams['managedFilter']; } +export interface ResolveWorkflowTriggerMatchesResult { + matched: WorkflowDetailDto[]; + invalidConditionWorkflows: Array<{ id: string; name: string }>; +} + +export interface ExecuteWorkflowSynchronouslyParams { + workflowId: string; + /** + * Pre-fetched workflow DTO (e.g. from trigger resolution). When supplied, the ES re-fetch + * inside `executeWorkflowSynchronously` is skipped — eliminates a redundant read on the + * inference hot path. The DTO is still validated (enabled, valid, definition) before execution. + */ + workflow?: WorkflowDetailDto; + context: Record; + spaceId: string; + request: KibanaRequest; + capabilities?: WorkflowExecutionCapabilities; + abortSignal?: AbortSignal; +} + export interface DeleteWorkflowsResponse { total: number; deleted: number; @@ -310,6 +346,63 @@ export class WorkflowsManagementApi { return this.workflowsService.getWorkflowsExecutionEngine(); } + /** + * Enforces trigger-declared exclusivity before a workflow becomes enabled. + * + * A trigger definition may declare `exclusivity: 'per-space'`, meaning at most + * one *enabled* workflow per space may subscribe to it. The rule is entirely + * registration-driven: this method never names a specific trigger, and a + * trigger that does not declare exclusivity costs nothing (no search is issued). + * + * Best-effort: two concurrent enables can interleave between the search and the + * write. There is no unique key in the workflow index to make this atomic; the + * guard exists to return a clear 409 rather than silently nondeterministic dispatch. + * + * @param triggerTypes Trigger ids the workflow will subscribe to once written. + * @param spaceId Space the workflow is being written into. + * @param excludeWorkflowId The workflow being updated, skipped as its own conflict. + * Omitted on create, where there is no self to skip yet. + */ + private async assertExclusiveTriggersAvailable({ + triggerTypes, + spaceId, + excludeWorkflowId, + }: { + triggerTypes: string[]; + spaceId: string; + excludeWorkflowId?: string; + }): Promise { + if (triggerTypes.length === 0) { + return; + } + + const workflowsExtensions = await this.workflowsService.getWorkflowsExtensions(); + + // Dedupe: a workflow may declare the same trigger twice (different conditions), + // so one exclusive trigger costs at most one search. Sequential by design — + // the common case is zero exclusive triggers and issues no search at all. + for (const triggerId of new Set(triggerTypes)) { + if (workflowsExtensions.getTriggerDefinition(triggerId)?.exclusivity === 'per-space') { + // Already filtered to enabled: true and includes global (spaceId: '*') subscribers. + const enabledSubscribers = await this.getWorkflowsSubscribedToTrigger(triggerId, spaceId); + const conflict = enabledSubscribers.find((w) => w.id !== excludeWorkflowId); + if (conflict) { + throw new WorkflowConflictError( + i18n.translate('workflowsManagement.exclusiveTriggerConflictError', { + defaultMessage: + 'Cannot enable: workflow "{conflictingWorkflowName}" is already enabled for the {triggerId} trigger. Disable it first.', + values: { + conflictingWorkflowName: conflict.name ?? conflict.id, + triggerId, + }, + }), + conflict.id + ); + } + } + } + } + public setSmlIndexAttachment(fn: SmlIndexAttachmentFn, logger: Logger): void { this.smlIndexAttachment = fn; this.smlLogger = logger; @@ -350,6 +443,50 @@ export class WorkflowsManagementApi { return this.workflowsService.getWorkflowsSubscribedToTrigger(triggerId, spaceId); } + public async resolveWorkflowTriggerMatches( + triggerId: string, + event: Record, + spaceId: string + ): Promise { + const subscribed = await this.getWorkflowsSubscribedToTrigger(triggerId, spaceId); + const matched: WorkflowDetailDto[] = []; + const invalidConditionWorkflows: Array<{ id: string; name: string }> = []; + + subscribed.forEach((workflow) => { + const outcome = classifyWorkflowTriggerMatch(workflow, triggerId, event); + if (outcome === 'matched') { + matched.push(workflow); + } else if (outcome === 'kql_error') { + invalidConditionWorkflows.push({ + id: workflow.id, + name: workflow.name ?? workflow.definition?.name ?? workflow.id, + }); + } + }); + + return { matched, invalidConditionWorkflows }; + } + + public async executeWorkflowSynchronously({ + workflowId, + workflow: prefetchedWorkflow, + context, + spaceId, + request, + capabilities, + abortSignal, + }: ExecuteWorkflowSynchronouslyParams): Promise { + const model = prefetchedWorkflow + ? this.validateAndBuildWorkflowModel(prefetchedWorkflow) + : await this.getSavedWorkflowExecutionModel(workflowId, spaceId); + const workflowsExecutionEngine = await this.getWorkflowsExecutionEngine(); + return workflowsExecutionEngine.executeWorkflow(model, context, request, { + executionMode: 'sync', + capabilities, + abortSignal, + }); + } + public async getWorkflow(id: string, spaceId: string): Promise { return this.workflowsService.getWorkflow(id, spaceId); } @@ -381,9 +518,19 @@ export class WorkflowsManagementApi { public async createWorkflow( workflow: CreateWorkflowCommand, spaceId: string, - request: KibanaRequest + request: KibanaRequest, + options?: { originManagedWorkflowId?: string } ): Promise { - const result = await this.workflowsService.createWorkflow(workflow, spaceId, request); + // A workflow can be created already enabled when the YAML declares `enabled: true`, + // so the exclusivity guard must run here, not only on the enable transition. + if (workflowYamlDeclaresEnabled(workflow.yaml)) { + await this.assertExclusiveTriggersAvailable({ + triggerTypes: getTriggerTypesFromYaml(workflow.yaml), + spaceId, + }); + } + + const result = await this.workflowsService.createWorkflow(workflow, spaceId, request, options); this.notifySml(result.id, 'create', request); return result; } @@ -419,6 +566,15 @@ export class WorkflowsManagementApi { })}`; const clonedYaml = updateWorkflowYamlFields(workflow.yaml, { name: cloneName }); + // A clone inherits the source's `enabled` state from the YAML. If the source is + // enabled for an exclusive trigger, a second enabled subscriber would form — block it. + if (workflowYamlDeclaresEnabled(clonedYaml)) { + await this.assertExclusiveTriggersAvailable({ + triggerTypes: getTriggerTypesFromYaml(clonedYaml), + spaceId, + }); + } + // `updateWorkflowYamlFields` cannot inject a `name` key when the YAML root is not a // mapping (a scalar or sequence), so it returns the YAML unchanged in that case. Pass // `cloneName` as an explicit fallback so the clone is still named " Copy" instead @@ -452,6 +608,36 @@ export class WorkflowsManagementApi { ) { throw new ManagedWorkflowUpdateForbiddenError(); } + + // Enforce trigger-declared exclusivity before any path that results in this + // workflow becoming enabled. The condition covers three routes to enablement: + // + // 1. Field-only enable (`workflow.enabled === true`): trigger types unchanged, + // so use the existing definition. + // 2. YAML update that declares top-level `enabled: true`: the YAML value wins + // over the `enabled` field (see workflow_crud_service.ts for the precedence + // rule), so check the NEW trigger types from the incoming YAML. + // 3. YAML update that adds an exclusive trigger to an already-enabled workflow + // (`workflow.enabled` absent but new triggers in YAML): same as case 2. + // + // Known gap: restoreWorkflowVersion — the restored snapshot YAML is fetched + // inside the crud service and is unavailable here. Track as a follow-up. + if (workflow.yaml) { + if (workflowYamlDeclaresEnabled(workflow.yaml) || workflow.enabled === true) { + await this.assertExclusiveTriggersAvailable({ + triggerTypes: getTriggerTypesFromYaml(workflow.yaml), + spaceId, + excludeWorkflowId: id, + }); + } + } else if (workflow.enabled === true) { + await this.assertExclusiveTriggersAvailable({ + triggerTypes: getTriggerTypesFromDefinition(originalWorkflow.definition), + spaceId, + excludeWorkflowId: id, + }); + } + const result = await this.workflowsService.updateWorkflow(id, workflow, spaceId, request); this.notifySml(id, 'update', request); return result; @@ -653,26 +839,33 @@ export class WorkflowsManagementApi { }; } + /** + * Validates an already-fetched workflow DTO and converts it to an execution model. + * Used by `executeWorkflowSynchronously` when the caller passes a pre-fetched workflow, + * and by `getSavedWorkflowExecutionModel` after the ES fetch. + */ + private validateAndBuildWorkflowModel(workflow: WorkflowDetailDto): WorkflowExecutionEngineModel { + if (!workflow.enabled) { + throw new Error(`Workflow '${workflow.id}' is disabled and cannot be executed.`); + } + if (!workflow.valid) { + throw new Error(`Workflow '${workflow.id}' has validation errors and cannot be executed.`); + } + if (!workflow.definition) { + throw new Error(`Workflow '${workflow.id}' has no definition and cannot be executed.`); + } + return toWorkflowExecutionEngineModel(workflow); + } + private async getSavedWorkflowExecutionModel( workflowId: string, spaceId: string ): Promise { const workflow = await this.getWorkflow(workflowId, spaceId); - if (!workflow) { throw new WorkflowNotFoundError(workflowId); } - if (!workflow.enabled) { - throw new Error(`Workflow '${workflowId}' is disabled and cannot be executed.`); - } - if (!workflow.valid) { - throw new Error(`Workflow '${workflowId}' has validation errors and cannot be executed.`); - } - if (!workflow.definition) { - throw new Error(`Workflow '${workflowId}' has no definition and cannot be executed.`); - } - - return toWorkflowExecutionEngineModel(workflow); + return this.validateAndBuildWorkflowModel(workflow); } private async waitForWorkflowExecution({ diff --git a/src/platform/plugins/shared/workflows_management/server/api/workflows_management_service.ts b/src/platform/plugins/shared/workflows_management/server/api/workflows_management_service.ts index 578cd6269e29d..28c205c41f480 100644 --- a/src/platform/plugins/shared/workflows_management/server/api/workflows_management_service.ts +++ b/src/platform/plugins/shared/workflows_management/server/api/workflows_management_service.ts @@ -408,7 +408,7 @@ export class WorkflowsService { workflow: CreateWorkflowCommand, spaceId: string, request: KibanaRequest, - options?: { nameFallback?: string } + options?: { nameFallback?: string; originManagedWorkflowId?: string } ): Promise { await this.ensureInitialized(); return this.crudService.createWorkflow(workflow, spaceId, request, options); diff --git a/src/platform/plugins/shared/workflows_management/server/services/workflow_crud_service.ts b/src/platform/plugins/shared/workflows_management/server/services/workflow_crud_service.ts index 81fdb1a543006..861f1d84c1623 100644 --- a/src/platform/plugins/shared/workflows_management/server/services/workflow_crud_service.ts +++ b/src/platform/plugins/shared/workflows_management/server/services/workflow_crud_service.ts @@ -499,7 +499,7 @@ export class WorkflowCrudService { workflow: CreateWorkflowCommand, spaceId: string, request: KibanaRequest, - options?: { nameFallback?: string } + options?: { nameFallback?: string; originManagedWorkflowId?: string } ): Promise { if (workflow.id) { validateWorkflowId(workflow.id); @@ -527,6 +527,7 @@ export class WorkflowCrudService { spaceId, triggerDefinitions, nameFallback: options?.nameFallback, + originManagedWorkflowId: options?.originManagedWorkflowId, }); let id = baseId; diff --git a/src/platform/plugins/shared/workflows_management/server/services/workflow_search_service.ts b/src/platform/plugins/shared/workflows_management/server/services/workflow_search_service.ts index 2145c074b8133..3eb0978aaaf48 100644 --- a/src/platform/plugins/shared/workflows_management/server/services/workflow_search_service.ts +++ b/src/platform/plugins/shared/workflows_management/server/services/workflow_search_service.ts @@ -104,6 +104,15 @@ export class WorkflowSearchService { 'valid', 'created_at', 'updated_at', + // Required for executeWorkflowSynchronously when the DTO is reused directly: + // 'version' is surfaced to Liquid templates as {{ workflow.version }} in sync mode; + // the managed fields feed execution telemetry. + 'version', + 'managed', + 'managedBy', + 'billable', + 'originManagedWorkflowId', + 'managedVersion', ]; const pitResponse = await esClient.openPointInTime({ diff --git a/x-pack/platform/plugins/shared/inference/server/chat_complete/api.test.ts b/x-pack/platform/plugins/shared/inference/server/chat_complete/api.test.ts index 8e461cf8cc5f7..d3004b03052d1 100644 --- a/x-pack/platform/plugins/shared/inference/server/chat_complete/api.test.ts +++ b/x-pack/platform/plugins/shared/inference/server/chat_complete/api.test.ts @@ -13,7 +13,16 @@ import { inferenceEndpointAdapterMock, } from './api.test.mocks'; -import { of, Subject, isObservable, toArray, firstValueFrom, filter } from 'rxjs'; +import { + concat, + filter, + firstValueFrom, + isObservable, + of, + Subject, + throwError, + toArray, +} from 'rxjs'; import { loggerMock, type MockedLogger } from '@kbn/logging-mocks'; import { httpServerMock } from '@kbn/core/server/mocks'; import { actionsMock } from '@kbn/actions-plugin/server/mocks'; @@ -31,12 +40,14 @@ import { createInferenceConnectorMock, createInferenceExecutorMock, createRegexWorkerServiceMock, + createPiiRegexWorkerServiceMock, chunkEvent, tokensEvent, } from '../test_utils'; import { createChatCompleteApi } from './api'; import { createChatCompleteCallbackApi } from './callback_api'; import { InferenceEndpointIdCache } from '../util/inference_endpoint_id_cache'; +import type { WorkflowAnonymizationProvider } from '../workflow_anonymization_provider'; describe('createChatCompleteApi', () => { let request: ReturnType; @@ -161,6 +172,106 @@ describe('createChatCompleteApi', () => { }); }); + it('uses original event data and workflow-transformed connector input in workflow mode', async () => { + const protectedMessages = [{ role: MessageRole.User, content: 'protected question' }] as const; + const provider: WorkflowAnonymizationProvider = { + supportsSynchronousExecution: true, + execute: jest.fn(async ({ event, namespace, proceed }) => { + expect(event).toEqual({ + system: 'original system', + messages: [{ role: MessageRole.User, content: 'original question' }], + sessionId: 'session-a', + agentId: 'agent-a', + }); + expect(namespace).toBe('space-a'); + await proceed.invoke({ messages: protectedMessages, tokenMap: {} }); + return { matched: true, content: 'workflow restored' }; + }), + }; + const workflowCallbackApi = createChatCompleteCallbackApi({ + request, + namespace: 'space-a', + actions, + logger, + // Workflow mode must not wait for or consume the legacy anonymization rule source. + anonymizationRulesPromise: new Promise(() => {}), + regexWorker, + esClient: mockEsClient, + endpointIdCache, + anonymization: { saltPromise: Promise.resolve('server-managed-salt') }, + workflowAnonymization: { + provider, + failureMode: 'block', + preLLMTimeoutMs: 0, + piiRegexWorker: createPiiRegexWorkerServiceMock(), + }, + }); + const workflowChatComplete = createChatCompleteApi({ callbackApi: workflowCallbackApi }); + + await expect( + workflowChatComplete({ + connectorId: 'connectorId', + system: 'original system', + messages: [{ role: MessageRole.User, content: 'original question' }], + sessionId: 'session-a', + metadata: { agentId: 'agent-a' }, + maxRetries: 0, + }) + ).resolves.toEqual(expect.objectContaining({ content: 'workflow restored' })); + expect(inferenceAdapter.chatComplete).toHaveBeenCalledWith( + expect.objectContaining({ messages: protectedMessages }) + ); + }); + + it('does not retry a workflow connector call after streaming has started', async () => { + const providerError = createInferenceProviderError('stream failed', { status: 500 }); + inferenceAdapter.chatComplete.mockReturnValue( + concat( + of(chunkEvent('partial')), + throwError(() => providerError) + ) + ); + const provider: WorkflowAnonymizationProvider = { + supportsSynchronousExecution: true, + execute: jest.fn(async ({ proceed }) => { + await proceed.invoke({ + messages: [{ role: MessageRole.User, content: 'protected question' }], + tokenMap: {}, + }); + return { matched: true, content: 'unreachable' }; + }), + }; + const workflowChatComplete = createChatCompleteApi({ + callbackApi: createChatCompleteCallbackApi({ + request, + namespace: 'space-a', + actions, + logger, + anonymizationRulesPromise: Promise.resolve([]), + regexWorker, + esClient: mockEsClient, + endpointIdCache, + anonymization: { saltPromise: Promise.resolve('server-managed-salt') }, + workflowAnonymization: { + provider, + failureMode: 'allow_unsafe', + preLLMTimeoutMs: 0, + piiRegexWorker: createPiiRegexWorkerServiceMock(), + }, + }), + }); + + await expect( + workflowChatComplete({ + connectorId: 'connectorId', + messages: [{ role: MessageRole.User, content: 'original question' }], + maxRetries: 1, + retryConfiguration: { retryOn: 'all', initialDelay: 0 }, + }) + ).rejects.toBe(providerError); + expect(inferenceAdapter.chatComplete).toHaveBeenCalledTimes(1); + }); + it('forwards `maxContentLength` down to `inferenceAdapter.chatComplete`', async () => { await chatComplete({ connectorId: 'connectorId', diff --git a/x-pack/platform/plugins/shared/inference/server/chat_complete/callback_api.ts b/x-pack/platform/plugins/shared/inference/server/chat_complete/callback_api.ts index 5f3f06632142e..e6e5dff8de6db 100644 --- a/x-pack/platform/plugins/shared/inference/server/chat_complete/callback_api.ts +++ b/x-pack/platform/plugins/shared/inference/server/chat_complete/callback_api.ts @@ -16,11 +16,12 @@ import { getConnectorPlatform, getConnectorDefaultModel, type ChatCompleteCompositeResponse, + type ChatCompletionEvent, MessageRole, } from '@kbn/inference-common'; import type { Logger } from '@kbn/logging'; import type { Observable } from 'rxjs'; -import { defer, forkJoin, from, identity, share, switchMap, catchError, throwError } from 'rxjs'; +import { catchError, defer, from, identity, share, switchMap, tap, throwError } from 'rxjs'; import { withChatCompleteSpan } from '@kbn/inference-tracing'; import type { ElasticsearchClient } from '@kbn/core/server'; import { omit } from 'lodash'; @@ -41,6 +42,7 @@ import { retryHoldingTokenCountEvents, streamToResponse, } from './utils'; +import { retryWithExponentialBackoff } from '../../common/utils/retry_with_exponential_backoff'; import type { InferenceCallbackManager } from '../inference_client/callback_manager'; import { getRetryFilter } from '../../common/utils/error_retry_filter'; import { deanonymizeMessage } from './anonymization/deanonymize_message'; @@ -51,6 +53,11 @@ import type { InferenceEndpointIdCache } from '../util/inference_endpoint_id_cac import { prepareAnonymization } from './prepare_anonymization'; import type { TokenUsageLogger } from '../token_usage'; import { handleTokenUsageLogging, buildTokenUsageContext } from '../token_usage'; +import type { WorkflowAnonymizationOptions } from '../inference_client/workflow_anonymization_options'; +import { + createWorkflowAnonymizationPipeline, + type WorkflowInvocationState, +} from './workflow_anonymization_pipeline'; interface CreateChatCompleteApiOptions { request: KibanaRequest; @@ -61,6 +68,7 @@ interface CreateChatCompleteApiOptions { regexWorker: RegexWorkerService; esClient: ElasticsearchClient; anonymization?: InferenceAnonymizationOptions; + workflowAnonymization?: WorkflowAnonymizationOptions; endpointIdCache: InferenceEndpointIdCache; callbackManager?: InferenceCallbackManager; tokenUsageLogger?: TokenUsageLogger; @@ -118,6 +126,7 @@ export function createChatCompleteCallbackApi({ regexWorker, esClient, anonymization, + workflowAnonymization, endpointIdCache, callbackManager, tokenUsageLogger, @@ -135,6 +144,8 @@ export function createChatCompleteCallbackApi({ }: ChatCompleteApiWithCallbackInitOptions, callback: ChatCompleteApiWithCallbackCallback ) => { + const workflowInvocationState: WorkflowInvocationState = { connectorInvoked: false }; + const retryFilter = getRetryFilter(retryConfiguration.retryOn); const inference$ = defer(() => resolveAndCreatePipeline({ connectorId, @@ -150,6 +161,14 @@ export function createChatCompleteCallbackApi({ stream, namespace, anonymization, + workflowAnonymization, + workflowInvocationState, + connectorRetry: { + maxRetries, + backoffMultiplier: retryConfiguration.backoffMultiplier, + initialDelay: retryConfiguration.initialDelay, + retryFilter, + }, tokenUsageLogger, isTokenUsageTrackingEnabled, isDefaultConnectorOnly, @@ -160,7 +179,7 @@ export function createChatCompleteCallbackApi({ maxRetry: maxRetries, backoffMultiplier: retryConfiguration.backoffMultiplier, initialDelay: retryConfiguration.initialDelay, - errorFilter: getRetryFilter(retryConfiguration.retryOn), + errorFilter: (error) => !workflowInvocationState.connectorInvoked && retryFilter(error), }), callbackManager ? handleLifecycleCallbacks({ callbackManager }) : identity, abortSignal ? handleCancellation(abortSignal) : identity @@ -176,6 +195,7 @@ export function createChatCompleteCallbackApi({ function createChatCompletePipeline({ resolve, + request, esClient, logger, anonymizationRulesPromise, @@ -185,11 +205,15 @@ function createChatCompletePipeline({ stream, namespace, anonymization, + workflowAnonymization, + workflowInvocationState, + connectorRetry, connectorId, tokenUsageLogger, isTokenUsageTrackingEnabled, }: { resolve: () => Promise; + request: KibanaRequest; esClient: ElasticsearchClient; logger: Logger; anonymizationRulesPromise: Promise; @@ -199,15 +223,20 @@ function createChatCompletePipeline({ stream?: boolean; namespace: string; anonymization?: InferenceAnonymizationOptions; + workflowAnonymization?: WorkflowAnonymizationOptions; + workflowInvocationState: WorkflowInvocationState; + connectorRetry: { + maxRetries: number; + backoffMultiplier?: number; + initialDelay?: number; + retryFilter: (error: Error) => boolean; + }; connectorId: string; tokenUsageLogger?: TokenUsageLogger; isTokenUsageTrackingEnabled?: () => Promise; }) { - return forkJoin({ - context: from(resolve()), - anonymizationRules: from(anonymizationRulesPromise), - }).pipe( - switchMap(({ context, anonymizationRules }) => { + return from(resolve()).pipe( + switchMap((context) => { const { callbackContext, getSpanModel, chatComplete } = context; const { @@ -228,84 +257,133 @@ function createChatCompletePipeline({ const messages = sanitizeMessages(givenMessages); - return from( - prepareAnonymization({ - namespace, - logger, - anonymizationRules, - regexWorker, - esClient, - replacementsEsClient: anonymization?.replacements?.esClient, - replacementsEncryptionKeyPromise: anonymization?.replacements?.encryptionKeyPromise, - usePersistentReplacements: anonymization?.replacements?.usePersistentReplacements, - requireReplacementsEncryptionKey: anonymization?.replacements?.requireEncryptionKey, - saltPromise: anonymization?.saltPromise, - resolveEffectivePolicy: anonymization?.resolveEffectivePolicy, - metadata, - system, - messages, - }) - ).pipe( - switchMap(({ anonymization: preparedAnonymization, replacementsId, effectivePolicy }) => { - const systemWithAnonymizationInstructions = preparedAnonymization.system - ? addAnonymizationInstruction( - preparedAnonymization.system, - anonymizationRules, - effectivePolicy - ) - : system; - - const spanModel = getSpanModel(modelName); - - return withChatCompleteSpan( - { - system: systemWithAnonymizationInstructions, - messages: preparedAnonymization.messages, - tools, + const invokeConnector = ({ + system: connectorSystem, + messages: connectorMessages, + abortSignal: connectorAbortSignal, + }: { + system?: string; + messages: readonly (typeof messages)[number][]; + abortSignal?: AbortSignal; + }): Observable => { + const spanModel = getSpanModel(modelName); + let emittedEvent = false; + return withChatCompleteSpan( + { + system: connectorSystem, + messages: [...connectorMessages], + tools, + toolChoice, + ...(spanModel ? { model: spanModel } : {}), + ...metadata?.attributes, + }, + () => + chatComplete({ + system: connectorSystem, + messages: [...connectorMessages], toolChoice, + tools, + temperature, + reasoning, + logger, + functionCalling, + modelName, + abortSignal: connectorAbortSignal, + metadata, + timeout, + maxContentLength, cacheControl, sessionId, - reasoning, - ...(spanModel ? { model: spanModel } : {}), - ...metadata?.attributes, - }, - () => { - return chatComplete({ - system: systemWithAnonymizationInstructions, - messages: preparedAnonymization.messages, - toolChoice, - tools, - temperature, - reasoning, - logger, - functionCalling, + stream, + }).pipe(chunksIntoMessage({ toolOptions: { toolChoice, tools }, logger })) + ).pipe( + tap(() => { + emittedEvent = true; + }), + workflowAnonymization + ? retryWithExponentialBackoff({ + maxRetry: connectorRetry.maxRetries, + backoffMultiplier: connectorRetry.backoffMultiplier, + initialDelay: connectorRetry.initialDelay, + errorFilter: (error: Error) => !emittedEvent && connectorRetry.retryFilter(error), + }) + : identity + ); + }; + + const tokenUsageOperator = tokenUsageLogger + ? handleTokenUsageLogging({ + tokenUsageLogger, + getContext: () => + buildTokenUsageContext({ + connectorId, + model: callbackContext.model, modelName, - abortSignal, - metadata, - timeout, - maxContentLength, - cacheControl, - sessionId, - stream, - }).pipe(chunksIntoMessage({ toolOptions: { toolChoice, tools }, logger })); - } - ).pipe(deanonymizeMessage({ ...preparedAnonymization, replacementsId })); - }), - tokenUsageLogger - ? handleTokenUsageLogging({ - tokenUsageLogger, - getContext: () => - buildTokenUsageContext({ - connectorId, - model: callbackContext.model, - modelName, - featureId: metadata?.connectorTelemetry?.pluginId, - parentFeatureId: metadata?.connectorTelemetry?.aggregateBy, - }), + featureId: metadata?.connectorTelemetry?.pluginId, + parentFeatureId: metadata?.connectorTelemetry?.aggregateBy, + }), + logger, + isEnabled: isTokenUsageTrackingEnabled, + }) + : identity; + + if (workflowAnonymization) { + return createWorkflowAnonymizationPipeline({ + request, + namespace, + system, + messages, + sessionId, + agentId: metadata?.agentId, + abortSignal, + regexWorker: workflowAnonymization.piiRegexWorker, + logger, + workflowAnonymization, + invocationState: workflowInvocationState, + invokeConnector, + }).pipe(tokenUsageOperator); + } + + return from(anonymizationRulesPromise).pipe( + switchMap((anonymizationRules) => + from( + prepareAnonymization({ + namespace, logger, - isEnabled: isTokenUsageTrackingEnabled, + anonymizationRules, + regexWorker, + esClient, + replacementsEsClient: anonymization?.replacements?.esClient, + replacementsEncryptionKeyPromise: anonymization?.replacements?.encryptionKeyPromise, + usePersistentReplacements: anonymization?.replacements?.usePersistentReplacements, + requireReplacementsEncryptionKey: anonymization?.replacements?.requireEncryptionKey, + saltPromise: anonymization?.saltPromise, + resolveEffectivePolicy: anonymization?.resolveEffectivePolicy, + metadata, + system, + messages, }) - : identity + ).pipe( + switchMap( + ({ anonymization: preparedAnonymization, replacementsId, effectivePolicy }) => { + const systemWithAnonymizationInstructions = preparedAnonymization.system + ? addAnonymizationInstruction( + preparedAnonymization.system, + anonymizationRules, + effectivePolicy + ) + : system; + + return invokeConnector({ + system: systemWithAnonymizationInstructions, + messages: preparedAnonymization.messages, + abortSignal, + }).pipe(deanonymizeMessage({ ...preparedAnonymization, replacementsId })); + } + ) + ) + ), + tokenUsageOperator ); }) ); @@ -325,6 +403,9 @@ function resolveAndCreatePipeline({ stream, namespace, anonymization, + workflowAnonymization, + workflowInvocationState, + connectorRetry, tokenUsageLogger, isTokenUsageTrackingEnabled, isDefaultConnectorOnly, @@ -343,6 +424,14 @@ function resolveAndCreatePipeline({ stream?: boolean; namespace: string; anonymization?: InferenceAnonymizationOptions; + workflowAnonymization?: WorkflowAnonymizationOptions; + workflowInvocationState: WorkflowInvocationState; + connectorRetry: { + maxRetries: number; + backoffMultiplier?: number; + initialDelay?: number; + retryFilter: (error: Error) => boolean; + }; tokenUsageLogger?: TokenUsageLogger; isTokenUsageTrackingEnabled?: () => Promise; isDefaultConnectorOnly?: () => Promise; @@ -463,6 +552,7 @@ function resolveAndCreatePipeline({ return createChatCompletePipeline({ resolve, + request, esClient, logger, anonymizationRulesPromise, @@ -472,6 +562,9 @@ function resolveAndCreatePipeline({ stream, namespace, anonymization, + workflowAnonymization, + workflowInvocationState, + connectorRetry, connectorId, tokenUsageLogger, isTokenUsageTrackingEnabled, diff --git a/x-pack/platform/plugins/shared/inference/server/chat_complete/workflow_anonymization_metrics.ts b/x-pack/platform/plugins/shared/inference/server/chat_complete/workflow_anonymization_metrics.ts new file mode 100644 index 0000000000000..2977fb16a61f6 --- /dev/null +++ b/x-pack/platform/plugins/shared/inference/server/chat_complete/workflow_anonymization_metrics.ts @@ -0,0 +1,29 @@ +/* + * 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 { metrics, ValueType } from '@opentelemetry/api'; + +const meter = metrics.getMeter('kibana.inference.anonymization'); + +export const pipelineRequestsCounter = meter.createCounter( + 'kibana.inference.anonymization.pipeline.requests', + { + description: 'Number of around-completion anonymization pipeline executions', + unit: '{request}', + valueType: ValueType.INT, + } +); + +export const pipelineFirstChunkDurationHistogram = meter.createHistogram( + 'kibana.inference.anonymization.pipeline.first_chunk.duration', + { + description: + 'Duration from connector invocation to first restored chunk emitted downstream, in milliseconds', + unit: 'ms', + valueType: ValueType.DOUBLE, + } +); diff --git a/x-pack/platform/plugins/shared/inference/server/chat_complete/workflow_anonymization_pipeline.test.ts b/x-pack/platform/plugins/shared/inference/server/chat_complete/workflow_anonymization_pipeline.test.ts new file mode 100644 index 0000000000000..d6301380699fc --- /dev/null +++ b/x-pack/platform/plugins/shared/inference/server/chat_complete/workflow_anonymization_pipeline.test.ts @@ -0,0 +1,487 @@ +/* + * 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 { httpServerMock } from '@kbn/core-http-server-mocks'; +import { ChatCompletionEventType, MessageRole } from '@kbn/inference-common'; +import type { WorkflowAnonymizationProvider } from '../workflow_anonymization_provider'; +import { loggerMock } from '@kbn/logging-mocks'; +import { firstValueFrom, Observable, of, throwError, toArray } from 'rxjs'; +import { + chunkEvent, + createPiiRegexWorkerServiceMock, + messageEvent, + tokensEvent, +} from '../test_utils'; +import { createWorkflowAnonymizationPipeline } from './workflow_anonymization_pipeline'; + +const token = 'EMAIL_0123456789abcdef0123456789abcdef'; +const tokenMap = { + [token]: { original: 'person@example.com', entityClass: 'EMAIL' }, +}; +const originalMessages = [{ role: MessageRole.User, content: 'original' }] as const; +const protectedMessages = [{ role: MessageRole.User, content: `Contact ${token}` }] as const; + +const createOptions = ({ + provider, + invokeConnector, + failureMode = 'block', + abortSignal, + preLLMTimeoutMs = 0, +}: { + provider: WorkflowAnonymizationProvider; + invokeConnector: jest.Mock; + failureMode?: 'block' | 'allow_unsafe'; + abortSignal?: AbortSignal; + preLLMTimeoutMs?: number; +}) => ({ + request: httpServerMock.createKibanaRequest(), + namespace: 'space-a', + system: 'original system', + messages: originalMessages, + sessionId: 'session-a', + agentId: 'agent-a', + abortSignal, + regexWorker: createPiiRegexWorkerServiceMock(), + logger: loggerMock.create(), + workflowAnonymization: { + provider, + failureMode, + preLLMTimeoutMs, + encryptionKey: 'server-managed-salt', + }, + invocationState: { connectorInvoked: false }, + invokeConnector, +}); + +describe('createWorkflowAnonymizationPipeline', () => { + it('passes the caller abort signal to workflow execution and the connector', async () => { + const abortSignal = new AbortController().signal; + const invokeConnector = jest.fn().mockReturnValue(of(messageEvent('direct'))); + const provider: WorkflowAnonymizationProvider = { + supportsSynchronousExecution: true, + execute: jest.fn(async (options) => { + expect(options.abortSignal).toBe(abortSignal); + return { matched: false } as const; + }), + }; + + await firstValueFrom( + createWorkflowAnonymizationPipeline( + createOptions({ provider, invokeConnector, abortSignal }) + ).pipe(toArray()) + ); + + expect(invokeConnector).toHaveBeenCalledWith({ + system: 'original system', + messages: originalMessages, + abortSignal, + }); + }); + + it('stops an active connector stream when the caller aborts', async () => { + const abortController = new AbortController(); + let markConnectorStarted: () => void = () => undefined; + const connectorStarted = new Promise((resolve) => { + markConnectorStarted = resolve; + }); + const connectorStopped = jest.fn(); + const invokeConnector = jest.fn( + ({ abortSignal }: { abortSignal?: AbortSignal }) => + new Observable((subscriber) => { + const stop = () => subscriber.error(abortSignal?.reason ?? new Error('aborted')); + abortSignal?.addEventListener('abort', stop, { once: true }); + markConnectorStarted(); + return () => { + abortSignal?.removeEventListener('abort', stop); + connectorStopped(); + }; + }) + ); + const provider: WorkflowAnonymizationProvider = { + supportsSynchronousExecution: true, + execute: jest.fn(async ({ proceed }) => { + await proceed.invoke({ messages: protectedMessages, tokenMap: {} }); + return { matched: true, content: 'unreachable' }; + }), + }; + const completion = firstValueFrom( + createWorkflowAnonymizationPipeline( + createOptions({ provider, invokeConnector, abortSignal: abortController.signal }) + ).pipe(toArray()) + ); + + await connectorStarted; + abortController.abort(new Error('cancelled')); + + await expect(completion).rejects.toThrow('cancelled'); + expect(connectorStopped).toHaveBeenCalledTimes(1); + }); + + it('relays restored chunks and emits one workflow-authoritative terminal message', async () => { + const invokeConnector = jest.fn().mockReturnValue( + of( + chunkEvent(`Contact ${token.slice(0, 14)}`, [ + { + index: 0, + toolCallId: 'tool-call-1', + function: { + name: 'send_email', + arguments: `{"recipient":"${token.slice(0, 8)}`, + }, + }, + ]), + chunkEvent(`${token.slice(14)} now`), + tokensEvent(), + { + ...messageEvent(`Contact ${token} now`, [ + { + toolCallId: 'tool-call-1', + function: { name: 'send_email', arguments: { recipient: token } }, + }, + ]), + refusal: 'preserved refusal metadata', + } + ) + ); + const provider: WorkflowAnonymizationProvider = { + supportsSynchronousExecution: true, + execute: jest.fn(async ({ event, namespace, proceed }) => { + expect(event).toEqual({ + system: 'original system', + messages: originalMessages, + sessionId: 'session-a', + agentId: 'agent-a', + }); + expect(namespace).toBe('space-a'); + await expect( + proceed.invoke({ system: 'protected system', messages: protectedMessages, tokenMap }) + ).resolves.toEqual({ rawContent: `Contact ${token} now` }); + return { matched: true, content: 'workflow-restored final' }; + }), + }; + + const events = await firstValueFrom( + createWorkflowAnonymizationPipeline(createOptions({ provider, invokeConnector })).pipe( + toArray() + ) + ); + const chunks = events.filter( + (event) => event.type === ChatCompletionEventType.ChatCompletionChunk + ); + const terminalMessages = events.filter( + (event) => event.type === ChatCompletionEventType.ChatCompletionMessage + ); + + expect(chunks.map(({ content }) => content).join('')).toBe('Contact person@example.com now'); + expect(chunks.flatMap(({ tool_calls: toolCalls }) => toolCalls)).toEqual([ + { + index: 0, + toolCallId: 'tool-call-1', + function: { + name: 'send_email', + arguments: JSON.stringify({ recipient: 'person@example.com' }), + }, + }, + ]); + expect(JSON.stringify(events)).not.toContain(token); + expect( + events.filter((event) => event.type === ChatCompletionEventType.ChatCompletionTokenCount) + ).toHaveLength(1); + expect(terminalMessages).toEqual([ + expect.objectContaining({ + content: 'workflow-restored final', + refusal: 'preserved refusal metadata', + toolCalls: [ + expect.objectContaining({ + function: { + name: 'send_email', + arguments: { recipient: 'person@example.com' }, + }, + }), + ], + }), + ]); + expect(invokeConnector).toHaveBeenCalledTimes(1); + expect(invokeConnector).toHaveBeenCalledWith({ + system: expect.stringContaining('protected system'), + messages: protectedMessages, + abortSignal: undefined, + }); + const calledSystem: string = invokeConnector.mock.calls[0][0].system; + expect(calledSystem).toContain('[Anonymization context]'); + expect(calledSystem).toContain(`: EMAIL`); + }); + + it('uses the original direct path when no workflow matches', async () => { + const invokeConnector = jest + .fn() + .mockReturnValue(of(chunkEvent('direct'), messageEvent('direct'))); + const provider: WorkflowAnonymizationProvider = { + supportsSynchronousExecution: true, + execute: jest.fn().mockResolvedValue({ matched: false }), + }; + + await expect( + firstValueFrom( + createWorkflowAnonymizationPipeline(createOptions({ provider, invokeConnector })).pipe( + toArray() + ) + ) + ).resolves.toEqual([chunkEvent('direct'), messageEvent('direct')]); + expect(invokeConnector).toHaveBeenCalledWith({ + system: 'original system', + messages: originalMessages, + abortSignal: undefined, + }); + }); + + it('allows an unsafe direct call only when protection fails before connector invocation', async () => { + const invokeConnector = jest + .fn() + .mockReturnValue(of(chunkEvent('unsafe direct'), messageEvent('unsafe direct'))); + const provider: WorkflowAnonymizationProvider = { + supportsSynchronousExecution: true, + execute: jest.fn().mockRejectedValue(new Error('selection failed')), + }; + const options = createOptions({ provider, invokeConnector, failureMode: 'allow_unsafe' }); + + await expect( + firstValueFrom(createWorkflowAnonymizationPipeline(options).pipe(toArray())) + ).resolves.toEqual([chunkEvent('unsafe direct'), messageEvent('unsafe direct')]); + expect(invokeConnector).toHaveBeenCalledTimes(1); + expect(options.logger.warn).toHaveBeenCalledWith(expect.stringContaining('allow_unsafe')); + }); + + it('blocks before connector invocation by default when protection fails', async () => { + const invokeConnector = jest.fn(); + const provider: WorkflowAnonymizationProvider = { + supportsSynchronousExecution: true, + execute: jest.fn().mockRejectedValue(new Error('selection failed')), + }; + + await expect( + firstValueFrom( + createWorkflowAnonymizationPipeline( + createOptions({ provider, invokeConnector, failureMode: 'block' }) + ).pipe(toArray()) + ) + ).rejects.toThrow('selection failed'); + expect(invokeConnector).not.toHaveBeenCalled(); + }); + + it('never starts a second connector call when the workflow fails after proceed', async () => { + const invokeConnector = jest + .fn() + .mockReturnValue(of(chunkEvent('called'), messageEvent('called'))); + const provider: WorkflowAnonymizationProvider = { + supportsSynchronousExecution: true, + execute: jest.fn(async ({ proceed }) => { + await proceed.invoke({ messages: protectedMessages, tokenMap: {} }); + throw new Error('post-proceed failure'); + }), + }; + const options = createOptions({ provider, invokeConnector, failureMode: 'allow_unsafe' }); + + await expect( + firstValueFrom(createWorkflowAnonymizationPipeline(options).pipe(toArray())) + ).rejects.toThrow('post-proceed failure'); + expect(invokeConnector).toHaveBeenCalledTimes(1); + expect(options.logger.warn).not.toHaveBeenCalled(); + }); + + it('propagates a connector error after proceed without an allow_unsafe retry', async () => { + const connectorError = new Error('connector failed'); + const invokeConnector = jest.fn().mockReturnValue(throwError(() => connectorError)); + const provider: WorkflowAnonymizationProvider = { + supportsSynchronousExecution: true, + execute: jest.fn(async ({ proceed }) => { + await proceed.invoke({ messages: protectedMessages, tokenMap: {} }); + return { matched: true, content: 'unreachable' }; + }), + }; + const options = createOptions({ provider, invokeConnector, failureMode: 'allow_unsafe' }); + + await expect( + firstValueFrom(createWorkflowAnonymizationPipeline(options).pipe(toArray())) + ).rejects.toThrow('connector failed'); + expect(invokeConnector).toHaveBeenCalledTimes(1); + expect(options.logger.warn).not.toHaveBeenCalled(); + }); + + it('fails when the connector does not emit a terminal message', async () => { + const invokeConnector = jest.fn().mockReturnValue(of(chunkEvent('partial'))); + const provider: WorkflowAnonymizationProvider = { + supportsSynchronousExecution: true, + execute: jest.fn(async ({ proceed }) => { + await proceed.invoke({ messages: protectedMessages, tokenMap: {} }); + return { matched: true, content: 'unreachable' }; + }), + }; + + await expect( + firstValueFrom( + createWorkflowAnonymizationPipeline(createOptions({ provider, invokeConnector })).pipe( + toArray() + ) + ) + ).rejects.toThrow('without a terminal message'); + }); + + it('enforces runtime single-use proceed semantics', async () => { + const invokeConnector = jest.fn().mockReturnValue(of(messageEvent('first'))); + const provider: WorkflowAnonymizationProvider = { + supportsSynchronousExecution: true, + execute: jest.fn(async ({ proceed }) => { + await proceed.invoke({ messages: protectedMessages, tokenMap: {} }); + await proceed.invoke({ messages: protectedMessages, tokenMap: {} }); + return { matched: true, content: 'unreachable' }; + }), + }; + + await expect( + firstValueFrom( + createWorkflowAnonymizationPipeline(createOptions({ provider, invokeConnector })).pipe( + toArray() + ) + ) + ).rejects.toThrow('may only be invoked once'); + expect(invokeConnector).toHaveBeenCalledTimes(1); + }); + + describe('pre-LLM timeout', () => { + const makeStalledProvider = (): WorkflowAnonymizationProvider => ({ + supportsSynchronousExecution: true, + execute: jest.fn( + ({ abortSignal: signal }) => + new Promise((_resolve, reject) => { + if (signal?.aborted) { + reject(signal.reason); + return; + } + signal?.addEventListener('abort', () => reject(signal!.reason), { once: true }); + }) + ), + }); + + it('clears the timeout when proceed.invoke is called before it fires', async () => { + const clearTimeoutSpy = jest.spyOn(global, 'clearTimeout'); + const invokeConnector = jest.fn().mockReturnValue(of(messageEvent('ok'))); + const provider: WorkflowAnonymizationProvider = { + supportsSynchronousExecution: true, + execute: jest.fn(async ({ proceed }) => { + await proceed.invoke({ messages: protectedMessages, tokenMap: {} }); + return { matched: true, content: 'ok' }; + }), + }; + + await firstValueFrom( + createWorkflowAnonymizationPipeline( + createOptions({ provider, invokeConnector, preLLMTimeoutMs: 5000 }) + ).pipe(toArray()) + ); + + expect(clearTimeoutSpy).toHaveBeenCalled(); + clearTimeoutSpy.mockRestore(); + }); + + it('falls through to allow_unsafe direct path when timeout fires before proceed.invoke', async () => { + jest.useFakeTimers(); + try { + const invokeConnector = jest + .fn() + .mockReturnValue(of(chunkEvent('fallback'), messageEvent('fallback'))); + const options = createOptions({ + provider: makeStalledProvider(), + invokeConnector, + failureMode: 'allow_unsafe', + preLLMTimeoutMs: 5000, + }); + + const completionPromise = firstValueFrom( + createWorkflowAnonymizationPipeline(options).pipe(toArray()) + ); + // Attach the handler before firing timers to avoid unhandled-rejection warnings. + const assertion = expect(completionPromise).resolves.toEqual([ + chunkEvent('fallback'), + messageEvent('fallback'), + ]); + await jest.runAllTimersAsync(); + await assertion; + + expect(invokeConnector).toHaveBeenCalledTimes(1); + expect(options.logger.warn).toHaveBeenCalledWith(expect.stringContaining('allow_unsafe')); + } finally { + jest.useRealTimers(); + } + }); + + it('propagates the timeout error when failureMode is block', async () => { + jest.useFakeTimers(); + try { + const invokeConnector = jest.fn(); + const options = createOptions({ + provider: makeStalledProvider(), + invokeConnector, + failureMode: 'block', + preLLMTimeoutMs: 5000, + }); + + const completionPromise = firstValueFrom( + createWorkflowAnonymizationPipeline(options).pipe(toArray()) + ); + // Attach the handler before firing timers to avoid unhandled-rejection warnings. + const assertion = expect(completionPromise).rejects.toThrow('timed out'); + await jest.runAllTimersAsync(); + await assertion; + + expect(invokeConnector).not.toHaveBeenCalled(); + } finally { + jest.useRealTimers(); + } + }); + }); + + it('keeps relay and terminal state isolated across concurrent executions', async () => { + const secondToken = 'USER_NAME_fedcba9876543210fedcba9876543210'; + const run = async ({ + currentToken, + original, + finalContent, + }: { + currentToken: string; + original: string; + finalContent: string; + }) => { + const currentTokenMap = { + [currentToken]: { original, entityClass: 'USER_NAME' }, + }; + const invokeConnector = jest + .fn() + .mockReturnValue(of(chunkEvent(currentToken), messageEvent(currentToken))); + const provider: WorkflowAnonymizationProvider = { + supportsSynchronousExecution: true, + execute: jest.fn(async ({ proceed }) => { + await proceed.invoke({ messages: protectedMessages, tokenMap: currentTokenMap }); + return { matched: true, content: finalContent }; + }), + }; + return firstValueFrom( + createWorkflowAnonymizationPipeline(createOptions({ provider, invokeConnector })).pipe( + toArray() + ) + ); + }; + + const [firstEvents, secondEvents] = await Promise.all([ + run({ currentToken: token, original: 'first@example.com', finalContent: 'first final' }), + run({ currentToken: secondToken, original: 'second-user', finalContent: 'second final' }), + ]); + + expect(firstEvents).toEqual([chunkEvent('first@example.com'), messageEvent('first final')]); + expect(secondEvents).toEqual([chunkEvent('second-user'), messageEvent('second final')]); + }); +}); diff --git a/x-pack/platform/plugins/shared/inference/server/chat_complete/workflow_anonymization_pipeline.ts b/x-pack/platform/plugins/shared/inference/server/chat_complete/workflow_anonymization_pipeline.ts new file mode 100644 index 0000000000000..1d889709fa921 --- /dev/null +++ b/x-pack/platform/plugins/shared/inference/server/chat_complete/workflow_anonymization_pipeline.ts @@ -0,0 +1,358 @@ +/* + * 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 { performance } from 'perf_hooks'; +import type { KibanaRequest } from '@kbn/core-http-server'; +import { + ChatCompletionEventType, + type ChatCompletionEvent, + type ChatCompletionChunkEvent, + type ChatCompletionMessageEvent, + type Message, +} from '@kbn/inference-common'; +import type { Logger } from '@kbn/logging'; +import type { Observable } from 'rxjs'; +import { catchError, defer, finalize, merge, of, Subject, switchMap, throwError } from 'rxjs'; +import { createPiiDetectionContext } from '../workflow_anonymization/create_pii_detection_context'; +import { createPiiTokenizationContext } from '../workflow_anonymization/create_pii_tokenization_context'; +import type { PiiRegexWorkerService } from '../workflow_anonymization/detection/regex_worker_service'; +import type { WorkflowAnonymizationOptions } from '../inference_client/workflow_anonymization_options'; +import type { + InferenceProceedInput, + InferenceTokenMapEntry, +} from '../workflow_anonymization_capabilities'; +import { + createStreamingContentRestorer, + restoreTokenizedString, + restoreTokenizedValue, +} from './workflow_anonymization_restoration'; +import { + pipelineFirstChunkDurationHistogram, + pipelineRequestsCounter, +} from './workflow_anonymization_metrics'; + +const buildAnonymizationInstruction = ( + tokenMap: Readonly> +): string => { + const tokenLines = Object.entries(tokenMap).map( + ([token, entry]) => `- ${token}: ${entry.entityClass}` + ); + return [ + '[Anonymization context]', + 'Some values in this conversation have been replaced with privacy tokens to protect sensitive information.', + 'Token registry (treat each token as a real value of its entity type):', + ...tokenLines, + 'Rules:', + '- Use these tokens exactly as given; do not modify, expand, or remove them.', + '- Do not guess, infer, or reveal the original values behind any token.', + '- When writing tool call arguments, preserve token strings verbatim.', + ].join('\n'); +}; + +export interface WorkflowInvocationState { + connectorInvoked: boolean; +} + +interface CreateWorkflowAnonymizationPipelineOptions { + readonly request: KibanaRequest; + readonly namespace: string; + readonly system?: string; + readonly messages: readonly Message[]; + readonly sessionId?: string; + readonly agentId?: string; + readonly abortSignal?: AbortSignal; + readonly regexWorker: PiiRegexWorkerService; + readonly logger: Logger; + readonly workflowAnonymization: WorkflowAnonymizationOptions; + readonly invocationState: WorkflowInvocationState; + readonly invokeConnector: (input: { + system?: string; + messages: readonly Message[]; + abortSignal?: AbortSignal; + }) => Observable; +} + +const restoreTerminalMessage = ( + event: ChatCompletionMessageEvent, + tokenMap: Readonly>, + restoreContent = true, + restoreToolCallArgs = true +): ChatCompletionMessageEvent => ({ + ...event, + content: restoreContent ? restoreTokenizedString(event.content, tokenMap) : event.content, + toolCalls: event.toolCalls.map((toolCall) => { + if (!restoreToolCallArgs) { + return toolCall; + } + const restoredArguments = restoreTokenizedValue(toolCall.function.arguments, tokenMap); + if ( + !restoredArguments || + typeof restoredArguments !== 'object' || + Array.isArray(restoredArguments) + ) { + throw new Error('Workflow token restoration produced invalid tool-call arguments'); + } + return { + ...toolCall, + function: { ...toolCall.function, arguments: restoredArguments }, + }; + }), +}); + +const createRestoredToolCallChunk = ( + event: ChatCompletionMessageEvent +): ChatCompletionChunkEvent | undefined => { + if (event.toolCalls.length === 0) { + return undefined; + } + return { + type: ChatCompletionEventType.ChatCompletionChunk, + content: '', + tool_calls: event.toolCalls.map((toolCall, index) => ({ + index, + toolCallId: toolCall.toolCallId, + function: { + name: toolCall.function.name, + arguments: JSON.stringify(toolCall.function.arguments), + }, + })), + }; +}; + +export const createWorkflowAnonymizationPipeline = ({ + request, + namespace, + system, + messages, + sessionId, + agentId, + abortSignal, + regexWorker, + logger, + workflowAnonymization, + invocationState, + invokeConnector, +}: CreateWorkflowAnonymizationPipelineOptions): Observable => { + const relay$ = new Subject(); + let restoredTerminalMessage: ChatCompletionMessageEvent | undefined; + let proceedInvoked = false; + let preLLMTimer: ReturnType | undefined; + let timedOut = false; + + const clearPreLLMTimer = () => { + if (preLLMTimer !== undefined) { + clearTimeout(preLLMTimer); + preLLMTimer = undefined; + } + }; + + const proceed = { + invoke: async (input: InferenceProceedInput): Promise<{ rawContent: string }> => { + clearPreLLMTimer(); + if (proceedInvoked) { + throw new Error('The workflow inference proceed capability may only be invoked once'); + } + proceedInvoked = true; + invocationState.connectorInvoked = true; + const tokenMap = input.tokenMap ?? {}; + const contentRestorer = createStreamingContentRestorer(tokenMap); + + const tokenCount = Object.keys(tokenMap).length; + const augmentedSystem = + tokenCount > 0 + ? [input.system, buildAnonymizationInstruction(tokenMap)].filter(Boolean).join('\n\n') + : input.system; + + if (input.dryRun) { + logger.debug( + `Dry run: returning anonymized payload without calling inference connector (${tokenCount} tokens, ${input.messages.length} messages)` + ); + const dryRunPayload = JSON.stringify( + { system: augmentedSystem, messages: input.messages }, + null, + 2 + ); + relay$.next({ + type: ChatCompletionEventType.ChatCompletionChunk, + content: dryRunPayload, + tool_calls: [], + }); + restoredTerminalMessage = { + type: ChatCompletionEventType.ChatCompletionMessage, + content: dryRunPayload, + toolCalls: [], + }; + return { rawContent: dryRunPayload }; + } + + logger.debug( + `Sending anonymized request to inference connector (${tokenCount} tokens, ${input.messages.length} messages)` + ); + + return new Promise<{ rawContent: string }>((resolve, reject) => { + let rawContent: string | undefined; + const connectorStartTime = performance.now(); + let firstChunkRecorded = false; + invokeConnector({ + system: augmentedSystem, + messages: input.messages, + abortSignal: input.abortSignal ?? abortSignal, + }).subscribe({ + next: (event) => { + const restoreStreaming = input.restoreStreamingContent !== false; + const restoreToolCallArgs = input.restoreToolCallArguments !== false; + if (event.type === ChatCompletionEventType.ChatCompletionChunk) { + if (restoreStreaming) { + const restoredContent = contentRestorer.push(event.content); + if (restoredContent) { + if (!firstChunkRecorded) { + firstChunkRecorded = true; + pipelineFirstChunkDurationHistogram.record( + performance.now() - connectorStartTime + ); + } + relay$.next({ + ...event, + content: restoredContent, + // Tool-call arguments can be split at arbitrary JSON boundaries. Suppress their + // deltas until the assembled terminal message can be restored structurally. + tool_calls: [], + }); + } + } else { + if (event.content) { + if (!firstChunkRecorded) { + firstChunkRecorded = true; + pipelineFirstChunkDurationHistogram.record( + performance.now() - connectorStartTime + ); + } + relay$.next({ ...event, tool_calls: [] }); + } + } + return; + } + if (event.type === ChatCompletionEventType.ChatCompletionMessage) { + if (restoreStreaming) { + const remainingContent = contentRestorer.flush(); + if (remainingContent) { + relay$.next({ + type: ChatCompletionEventType.ChatCompletionChunk, + content: remainingContent, + tool_calls: [], + }); + } + } + rawContent = event.content; + restoredTerminalMessage = restoreTerminalMessage( + event, + tokenMap, + restoreStreaming, + restoreToolCallArgs + ); + return; + } + relay$.next(event); + }, + error: (error) => { + // Reject through the workflow so it unwinds before the merged stream terminates. + // connectorInvoked is already true, so allow_unsafe can never retry this call. + reject(error); + }, + complete: () => { + if (rawContent === undefined || !restoredTerminalMessage) { + reject(new Error('Inference connector completed without a terminal message')); + return; + } + resolve({ rawContent }); + }, + }); + }); + }, + }; + + const around$ = defer(async () => { + const serverSalt = workflowAnonymization.encryptionKey; + + let effectiveAbortSignal = abortSignal; + if (workflowAnonymization.preLLMTimeoutMs > 0) { + const preLLMController = new AbortController(); + preLLMTimer = setTimeout(() => { + timedOut = true; + preLLMController.abort( + new Error( + `Pre-LLM anonymization timed out after ${workflowAnonymization.preLLMTimeoutMs}ms` + ) + ); + }, workflowAnonymization.preLLMTimeoutMs); + effectiveAbortSignal = abortSignal + ? AbortSignal.any([preLLMController.signal, abortSignal]) + : preLLMController.signal; + } + + return workflowAnonymization.provider.execute({ + event: { system, messages, sessionId, agentId }, + namespace, + request, + pii: createPiiTokenizationContext({ + detectionContext: createPiiDetectionContext({ regexWorker }), + serverSalt, + sessionId, + }), + proceed, + abortSignal: effectiveAbortSignal, + }); + }).pipe( + switchMap((result) => { + if (!result.matched) { + // Close the unused relay before subscribing to the direct stream. merge() keeps the + // around branch active, so no relay events can interleave with the unmatched response. + relay$.complete(); + clearPreLLMTimer(); + invocationState.connectorInvoked = true; + pipelineRequestsCounter.add(1, { outcome: 'unmatched' }); + return invokeConnector({ system, messages, abortSignal }); + } + if (!restoredTerminalMessage) { + relay$.complete(); + pipelineRequestsCounter.add(1, { outcome: 'error' }); + return throwError( + () => new Error('Workflow completed without invoking the inference connector') + ); + } + // All relayed connector events precede the workflow-authoritative terminal output. + relay$.complete(); + pipelineRequestsCounter.add(1, { outcome: 'matched' }); + const restoredToolCallChunk = createRestoredToolCallChunk(restoredTerminalMessage); + const terminalMessage = { ...restoredTerminalMessage, content: result.content }; + return restoredToolCallChunk + ? of(restoredToolCallChunk, terminalMessage) + : of(terminalMessage); + }), + catchError((error) => { + clearPreLLMTimer(); + relay$.complete(); + if ( + workflowAnonymization.failureMode === 'allow_unsafe' && + !invocationState.connectorInvoked + ) { + logger.warn( + 'Workflow-driven anonymization failed before connector invocation; using the direct inference path because allow_unsafe is configured' + ); + invocationState.connectorInvoked = true; + pipelineRequestsCounter.add(1, { outcome: timedOut ? 'fallback_timeout' : 'fallback' }); + return invokeConnector({ system, messages, abortSignal }); + } + pipelineRequestsCounter.add(1, { outcome: 'error' }); + return throwError(() => error); + }), + finalize(() => clearPreLLMTimer()) + ); + + // Subscribe the relay first so synchronous connector emissions cannot be lost. + return merge(relay$, around$); +}; diff --git a/x-pack/platform/plugins/shared/inference/server/chat_complete/workflow_anonymization_restoration.test.ts b/x-pack/platform/plugins/shared/inference/server/chat_complete/workflow_anonymization_restoration.test.ts new file mode 100644 index 0000000000000..ae785bc82055d --- /dev/null +++ b/x-pack/platform/plugins/shared/inference/server/chat_complete/workflow_anonymization_restoration.test.ts @@ -0,0 +1,86 @@ +/* + * 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 { + createStreamingContentRestorer, + restoreTokenizedString, + restoreTokenizedValue, + type InferenceTokenMap, +} from './workflow_anonymization_restoration'; + +const token = 'EMAIL_0123456789abcdef0123456789abcdef'; +const tokenMap: InferenceTokenMap = { + [token]: { original: 'person@example.com', entityClass: 'EMAIL' }, +}; + +describe('workflow anonymization restoration', () => { + it('restores complete tokens and nested structured values', () => { + expect(restoreTokenizedString(`Contact ${token}`, tokenMap)).toBe('Contact person@example.com'); + expect( + restoreTokenizedValue( + { recipients: [token], encoded: JSON.stringify({ email: token }) }, + tokenMap + ) + ).toEqual({ + recipients: ['person@example.com'], + encoded: JSON.stringify({ email: 'person@example.com' }), + }); + }); + + it('holds a token split across chunks until it can be restored safely', () => { + const restorer = createStreamingContentRestorer(tokenMap); + + expect(restorer.push(`Contact ${token.slice(0, 12)}`)).toBe('Contact '); + expect(restorer.push(`${token.slice(12)} now`)).toBe('person@example.com now'); + expect(restorer.flush()).toBe(''); + }); + + it('does not hold back a single-character token prefix to avoid gaps on common letters', () => { + // MIN_PREFIX_HOLDBACK = 2: tokens split at their very first character are emitted immediately + // rather than causing streaming gaps on single letters like 'E' (EMAIL) or 'I' (IP). + // Accepted tradeoff: restoration does not happen when the token is split at position 1. + const restorer = createStreamingContentRestorer(tokenMap); + const first = restorer.push(`Contact ${token.slice(0, 1)}`); + const second = restorer.push(`${token.slice(1)} now`); + + expect(first).toBe(`Contact ${token.slice(0, 1)}`); + // The two halves are processed as separate chunks, so the token is not restored. + expect(`${first}${second}${restorer.flush()}`).toBe(`Contact ${token} now`); + }); + + for (const splitAt of Array.from({ length: token.length - 2 }, (_, index) => index + 2)) { + it(`restores a token split at character ${splitAt} without emitting its prefix`, () => { + const restorer = createStreamingContentRestorer(tokenMap); + const first = restorer.push(`Contact ${token.slice(0, splitAt)}`); + const second = restorer.push(`${token.slice(splitAt)} now`); + + expect(first).toBe('Contact '); + expect(`${first}${second}${restorer.flush()}`).toBe('Contact person@example.com now'); + }); + } + + it('restores a complete token at a chunk boundary without delaying it', () => { + const restorer = createStreamingContentRestorer(tokenMap); + + expect(restorer.push(token)).toBe('person@example.com'); + expect(restorer.flush()).toBe(''); + }); + + it('does not delay unrelated uppercase text that merely resembles a token', () => { + const restorer = createStreamingContentRestorer(tokenMap); + + expect(restorer.push('Status EMAIL_PENDING')).toBe('Status EMAIL_PENDING'); + expect(restorer.flush()).toBe(''); + }); + + it('flushes an incomplete token-shaped suffix without dropping content', () => { + const restorer = createStreamingContentRestorer(tokenMap); + + expect(restorer.push('Value EMAIL_0123')).toBe('Value '); + expect(restorer.flush()).toBe('EMAIL_0123'); + }); +}); diff --git a/x-pack/platform/plugins/shared/inference/server/chat_complete/workflow_anonymization_restoration.ts b/x-pack/platform/plugins/shared/inference/server/chat_complete/workflow_anonymization_restoration.ts new file mode 100644 index 0000000000000..24c904bb66900 --- /dev/null +++ b/x-pack/platform/plugins/shared/inference/server/chat_complete/workflow_anonymization_restoration.ts @@ -0,0 +1,118 @@ +/* + * 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 { InferenceTokenMapEntry } from '../workflow_anonymization_capabilities'; + +export type InferenceTokenMap = Readonly>; + +const orderedTokens = (tokenMap: InferenceTokenMap): string[] => + Object.keys(tokenMap) + .filter(Boolean) + .sort((left, right) => right.length - left.length); + +const escapeRegExp = (value: string): string => value.replace(/[.*+?^${}()|[\]\\]/g, '\\$&'); + +const replaceStringValues = ( + text: string, + replacements: Readonly> +): string => { + const keys = Object.keys(replacements) + .filter((key) => key.length > 0) + .sort((left, right) => right.length - left.length); + + if (keys.length === 0) { + return text; + } + + const pattern = new RegExp(keys.map(escapeRegExp).join('|'), 'g'); + return text.replace(pattern, (matched) => replacements[matched]); +}; + +export const restoreTokenizedString = (value: string, tokenMap: InferenceTokenMap): string => + replaceStringValues( + value, + Object.fromEntries(Object.entries(tokenMap).map(([token, entry]) => [token, entry.original])) + ); + +export const restoreTokenizedValue = (value: unknown, tokenMap: InferenceTokenMap): unknown => { + if (typeof value === 'string') { + return restoreTokenizedString(value, tokenMap); + } + if (Array.isArray(value)) { + return value.map((item) => restoreTokenizedValue(item, tokenMap)); + } + if (!value || typeof value !== 'object') { + return value; + } + return Object.fromEntries( + Object.entries(value).map(([key, item]) => [key, restoreTokenizedValue(item, tokenMap)]) + ); +}; + +const createProperTokenPrefixes = ( + tokens: readonly string[] +): { prefixes: ReadonlySet; maximumLength: number } => { + const prefixes = new Set(); + let maximumLength = 0; + tokens.forEach((token) => { + for (let length = 1; length < token.length; length += 1) { + prefixes.add(token.slice(0, length)); + maximumLength = Math.max(maximumLength, length); + } + }); + return { prefixes, maximumLength }; +}; + +// Minimum tail length before we consider it a potential split token. Single-character +// matches (e.g. 'E' for EMAIL_, 'I' for IP_) are too common in English text and cause +// visible streaming gaps by suppressing relay events for those chunks. +const MIN_PREFIX_HOLDBACK = 2; + +const longestPossibleTokenPrefix = ( + value: string, + prefixes: ReadonlySet, + maximumLength: number +): number => { + for ( + let length = Math.min(value.length, maximumLength); + length >= MIN_PREFIX_HOLDBACK; + length -= 1 + ) { + if (prefixes.has(value.slice(-length))) { + return length; + } + } + return 0; +}; + +export interface StreamingContentRestorer { + push(content: string): string; + flush(): string; +} + +export const createStreamingContentRestorer = ( + tokenMap: InferenceTokenMap +): StreamingContentRestorer => { + const tokens = orderedTokens(tokenMap); + const { prefixes, maximumLength } = createProperTokenPrefixes(tokens); + let heldContent = ''; + + return { + push: (content) => { + const buffered = `${heldContent}${content}`; + const heldLength = longestPossibleTokenPrefix(buffered, prefixes, maximumLength); + const safeLength = buffered.length - heldLength; + heldContent = buffered.slice(safeLength); + return restoreTokenizedString(buffered.slice(0, safeLength), tokenMap); + }, + flush: () => { + const restored = restoreTokenizedString(heldContent, tokenMap); + heldContent = ''; + return restored; + }, + }; +}; diff --git a/x-pack/platform/plugins/shared/inference/server/config.test.ts b/x-pack/platform/plugins/shared/inference/server/config.test.ts index fb498d45455e4..957c647055b21 100644 --- a/x-pack/platform/plugins/shared/inference/server/config.test.ts +++ b/x-pack/platform/plugins/shared/inference/server/config.test.ts @@ -11,6 +11,32 @@ describe('inference config schema', () => { it('validates with defaults', () => { expect(configSchema.validate({})).toMatchObject({ enabled: true, + anonymization: { + workflowDriven: false, + failureMode: 'block', + triggerCacheTtlSeconds: 30, + }, }); }); + + it('accepts the explicit unsafe failure mode', () => { + expect( + configSchema.validate({ + anonymization: { + workflowDriven: true, + failureMode: 'allow_unsafe', + }, + }).anonymization + ).toMatchObject({ + workflowDriven: true, + failureMode: 'allow_unsafe', + }); + }); + + it('accepts triggerCacheTtlSeconds of 0 to disable caching', () => { + expect( + configSchema.validate({ anonymization: { triggerCacheTtlSeconds: 0 } }).anonymization + .triggerCacheTtlSeconds + ).toBe(0); + }); }); diff --git a/x-pack/platform/plugins/shared/inference/server/config.ts b/x-pack/platform/plugins/shared/inference/server/config.ts index 795f2ceb06e4c..25f3b399cea9f 100644 --- a/x-pack/platform/plugins/shared/inference/server/config.ts +++ b/x-pack/platform/plugins/shared/inference/server/config.ts @@ -9,6 +9,23 @@ import { schema, type TypeOf } from '@kbn/config-schema'; export const configSchema = schema.object({ enabled: schema.boolean({ defaultValue: true }), + anonymization: schema.object({ + workflowDriven: schema.boolean({ defaultValue: false }), + encryptionKey: schema.maybe(schema.string({ maxLength: 512 })), + failureMode: schema.oneOf([schema.literal('block'), schema.literal('allow_unsafe')], { + defaultValue: 'block', + }), + // How long (in seconds) the resolved trigger-match result is cached per (space, agentId). + // A workflow configuration change (enable/disable/delete) takes up to this many seconds to + // take effect. Set to 0 to disable caching entirely at the cost of an ES lookup on every + // anonymization-eligible inference call. + triggerCacheTtlSeconds: schema.number({ defaultValue: 30, min: 0 }), + // Maximum time (ms) the anonymization workflow may run before the LLM connector is invoked. + // The clock starts after saltPromise resolves; salt-resolution time is excluded from this + // budget. Bounds only the anonymization execution overhead, not LLM response time. + // Set to 0 to disable the timeout. + preLLMTimeoutMs: schema.number({ defaultValue: 5000, min: 0 }), + }), workers: schema.object({ anonymization: schema.object({ enabled: schema.boolean({ defaultValue: true }), @@ -18,9 +35,22 @@ export const configSchema = schema.object({ idleTimeout: schema.duration({ defaultValue: '30s' }), taskTimeout: schema.duration({ defaultValue: '15s' }), }), + workflowAnonymization: schema.object({ + enabled: schema.boolean({ defaultValue: true }), + // Defaults to maxThreads to keep workers pre-warmed; the workflow-driven path runs + // synchronously on the request hot-path so cold-start latency is unacceptable. + // Lower this to allow partial thread scaling at the cost of occasional cold starts. + minThreads: schema.number({ defaultValue: 3, min: 0 }), + maxThreads: schema.number({ defaultValue: 3, min: 1 }), + maxQueue: schema.number({ defaultValue: 20, min: 1 }), + idleTimeout: schema.duration({ defaultValue: '30s' }), + taskTimeout: schema.duration({ defaultValue: '15s' }), + }), }), }); export type InferenceConfig = TypeOf; export type AnonymizationWorkerConfig = InferenceConfig['workers']['anonymization']; +export type WorkflowAnonymizationWorkerConfig = InferenceConfig['workers']['workflowAnonymization']; +export type WorkflowAnonymizationFailureMode = InferenceConfig['anonymization']['failureMode']; diff --git a/x-pack/platform/plugins/shared/inference/server/index.ts b/x-pack/platform/plugins/shared/inference/server/index.ts index 16a23c47e005b..e17e1d8ee40d0 100644 --- a/x-pack/platform/plugins/shared/inference/server/index.ts +++ b/x-pack/platform/plugins/shared/inference/server/index.ts @@ -21,6 +21,15 @@ import type { export type { InferenceServerSetup, InferenceServerStart }; export type { InferenceEndpoint } from './util/get_inference_endpoints'; +export type { WorkflowAnonymizationProvider } from './workflow_anonymization_provider'; +export { + createPiiTokenizationCapabilityValue, + resolvePiiTokenizationCapabilityValue, + createInferenceProceedCapabilityValue, + resolveInferenceProceedCapabilityValue, + INFERENCE_PROCEED_CAPABILITY_ID, + PII_TOKENIZATION_CAPABILITY_ID, +} from './workflow_anonymization_capabilities'; export { naturalLanguageToEsql, diff --git a/x-pack/platform/plugins/shared/inference/server/inference_client/create_chat_model.test.ts b/x-pack/platform/plugins/shared/inference/server/inference_client/create_chat_model.test.ts index 66e416350cfd2..30804d7735498 100644 --- a/x-pack/platform/plugins/shared/inference/server/inference_client/create_chat_model.test.ts +++ b/x-pack/platform/plugins/shared/inference/server/inference_client/create_chat_model.test.ts @@ -57,6 +57,7 @@ describe('createChatModel', () => { it('calls createClient with the right parameters', async () => { await createChatModel({ request, + namespace: 'space-a', connectorId: '.my-connector', actions, logger, @@ -73,6 +74,7 @@ describe('createChatModel', () => { expect(createClientMock).toHaveBeenCalledWith({ actions, request, + namespace: 'space-a', logger, esClient: mockEsClient, anonymizationRulesPromise: Promise.resolve([]), @@ -84,6 +86,7 @@ describe('createChatModel', () => { it('calls getConnectorById with the right parameters', async () => { await createChatModel({ request, + namespace: 'space-a', connectorId: '.my-connector', actions, logger, @@ -117,6 +120,7 @@ describe('createChatModel', () => { await createChatModel({ request, + namespace: 'space-a', connectorId: '.my-connector', actions, logger, diff --git a/x-pack/platform/plugins/shared/inference/server/inference_client/create_chat_model.ts b/x-pack/platform/plugins/shared/inference/server/inference_client/create_chat_model.ts index 4d6c136c0ac60..4b19d9a5b3c23 100644 --- a/x-pack/platform/plugins/shared/inference/server/inference_client/create_chat_model.ts +++ b/x-pack/platform/plugins/shared/inference/server/inference_client/create_chat_model.ts @@ -17,9 +17,11 @@ import type { RegexWorkerService } from '../chat_complete/anonymization/regex_wo import type { InferenceAnonymizationOptions } from './anonymization_options'; import type { InferenceEndpointIdCache } from '../util/inference_endpoint_id_cache'; import type { TokenUsageLogger } from '../token_usage'; +import type { WorkflowAnonymizationOptions } from './workflow_anonymization_options'; export interface CreateChatModelOptions { request: KibanaRequest; + namespace: string; connectorId: string; actions: ActionsClientProvider; logger: Logger; @@ -31,6 +33,7 @@ export interface CreateChatModelOptions { endpointIdCache: InferenceEndpointIdCache; callbacks?: InferenceCallbacks; anonymization?: InferenceAnonymizationOptions; + workflowAnonymization?: WorkflowAnonymizationOptions; tokenUsageLogger?: TokenUsageLogger; isTokenUsageTrackingEnabled?: () => Promise; isDefaultConnectorOnly?: () => Promise; @@ -39,6 +42,7 @@ export interface CreateChatModelOptions { export const createChatModel = async ({ request, + namespace, connectorId, actions, logger, @@ -50,6 +54,7 @@ export const createChatModel = async ({ endpointIdCache, callbacks, anonymization, + workflowAnonymization, tokenUsageLogger, isTokenUsageTrackingEnabled, isDefaultConnectorOnly, @@ -58,6 +63,7 @@ export const createChatModel = async ({ const client = createClient({ actions, request, + namespace, anonymizationRulesPromise, regexWorker, esClient, @@ -66,6 +72,7 @@ export const createChatModel = async ({ logger, callbacks, anonymization, + workflowAnonymization, tokenUsageLogger, isTokenUsageTrackingEnabled, isDefaultConnectorOnly, diff --git a/x-pack/platform/plugins/shared/inference/server/inference_client/create_client.ts b/x-pack/platform/plugins/shared/inference/server/inference_client/create_client.ts index 5745283165b51..7388fcf412781 100644 --- a/x-pack/platform/plugins/shared/inference/server/inference_client/create_client.ts +++ b/x-pack/platform/plugins/shared/inference/server/inference_client/create_client.ts @@ -24,6 +24,7 @@ import type { RegexWorkerService } from '../chat_complete/anonymization/regex_wo import type { InferenceAnonymizationOptions } from './anonymization_options'; import type { InferenceEndpointIdCache } from '../util/inference_endpoint_id_cache'; import type { TokenUsageLogger } from '../token_usage'; +import type { WorkflowAnonymizationOptions } from './workflow_anonymization_options'; interface CreateClientOptions { request: KibanaRequest; @@ -37,6 +38,7 @@ interface CreateClientOptions { endpointIdCache: InferenceEndpointIdCache; callbacks?: InferenceCallbacks; anonymization?: InferenceAnonymizationOptions; + workflowAnonymization?: WorkflowAnonymizationOptions; tokenUsageLogger?: TokenUsageLogger; isTokenUsageTrackingEnabled?: () => Promise; isDefaultConnectorOnly?: () => Promise; @@ -64,6 +66,7 @@ export function createClient( endpointIdCache, callbacks, anonymization, + workflowAnonymization, tokenUsageLogger, isTokenUsageTrackingEnabled, isDefaultConnectorOnly, @@ -81,6 +84,7 @@ export function createClient( endpointIdCache, callbacks, anonymization, + workflowAnonymization, tokenUsageLogger, isTokenUsageTrackingEnabled, isDefaultConnectorOnly, diff --git a/x-pack/platform/plugins/shared/inference/server/inference_client/inference_client.ts b/x-pack/platform/plugins/shared/inference/server/inference_client/inference_client.ts index 815cf5a0defb1..7779e66e0955c 100644 --- a/x-pack/platform/plugins/shared/inference/server/inference_client/inference_client.ts +++ b/x-pack/platform/plugins/shared/inference/server/inference_client/inference_client.ts @@ -27,6 +27,7 @@ import { createCallbackManager } from './callback_manager'; import type { InferenceAnonymizationOptions } from './anonymization_options'; import type { InferenceEndpointIdCache } from '../util/inference_endpoint_id_cache'; import type { TokenUsageLogger } from '../token_usage'; +import type { WorkflowAnonymizationOptions } from './workflow_anonymization_options'; export function createInferenceClient({ request, @@ -40,6 +41,7 @@ export function createInferenceClient({ endpointIdCache, callbacks, anonymization, + workflowAnonymization, tokenUsageLogger, isTokenUsageTrackingEnabled, isDefaultConnectorOnly, @@ -56,6 +58,7 @@ export function createInferenceClient({ endpointIdCache: InferenceEndpointIdCache; callbacks?: InferenceCallbacks; anonymization?: InferenceAnonymizationOptions; + workflowAnonymization?: WorkflowAnonymizationOptions; tokenUsageLogger?: TokenUsageLogger; isTokenUsageTrackingEnabled?: () => Promise; isDefaultConnectorOnly?: () => Promise; @@ -80,6 +83,7 @@ export function createInferenceClient({ ...(replacementsEsClient ? { esClient: replacementsEsClient } : {}), }, }, + workflowAnonymization, tokenUsageLogger, isTokenUsageTrackingEnabled, isDefaultConnectorOnly, diff --git a/x-pack/platform/plugins/shared/inference/server/inference_client/workflow_anonymization_options.ts b/x-pack/platform/plugins/shared/inference/server/inference_client/workflow_anonymization_options.ts new file mode 100644 index 0000000000000..28db97621bb31 --- /dev/null +++ b/x-pack/platform/plugins/shared/inference/server/inference_client/workflow_anonymization_options.ts @@ -0,0 +1,19 @@ +/* + * 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 { WorkflowAnonymizationFailureMode } from '../config'; +import type { WorkflowAnonymizationProvider } from '../workflow_anonymization_provider'; +import type { PiiRegexWorkerService } from '../workflow_anonymization/detection'; + +export interface WorkflowAnonymizationOptions { + readonly provider: WorkflowAnonymizationProvider; + readonly failureMode: WorkflowAnonymizationFailureMode; + readonly preLLMTimeoutMs: number; + /** HMAC server salt derived from xpack.inference.anonymization.encryptionKey. Undefined when the key is not configured; tokens are session-ID-derived only. */ + readonly encryptionKey?: string; + readonly piiRegexWorker: PiiRegexWorkerService; +} diff --git a/x-pack/platform/plugins/shared/inference/server/plugin.test.ts b/x-pack/platform/plugins/shared/inference/server/plugin.test.ts index 21a3181593219..908d29f959383 100644 --- a/x-pack/platform/plugins/shared/inference/server/plugin.test.ts +++ b/x-pack/platform/plugins/shared/inference/server/plugin.test.ts @@ -5,7 +5,9 @@ * 2.0. */ -import { resolveReplacementsEncryptionKey } from './plugin'; +import type { WorkflowAnonymizationProvider } from './workflow_anonymization_provider'; +import { resolveReplacementsEncryptionKey, resolveWorkflowAnonymizationOptions } from './plugin'; +import { createPiiRegexWorkerServiceMock } from './test_utils'; describe('resolveReplacementsEncryptionKey', () => { it('returns undefined when anonymization is disabled', async () => { @@ -38,3 +40,110 @@ describe('resolveReplacementsEncryptionKey', () => { ).resolves.toBeUndefined(); }); }); + +describe('resolveWorkflowAnonymizationOptions', () => { + const provider: WorkflowAnonymizationProvider = { + supportsSynchronousExecution: true, + execute: jest.fn(), + }; + const piiRegexWorker = createPiiRegexWorkerServiceMock(); + + it('does not enable or log when workflow mode is disabled', () => { + const logger = { error: jest.fn() }; + + expect( + resolveWorkflowAnonymizationOptions({ + enabled: false, + failureMode: 'block', + preLLMTimeoutMs: 5000, + provider, + piiRegexWorker, + logger, + }) + ).toBeUndefined(); + expect(logger.error).not.toHaveBeenCalled(); + }); + + it('enables the registered synchronous provider', () => { + const logger = { error: jest.fn() }; + + expect( + resolveWorkflowAnonymizationOptions({ + enabled: true, + failureMode: 'allow_unsafe', + preLLMTimeoutMs: 3000, + provider, + piiRegexWorker, + logger, + }) + ).toEqual({ + provider, + failureMode: 'allow_unsafe', + preLLMTimeoutMs: 3000, + encryptionKey: undefined, + piiRegexWorker, + }); + expect(logger.error).not.toHaveBeenCalled(); + }); + + it('passes encryptionKey through to the returned options', () => { + const logger = { error: jest.fn() }; + + expect( + resolveWorkflowAnonymizationOptions({ + enabled: true, + failureMode: 'block', + preLLMTimeoutMs: 5000, + encryptionKey: 'my-hmac-key', + provider, + piiRegexWorker, + logger, + }) + ).toEqual({ + provider, + failureMode: 'block', + preLLMTimeoutMs: 5000, + encryptionKey: 'my-hmac-key', + piiRegexWorker, + }); + }); + + it('logs once and retains legacy behavior when the provider is unavailable', () => { + const logger = { error: jest.fn() }; + + expect( + resolveWorkflowAnonymizationOptions({ + enabled: true, + failureMode: 'block', + preLLMTimeoutMs: 5000, + piiRegexWorker, + logger, + }) + ).toBeUndefined(); + expect(logger.error).toHaveBeenCalledWith( + expect.stringContaining('retaining legacy anonymization') + ); + }); + + it('logs and falls back to legacy when provider does not support synchronous execution', () => { + const logger = { error: jest.fn() }; + const asyncOnlyProvider: WorkflowAnonymizationProvider = { + supportsSynchronousExecution: false, + execute: jest.fn(), + }; + + expect( + resolveWorkflowAnonymizationOptions({ + enabled: true, + failureMode: 'block', + preLLMTimeoutMs: 5000, + provider: asyncOnlyProvider, + piiRegexWorker, + logger, + }) + ).toBeUndefined(); + expect(logger.error).toHaveBeenCalledWith( + expect.stringContaining('retaining legacy anonymization') + ); + }); +}); diff --git a/x-pack/platform/plugins/shared/inference/server/plugin.ts b/x-pack/platform/plugins/shared/inference/server/plugin.ts index 1a73d55360c71..a257a9402fed2 100644 --- a/x-pack/platform/plugins/shared/inference/server/plugin.ts +++ b/x-pack/platform/plugins/shared/inference/server/plugin.ts @@ -28,6 +28,7 @@ import { createChatModel, } from './inference_client'; import { RegexWorkerService } from './chat_complete/anonymization/regex_worker_service'; +import { PiiRegexWorkerService } from './workflow_anonymization/detection'; import { registerRoutes } from './routes'; import type { InferenceConfig } from './config'; import type { @@ -47,6 +48,8 @@ import { getInferenceEndpointById } from './util/get_inference_endpoint_by_id'; import { InferenceEndpointIdCache } from './util/inference_endpoint_id_cache'; import { TokenUsageLogger } from './token_usage'; import { installTokenUsageDashboard } from './dashboard'; +import type { WorkflowAnonymizationProvider } from './workflow_anonymization_provider'; +import type { WorkflowAnonymizationOptions } from './inference_client/workflow_anonymization_options'; const parseLegacyAnonymizationRules = (value: unknown): AnonymizationRule[] => { let parsed: unknown = value; @@ -90,6 +93,35 @@ export const resolveReplacementsEncryptionKey = async ({ return policyService.getReplacementsEncryptionKey(namespace); }; +export const resolveWorkflowAnonymizationOptions = ({ + enabled, + failureMode, + preLLMTimeoutMs, + encryptionKey, + provider, + piiRegexWorker, + logger, +}: { + enabled: boolean; + failureMode: WorkflowAnonymizationOptions['failureMode']; + preLLMTimeoutMs: number; + encryptionKey?: string; + provider?: WorkflowAnonymizationProvider; + piiRegexWorker: PiiRegexWorkerService; + logger: Pick; +}): WorkflowAnonymizationOptions | undefined => { + if (!enabled) { + return undefined; + } + if (!provider?.supportsSynchronousExecution) { + logger.error( + 'Workflow-driven inference anonymization is configured but synchronous workflow support is unavailable; retaining legacy anonymization' + ); + return undefined; + } + return { provider, failureMode, preLLMTimeoutMs, encryptionKey, piiRegexWorker }; +}; + export class InferencePlugin implements Plugin< @@ -102,8 +134,10 @@ export class InferencePlugin private logger: Logger; private config: InferenceConfig; private regexWorker?: RegexWorkerService; + private piiRegexWorker?: PiiRegexWorkerService; private endpointIdCache: InferenceEndpointIdCache; private tokenUsageLogger: TokenUsageLogger; + private workflowAnonymizationProvider?: WorkflowAnonymizationProvider; constructor(context: PluginInitializerContext) { this.logger = context.logger.get(); @@ -124,7 +158,18 @@ export class InferencePlugin logger: this.logger, }); - return {}; + return { + registerWorkflowAnonymizationProvider: (provider) => { + if (this.workflowAnonymizationProvider) { + throw new Error('A workflow anonymization provider is already registered'); + } + this.workflowAnonymizationProvider = provider; + }, + anonymizationConfig: { + triggerCacheTtlMs: Math.round(this.config.anonymization.triggerCacheTtlSeconds * 1000), + workflowDrivenEnabled: this.config.anonymization.workflowDriven, + }, + }; } start(core: CoreStart, pluginsStart: InferenceStartDependencies): InferenceServerStart { @@ -154,10 +199,34 @@ export class InferencePlugin ); } + if (this.config.anonymization.workflowDriven && !this.config.anonymization.encryptionKey) { + this.logger.warn( + 'xpack.inference.anonymization.encryptionKey is not configured; tokens will be derived from the session ID alone and are not server-hardened. Configure this key for HMAC-backed tokens.' + ); + } + this.regexWorker = new RegexWorkerService( this.config.workers.anonymization, this.logger.get('regex_worker') ); + if (this.config.anonymization.workflowDriven) { + this.piiRegexWorker = new PiiRegexWorkerService( + this.config.workers.workflowAnonymization, + this.logger.get('pii_regex_worker') + ); + } + + const workflowAnonymization = this.piiRegexWorker + ? resolveWorkflowAnonymizationOptions({ + enabled: this.config.anonymization.workflowDriven, + failureMode: this.config.anonymization.failureMode, + preLLMTimeoutMs: this.config.anonymization.preLLMTimeoutMs, + encryptionKey: this.config.anonymization.encryptionKey, + provider: this.workflowAnonymizationProvider, + piiRegexWorker: this.piiRegexWorker, + logger: this.logger, + }) + : undefined; const createAnonymizationRulesPromise = async (request: KibanaRequest) => { const namespace = @@ -220,6 +289,8 @@ export class InferencePlugin })(), esClient: core.elasticsearch.client.asScoped(request).asCurrentUser, anonymization: { + // Legacy salt path — always undefined today (ANONYMIZATION_FEATURE_ACTIVE is hardcoded + // false). Will be removed when the anonymization plugin is deleted. saltPromise: anonymizationEnabled ? policyService?.getSalt(namespace) : undefined, resolveEffectivePolicy: async (target?: ChatCompleteAnonymizationTarget) => { if (!anonymizationEnabled || !policyService || !target) { @@ -288,6 +359,7 @@ export class InferencePlugin esClient: core.elasticsearch.client.asScoped(options.request).asCurrentUser, endpointIdCache: this.endpointIdCache, tokenUsageLogger: this.tokenUsageLogger, + workflowAnonymization, isTokenUsageTrackingEnabled: createTokenUsageTrackingEnabledCheck(options.request), isDefaultConnectorOnly: createDefaultConnectorOnlyCheck(options.request), getDefaultConnectorId: createDefaultConnectorIdGetter(options.request), @@ -308,6 +380,7 @@ export class InferencePlugin endpointIdCache: this.endpointIdCache, logger: this.logger, tokenUsageLogger: this.tokenUsageLogger, + workflowAnonymization, isTokenUsageTrackingEnabled: createTokenUsageTrackingEnabledCheck(options.request), isDefaultConnectorOnly: createDefaultConnectorOnlyCheck(options.request), getDefaultConnectorId: createDefaultConnectorIdGetter(options.request), @@ -381,5 +454,6 @@ export class InferencePlugin async stop() { await this.regexWorker?.stop(); + await this.piiRegexWorker?.stop(); } } diff --git a/x-pack/platform/plugins/shared/inference/server/test_utils/index.ts b/x-pack/platform/plugins/shared/inference/server/test_utils/index.ts index 2cb236f1e4fa2..a18fa73d1d424 100644 --- a/x-pack/platform/plugins/shared/inference/server/test_utils/index.ts +++ b/x-pack/platform/plugins/shared/inference/server/test_utils/index.ts @@ -11,3 +11,4 @@ export { createInferenceConnectorMock } from './inference_connector'; export { createInferenceConnectorAdapterMock } from './inference_connector_adapter'; export { createInferenceExecutorMock } from './inference_executor'; export { createRegexWorkerServiceMock } from './regex_worker_service.mock'; +export { createPiiRegexWorkerServiceMock } from './pii_regex_worker_service.mock'; diff --git a/x-pack/platform/plugins/shared/inference/server/test_utils/pii_regex_worker_service.mock.ts b/x-pack/platform/plugins/shared/inference/server/test_utils/pii_regex_worker_service.mock.ts new file mode 100644 index 0000000000000..390088e788972 --- /dev/null +++ b/x-pack/platform/plugins/shared/inference/server/test_utils/pii_regex_worker_service.mock.ts @@ -0,0 +1,17 @@ +/* + * 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 { PiiRegexMatch } from '../workflow_anonymization/detection/types'; +import type { PiiRegexWorkerService } from '../workflow_anonymization/detection/regex_worker_service'; + +export const createPiiRegexWorkerServiceMock = () => { + const mock = { + run: jest.fn((): Promise => Promise.resolve([])), + stop: jest.fn().mockResolvedValue(undefined), + }; + return mock as unknown as PiiRegexWorkerService; +}; diff --git a/x-pack/platform/plugins/shared/inference/server/types.ts b/x-pack/platform/plugins/shared/inference/server/types.ts index 0d53d639538d0..3d7093c6760c8 100644 --- a/x-pack/platform/plugins/shared/inference/server/types.ts +++ b/x-pack/platform/plugins/shared/inference/server/types.ts @@ -36,8 +36,7 @@ import type { AnonymizationPluginSetup, } from '@kbn/anonymization-plugin/server'; import type { InferenceEndpoint } from './util/get_inference_endpoints'; - -/* eslint-disable @typescript-eslint/no-empty-interface*/ +import type { WorkflowAnonymizationProvider } from './workflow_anonymization_provider'; export interface InferenceSetupDependencies { actions: ActionsPluginSetup; @@ -52,7 +51,19 @@ export interface InferenceStartDependencies { /** * Setup contract of the inference plugin. */ -export interface InferenceServerSetup {} +export interface InferenceServerSetup { + registerWorkflowAnonymizationProvider(provider: WorkflowAnonymizationProvider): void; + /** Operational config values derived from the inference plugin's config, exposed so that + * consumer plugins (e.g. inference_workflows) can share the same settings without + * duplicating config keys. */ + anonymizationConfig: { + /** Trigger-resolution cache TTL in milliseconds. 0 means caching is disabled. */ + triggerCacheTtlMs: number; + /** Whether workflow-driven anonymization is enabled. Consumer plugins should use this + * to gate any unconditional startup work (e.g. managed workflow installation). */ + workflowDrivenEnabled: boolean; + }; +} /** * Options to create an inference client using the {@link InferenceServerStart.getClient} API. diff --git a/x-pack/platform/plugins/shared/inference/server/workflow_anonymization/create_pii_detection_context.ts b/x-pack/platform/plugins/shared/inference/server/workflow_anonymization/create_pii_detection_context.ts new file mode 100644 index 0000000000000..1d90936af60eb --- /dev/null +++ b/x-pack/platform/plugins/shared/inference/server/workflow_anonymization/create_pii_detection_context.ts @@ -0,0 +1,62 @@ +/* + * 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 { AnonymizationRule } from '@kbn/inference-common'; +import type { + DetectedPiiEntity, + PiiDetectionContext, + PiiTextRecord, +} from './pii_detection_context'; +import type { PiiRegexWorkerService } from './detection/regex_worker_service'; + +export const createPiiDetectionContext = ({ + regexWorker, +}: { + regexWorker: PiiRegexWorkerService; +}): PiiDetectionContext => ({ + detectEntities: async ({ records, rules, abortSignal }) => { + abortSignal?.throwIfAborted(); + + if (hasEnabledNerRule(rules)) { + throw new Error('NER detection is not supported by workflow-driven anonymization'); + } + + const piiRules = rules + .filter( + (rule): rule is Extract => + rule.enabled && rule.type === 'RegExp' + ) + .map((rule) => ({ entityClass: rule.entityClass, pattern: rule.pattern })); + + if (piiRules.length === 0 || records.length === 0) { + return []; + } + + const workerRecords = records.map(({ id, text }) => ({ [id]: text })); + const matches = await regexWorker.run({ rules: piiRules, records: workerRecords }); + abortSignal?.throwIfAborted(); + + return matches.map((match) => ({ + recordId: getRecordId(records, match.recordIndex), + start: match.start, + end: match.end, + value: match.matchValue, + entityClass: match.entityClass, + })); + }, +}); + +const hasEnabledNerRule = (rules: readonly AnonymizationRule[]): boolean => + rules.some((rule) => rule.enabled && rule.type === 'NER'); + +const getRecordId = (records: readonly PiiTextRecord[], recordIndex: number): string => { + const record = records[recordIndex]; + if (!record) { + throw new Error(`PII detector returned an invalid record index: ${recordIndex}`); + } + return record.id; +}; diff --git a/x-pack/platform/plugins/shared/inference/server/workflow_anonymization/create_pii_tokenization_context.ts b/x-pack/platform/plugins/shared/inference/server/workflow_anonymization/create_pii_tokenization_context.ts new file mode 100644 index 0000000000000..cd4e36bfa490c --- /dev/null +++ b/x-pack/platform/plugins/shared/inference/server/workflow_anonymization/create_pii_tokenization_context.ts @@ -0,0 +1,47 @@ +/* + * 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 { createHmac, randomBytes } from 'crypto'; +import type { PiiTokenizationContext } from '../workflow_anonymization_capabilities'; +import type { PiiDetectionContext } from './pii_detection_context'; +import { generateEntityToken } from './detection/entity_mask'; + +const deriveExecutionScope = ({ + serverSalt, + sessionId, +}: { + serverSalt?: string; + sessionId?: string; +}): string => { + if (!sessionId) { + return randomBytes(32).toString('hex'); + } + // Without a server key, use the session ID directly as salt — tokens are within-session stable + // but not server-hardened. Configure xpack.inference.anonymization.encryptionKey for HMAC-backed + // tokens that are opaque even if the session ID is known. + if (!serverSalt) { + return sessionId; + } + return createHmac('sha256', serverSalt).update(`session:${sessionId}`).digest('hex'); +}; + +export const createPiiTokenizationContext = ({ + detectionContext, + serverSalt, + sessionId, +}: { + detectionContext: PiiDetectionContext; + serverSalt?: string; + sessionId?: string; +}): PiiTokenizationContext => { + const executionScope = deriveExecutionScope({ serverSalt, sessionId }); + + return { + detectEntities: (options) => detectionContext.detectEntities(options), + tokenize: (entityClass, value) => generateEntityToken(executionScope, entityClass, value), + }; +}; diff --git a/x-pack/platform/plugins/shared/inference/server/workflow_anonymization/detection/assert_re2_compilable.test.ts b/x-pack/platform/plugins/shared/inference/server/workflow_anonymization/detection/assert_re2_compilable.test.ts new file mode 100644 index 0000000000000..e83e4928534d0 --- /dev/null +++ b/x-pack/platform/plugins/shared/inference/server/workflow_anonymization/detection/assert_re2_compilable.test.ts @@ -0,0 +1,61 @@ +/* + * 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 { assertRe2Compilable } from './assert_re2_compilable'; + +describe('assertRe2Compilable', () => { + describe('valid RE2 patterns', () => { + it('does not throw for a simple literal pattern', () => { + expect(() => assertRe2Compilable('hello')).not.toThrow(); + }); + + it('does not throw for a character class pattern', () => { + expect(() => assertRe2Compilable('[a-z0-9]+')).not.toThrow(); + }); + + it('does not throw for an IP-address pattern', () => { + expect(() => + assertRe2Compilable('\\b\\d{1,3}\\.\\d{1,3}\\.\\d{1,3}\\.\\d{1,3}\\b') + ).not.toThrow(); + }); + }); + + describe('RE2-unsupported constructs throw', () => { + it('throws for a positive lookahead', () => { + expect(() => assertRe2Compilable('(?=a)b')).toThrow(); + }); + + it('throws for a negative lookahead', () => { + expect(() => assertRe2Compilable('(?!a)b')).toThrow(); + }); + + it('throws for a lookbehind', () => { + expect(() => assertRe2Compilable('(?<=@)\\w+')).toThrow(); + }); + + it('throws for a backreference', () => { + expect(() => assertRe2Compilable('(\\w+)\\s+\\1')).toThrow(); + }); + }); + + describe('invalid patterns throw', () => { + it('throws for an unclosed group', () => { + expect(() => assertRe2Compilable('(unclosed')).toThrow(); + }); + }); + + describe('error message quality', () => { + it('includes the offending pattern in the error message', () => { + const pattern = '(?=a)b'; + expect(() => assertRe2Compilable(pattern)).toThrow(pattern); + }); + + it('mentions lookahead / lookbehind / backreferences in the message', () => { + expect(() => assertRe2Compilable('(?=a)b')).toThrow('lookahead'); + }); + }); +}); diff --git a/x-pack/platform/plugins/shared/inference/server/workflow_anonymization/detection/assert_re2_compilable.ts b/x-pack/platform/plugins/shared/inference/server/workflow_anonymization/detection/assert_re2_compilable.ts new file mode 100644 index 0000000000000..2675805e1e40a --- /dev/null +++ b/x-pack/platform/plugins/shared/inference/server/workflow_anonymization/detection/assert_re2_compilable.ts @@ -0,0 +1,27 @@ +/* + * 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 { RE2JS } from 're2js'; + +/** + * Asserts that `pattern` is valid RE2 syntax. + * + * Throws a descriptive Error when the pattern contains constructs RE2 does not + * support — most commonly lookahead (`(?=…)`), lookbehind (`(?<=…)`), or + * backreferences (`\1`). Call this at workflow-definition-save time so a bad + * pattern is caught before it reaches the regex worker at inference time. + */ +export const assertRe2Compilable = (pattern: string): void => { + try { + RE2JS.compile(pattern); + } catch (err) { + throw new Error( + `Pattern is not valid RE2 syntax (lookahead, lookbehind and backreferences are ` + + `not supported): ${pattern}\nCause: ${err instanceof Error ? err.message : String(err)}` + ); + } +}; diff --git a/x-pack/platform/plugins/shared/inference/server/workflow_anonymization/detection/entity_mask.test.ts b/x-pack/platform/plugins/shared/inference/server/workflow_anonymization/detection/entity_mask.test.ts new file mode 100644 index 0000000000000..3b89cc42060de --- /dev/null +++ b/x-pack/platform/plugins/shared/inference/server/workflow_anonymization/detection/entity_mask.test.ts @@ -0,0 +1,103 @@ +/* + * 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 { generateEntityToken } from './entity_mask'; + +describe('generateEntityToken', () => { + describe('output format', () => { + it('returns a string of the form _', () => { + const token = generateEntityToken('scope', 'EMAIL', 'user@example.com'); + expect(token).toMatch(/^EMAIL_[0-9a-f]+$/); + }); + + it('default hash segment is 32 hex characters', () => { + const token = generateEntityToken('scope', 'EMAIL', 'user@example.com'); + const hash = token.slice('EMAIL_'.length); + expect(hash).toHaveLength(32); + }); + + it('hashLength parameter controls the hash segment length', () => { + const token = generateEntityToken('scope', 'IP', '10.0.0.1', 16); + const hash = token.slice('IP_'.length); + expect(hash).toHaveLength(16); + }); + + it('clamps hashLength to MAX_HASH_LENGTH (64)', () => { + const token = generateEntityToken('scope', 'IP', '10.0.0.1', 200); + const hash = token.slice('IP_'.length); + expect(hash).toHaveLength(64); + }); + + it('floors a fractional hashLength', () => { + const token = generateEntityToken('scope', 'IP', '10.0.0.1', 8.9); + const hash = token.slice('IP_'.length); + expect(hash).toHaveLength(8); + }); + + it('falls back to the default (32) for non-finite hashLength', () => { + const withNaN = generateEntityToken('scope', 'IP', '10.0.0.1', NaN); + const withInf = generateEntityToken('scope', 'IP', '10.0.0.1', Infinity); + expect(withNaN.slice('IP_'.length)).toHaveLength(32); + expect(withInf.slice('IP_'.length)).toHaveLength(32); + }); + + it('falls back to the default (32) for zero or negative hashLength', () => { + const withZero = generateEntityToken('scope', 'IP', '10.0.0.1', 0); + const withNeg = generateEntityToken('scope', 'IP', '10.0.0.1', -1); + expect(withZero.slice('IP_'.length)).toHaveLength(32); + expect(withNeg.slice('IP_'.length)).toHaveLength(32); + }); + }); + + describe('determinism', () => { + it('returns the same token for the same inputs', () => { + const a = generateEntityToken('scope-abc', 'EMAIL', 'user@example.com'); + const b = generateEntityToken('scope-abc', 'EMAIL', 'user@example.com'); + expect(a).toBe(b); + }); + }); + + describe('delimiter collision protection', () => { + // The HMAC input uses length-prefixed components so that a value containing ":" + // cannot collide with a different entityClass/value pair. + it('entityClass containing ":" does not collide with a split across entityClass/value', () => { + // Without length-prefixing, ("A:B", "C") and ("A", "B:C") could produce the + // same HMAC input string. + const a = generateEntityToken('scope', 'A:B', 'C'); + const b = generateEntityToken('scope', 'A', 'B:C'); + // Compare hash portions only — the entityClass prefixes differ by construction, + // so comparing the full token would pass even if the HMAC inputs collided. + expect(a.slice(a.lastIndexOf('_') + 1)).not.toBe(b.slice(b.lastIndexOf('_') + 1)); + }); + }); + + describe('invalid inputs', () => { + it('throws for an empty entityClass', () => { + expect(() => generateEntityToken('scope', '', 'value')).toThrow('entityClass'); + }); + }); + + describe('sensitivity to each input', () => { + it('produces a different token when value differs', () => { + const a = generateEntityToken('scope', 'EMAIL', 'a@example.com'); + const b = generateEntityToken('scope', 'EMAIL', 'b@example.com'); + expect(a).not.toBe(b); + }); + + it('produces a different token when entityClass differs', () => { + const a = generateEntityToken('scope', 'EMAIL', 'x@example.com'); + const b = generateEntityToken('scope', 'IP', 'x@example.com'); + expect(a).not.toBe(b); + }); + + it('produces a different token when executionScope differs', () => { + const a = generateEntityToken('scope-1', 'EMAIL', 'user@example.com'); + const b = generateEntityToken('scope-2', 'EMAIL', 'user@example.com'); + expect(a).not.toBe(b); + }); + }); +}); diff --git a/x-pack/platform/plugins/shared/inference/server/workflow_anonymization/detection/entity_mask.ts b/x-pack/platform/plugins/shared/inference/server/workflow_anonymization/detection/entity_mask.ts new file mode 100644 index 0000000000000..ec174d35e8249 --- /dev/null +++ b/x-pack/platform/plugins/shared/inference/server/workflow_anonymization/detection/entity_mask.ts @@ -0,0 +1,46 @@ +/* + * 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 { createHmac } from 'crypto'; + +const DEFAULT_HASH_LENGTH = 32; +const MAX_HASH_LENGTH = 64; + +/** + * Generates a deterministic anonymization token for a PII value. + * + * Format: `_` where the hash is the first `hashLength` hex + * characters of HMAC-SHA256 keyed with `executionScope`. + * + * The `executionScope` is per-execution (derived from the server salt + session ID + * in `createPiiTokenizationContext`), so tokens are stable within an execution but + * opaque across executions or spaces. + * + * Deliberately independent of `@kbn/anonymization-common` — that package belongs + * to the abandoned third-effort anonymization plugin. + */ +export const generateEntityToken = ( + executionScope: string, + entityClass: string, + value: string, + hashLength = DEFAULT_HASH_LENGTH +): string => { + if (!entityClass) { + throw new Error('entityClass must be a non-empty string'); + } + + const clampedLen = + Number.isFinite(hashLength) && hashLength > 0 + ? Math.min(Math.floor(hashLength), MAX_HASH_LENGTH) + : DEFAULT_HASH_LENGTH; + + // Length-prefixed format prevents delimiter collisions when components + // contain the separator character. + const hmacInput = `${entityClass.length}:${entityClass}:${value.length}:${value}`; + const hash = createHmac('sha256', executionScope).update(hmacInput).digest('hex'); + return `${entityClass}_${hash.substring(0, clampedLen)}`; +}; diff --git a/x-pack/platform/plugins/shared/inference/server/workflow_anonymization/detection/execute_regex_rules.test.ts b/x-pack/platform/plugins/shared/inference/server/workflow_anonymization/detection/execute_regex_rules.test.ts new file mode 100644 index 0000000000000..d4982adf34d18 --- /dev/null +++ b/x-pack/platform/plugins/shared/inference/server/workflow_anonymization/detection/execute_regex_rules.test.ts @@ -0,0 +1,173 @@ +/* + * 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 { executeRegexRules } from './execute_regex_rules'; +import type { PiiRegexRule } from './types'; + +const r = (entityClass: string, pattern: string, maxMatchLength?: number): PiiRegexRule => ({ + entityClass, + pattern, + maxMatchLength, +}); + +describe('executeRegexRules', () => { + describe('basic matching', () => { + it('finds a simple IP address match', () => { + const rules = [r('IP', '\\b\\d{1,3}\\.\\d{1,3}\\.\\d{1,3}\\.\\d{1,3}\\b')]; + const records = [{ content: 'connect to 10.0.0.1 now' }]; + const results = executeRegexRules({ rules, records }); + + expect(results).toHaveLength(1); + expect(results[0]).toMatchObject({ + matchValue: '10.0.0.1', + entityClass: 'IP', + start: 11, + end: 19, + ruleIndex: 0, + recordIndex: 0, + recordKey: 'content', + }); + }); + + it('finds multiple matches in one field', () => { + const rules = [r('IP', '\\b\\d{1,3}\\.\\d{1,3}\\.\\d{1,3}\\.\\d{1,3}\\b')]; + const records = [{ content: '10.0.0.1 and 192.168.1.1' }]; + const results = executeRegexRules({ rules, records }); + + expect(results).toHaveLength(2); + expect(results[0].matchValue).toBe('10.0.0.1'); + expect(results[1].matchValue).toBe('192.168.1.1'); + }); + + it('returns no matches when nothing in the field satisfies the rule', () => { + const rules = [r('EMAIL', '[a-zA-Z0-9._%+\\-]+@[a-zA-Z0-9.\\-]+\\.[a-zA-Z]{2,}')]; + const records = [{ content: 'no email here' }]; + expect(executeRegexRules({ rules, records })).toHaveLength(0); + }); + + it('preserves ruleIndex and recordIndex for multi-rule multi-record inputs', () => { + const rules = [ + r('IP', '\\b\\d{1,3}\\.\\d{1,3}\\.\\d{1,3}\\.\\d{1,3}\\b'), + r('EMAIL', '[a-zA-Z0-9._%+\\-]+@[a-zA-Z0-9.\\-]+\\.[a-zA-Z]{2,}'), + ]; + const records = [{ content: 'no match here' }, { content: 'reach me at user@example.com' }]; + const results = executeRegexRules({ rules, records }); + + // Only the email rule fires on record[1] + expect(results).toHaveLength(1); + expect(results[0]).toMatchObject({ ruleIndex: 1, recordIndex: 1, entityClass: 'EMAIL' }); + }); + }); + + describe('zero-length match handling', () => { + // The snapshot used `break` on zero-length matches, which abandoned the rest of the + // field — this verifies the correct `continue` (advance-one-char) behavior. + + it('finds ALL non-zero occurrences of a pattern that can match zero characters', () => { + // a* matches "" (zero-length) between non-'a' chars AND "aaa" as a run. + // The advance-one-char guard must skip the zero-length hits so both "aaa" + // runs are still found. + const rules = [r('A_RUN', 'a*')]; + const records = [{ content: 'aaa hello aaa world' }]; + const results = executeRegexRules({ rules, records }); + + const matchValues = results.map((m) => m.matchValue); + expect(matchValues).toContain('aaa'); + expect(matchValues.filter((v) => v === 'aaa')).toHaveLength(2); + }); + + it('does not infinite-loop on a pattern that always matches empty string', () => { + const rules = [r('EMPTY', 'x*')]; // matches "" everywhere + const records = [{ content: 'abc' }]; + // Should terminate without hanging; zero-length matches are skipped + const results = executeRegexRules({ rules, records }); + // x* matches "" at every position — all skipped; "x" not present so no non-empty match + expect(results).toHaveLength(0); + }); + }); + + describe('native RegExp fallback for RE2-unsupported constructs', () => { + // RE2 does not support lookahead, lookbehind, or backreferences. + // Those patterns fall back to native RegExp and must still produce matches. + + it('matches a lookahead pattern via native RegExp fallback', () => { + // (?=[a-z]) is valid PCRE/native but not RE2 + const rules = [r('HOST', '(?=[a-z])\\w+')]; + const records = [{ content: 'somehost' }]; + const results = executeRegexRules({ rules, records }); + expect(results.length).toBeGreaterThan(0); + expect(results[0].matchValue).toBe('somehost'); + }); + + it('matches a lookbehind pattern via native RegExp fallback', () => { + const rules = [r('AFTER_AT', '(?<=@)\\w+')]; + const records = [{ content: 'user@example' }]; + const results = executeRegexRules({ rules, records }); + expect(results.length).toBe(1); + expect(results[0].matchValue).toBe('example'); + }); + + it('matches a backreference pattern via native RegExp fallback', () => { + const rules = [r('REPEATED', '(\\w+)\\s+\\1')]; + const records = [{ content: 'hello hello' }]; + const results = executeRegexRules({ rules, records }); + expect(results.length).toBe(1); + expect(results[0].matchValue).toBe('hello hello'); + }); + }); + + describe('truly invalid patterns throw (fail closed)', () => { + // Patterns that are invalid in both RE2 and native RegExp must still throw so + // the caller can apply its own failure-mode policy. + + it('throws for an unclosed group', () => { + const rules = [r('BAD', '(unclosed')]; + const records = [{ content: 'test' }]; + expect(() => executeRegexRules({ rules, records })).toThrow(); + }); + }); + + describe('maxMatchLength', () => { + it('discards matches that exceed maxMatchLength', () => { + // Use a simple unbounded pattern and a short limit to verify the filter + const maxLen = 5; + const rules = [r('WORD', '[a-z]+', maxLen)]; + // 'short' (5 chars, exactly at limit) passes; 'toolong' (7 chars) does not + const records = [{ content: 'short toolong' }]; + const results = executeRegexRules({ rules, records }); + + expect(results.map((m) => m.matchValue)).toEqual(['short']); + }); + + it('accepts matches exactly at maxMatchLength', () => { + const rules = [r('WORD', '[a-z]+', 5)]; + const records = [{ content: 'exact' }]; // 5 chars + const results = executeRegexRules({ rules, records }); + expect(results).toHaveLength(1); + expect(results[0].matchValue).toBe('exact'); + }); + + it('does not filter matches when maxMatchLength is not set', () => { + const longWord = 'a'.repeat(300); + const rules = [r('WORD', '[a-z]+')]; // no maxMatchLength + const records = [{ content: longWord }]; + const results = executeRegexRules({ rules, records }); + expect(results).toHaveLength(1); + expect(results[0].matchValue).toBe(longWord); + }); + }); + + describe('skips empty string fields', () => { + it('ignores empty string fields', () => { + const rules = [r('IP', '\\b\\d+\\.\\d+\\.\\d+\\.\\d+\\b')]; + const records = [{ content: '', other: '10.0.0.1' }]; + const results = executeRegexRules({ rules, records }); + expect(results).toHaveLength(1); + expect(results[0].recordKey).toBe('other'); + }); + }); +}); diff --git a/x-pack/platform/plugins/shared/inference/server/workflow_anonymization/detection/execute_regex_rules.ts b/x-pack/platform/plugins/shared/inference/server/workflow_anonymization/detection/execute_regex_rules.ts new file mode 100644 index 0000000000000..ce00f79d01f82 --- /dev/null +++ b/x-pack/platform/plugins/shared/inference/server/workflow_anonymization/detection/execute_regex_rules.ts @@ -0,0 +1,135 @@ +/* + * 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 { RE2JS } from 're2js'; +import type { PiiRegexRule, PiiRegexMatch, PiiRegexWorkerTaskPayload } from './types'; + +type CompiledRule = + | { engine: 're2'; pattern: ReturnType } + | { engine: 'native'; pattern: RegExp }; + +export function compileRule(rawPattern: string, re2Only = false): CompiledRule { + try { + return { engine: 're2', pattern: RE2JS.compile(rawPattern) }; + } catch (err) { + if (re2Only) { + throw err instanceof Error ? err : new Error(String(err)); + } + // RE2 does not support lookahead, lookbehind, or backreferences. Fall back to + // native RegExp. ReDoS protection is provided by the Piscina worker timeout and + // per-task abort via AbortSignal. + return { engine: 'native', pattern: new RegExp(rawPattern, 'g') }; + } +} + +function findSpans( + compiled: CompiledRule, + value: string +): Array<{ start: number; end: number; matchValue: string }> { + const spans: Array<{ start: number; end: number; matchValue: string }> = []; + + if (compiled.engine === 're2') { + const matcher = compiled.pattern.matcher(value); + let pos: number | null = null; + + while (true) { + const found = pos === null ? matcher.find() : matcher.find(pos); + pos = null; + if (!found) break; + + const start = matcher.start(); + const end = matcher.end(); + + if (end <= start) { + const next = start + 1; + if (next > value.length) break; + pos = next; + continue; + } + + const matchValue = matcher.group(); + if (matchValue !== null) { + spans.push({ start, end, matchValue }); + } + } + } else { + compiled.pattern.lastIndex = 0; + let match: RegExpExecArray | null; + + while ((match = compiled.pattern.exec(value)) !== null) { + const start = match.index; + const matchValue = match[0]; + const end = start + matchValue.length; + + if (end <= start) { + // Zero-length match: advance one character to avoid an infinite loop + compiled.pattern.lastIndex = start + 1; + continue; + } + + spans.push({ start, end, matchValue }); + } + } + + return spans; +} + +/** + * Executes a set of regex rules against a batch of text records. + * + * RE2JS is tried first for each pattern. Patterns that contain constructs RE2 does + * not support (lookahead, lookbehind, backreferences) fall back to native RegExp. + * The Piscina worker timeout and per-task AbortSignal provide ReDoS protection for + * native RegExp patterns. + * + * When `re2Only` is true, native RegExp fallback is disabled and any non-RE2 pattern + * throws immediately. Use this on the synchronous (non-worker) path where there is no + * timeout to contain catastrophic backtracking. + * + * Zero-length matches advance one character and continue scanning; they do not + * terminate the search for that field. + * + * Throws only when a pattern is invalid in both RE2 and native RegExp syntax (or in + * RE2 alone when `re2Only` is true). + */ +export const executeRegexRules = ( + { rules, records }: PiiRegexWorkerTaskPayload, + { re2Only = false }: { re2Only?: boolean } = {} +): PiiRegexMatch[] => { + const compiled = rules.map((rule) => compileRule(rule.pattern, re2Only)); + const results: PiiRegexMatch[] = []; + + for (let ruleIndex = 0; ruleIndex < rules.length; ruleIndex++) { + const rule = rules[ruleIndex] as PiiRegexRule; + + for (let recordIndex = 0; recordIndex < records.length; recordIndex++) { + const record = records[recordIndex]; + for (const [recordKey, value] of Object.entries(record)) { + if (value.length === 0) { + continue; + } + + for (const { start, end, matchValue } of findSpans(compiled[ruleIndex], value)) { + if (rule.maxMatchLength !== undefined && matchValue.length > rule.maxMatchLength) { + continue; + } + results.push({ + ruleIndex, + recordIndex, + recordKey, + start, + end, + matchValue, + entityClass: rule.entityClass, + }); + } + } + } + } + + return results; +}; diff --git a/x-pack/platform/plugins/shared/inference/server/workflow_anonymization/detection/index.ts b/x-pack/platform/plugins/shared/inference/server/workflow_anonymization/detection/index.ts new file mode 100644 index 0000000000000..5630389eda3c8 --- /dev/null +++ b/x-pack/platform/plugins/shared/inference/server/workflow_anonymization/detection/index.ts @@ -0,0 +1,16 @@ +/* + * 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 { assertRe2Compilable } from './assert_re2_compilable'; +export { generateEntityToken } from './entity_mask'; +export { PiiRegexWorkerService } from './regex_worker_service'; +export type { + PiiRegexRule, + PiiRegexMatch, + PiiRegexWorkerTaskPayload, + PiiDetectionFailureMode, +} from './types'; diff --git a/x-pack/platform/plugins/shared/inference/server/workflow_anonymization/detection/regex_worker_service.test.ts b/x-pack/platform/plugins/shared/inference/server/workflow_anonymization/detection/regex_worker_service.test.ts new file mode 100644 index 0000000000000..97f75c6aa3444 --- /dev/null +++ b/x-pack/platform/plugins/shared/inference/server/workflow_anonymization/detection/regex_worker_service.test.ts @@ -0,0 +1,176 @@ +/* + * 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 { loggerMock, type MockedLogger } from '@kbn/logging-mocks'; +import { PiiRegexWorkerService } from './regex_worker_service'; +import type { WorkflowAnonymizationWorkerConfig } from '../../config'; +import type { PiiRegexWorkerTaskPayload } from './types'; + +function createTestConfig( + overrides: Partial = {} +): WorkflowAnonymizationWorkerConfig { + return { + enabled: true, + minThreads: 1, + maxThreads: 2, + maxQueue: 20, + idleTimeout: { asMilliseconds: () => 30_000 }, + taskTimeout: { asMilliseconds: () => 15_000 }, + ...overrides, + } as WorkflowAnonymizationWorkerConfig; +} + +const IP_PAYLOAD: PiiRegexWorkerTaskPayload = { + rules: [{ entityClass: 'IP', pattern: '\\b\\d{1,3}\\.\\d{1,3}\\.\\d{1,3}\\.\\d{1,3}\\b' }], + records: [{ content: 'connect to 10.0.0.1' }], +}; + +describe('PiiRegexWorkerService', () => { + let logger: MockedLogger; + let service: PiiRegexWorkerService; + + beforeEach(() => { + jest.resetAllMocks(); + logger = loggerMock.create(); + }); + + afterEach(async () => { + await service?.stop(); + }); + + it('executes rules through the worker pool and returns matches', async () => { + service = new PiiRegexWorkerService(createTestConfig(), logger); + const results = await service.run(IP_PAYLOAD); + + expect(results).toHaveLength(1); + expect(results[0]).toMatchObject({ + entityClass: 'IP', + matchValue: '10.0.0.1', + recordIndex: 0, + recordKey: 'content', + ruleIndex: 0, + start: expect.any(Number), + end: expect.any(Number), + }); + }); + + it('runs synchronously when the worker pool is disabled', async () => { + service = new PiiRegexWorkerService(createTestConfig({ enabled: false }), logger); + const results = await service.run(IP_PAYLOAD); + + expect((service as any).worker).toBeUndefined(); + expect(results).toHaveLength(1); + expect(results[0].matchValue).toBe('10.0.0.1'); + }); + + it('throws on the sync path for patterns that require native RegExp (lookahead/lookbehind/backrefs)', async () => { + service = new PiiRegexWorkerService(createTestConfig({ enabled: false }), logger); + // (?=a) is a positive lookahead — RE2 rejects it; native RegExp can backtrack catastrophically + await expect( + service.run({ + rules: [{ entityClass: 'MISC', pattern: '(?=a)a+' }], + records: [{ content: 'aaa' }], + }) + ).rejects.toThrow(); + }); + + it('aborts the timed-out task and throws when taskTimeout elapses', async () => { + service = new PiiRegexWorkerService( + createTestConfig({ taskTimeout: { asMilliseconds: () => 1 } } as any), + logger + ); + + // (?=a)(a+)+$ falls back to native RegExp (RE2 rejects the lookahead) and + // backtracks catastrophically on a long all-'a' string — guaranteed timeout. + await expect( + service.run({ + rules: [{ entityClass: 'MISC', pattern: '(?=a)(a+)+$' }], + records: [{ content: 'a'.repeat(10_000) + 'b' }], + }) + ).rejects.toThrow('timed out'); + }); + + describe('allow_unsafe failure mode', () => { + it('skips the invalid rule and still returns matches from valid rules', async () => { + service = new PiiRegexWorkerService(createTestConfig(), logger); + const results = await service.run( + { + rules: [ + { entityClass: 'BAD', pattern: '(unclosed' }, + { entityClass: 'IP', pattern: '\\b\\d{1,3}\\.\\d{1,3}\\.\\d{1,3}\\.\\d{1,3}\\b' }, + ], + records: [{ content: 'connect to 10.0.0.1' }], + }, + 'allow_unsafe' + ); + + // The bad rule is skipped; the good IP rule still fires + expect(results).toHaveLength(1); + expect(results[0]).toMatchObject({ entityClass: 'IP', matchValue: '10.0.0.1' }); + expect(logger.warn).toHaveBeenCalledWith( + 'PII regex rule skipped: pattern could not be compiled', + expect.objectContaining({ entityClass: 'BAD' }) + ); + expect(logger.error).not.toHaveBeenCalled(); + }); + + it('returns [] when every rule is invalid', async () => { + service = new PiiRegexWorkerService(createTestConfig(), logger); + const results = await service.run( + { + rules: [{ entityClass: 'BAD', pattern: '(unclosed' }], + records: [{ content: 'test' }], + }, + 'allow_unsafe' + ); + + expect(results).toEqual([]); + expect(logger.warn).toHaveBeenCalledWith( + 'PII regex rule skipped: pattern could not be compiled', + expect.objectContaining({ entityClass: 'BAD' }) + ); + expect(logger.error).not.toHaveBeenCalled(); + }); + }); + + it('throws when failureMode is block (default)', async () => { + service = new PiiRegexWorkerService(createTestConfig(), logger); + const badPayload: PiiRegexWorkerTaskPayload = { + rules: [{ entityClass: 'BAD', pattern: '(unclosed' }], + records: [{ content: 'test' }], + }; + + await expect(service.run(badPayload)).rejects.toThrow(); + expect(logger.error).not.toHaveBeenCalled(); + }); + + describe('worker queue at capacity', () => { + it('throws a distinct "queue at capacity" error distinguishable from rule errors', async () => { + service = new PiiRegexWorkerService(createTestConfig(), logger); + jest + .spyOn((service as any).worker, 'run') + .mockRejectedValueOnce(new Error('Task queue is at limit')); + + await expect(service.run(IP_PAYLOAD)).rejects.toThrow('queue at capacity'); + }); + + it('logs and returns [] in allow_unsafe mode when queue is at capacity', async () => { + service = new PiiRegexWorkerService(createTestConfig(), logger); + jest + .spyOn((service as any).worker, 'run') + .mockRejectedValueOnce(new Error('Task queue is at limit')); + + const results = await service.run(IP_PAYLOAD, 'allow_unsafe'); + + expect(results).toEqual([]); + expect(logger.error).toHaveBeenCalledWith( + 'PII regex detection failed; proceeding without anonymization', + expect.objectContaining({ error: expect.anything() }) + ); + }); + }); +}); diff --git a/x-pack/platform/plugins/shared/inference/server/workflow_anonymization/detection/regex_worker_service.ts b/x-pack/platform/plugins/shared/inference/server/workflow_anonymization/detection/regex_worker_service.ts new file mode 100644 index 0000000000000..52f3f585c86bb --- /dev/null +++ b/x-pack/platform/plugins/shared/inference/server/workflow_anonymization/detection/regex_worker_service.ts @@ -0,0 +1,143 @@ +/* + * 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 Piscina from 'piscina'; +import type { Logger } from '@kbn/logging'; +import type { WorkflowAnonymizationWorkerConfig } from '../../config'; +import type { PiiRegexWorkerTaskPayload, PiiRegexMatch, PiiDetectionFailureMode } from './types'; +import { compileRule, executeRegexRules } from './execute_regex_rules'; + +function runSync(payload: PiiRegexWorkerTaskPayload): PiiRegexMatch[] { + // On the sync path there is no timeout to contain catastrophic backtracking, so + // enforce RE2-only patterns. executeRegexRules throws from compileRule on first + // non-RE2 pattern when re2Only is true. + return executeRegexRules(payload, { re2Only: true }); +} + +/** + * Manages the Piscina worker pool for the workflow-driven PII regex executor. + * + * Separate from the o11y `RegexWorkerService` in `chat_complete/anonymization/`. + * This pool is exclusively used by the workflow-driven anonymization path. + */ +export class PiiRegexWorkerService { + private readonly enabled: boolean; + private worker?: Piscina; + private readonly config: WorkflowAnonymizationWorkerConfig; + + constructor(config: WorkflowAnonymizationWorkerConfig, private readonly logger: Logger) { + this.config = config; + this.enabled = config.enabled; + + if (this.enabled) { + this.worker = this.createWorkerPool(); + } + } + + private createWorkerPool(): Piscina { + this.logger.debug( + `Initializing PII regex worker pool (min=${this.config.minThreads} | max=${ + this.config.maxThreads + } | idle=${this.config.idleTimeout.asMilliseconds()}ms)` + ); + + return new Piscina({ + filename: require.resolve('./regex_worker_wrapper.js'), + minThreads: this.config.minThreads, + maxThreads: this.config.maxThreads, + maxQueue: this.config.maxQueue, + idleTimeout: this.config.idleTimeout.asMilliseconds(), + }); + } + + /** + * Executes PII regex rules against records. + * + * Throws when a rule has an invalid pattern and `failureMode` is `'block'` (the default). + * With `'allow_unsafe'`, invalid rules are logged individually and skipped; the remaining + * rules still run and return their matches. Infrastructure failures (timeout, queue + * saturation) return no matches for the entire payload regardless of `failureMode`. + * + * When the worker pool is disabled, runs synchronously on the main event loop. + * In that mode only RE2-compilable patterns are accepted; patterns that require + * native RegExp (lookahead / lookbehind / backreferences) are also skipped (or rejected + * in `'block'` mode) to prevent unbounded backtracking on the event loop. + */ + async run( + payload: PiiRegexWorkerTaskPayload, + failureMode: PiiDetectionFailureMode = 'block' + ): Promise { + const re2Only = !this.enabled; + const effectivePayload = + failureMode === 'allow_unsafe' ? this.filterRules(payload, re2Only) : payload; + + try { + if (!this.enabled) { + return runSync(effectivePayload); + } + if (!this.worker) { + throw new Error('PII regex worker pool was not initialized'); + } + + const controller = new AbortController(); + const timer = setTimeout(() => controller.abort(), this.config.taskTimeout.asMilliseconds()); + + try { + return await this.worker.run(effectivePayload, { signal: controller.signal }); + } catch (err) { + if (err instanceof Error && err.name === 'AbortError') { + throw new Error( + `PII regex detection task timed out after ${this.config.taskTimeout.asMilliseconds()}ms` + ); + } + // Piscina does not expose a stable error code; match on the message + // (verified against piscina@5.3.1 dist/errors.js TaskQueueAtLimit). + if (err instanceof Error && err.message === 'Task queue is at limit') { + throw new Error( + `PII regex detection rejected: worker queue at capacity (maxQueue=${this.config.maxQueue})` + ); + } + throw err; + } finally { + clearTimeout(timer); + } + } catch (err) { + if (failureMode === 'allow_unsafe') { + // Only infrastructure failures reach here (timeout, saturation, worker crash). + // Rule-level errors were already handled by filterRules above. + this.logger.error('PII regex detection failed; proceeding without anonymization', { + error: err, + }); + return []; + } + throw err; + } + } + + private filterRules( + payload: PiiRegexWorkerTaskPayload, + re2Only: boolean + ): PiiRegexWorkerTaskPayload { + const safeRules = payload.rules.filter((rule) => { + try { + compileRule(rule.pattern, re2Only); + return true; + } catch (err) { + this.logger.warn('PII regex rule skipped: pattern could not be compiled', { + entityClass: rule.entityClass, + error: err, + }); + return false; + } + }); + return { ...payload, rules: safeRules }; + } + + async stop(): Promise { + await this.worker?.destroy(); + } +} diff --git a/x-pack/platform/plugins/shared/inference/server/workflow_anonymization/detection/regex_worker_task.ts b/x-pack/platform/plugins/shared/inference/server/workflow_anonymization/detection/regex_worker_task.ts new file mode 100644 index 0000000000000..6cba3080c3a9b --- /dev/null +++ b/x-pack/platform/plugins/shared/inference/server/workflow_anonymization/detection/regex_worker_task.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. + */ + +import type { PiiRegexWorkerTaskPayload } from './types'; +import { executeRegexRules } from './execute_regex_rules'; + +// eslint-disable-next-line import/no-default-export +export default function (payload: PiiRegexWorkerTaskPayload) { + return executeRegexRules(payload); +} diff --git a/x-pack/platform/plugins/shared/inference/server/workflow_anonymization/detection/regex_worker_wrapper.js b/x-pack/platform/plugins/shared/inference/server/workflow_anonymization/detection/regex_worker_wrapper.js new file mode 100644 index 0000000000000..009a8c47701cb --- /dev/null +++ b/x-pack/platform/plugins/shared/inference/server/workflow_anonymization/detection/regex_worker_wrapper.js @@ -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. + */ +if (process.env.NODE_ENV !== 'production') { + require('@kbn/setup-node-env'); +} else { + require('@kbn/setup-node-env/dist'); +} + +module.exports = require('./regex_worker_task'); diff --git a/x-pack/platform/plugins/shared/inference/server/workflow_anonymization/detection/types.ts b/x-pack/platform/plugins/shared/inference/server/workflow_anonymization/detection/types.ts new file mode 100644 index 0000000000000..043603365a1f1 --- /dev/null +++ b/x-pack/platform/plugins/shared/inference/server/workflow_anonymization/detection/types.ts @@ -0,0 +1,60 @@ +/* + * 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. + */ + +/** + * A single regex PII detection rule, as authored in a workflow definition. + * + * Deliberately independent of `RegexAnonymizationRule` in `@kbn/inference-common`: + * that type belongs to the legacy text-level anonymization path and carries fields + * (NER variants, mask types) this runtime does not implement. Rules here arrive from + * workflow YAML, so the workflow step schema is the only contract that shapes them. + */ +export interface PiiRegexRule { + /** Token prefix for matches of this rule, e.g. `EMAIL`, `HOST_NAME`. */ + entityClass: string; + /** Regex pattern. RE2JS is tried first; native RegExp is used as a fallback for constructs RE2 does not support (lookahead, lookbehind, backreferences). */ + pattern: string; + /** + * Maximum length of an accepted match, in characters. Matches longer than this are + * discarded. Use it to express length bounds that RE2 cannot encode as lookahead. + */ + maxMatchLength?: number; +} + +/** A single detected PII span, positioned against the unmodified input text. */ +export interface PiiRegexMatch { + /** Index of the rule that produced this match, preserving rule precedence. */ + ruleIndex: number; + /** Index of the record within the payload's `records` array. */ + recordIndex: number; + /** Key within the record whose value produced this match. */ + recordKey: string; + /** Inclusive start offset within the original field value. */ + start: number; + /** Exclusive end offset within the original field value. */ + end: number; + /** The matched text, verbatim. */ + matchValue: string; + /** `entityClass` of the rule that produced this match. */ + entityClass: string; +} + +/** Payload handed to the regex worker pool. */ +export interface PiiRegexWorkerTaskPayload { + rules: readonly PiiRegexRule[]; + records: ReadonlyArray>; +} + +/** + * How the runtime reacts when a rule cannot be compiled or executed. + * + * `block` refuses to produce a partial result, so a broken rule fails the LLM call + * rather than letting the PII class it was meant to catch through unmasked. + * `allow_unsafe` logs and skips each invalid rule; valid rules still run and return their matches. + * Infrastructure failures (timeout, queue saturation) still return no matches for the entire payload. + */ +export type PiiDetectionFailureMode = 'block' | 'allow_unsafe'; diff --git a/x-pack/platform/plugins/shared/inference/server/workflow_anonymization/pii_detection_context.ts b/x-pack/platform/plugins/shared/inference/server/workflow_anonymization/pii_detection_context.ts new file mode 100644 index 0000000000000..c2d24201e55ef --- /dev/null +++ b/x-pack/platform/plugins/shared/inference/server/workflow_anonymization/pii_detection_context.ts @@ -0,0 +1,19 @@ +/* + * 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 { AnonymizationRule } from '@kbn/inference-common'; +import type { DetectedPiiEntity, PiiTextRecord } from '../workflow_anonymization_capabilities'; + +export type { DetectedPiiEntity, PiiTextRecord }; + +export interface PiiDetectionContext { + detectEntities(options: { + records: readonly PiiTextRecord[]; + rules: readonly AnonymizationRule[]; + abortSignal?: AbortSignal; + }): Promise; +} diff --git a/x-pack/platform/plugins/shared/inference/server/workflow_anonymization_capabilities.test.ts b/x-pack/platform/plugins/shared/inference/server/workflow_anonymization_capabilities.test.ts new file mode 100644 index 0000000000000..175857271addd --- /dev/null +++ b/x-pack/platform/plugins/shared/inference/server/workflow_anonymization_capabilities.test.ts @@ -0,0 +1,79 @@ +/* + * 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 { + createPiiTokenizationCapabilityValue, + resolvePiiTokenizationCapabilityValue, + createInferenceProceedCapabilityValue, + resolveInferenceProceedCapabilityValue, +} from './workflow_anonymization_capabilities'; +import type { + PiiTokenizationContext, + InferenceProceedCapability, +} from './workflow_anonymization_capabilities'; + +const makePiiContext = (): PiiTokenizationContext => ({ + detectEntities: jest.fn(), + tokenize: jest.fn(), +}); + +const makeProceedCapability = (): InferenceProceedCapability => ({ + invoke: jest.fn(), +}); + +describe('PiiTokenizationCapability', () => { + it('round-trips: resolved value matches registered capability', () => { + const capability = makePiiContext(); + const token = createPiiTokenizationCapabilityValue(capability); + expect(resolvePiiTokenizationCapabilityValue(token)).toBe(capability); + }); + + it('returns undefined for a plain object without the marker', () => { + expect(resolvePiiTokenizationCapabilityValue({})).toBeUndefined(); + }); + + it('returns undefined for a frozen object without the marker', () => { + expect(resolvePiiTokenizationCapabilityValue(Object.freeze({}))).toBeUndefined(); + }); + + it('produces an opaque token — the value object exposes no capability methods', () => { + const token = createPiiTokenizationCapabilityValue(makePiiContext()); + expect((token as Record).detectEntities).toBeUndefined(); + expect((token as Record).tokenize).toBeUndefined(); + }); +}); + +describe('InferenceProceedCapability', () => { + it('round-trips: resolved value matches registered capability', () => { + const capability = makeProceedCapability(); + const token = createInferenceProceedCapabilityValue(capability); + expect(resolveInferenceProceedCapabilityValue(token)).toBe(capability); + }); + + it('returns undefined for a plain object without the marker', () => { + expect(resolveInferenceProceedCapabilityValue({})).toBeUndefined(); + }); + + it('returns undefined for a frozen object without the marker', () => { + expect(resolveInferenceProceedCapabilityValue(Object.freeze({}))).toBeUndefined(); + }); + + it('produces an opaque token — the value object exposes no capability methods', () => { + const token = createInferenceProceedCapabilityValue(makeProceedCapability()); + expect((token as Record).invoke).toBeUndefined(); + }); + + it('different capability instances produce distinct tokens that each resolve correctly', () => { + const capA = makeProceedCapability(); + const capB = makeProceedCapability(); + const tokenA = createInferenceProceedCapabilityValue(capA); + const tokenB = createInferenceProceedCapabilityValue(capB); + expect(resolveInferenceProceedCapabilityValue(tokenA)).toBe(capA); + expect(resolveInferenceProceedCapabilityValue(tokenB)).toBe(capB); + expect(resolveInferenceProceedCapabilityValue(tokenA)).not.toBe(capB); + }); +}); diff --git a/x-pack/platform/plugins/shared/inference/server/workflow_anonymization_capabilities.ts b/x-pack/platform/plugins/shared/inference/server/workflow_anonymization_capabilities.ts new file mode 100644 index 0000000000000..320086a714a38 --- /dev/null +++ b/x-pack/platform/plugins/shared/inference/server/workflow_anonymization_capabilities.ts @@ -0,0 +1,110 @@ +/* + * 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 { AnonymizationRule, Message } from '@kbn/inference-common'; + +export const PII_TOKENIZATION_CAPABILITY_ID = 'inference.pii_tokenization'; +export const INFERENCE_PROCEED_CAPABILITY_ID = 'inference.proceed'; + +export interface PiiTextRecord { + readonly id: string; + readonly text: string; +} + +export interface DetectedPiiEntity { + readonly recordId: string; + readonly start: number; + readonly end: number; + readonly value: string; + readonly entityClass: string; +} + +export interface PiiTokenizationContext { + detectEntities(options: { + records: readonly PiiTextRecord[]; + rules: readonly AnonymizationRule[]; + abortSignal?: AbortSignal; + }): Promise; + + /** + * Generates a replacement token for the given entity class and original value. + * + * **MUST be deterministic**: calling `tokenize` with the same `(entityClass, value)` pair + * within a single call MUST return the same token. The PII-application logic uses this + * guarantee to deduplicate repeated occurrences into a single token-map entry and to detect + * collisions. A non-deterministic implementation would silently corrupt the token map and + * break PII restoration. + */ + tokenize(entityClass: string, value: string): string; +} + +export interface InferenceTokenMapEntry { + readonly original: string; + readonly entityClass: string; +} + +export interface InferenceProceedInput { + readonly system?: string; + readonly messages: readonly Message[]; + readonly tokenMap?: Readonly>; + readonly restoreStreamingContent?: boolean; + readonly restoreToolCallArguments?: boolean; + readonly dryRun?: boolean; + readonly abortSignal?: AbortSignal; +} + +export interface InferenceProceedCapability { + invoke(input: InferenceProceedInput): Promise<{ rawContent: string }>; +} + +// Non-enumerable Symbols used as sentinels on capability value objects. They survive +// Object.freeze() but are stripped by JSON.stringify, so a serialized-and-reparsed value +// will be missing the sentinel and fail the identity check with a diagnostic error rather +// than a silent WeakMap miss. +const PII_TOKENIZATION_CAPABILITY_MARKER = Symbol('inference.pii_tokenization.capability'); +const INFERENCE_PROCEED_CAPABILITY_MARKER = Symbol('inference.proceed.capability'); + +const piiTokenizationCapabilities = new WeakMap(); +const inferenceProceedCapabilities = new WeakMap(); + +export const createPiiTokenizationCapabilityValue = ( + capability: PiiTokenizationContext +): object => { + const value: Record = {}; + value[PII_TOKENIZATION_CAPABILITY_MARKER] = true; + Object.freeze(value); + piiTokenizationCapabilities.set(value, capability); + return value; +}; + +export const resolvePiiTokenizationCapabilityValue = ( + value: object +): PiiTokenizationContext | undefined => { + if (!(PII_TOKENIZATION_CAPABILITY_MARKER in value)) { + return undefined; + } + return piiTokenizationCapabilities.get(value); +}; + +export const createInferenceProceedCapabilityValue = ( + capability: InferenceProceedCapability +): object => { + const value: Record = {}; + value[INFERENCE_PROCEED_CAPABILITY_MARKER] = true; + Object.freeze(value); + inferenceProceedCapabilities.set(value, capability); + return value; +}; + +export const resolveInferenceProceedCapabilityValue = ( + value: object +): InferenceProceedCapability | undefined => { + if (!(INFERENCE_PROCEED_CAPABILITY_MARKER in value)) { + return undefined; + } + return inferenceProceedCapabilities.get(value); +}; diff --git a/x-pack/platform/plugins/shared/inference/server/workflow_anonymization_provider.ts b/x-pack/platform/plugins/shared/inference/server/workflow_anonymization_provider.ts new file mode 100644 index 0000000000000..b5a20daf0be89 --- /dev/null +++ b/x-pack/platform/plugins/shared/inference/server/workflow_anonymization_provider.ts @@ -0,0 +1,41 @@ +/* + * 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 { KibanaRequest } from '@kbn/core/server'; +import type { Message } from '@kbn/inference-common'; +import type { + InferenceProceedCapability, + PiiTokenizationContext, +} from './workflow_anonymization_capabilities'; + +export interface AroundCompletionEvent { + readonly system?: string; + readonly messages: readonly Message[]; + readonly sessionId?: string; + readonly agentId?: string; +} + +export type WorkflowAroundCompletionResult = + | { matched: false } + | { matched: true; content: string }; + +export interface WorkflowAnonymizationProvider { + /** + * Whether this provider supports synchronous (within-request) workflow execution. + * Checked at runtime before the hook path is entered; must be `true` for + * providers that implement the around-completion flow. + */ + readonly supportsSynchronousExecution: boolean; + execute(options: { + event: AroundCompletionEvent; + namespace: string; + request: KibanaRequest; + pii: PiiTokenizationContext; + proceed: InferenceProceedCapability; + abortSignal?: AbortSignal; + }): Promise; +} diff --git a/x-pack/platform/plugins/shared/inference_workflows/common/workflow_anonymization.test.ts b/x-pack/platform/plugins/shared/inference_workflows/common/workflow_anonymization.test.ts new file mode 100644 index 0000000000000..ff80ff7cf86ea --- /dev/null +++ b/x-pack/platform/plugins/shared/inference_workflows/common/workflow_anonymization.test.ts @@ -0,0 +1,122 @@ +/* + * 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 { MessageRole, type Message } from '@kbn/inference-common'; +import { + aiPiiInputSchema, + anonymizationRuleSchema, + tokenMapSchema, + workflowChatMessageSchema, +} from './workflow_anonymization'; + +describe('workflow anonymization schemas', () => { + const supportedEntityClasses = [ + 'PER', + 'ORG', + 'LOC', + 'MISC', + 'HOST_NAME', + 'USER_NAME', + 'IP', + 'URL', + 'EMAIL', + 'CLOUD_ACCOUNT', + 'ENTITY_NAME', + 'RESOURCE_NAME', + 'RESOURCE_ID', + ] as const; + const messages: Message[] = [ + { + role: MessageRole.User, + content: [ + { type: 'text', text: 'hello' }, + { type: 'image', source: { data: 'base64', mimeType: 'image/png' } }, + ], + }, + { + role: MessageRole.Assistant, + content: null, + refusal: null, + toolCalls: [ + { + toolCallId: 'call-1', + function: { name: 'lookup', arguments: { email: 'person@example.com' } }, + }, + ], + }, + { + role: MessageRole.Tool, + name: 'lookup', + toolCallId: 'call-1', + response: { nested: ['value'] }, + data: { source: 'test' }, + }, + ]; + + it.each(messages)('accepts a valid inference Message %#', (message) => { + expect(workflowChatMessageSchema.parse(message)).toEqual(message); + }); + + it('preserves opaque image extensions for forward compatibility', () => { + const message = { + role: MessageRole.User, + content: [ + { + type: 'image', + source: { data: 'base64', mimeType: 'image/png', sourceMetadata: 'preserved' }, + altText: 'preserved', + }, + ], + }; + + expect(workflowChatMessageSchema.parse(message)).toEqual(message); + }); + + it('rejects unknown message, rule, and token-map fields', () => { + expect(() => + workflowChatMessageSchema.parse({ + role: MessageRole.User, + content: 'hello', + unexpected: 'value', + }) + ).toThrow(); + expect(() => + anonymizationRuleSchema.parse({ + type: 'RegExp', + enabled: true, + pattern: 'secret', + entityClass: 'ENTITY_NAME', + unexpected: 'value', + }) + ).toThrow(); + expect(() => + tokenMapSchema.parse({ + TOKEN: { original: 'secret', entityClass: 'ENTITY_NAME', unexpected: 'value' }, + }) + ).toThrow(); + }); + + it.each(supportedEntityClasses)('accepts the supported %s entity class', (entityClass) => { + expect( + anonymizationRuleSchema.parse({ + type: 'RegExp', + enabled: true, + pattern: 'value', + entityClass, + }) + ).toEqual({ type: 'RegExp', enabled: true, pattern: 'value', entityClass }); + }); + + it('rejects an unsupported message role through the PII input boundary', () => { + expect(() => + aiPiiInputSchema.parse({ + messages: [{ role: 'system', content: 'not a Message role' }], + rules: [], + }) + ).toThrow(); + }); +}); diff --git a/x-pack/platform/plugins/shared/inference_workflows/common/workflow_anonymization.ts b/x-pack/platform/plugins/shared/inference_workflows/common/workflow_anonymization.ts new file mode 100644 index 0000000000000..34154afad14ca --- /dev/null +++ b/x-pack/platform/plugins/shared/inference_workflows/common/workflow_anonymization.ts @@ -0,0 +1,246 @@ +/* + * 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 { i18n } from '@kbn/i18n'; +import { MessageRole, type AnonymizationEntityClass, type Message } from '@kbn/inference-common'; +import { StepCategory } from '@kbn/workflows'; +import type { + CommonStepDefinition, + CommonTriggerDefinition, +} from '@kbn/workflows-extensions/common'; +import { z } from '@kbn/zod/v4'; + +export const INFERENCE_AROUND_COMPLETION_TRIGGER_ID = 'inference.aroundCompletion'; +export const AI_PII_STEP_ID = 'ai.pii'; +export const CALL_SITE_PROCEED_STEP_ID = 'call_site.proceed'; +export const TRANSFORM_PII_RESTORE_STEP_ID = 'transform.pii_restore'; + +const textContentSchema = z.object({ type: z.literal('text'), text: z.string() }).strict(); +const imageContentSchema = z + .object({ + type: z.literal('image'), + // Images are opaque to PII traversal. Preserve additional source metadata so valid message + // extensions do not become runtime failures before anonymization sees the textual content. + source: z.object({ data: z.string(), mimeType: z.string() }).passthrough(), + }) + .passthrough(); + +const userMessageSchema = z + .object({ + role: z.literal(MessageRole.User), + content: z.union([z.string(), z.array(z.union([textContentSchema, imageContentSchema]))]), + }) + .strict(); + +const toolCallSchema = z + .object({ + toolCallId: z.string(), + function: z + .object({ + name: z.string(), + arguments: z.record(z.string(), z.unknown()), + }) + .strict(), + }) + .strict(); + +const assistantMessageSchema = z + .object({ + role: z.literal(MessageRole.Assistant), + content: z.string().nullable(), + refusal: z.string().nullable().optional(), + toolCalls: z.array(toolCallSchema).optional(), + }) + .strict(); + +const toolMessageSchema = z + .object({ + role: z.literal(MessageRole.Tool), + name: z.string(), + toolCallId: z.string(), + response: z.union([z.string(), z.record(z.string(), z.unknown())]), + data: z.record(z.string(), z.unknown()).optional(), + }) + .strict(); + +const inferredWorkflowChatMessageSchema = z.discriminatedUnion('role', [ + userMessageSchema, + assistantMessageSchema, + toolMessageSchema, +]) satisfies z.ZodType; + +// Consumers need the stable Message output contract rather than the narrower inferred shape of +// passthrough image branches. The definition above remains compile-time checked without a cast. +export const workflowChatMessageSchema: z.ZodType = inferredWorkflowChatMessageSchema; + +export const tokenMapEntrySchema = z + .object({ + original: z.string(), + entityClass: z.string(), + }) + .strict(); + +export const tokenMapSchema = z.record(z.string(), tokenMapEntrySchema); + +const ANONYMIZATION_ENTITY_CLASSES = { + PER: 'PER', + ORG: 'ORG', + LOC: 'LOC', + MISC: 'MISC', + HOST_NAME: 'HOST_NAME', + USER_NAME: 'USER_NAME', + IP: 'IP', + URL: 'URL', + EMAIL: 'EMAIL', + CLOUD_ACCOUNT: 'CLOUD_ACCOUNT', + ENTITY_NAME: 'ENTITY_NAME', + RESOURCE_NAME: 'RESOURCE_NAME', + RESOURCE_ID: 'RESOURCE_ID', +} as const satisfies Record; + +const regexRuleSchema = z + .object({ + type: z.literal('RegExp'), + enabled: z.boolean(), + pattern: z.string(), + entityClass: z.enum(ANONYMIZATION_ENTITY_CLASSES), + }) + .strict(); + +const nerRuleSchema = z + .object({ + type: z.literal('NER'), + enabled: z.boolean(), + modelId: z.string().optional(), + timeoutSeconds: z.number().positive().optional(), + allowedEntityClasses: z.array(z.enum(['PER', 'ORG', 'LOC', 'MISC'])).optional(), + }) + .strict(); + +export const anonymizationRuleSchema = z.discriminatedUnion('type', [ + regexRuleSchema, + nerRuleSchema, +]); + +export const aroundCompletionEventSchema = z + .object({ + system: z.string().optional(), + messages: z.array(workflowChatMessageSchema), + sessionId: z.string().optional(), + agentId: z.string().optional(), + }) + .strict(); + +export const aiPiiInputSchema = z + .object({ + system: z.string().optional(), + messages: z.array(workflowChatMessageSchema), + rules: z.array(anonymizationRuleSchema), + }) + .strict(); + +export const anonymizedCompletionSchema = z + .object({ + system: z.string().optional(), + messages: z.array(workflowChatMessageSchema), + tokenMap: tokenMapSchema, + }) + .strict(); + +export const callSiteProceedInputSchema = z + .object({ + system: z.string().optional(), + messages: z.array(workflowChatMessageSchema), + tokenMap: tokenMapSchema.optional(), + restoreStreamingContent: z.boolean().optional(), + restoreToolCallArguments: z.boolean().optional(), + dryRun: z.boolean().optional(), + }) + .strict(); +export const callSiteProceedOutputSchema = z.object({ rawContent: z.string() }).strict(); + +export const piiRestoreInputSchema = z + .object({ rawContent: z.string(), tokenMap: tokenMapSchema }) + .strict(); +export const piiRestoreOutputSchema = z.object({ content: z.string() }).strict(); + +const emptyConfigSchema = z.object({}).strict(); + +export const aroundCompletionTriggerDefinition: CommonTriggerDefinition< + typeof aroundCompletionEventSchema +> = { + id: INFERENCE_AROUND_COMPLETION_TRIGGER_ID, + eventSchema: aroundCompletionEventSchema, + title: i18n.translate('xpack.inferenceWorkflows.aroundCompletionTrigger.title', { + defaultMessage: 'Around inference completion', + }), + description: i18n.translate('xpack.inferenceWorkflows.aroundCompletionTrigger.description', { + defaultMessage: 'Runs a workflow around an inference completion call.', + }), + stability: 'tech_preview', +}; + +export const aiPiiCommonDefinition: CommonStepDefinition< + typeof aiPiiInputSchema, + typeof anonymizedCompletionSchema, + typeof emptyConfigSchema +> = { + id: AI_PII_STEP_ID, + category: StepCategory.Ai, + label: i18n.translate('xpack.inferenceWorkflows.aiPiiStep.label', { + defaultMessage: 'Protect PII', + }), + description: i18n.translate('xpack.inferenceWorkflows.aiPiiStep.description', { + defaultMessage: 'Detects and replaces sensitive values for the current inference call.', + }), + inputSchema: aiPiiInputSchema, + outputSchema: anonymizedCompletionSchema, + configSchema: emptyConfigSchema, + stability: 'tech_preview', +}; + +export const callSiteProceedCommonDefinition: CommonStepDefinition< + typeof callSiteProceedInputSchema, + typeof callSiteProceedOutputSchema, + typeof emptyConfigSchema +> = { + id: CALL_SITE_PROCEED_STEP_ID, + category: StepCategory.Ai, + label: i18n.translate('xpack.inferenceWorkflows.callSiteProceedStep.label', { + defaultMessage: 'Call inference model', + }), + description: i18n.translate('xpack.inferenceWorkflows.callSiteProceedStep.description', { + defaultMessage: 'Calls the inference model once with workflow-transformed input.', + }), + inputSchema: callSiteProceedInputSchema, + outputSchema: callSiteProceedOutputSchema, + configSchema: emptyConfigSchema, + stability: 'tech_preview', +}; + +export const piiRestoreCommonDefinition: CommonStepDefinition< + typeof piiRestoreInputSchema, + typeof piiRestoreOutputSchema, + typeof emptyConfigSchema +> = { + id: TRANSFORM_PII_RESTORE_STEP_ID, + category: StepCategory.Ai, + label: i18n.translate('xpack.inferenceWorkflows.piiRestoreStep.label', { + defaultMessage: 'Restore PII', + }), + description: i18n.translate('xpack.inferenceWorkflows.piiRestoreStep.description', { + defaultMessage: 'Restores protected values in the final inference response.', + }), + inputSchema: piiRestoreInputSchema, + outputSchema: piiRestoreOutputSchema, + configSchema: emptyConfigSchema, + stability: 'tech_preview', +}; + +export type TokenMap = z.infer; +export type AiPiiInput = z.infer; +export type AnonymizedCompletion = z.infer; diff --git a/x-pack/platform/plugins/shared/inference_workflows/integration_tests/anonymization_workflow.integration.test.ts b/x-pack/platform/plugins/shared/inference_workflows/integration_tests/anonymization_workflow.integration.test.ts new file mode 100644 index 0000000000000..f8e72c4d0fd18 --- /dev/null +++ b/x-pack/platform/plugins/shared/inference_workflows/integration_tests/anonymization_workflow.integration.test.ts @@ -0,0 +1,300 @@ +/* + * 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. + */ + +/** + * End-to-end integration test for the PII anonymization workflow. + * + * Runs the real execution engine with real step handlers (ai.pii, call_site.proceed, + * transform.pii_restore) and controlled capability mocks. No existing test covers all + * three steps together through the real engine — this closes that gap. + * + * The central claim under test: PII in the prompt never reaches the LLM. The external + * model sees only opaque tokens; the caller sees original values restored. + */ + +import { ExecutionStatus } from '@kbn/workflows'; + +import { WorkflowRunFixture } from '@kbn/workflows-execution-engine/integration_tests/workflow_run_fixture'; + +import { aiPiiStepDefinition } from '../server/workflow_anonymization/ai_pii_step'; +import { callSiteProceedStepDefinition } from '../server/workflow_anonymization/call_site_proceed_step'; +import { piiRestoreStepDefinition } from '../server/workflow_anonymization/pii_restore_step'; +import { + AI_PII_STEP_ID, + CALL_SITE_PROCEED_STEP_ID, + TRANSFORM_PII_RESTORE_STEP_ID, +} from '../common/workflow_anonymization'; +import { + PII_TOKENIZATION_CAPABILITY_ID, + INFERENCE_PROCEED_CAPABILITY_ID, + createPiiTokenizationCapabilityValue, + createInferenceProceedCapabilityValue, + type PiiTokenizationContext, + type InferenceProceedCapability, + type DetectedPiiEntity, +} from '@kbn/inference-plugin/server/workflow_anonymization_capabilities'; + +// ─── Test fixtures ──────────────────────────────────────────────────────────── + +const TEST_EMAIL = 'john@acme.com'; +const TEST_IP = '10.0.5.42'; +const EMAIL_TOKEN = 'EMAIL_REDACTED'; +const IP_TOKEN = 'IP_REDACTED'; + +/** The managed YAML shipped in the codebase — used verbatim so the test covers the real definition. */ +const ANONYMIZATION_WORKFLOW_YAML = ` +name: "Protect sensitive inference data" +enabled: false +description: "Protects common identifiers around inference completion calls." +version: "1" +tags: + - inference + - anonymization + +triggers: + - type: inference.aroundCompletion + +outputs: + - name: content + type: string + required: true + +steps: + - name: anonymize_completion + type: ai.pii + with: + system: "\${{ event.system }}" + messages: "\${{ event.messages }}" + rules: + - type: RegExp + enabled: true + entityClass: EMAIL + pattern: '([a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\\.[a-zA-Z]{2,})' + - type: RegExp + enabled: true + entityClass: IP + pattern: '\\b(?:(?:25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\\.){3}(?:25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\\b' + + - name: invoke_inference + type: call_site.proceed + with: + system: "\${{ steps.anonymize_completion.output.system }}" + messages: "\${{ steps.anonymize_completion.output.messages }}" + tokenMap: "\${{ steps.anonymize_completion.output.tokenMap }}" + + - name: restore_completion + type: transform.pii_restore + with: + rawContent: "\${{ steps.invoke_inference.output.rawContent }}" + tokenMap: "\${{ steps.anonymize_completion.output.tokenMap }}" + + - name: emit_restored_completion + type: workflow.output + with: + content: "\${{ steps.restore_completion.output.content }}" +`; + +// ─── Helpers ────────────────────────────────────────────────────────────────── + +/** Builds a PiiTokenizationContext that finds TEST_EMAIL and TEST_IP by direct string search. */ +const buildPiiContext = (): PiiTokenizationContext => ({ + detectEntities: jest.fn().mockImplementation(async ({ records }) => { + const entities: DetectedPiiEntity[] = []; + for (const record of records) { + const emailIdx = record.text.indexOf(TEST_EMAIL); + if (emailIdx >= 0) { + entities.push({ + recordId: record.id, + start: emailIdx, + end: emailIdx + TEST_EMAIL.length, + value: TEST_EMAIL, + entityClass: 'EMAIL', + }); + } + const ipIdx = record.text.indexOf(TEST_IP); + if (ipIdx >= 0) { + entities.push({ + recordId: record.id, + start: ipIdx, + end: ipIdx + TEST_IP.length, + value: TEST_IP, + entityClass: 'IP', + }); + } + } + return entities; + }), + tokenize: jest.fn().mockImplementation((entityClass: string) => `${entityClass}_REDACTED`), +}); + +/** Wires step definitions and capabilities into the fixture. */ +const setupFixture = ( + fixture: WorkflowRunFixture, + proceedCapability: InferenceProceedCapability +) => { + const piiContext = buildPiiContext(); + + const stepMap: Record = { + [AI_PII_STEP_ID]: aiPiiStepDefinition, + [CALL_SITE_PROCEED_STEP_ID]: callSiteProceedStepDefinition, + [TRANSFORM_PII_RESTORE_STEP_ID]: piiRestoreStepDefinition, + }; + + fixture.dependencies.workflowsExtensions.hasStepDefinition.mockImplementation( + (id) => id in stepMap + ); + fixture.dependencies.workflowsExtensions.getStepDefinition.mockImplementation( + (id) => stepMap[id] + ); + + fixture.dependencies.capabilities = [ + { id: PII_TOKENIZATION_CAPABILITY_ID, value: createPiiTokenizationCapabilityValue(piiContext) }, + { + id: INFERENCE_PROCEED_CAPABILITY_ID, + value: createInferenceProceedCapabilityValue(proceedCapability), + }, + ]; + + return { piiContext }; +}; + +// ─── Tests ──────────────────────────────────────────────────────────────────── + +describe('PII anonymization workflow — end-to-end', () => { + let fixture: WorkflowRunFixture; + let mockProceed: InferenceProceedCapability; + + beforeEach(() => { + fixture = new WorkflowRunFixture(); + mockProceed = { + invoke: jest + .fn() + .mockResolvedValue({ rawContent: `I can help ${EMAIL_TOKEN} at ${IP_TOKEN}.` }), + }; + setupFixture(fixture, mockProceed); + }); + + describe('PII never reaches the LLM', () => { + it('sends only tokens to proceed.invoke — no raw PII in messages', async () => { + await fixture.runWorkflow({ + workflowYaml: ANONYMIZATION_WORKFLOW_YAML, + event: { + system: 'You are a helpful assistant.', + messages: [ + { role: 'user', content: `Please help with user ${TEST_EMAIL} at IP ${TEST_IP}.` }, + ], + }, + }); + + expect(mockProceed.invoke).toHaveBeenCalledTimes(1); + const calledWith = (mockProceed.invoke as jest.Mock).mock.calls[0][0]; + const serializedMessages = JSON.stringify(calledWith.messages); + + expect(serializedMessages).not.toContain(TEST_EMAIL); + expect(serializedMessages).not.toContain(TEST_IP); + expect(serializedMessages).toContain(EMAIL_TOKEN); + expect(serializedMessages).toContain(IP_TOKEN); + }); + + it('sends only tokens in the system prompt — no raw PII', async () => { + await fixture.runWorkflow({ + workflowYaml: ANONYMIZATION_WORKFLOW_YAML, + event: { + system: `System context includes ${TEST_EMAIL}.`, + messages: [{ role: 'user', content: 'Hello.' }], + }, + }); + + const calledWith = (mockProceed.invoke as jest.Mock).mock.calls[0][0]; + + expect(calledWith.system).not.toContain(TEST_EMAIL); + expect(calledWith.system).toContain(EMAIL_TOKEN); + }); + }); + + describe('PII is restored in the workflow output', () => { + it('restores tokens back to original values in content', async () => { + await fixture.runWorkflow({ + workflowYaml: ANONYMIZATION_WORKFLOW_YAML, + event: { + system: 'You are a helpful assistant.', + messages: [ + { role: 'user', content: `Please help with user ${TEST_EMAIL} at IP ${TEST_IP}.` }, + ], + }, + }); + + const execution = fixture.workflowExecutionRepositoryMock.workflowExecutions.get( + 'fake_workflow_execution_id' + ); + + expect(execution?.status).toBe(ExecutionStatus.COMPLETED); + + const content = execution?.context?.output?.content as string; + expect(content).toContain(TEST_EMAIL); + expect(content).toContain(TEST_IP); + expect(content).not.toContain(EMAIL_TOKEN); + expect(content).not.toContain(IP_TOKEN); + }); + }); + + describe('token map isolation between calls', () => { + it('does not bleed token maps across independent workflow executions', async () => { + const event = { + system: 'You are a helpful assistant.', + messages: [{ role: 'user', content: `Contact ${TEST_EMAIL}.` }], + }; + + await fixture.runWorkflow({ workflowYaml: ANONYMIZATION_WORKFLOW_YAML, event }); + const firstProceedCall = (mockProceed.invoke as jest.Mock).mock.calls[0][0]; + + // Second run on a fresh fixture — different token map + const fixture2 = new WorkflowRunFixture(); + const proceed2: InferenceProceedCapability = { + invoke: jest.fn().mockResolvedValue({ rawContent: `Response about ${EMAIL_TOKEN}.` }), + }; + setupFixture(fixture2, proceed2); + await fixture2.runWorkflow({ workflowYaml: ANONYMIZATION_WORKFLOW_YAML, event }); + const secondProceedCall = (proceed2.invoke as jest.Mock).mock.calls[0][0]; + + // Both runs tokenized independently — same token name for the same value is fine, + // but the tokenMap objects must be distinct instances + expect(firstProceedCall.tokenMap).not.toBe(secondProceedCall.tokenMap); + }); + }); + + describe('workflow completes successfully', () => { + it('reaches COMPLETED status with no error', async () => { + await fixture.runWorkflow({ + workflowYaml: ANONYMIZATION_WORKFLOW_YAML, + event: { + system: 'You are a helpful assistant.', + messages: [{ role: 'user', content: `Email ${TEST_EMAIL}, IP ${TEST_IP}.` }], + }, + }); + + const execution = fixture.workflowExecutionRepositoryMock.workflowExecutions.get( + 'fake_workflow_execution_id' + ); + + expect(execution?.status).toBe(ExecutionStatus.COMPLETED); + expect(execution?.error).toBeUndefined(); + }); + + it('calls proceed.invoke exactly once per workflow execution', async () => { + await fixture.runWorkflow({ + workflowYaml: ANONYMIZATION_WORKFLOW_YAML, + event: { + system: '', + messages: [{ role: 'user', content: `Email ${TEST_EMAIL}.` }], + }, + }); + + expect(mockProceed.invoke).toHaveBeenCalledTimes(1); + }); + }); +}); diff --git a/x-pack/platform/plugins/shared/inference_workflows/integration_tests/jest.integration.config.js b/x-pack/platform/plugins/shared/inference_workflows/integration_tests/jest.integration.config.js new file mode 100644 index 0000000000000..28bd0a9ca5cb2 --- /dev/null +++ b/x-pack/platform/plugins/shared/inference_workflows/integration_tests/jest.integration.config.js @@ -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. + */ + +module.exports = { + preset: '@kbn/test/jest_integration', + rootDir: '../../../../../../', + roots: ['/x-pack/platform/plugins/shared/inference_workflows/integration_tests'], + forceExit: true, +}; diff --git a/x-pack/platform/plugins/shared/inference_workflows/kibana.jsonc b/x-pack/platform/plugins/shared/inference_workflows/kibana.jsonc index bff5520895a75..0840a0b1d9cd6 100644 --- a/x-pack/platform/plugins/shared/inference_workflows/kibana.jsonc +++ b/x-pack/platform/plugins/shared/inference_workflows/kibana.jsonc @@ -10,7 +10,7 @@ "server": true, "browser": true, "configPath": ["xpack", "inferenceWorkflows"], - "requiredPlugins": ["inference", "workflowsExtensions"], + "requiredPlugins": ["inference", "spaces", "workflowsExtensions", "workflowsManagement"], "optionalPlugins": ["searchInferenceEndpoints"] } } diff --git a/x-pack/platform/plugins/shared/inference_workflows/moon.yml b/x-pack/platform/plugins/shared/inference_workflows/moon.yml index bf07c94a673a8..ad1756a2de0f1 100644 --- a/x-pack/platform/plugins/shared/inference_workflows/moon.yml +++ b/x-pack/platform/plugins/shared/inference_workflows/moon.yml @@ -22,9 +22,15 @@ dependsOn: - '@kbn/zod' - '@kbn/workflows' - '@kbn/workflows-extensions' + - '@kbn/workflows-management-plugin' - '@kbn/inference-plugin' - '@kbn/inference-common' - '@kbn/search-inference-endpoints' + - '@kbn/core-http-server-mocks' + - '@kbn/core-logging-server-mocks' + - '@kbn/core-spaces-common' + - '@kbn/spaces-plugin' + - '@kbn/config-schema' tags: - plugin - prod diff --git a/x-pack/platform/plugins/shared/inference_workflows/public/plugin.ts b/x-pack/platform/plugins/shared/inference_workflows/public/plugin.ts index 284fb0b04c850..cb1205c34aedf 100644 --- a/x-pack/platform/plugins/shared/inference_workflows/public/plugin.ts +++ b/x-pack/platform/plugins/shared/inference_workflows/public/plugin.ts @@ -25,6 +25,20 @@ export class InferenceWorkflowsPublicPlugin deps.workflowsExtensions.registerStepDefinition(() => import('./steps/ai/ai_classify_step').then((m) => m.AiClassifyStepDefinition) ); + deps.workflowsExtensions.registerStepDefinition(() => + import('./workflow_anonymization').then((module) => module.aiPiiStepDefinition) + ); + deps.workflowsExtensions.registerStepDefinition(() => + import('./workflow_anonymization').then((module) => module.callSiteProceedStepDefinition) + ); + deps.workflowsExtensions.registerStepDefinition(() => + import('./workflow_anonymization').then((module) => module.piiRestoreStepDefinition) + ); + deps.workflowsExtensions.registerTriggerDefinition(() => + import('./workflow_anonymization').then( + (module) => module.aroundCompletionPublicTriggerDefinition + ) + ); return {}; } diff --git a/x-pack/platform/plugins/shared/inference_workflows/public/workflow_anonymization.ts b/x-pack/platform/plugins/shared/inference_workflows/public/workflow_anonymization.ts new file mode 100644 index 0000000000000..d3ec19c54975e --- /dev/null +++ b/x-pack/platform/plugins/shared/inference_workflows/public/workflow_anonymization.ts @@ -0,0 +1,47 @@ +/* + * 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 React from 'react'; +import { + createPublicStepDefinition, + type PublicTriggerDefinition, +} from '@kbn/workflows-extensions/public'; +import { + aiPiiCommonDefinition, + aroundCompletionTriggerDefinition, + callSiteProceedCommonDefinition, + piiRestoreCommonDefinition, +} from '../common/workflow_anonymization'; + +const loadAgentIcon = () => + import('@elastic/eui/es/components/icon/assets/product_agent').then(({ icon }) => ({ + default: icon, + })); + +const AgentIcon = React.lazy(loadAgentIcon); + +export const aiPiiStepDefinition = createPublicStepDefinition({ + ...aiPiiCommonDefinition, + icon: AgentIcon, +}); + +export const callSiteProceedStepDefinition = createPublicStepDefinition({ + ...callSiteProceedCommonDefinition, + icon: AgentIcon, +}); + +export const piiRestoreStepDefinition = createPublicStepDefinition({ + ...piiRestoreCommonDefinition, + icon: AgentIcon, +}); + +export const aroundCompletionPublicTriggerDefinition: PublicTriggerDefinition< + typeof aroundCompletionTriggerDefinition.eventSchema +> = { + ...aroundCompletionTriggerDefinition, + icon: AgentIcon, +}; diff --git a/x-pack/platform/plugins/shared/inference_workflows/server/index.ts b/x-pack/platform/plugins/shared/inference_workflows/server/index.ts index 0d6235e175f0f..040f1cffddae8 100644 --- a/x-pack/platform/plugins/shared/inference_workflows/server/index.ts +++ b/x-pack/platform/plugins/shared/inference_workflows/server/index.ts @@ -9,5 +9,5 @@ import type { PluginInitializerContext } from '@kbn/core/server'; export async function plugin(initializerContext: PluginInitializerContext) { const { InferenceWorkflowsPlugin } = await import('./plugin'); - return new InferenceWorkflowsPlugin(); + return new InferenceWorkflowsPlugin(initializerContext); } diff --git a/x-pack/platform/plugins/shared/inference_workflows/server/plugin.ts b/x-pack/platform/plugins/shared/inference_workflows/server/plugin.ts index 2dc8da536dac4..bb56aa71c8451 100644 --- a/x-pack/platform/plugins/shared/inference_workflows/server/plugin.ts +++ b/x-pack/platform/plugins/shared/inference_workflows/server/plugin.ts @@ -5,20 +5,50 @@ * 2.0. */ -import type { CoreSetup, CoreStart, Plugin } from '@kbn/core/server'; +import type { + CoreSetup, + CoreStart, + Logger, + Plugin, + PluginInitializerContext, +} from '@kbn/core/server'; import type { InferenceWorkflowsSetupDeps, InferenceWorkflowsStartDeps } from './types'; import { aiPromptStepDefinition } from './steps/ai/ai_prompt_step/step'; import { aiSummarizeStepDefinition } from './steps/ai/ai_summarize_step/step'; import { aiClassifyStepDefinition } from './steps/ai/ai_classify_step/step'; import { registerInferenceFeatures } from './steps/ai/register_inference_features'; +import { aroundCompletionTriggerDefinition } from '../common/workflow_anonymization'; +import { aiPiiStepDefinition } from './workflow_anonymization/ai_pii_step'; +import { callSiteProceedStepDefinition } from './workflow_anonymization/call_site_proceed_step'; +import { piiRestoreStepDefinition } from './workflow_anonymization/pii_restore_step'; +import { createWorkflowAnonymizationProvider } from './workflow_anonymization/create_workflow_anonymization_provider'; export class InferenceWorkflowsPlugin implements Plugin<{}, {}, InferenceWorkflowsSetupDeps, InferenceWorkflowsStartDeps> { + private readonly logger: Logger; + + constructor(initializerContext: PluginInitializerContext) { + this.logger = initializerContext.logger.get(); + } + setup(core: CoreSetup, deps: InferenceWorkflowsSetupDeps) { deps.workflowsExtensions.registerStepDefinition(aiPromptStepDefinition(core)); deps.workflowsExtensions.registerStepDefinition(aiSummarizeStepDefinition(core)); deps.workflowsExtensions.registerStepDefinition(aiClassifyStepDefinition(core)); + deps.workflowsExtensions.registerStepDefinition(aiPiiStepDefinition); + deps.workflowsExtensions.registerStepDefinition(callSiteProceedStepDefinition); + deps.workflowsExtensions.registerStepDefinition(piiRestoreStepDefinition); + deps.workflowsExtensions.registerTriggerDefinition(aroundCompletionTriggerDefinition); + deps.inference.registerWorkflowAnonymizationProvider( + createWorkflowAnonymizationProvider({ + management: deps.workflowsManagement.management, + triggerCacheTtlMs: deps.inference.anonymizationConfig.triggerCacheTtlMs, + // Managed workflow installation is wired in PR 9. Until then the trigger resolution + // returns matched:false for all requests because no workflow is installed. + ensureManagedWorkflow: async () => {}, + }) + ); if (deps.searchInferenceEndpoints) { registerInferenceFeatures(deps.searchInferenceEndpoints); @@ -27,7 +57,7 @@ export class InferenceWorkflowsPlugin return {}; } - start(_core: CoreStart) { + start(_core: CoreStart, _deps: InferenceWorkflowsStartDeps) { return {}; } } diff --git a/x-pack/platform/plugins/shared/inference_workflows/server/types.ts b/x-pack/platform/plugins/shared/inference_workflows/server/types.ts index d0f07f925daea..a6472036c03da 100644 --- a/x-pack/platform/plugins/shared/inference_workflows/server/types.ts +++ b/x-pack/platform/plugins/shared/inference_workflows/server/types.ts @@ -5,19 +5,27 @@ * 2.0. */ -import type { InferenceServerStart } from '@kbn/inference-plugin/server'; +import type { InferenceServerSetup, InferenceServerStart } from '@kbn/inference-plugin/server'; +import type { SpacesPluginSetup, SpacesPluginStart } from '@kbn/spaces-plugin/server'; +import type { WorkflowsServerPluginSetup } from '@kbn/workflows-management-plugin/server'; import type { WorkflowsExtensionsServerPluginSetup } from '@kbn/workflows-extensions/server'; +import type { WorkflowsExtensionsServerPluginStart } from '@kbn/workflows-extensions/server'; import type { SearchInferenceEndpointsPluginSetup, SearchInferenceEndpointsPluginStart, } from '@kbn/search-inference-endpoints/server'; export interface InferenceWorkflowsSetupDeps { + inference: InferenceServerSetup; + spaces: SpacesPluginSetup; workflowsExtensions: WorkflowsExtensionsServerPluginSetup; + workflowsManagement: WorkflowsServerPluginSetup; searchInferenceEndpoints?: SearchInferenceEndpointsPluginSetup; } export interface InferenceWorkflowsStartDeps { inference: InferenceServerStart; + spaces: SpacesPluginStart; + workflowsExtensions: WorkflowsExtensionsServerPluginStart; searchInferenceEndpoints?: SearchInferenceEndpointsPluginStart; } diff --git a/x-pack/platform/plugins/shared/inference_workflows/server/workflow_anonymization/ai_pii_step.test.ts b/x-pack/platform/plugins/shared/inference_workflows/server/workflow_anonymization/ai_pii_step.test.ts new file mode 100644 index 0000000000000..a5792932f0e17 --- /dev/null +++ b/x-pack/platform/plugins/shared/inference_workflows/server/workflow_anonymization/ai_pii_step.test.ts @@ -0,0 +1,267 @@ +/* + * 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 { MessageRole } from '@kbn/inference-common'; +import { + createPiiTokenizationCapabilityValue, + PII_TOKENIZATION_CAPABILITY_ID, + type DetectedPiiEntity, + type PiiTextRecord, + type PiiTokenizationContext, +} from '@kbn/inference-plugin/server'; +import type { WorkflowExecutionCapabilities } from '@kbn/workflows-extensions/server'; +import type { AiPiiInput } from '../../common/workflow_anonymization'; +import { executePiiProtection } from './ai_pii_step'; + +const rule = { + type: 'RegExp', + enabled: true, + pattern: '[^\\s]+@[^\\s]+', + entityClass: 'EMAIL', +} as const; + +const createCapabilities = (pii: PiiTokenizationContext): WorkflowExecutionCapabilities => [ + { id: PII_TOKENIZATION_CAPABILITY_ID, value: createPiiTokenizationCapabilityValue(pii) }, +]; + +const createLogger = () => ({ warn: jest.fn() }); + +const input: AiPiiInput = { + system: 'Contact person@example.com', + messages: [ + { role: MessageRole.User, content: 'Email person@example.com' }, + { + role: MessageRole.Assistant, + content: null, + toolCalls: [ + { + toolCallId: 'call-1', + function: { name: 'lookup', arguments: { email: 'person@example.com' } }, + }, + ], + }, + ], + rules: [rule], +}; + +const detectEmailEntities = (records: readonly PiiTextRecord[]): DetectedPiiEntity[] => + records.flatMap((record) => { + const value = 'person@example.com'; + const start = record.text.indexOf(value); + return start === -1 + ? [] + : [ + { + recordId: record.id, + start, + end: start + value.length, + value, + entityClass: 'EMAIL', + }, + ]; + }); + +describe('executePiiProtection', () => { + it('protects system, message content, and structured tool arguments with one call-local map', async () => { + const detectEntities = jest.fn(({ records }) => Promise.resolve(detectEmailEntities(records))); + const capabilities = createCapabilities({ + detectEntities, + tokenize: () => 'EMAIL_TOKEN', + }); + + const output = await executePiiProtection({ + input, + capabilities, + abortSignal: new AbortController().signal, + logger: createLogger(), + }); + + expect(output.system).toBe('Contact EMAIL_TOKEN'); + expect(output.messages[0]).toEqual({ role: MessageRole.User, content: 'Email EMAIL_TOKEN' }); + expect(output.messages[1]).toEqual({ + role: MessageRole.Assistant, + content: null, + toolCalls: [ + { + toolCallId: 'call-1', + function: { name: 'lookup', arguments: { email: 'EMAIL_TOKEN' } }, + }, + ], + }); + expect(output.tokenMap).toEqual({ + EMAIL_TOKEN: { original: 'person@example.com', entityClass: 'EMAIL' }, + }); + expect(detectEntities).toHaveBeenCalledWith({ + records: expect.any(Array), + rules: [rule], + abortSignal: expect.any(AbortSignal), + }); + const detectedRecords = detectEntities.mock.calls[0][0].records; + expect(detectedRecords).not.toContainEqual( + expect.objectContaining({ id: expect.stringContaining('/function/name') }) + ); + }); + + it('fails closed when the detector returns an invalid range', async () => { + const capabilities = createCapabilities({ + detectEntities: jest + .fn() + .mockResolvedValue([ + { recordId: '/system', start: 0, end: 999, value: 'Contact', entityClass: 'EMAIL' }, + ]), + tokenize: jest.fn(), + }); + + await expect( + executePiiProtection({ + input, + capabilities, + abortSignal: new AbortController().signal, + logger: createLogger(), + }) + ).rejects.toThrow('PII detector returned an invalid range'); + }); + + it('honors abort signal that fires during detection and before entity application', async () => { + const abortController = new AbortController(); + const capabilities = createCapabilities({ + detectEntities: jest.fn().mockImplementation(async () => { + abortController.abort(); + return []; + }), + tokenize: jest.fn(), + }); + + await expect( + executePiiProtection({ + input, + capabilities, + abortSignal: abortController.signal, + logger: createLogger(), + }) + ).rejects.toMatchObject({ name: 'AbortError' }); + }); + + it('deduplicates the same detected value into a single token-map entry', async () => { + // Regression guard for the tokenize-determinism contract on PiiTokenizationContext. + // If tokenize() were non-deterministic, repeated occurrences would produce different tokens, + // triggering a false collision error or leaving orphaned token-map entries. + const detectEntities = jest.fn(({ records }) => Promise.resolve(detectEmailEntities(records))); + const capabilities = createCapabilities({ + detectEntities, + tokenize: () => 'EMAIL_TOKEN', + }); + + const output = await executePiiProtection({ + input: { + messages: [ + { role: MessageRole.User, content: 'First person@example.com mention' }, + { role: MessageRole.User, content: 'Second person@example.com mention' }, + ], + rules: [rule], + }, + capabilities, + abortSignal: new AbortController().signal, + logger: createLogger(), + }); + + expect(output.messages[0]).toEqual({ + role: MessageRole.User, + content: 'First EMAIL_TOKEN mention', + }); + expect(output.messages[1]).toEqual({ + role: MessageRole.User, + content: 'Second EMAIL_TOKEN mention', + }); + // Both occurrences share one entry — not two separate entries + expect(Object.keys(output.tokenMap)).toHaveLength(1); + expect(output.tokenMap).toEqual({ + EMAIL_TOKEN: { original: 'person@example.com', entityClass: 'EMAIL' }, + }); + }); + + it('requires a valid request-local PII capability', async () => { + await expect( + executePiiProtection({ + input, + capabilities: [], + abortSignal: new AbortController().signal, + logger: createLogger(), + }) + ).rejects.toThrow(PII_TOKENIZATION_CAPABILITY_ID); + }); + + it('warns when a later detected entity overlaps an earlier protected range', async () => { + const logger = createLogger(); + const capabilities = createCapabilities({ + detectEntities: jest.fn().mockResolvedValue([ + { recordId: '/system', start: 0, end: 7, value: 'Contact', entityClass: 'ENTITY_NAME' }, + { recordId: '/system', start: 4, end: 10, value: 'act pe', entityClass: 'ENTITY_NAME' }, + ]), + tokenize: () => 'ENTITY_TOKEN', + }); + + await executePiiProtection({ + input, + capabilities, + abortSignal: new AbortController().signal, + logger, + }); + + expect(logger.warn).toHaveBeenCalledWith( + 'PII detector returned overlapping entities; ignoring the later match', + { recordId: '/system', entityClass: 'ENTITY_NAME', start: 4, end: 10 } + ); + }); + + it('logs token collisions without recording sensitive values and then fails closed', async () => { + const logger = createLogger(); + // Two different emails both tokenize to the same token — within-call collision. + const capabilities = createCapabilities({ + detectEntities: jest.fn(({ records }) => + Promise.resolve( + records.flatMap((record) => { + const entities: DetectedPiiEntity[] = []; + for (const value of ['person@example.com', 'other@example.com']) { + const start = record.text.indexOf(value); + if (start !== -1) { + entities.push({ + recordId: record.id, + start, + end: start + value.length, + value, + entityClass: 'EMAIL', + }); + } + } + return entities; + }) + ) + ), + tokenize: () => 'COLLISION_TOKEN', + }); + + await expect( + executePiiProtection({ + input: { + system: 'person@example.com and other@example.com', + messages: [], + rules: [rule], + }, + capabilities, + abortSignal: new AbortController().signal, + logger, + }) + ).rejects.toThrow('PII token collision detected'); + expect(logger.warn).toHaveBeenCalledWith( + 'PII token collision detected; failing workflow protection', + { existingEntityClass: 'EMAIL', detectedEntityClass: 'EMAIL' } + ); + expect(JSON.stringify(logger.warn.mock.calls)).not.toContain('person@example.com'); + expect(JSON.stringify(logger.warn.mock.calls)).not.toContain('other@example.com'); + }); +}); diff --git a/x-pack/platform/plugins/shared/inference_workflows/server/workflow_anonymization/ai_pii_step.ts b/x-pack/platform/plugins/shared/inference_workflows/server/workflow_anonymization/ai_pii_step.ts new file mode 100644 index 0000000000000..84ecc3c89dbca --- /dev/null +++ b/x-pack/platform/plugins/shared/inference_workflows/server/workflow_anonymization/ai_pii_step.ts @@ -0,0 +1,166 @@ +/* + * 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 { DetectedPiiEntity, PiiTextRecord } from '@kbn/inference-plugin/server'; +import { + createServerStepDefinition, + type WorkflowExecutionCapabilities, +} from '@kbn/workflows-extensions/server'; +import { + aiPiiCommonDefinition, + type AiPiiInput, + type AnonymizedCompletion, + type TokenMap, +} from '../../common/workflow_anonymization'; +import { getPiiTokenizationContext } from './capabilities'; +import { createCompletionTextRecords } from './message_records'; +import { piiReplacementsCounter } from './anonymization_metrics'; + +interface IndexedEntity extends DetectedPiiEntity { + readonly detectionIndex: number; +} + +interface PiiProtectionLogger { + warn(message: string, meta?: object): void; +} + +const validateEntity = (entity: DetectedPiiEntity, record: PiiTextRecord | undefined): void => { + if (!record) { + throw new Error(`PII detector returned unknown record "${entity.recordId}"`); + } + if ( + !Number.isInteger(entity.start) || + !Number.isInteger(entity.end) || + entity.start < 0 || + entity.end <= entity.start || + entity.end > record.text.length || + record.text.slice(entity.start, entity.end) !== entity.value + ) { + throw new Error(`PII detector returned an invalid range for record "${entity.recordId}"`); + } +}; + +const applyDetectedEntities = ({ + records, + entities, + tokenize, + logger, +}: { + records: readonly PiiTextRecord[]; + entities: readonly DetectedPiiEntity[]; + tokenize: (entityClass: string, value: string) => string; + logger: PiiProtectionLogger; +}): { values: ReadonlyMap; tokenMap: TokenMap } => { + const recordsById = new Map(records.map((record) => [record.id, record])); + const entitiesByRecord = new Map(); + + entities.forEach((entity, detectionIndex) => { + validateEntity(entity, recordsById.get(entity.recordId)); + const recordEntities = entitiesByRecord.get(entity.recordId) ?? []; + recordEntities.push({ ...entity, detectionIndex }); + entitiesByRecord.set(entity.recordId, recordEntities); + }); + + const nextTokenMap: TokenMap = {}; + const values = new Map(); + + records.forEach((record) => { + const sortedEntities = [...(entitiesByRecord.get(record.id) ?? [])].sort( + (left, right) => left.start - right.start || left.detectionIndex - right.detectionIndex + ); + let cursor = 0; + let output = ''; + + sortedEntities.forEach((entity) => { + if (entity.start < cursor) { + logger.warn('PII detector returned overlapping entities; ignoring the later match', { + recordId: entity.recordId, + entityClass: entity.entityClass, + start: entity.start, + end: entity.end, + }); + return; + } + + const token = tokenize(entity.entityClass, entity.value); + const existing = nextTokenMap[token]; + if ( + existing && + (existing.original !== entity.value || existing.entityClass !== entity.entityClass) + ) { + logger.warn('PII token collision detected; failing workflow protection', { + existingEntityClass: existing.entityClass, + detectedEntityClass: entity.entityClass, + }); + throw new Error('PII token collision detected'); + } + + output += record.text.slice(cursor, entity.start); + output += token; + cursor = entity.end; + nextTokenMap[token] = { + original: entity.value, + entityClass: entity.entityClass, + }; + }); + + output += record.text.slice(cursor); + values.set(record.id, output); + }); + + return { values, tokenMap: nextTokenMap }; +}; + +export const executePiiProtection = async ({ + input, + capabilities, + abortSignal, + logger, +}: { + input: AiPiiInput; + capabilities: WorkflowExecutionCapabilities | undefined; + abortSignal: AbortSignal; + logger: PiiProtectionLogger; +}): Promise => { + const pii = getPiiTokenizationContext(capabilities); + const detectionRecords = createCompletionTextRecords(input); + const entities = await pii.detectEntities({ + records: detectionRecords.records, + rules: input.rules, + abortSignal, + }); + abortSignal.throwIfAborted(); + + const protectedRecords = applyDetectedEntities({ + records: detectionRecords.records, + entities, + tokenize: pii.tokenize, + logger, + }); + + const newTokensByClass = new Map(); + for (const [, entry] of Object.entries(protectedRecords.tokenMap)) { + newTokensByClass.set(entry.entityClass, (newTokensByClass.get(entry.entityClass) ?? 0) + 1); + } + for (const [entityClass, count] of newTokensByClass) { + piiReplacementsCounter.add(count, { entity_class: entityClass }); + } + + return { + ...detectionRecords.replace(protectedRecords.values), + tokenMap: protectedRecords.tokenMap, + }; +}; + +export const aiPiiStepDefinition = createServerStepDefinition({ + ...aiPiiCommonDefinition, + handler: async ({ input, capabilities, abortSignal, logger }) => { + return { + output: await executePiiProtection({ input, capabilities, abortSignal, logger }), + }; + }, +}); diff --git a/x-pack/platform/plugins/shared/inference_workflows/server/workflow_anonymization/anonymization_metrics.ts b/x-pack/platform/plugins/shared/inference_workflows/server/workflow_anonymization/anonymization_metrics.ts new file mode 100644 index 0000000000000..3c37ee20c04c9 --- /dev/null +++ b/x-pack/platform/plugins/shared/inference_workflows/server/workflow_anonymization/anonymization_metrics.ts @@ -0,0 +1,47 @@ +/* + * 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 { metrics, ValueType } from '@opentelemetry/api'; + +const meter = metrics.getMeter('kibana.inference.anonymization'); + +export const triggerEvaluationsCounter = meter.createCounter( + 'kibana.inference.anonymization.trigger.evaluations', + { + description: 'Number of around-completion workflow trigger evaluations', + unit: '{evaluation}', + valueType: ValueType.INT, + } +); + +export const piiReplacementsCounter = meter.createCounter( + 'kibana.inference.anonymization.pii.replacements', + { + description: + 'Number of new PII token replacements added to the token map per ai.pii step execution', + unit: '{token}', + valueType: ValueType.INT, + } +); + +export const managedWorkflowInstallationsCounter = meter.createCounter( + 'kibana.inference.anonymization.managed_workflow.installations', + { + description: 'Number of managed anonymization workflow space-level installations', + unit: '{installation}', + valueType: ValueType.INT, + } +); + +export const legacyMigrationRunsCounter = meter.createCounter( + 'kibana.inference.anonymization.legacy_migration.runs', + { + description: 'Number of legacy anonymization configuration migration runs', + unit: '{run}', + valueType: ValueType.INT, + } +); diff --git a/x-pack/platform/plugins/shared/inference_workflows/server/workflow_anonymization/call_site_proceed_step.test.ts b/x-pack/platform/plugins/shared/inference_workflows/server/workflow_anonymization/call_site_proceed_step.test.ts new file mode 100644 index 0000000000000..21673e39213df --- /dev/null +++ b/x-pack/platform/plugins/shared/inference_workflows/server/workflow_anonymization/call_site_proceed_step.test.ts @@ -0,0 +1,72 @@ +/* + * 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 { MessageRole } from '@kbn/inference-common'; +import { + createInferenceProceedCapabilityValue, + INFERENCE_PROCEED_CAPABILITY_ID, + type InferenceProceedInput, +} from '@kbn/inference-plugin/server'; +import { callSiteProceedStepDefinition } from './call_site_proceed_step'; +import { getInferenceProceedCapability } from './capabilities'; + +describe('callSiteProceedStepDefinition handler', () => { + it('forwards combined input and abortSignal to the proceed capability and returns rawContent', async () => { + const invoke = jest.fn().mockResolvedValue({ rawContent: 'RESULT' }); + const abortSignal = new AbortController().signal; + const input: Omit = { + system: 'protected system', + messages: [{ role: MessageRole.User, content: 'text TOKEN' }], + tokenMap: { TOKEN: { original: 'secret', entityClass: 'ENTITY_NAME' } }, + }; + const capabilities = [ + { + id: INFERENCE_PROCEED_CAPABILITY_ID, + value: createInferenceProceedCapabilityValue({ invoke }), + }, + ]; + + const result = await callSiteProceedStepDefinition.handler({ + input, + capabilities, + abortSignal, + }); + + expect(invoke).toHaveBeenCalledWith({ ...input, abortSignal }); + expect(result).toEqual({ output: { rawContent: 'RESULT' } }); + }); +}); + +describe('inference proceed capability', () => { + it('receives workflow-transformed input and its call-local token map', async () => { + const invoke = jest.fn().mockResolvedValue({ rawContent: 'TOKEN' }); + const proceed = getInferenceProceedCapability([ + { + id: INFERENCE_PROCEED_CAPABILITY_ID, + value: createInferenceProceedCapabilityValue({ invoke }), + }, + ]); + const abortSignal = new AbortController().signal; + const input: InferenceProceedInput = { + system: 'protected system', + messages: [{ role: MessageRole.User, content: 'protected TOKEN' }], + tokenMap: { TOKEN: { original: 'secret', entityClass: 'ENTITY_NAME' } }, + abortSignal, + }; + + await expect(proceed.invoke(input)).resolves.toEqual({ rawContent: 'TOKEN' }); + expect(invoke).toHaveBeenCalledWith(input); + }); + + it('rejects an unregistered look-alike capability', () => { + expect(() => + getInferenceProceedCapability([ + { id: INFERENCE_PROCEED_CAPABILITY_ID, value: { invoke: jest.fn() } }, + ]) + ).toThrow(`Workflow capability "${INFERENCE_PROCEED_CAPABILITY_ID}" is invalid`); + }); +}); diff --git a/x-pack/platform/plugins/shared/inference_workflows/server/workflow_anonymization/call_site_proceed_step.ts b/x-pack/platform/plugins/shared/inference_workflows/server/workflow_anonymization/call_site_proceed_step.ts new file mode 100644 index 0000000000000..65f69cae43832 --- /dev/null +++ b/x-pack/platform/plugins/shared/inference_workflows/server/workflow_anonymization/call_site_proceed_step.ts @@ -0,0 +1,19 @@ +/* + * 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 { createServerStepDefinition } from '@kbn/workflows-extensions/server'; +import { callSiteProceedCommonDefinition } from '../../common/workflow_anonymization'; +import { getInferenceProceedCapability } from './capabilities'; + +export const callSiteProceedStepDefinition = createServerStepDefinition({ + ...callSiteProceedCommonDefinition, + handler: async ({ input, capabilities, abortSignal }) => { + const proceed = getInferenceProceedCapability(capabilities); + const output = await proceed.invoke({ ...input, abortSignal }); + return { output }; + }, +}); diff --git a/x-pack/platform/plugins/shared/inference_workflows/server/workflow_anonymization/capabilities.ts b/x-pack/platform/plugins/shared/inference_workflows/server/workflow_anonymization/capabilities.ts new file mode 100644 index 0000000000000..6ed79f8f6e7dc --- /dev/null +++ b/x-pack/platform/plugins/shared/inference_workflows/server/workflow_anonymization/capabilities.ts @@ -0,0 +1,57 @@ +/* + * 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 { + INFERENCE_PROCEED_CAPABILITY_ID, + PII_TOKENIZATION_CAPABILITY_ID, + resolveInferenceProceedCapabilityValue, + resolvePiiTokenizationCapabilityValue, + type InferenceProceedCapability, + type PiiTokenizationContext, +} from '@kbn/inference-plugin/server'; +import type { WorkflowExecutionCapabilities } from '@kbn/workflows-extensions/server'; + +const getCapabilityValue = ( + capabilities: WorkflowExecutionCapabilities | undefined, + id: string +): object => { + const capability = capabilities?.find((entry) => entry.id === id); + if (!capability) { + throw new Error(`Workflow step requires the request-local "${id}" capability`); + } + return capability.value; +}; + +export const getPiiTokenizationContext = ( + capabilities: WorkflowExecutionCapabilities | undefined +): PiiTokenizationContext => { + const capability = resolvePiiTokenizationCapabilityValue( + getCapabilityValue(capabilities, PII_TOKENIZATION_CAPABILITY_ID) + ); + if (!capability) { + throw new Error( + `Workflow capability "${PII_TOKENIZATION_CAPABILITY_ID}" is invalid. ` + + `Capabilities are in-memory handles that cannot survive JSON serialization.` + ); + } + return capability; +}; + +export const getInferenceProceedCapability = ( + capabilities: WorkflowExecutionCapabilities | undefined +): InferenceProceedCapability => { + const capability = resolveInferenceProceedCapabilityValue( + getCapabilityValue(capabilities, INFERENCE_PROCEED_CAPABILITY_ID) + ); + if (!capability) { + throw new Error( + `Workflow capability "${INFERENCE_PROCEED_CAPABILITY_ID}" is invalid. ` + + `Capabilities are in-memory handles that cannot survive JSON serialization.` + ); + } + return capability; +}; diff --git a/x-pack/platform/plugins/shared/inference_workflows/server/workflow_anonymization/create_workflow_anonymization_provider.test.ts b/x-pack/platform/plugins/shared/inference_workflows/server/workflow_anonymization/create_workflow_anonymization_provider.test.ts new file mode 100644 index 0000000000000..b6c75ee9186ba --- /dev/null +++ b/x-pack/platform/plugins/shared/inference_workflows/server/workflow_anonymization/create_workflow_anonymization_provider.test.ts @@ -0,0 +1,361 @@ +/* + * 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 { httpServerMock } from '@kbn/core-http-server-mocks'; +import { MessageRole } from '@kbn/inference-common'; +import type { + AroundCompletionEvent, + InferenceProceedCapability, + PiiTokenizationContext, +} from '@kbn/inference-plugin/server'; +import { ExecutionStatus, type WorkflowDetailDto, type WorkflowYaml } from '@kbn/workflows'; +import type { WorkflowsManagementApi } from '@kbn/workflows-management-plugin/server'; +import { + INFERENCE_PROCEED_CAPABILITY_ID, + PII_TOKENIZATION_CAPABILITY_ID, + resolveInferenceProceedCapabilityValue, + resolvePiiTokenizationCapabilityValue, +} from '@kbn/inference-plugin/server'; +import { aroundCompletionEventSchema } from '../../common/workflow_anonymization'; +import { createWorkflowAnonymizationProvider } from './create_workflow_anonymization_provider'; + +type Management = Pick< + WorkflowsManagementApi, + 'resolveWorkflowTriggerMatches' | 'executeWorkflowSynchronously' +>; + +const createManagement = (): jest.Mocked => ({ + resolveWorkflowTriggerMatches: jest.fn(), + executeWorkflowSynchronously: jest.fn(), +}); + +const createProvider = ( + management: jest.Mocked, + ensureManagedWorkflow: jest.Mock = jest.fn().mockResolvedValue(undefined), + triggerCacheTtlMs = 30_000 +) => createWorkflowAnonymizationProvider({ management, ensureManagedWorkflow, triggerCacheTtlMs }); + +const createAroundCompletionTrigger = (): WorkflowYaml['triggers'][number] => { + // The static WorkflowYaml type models built-ins only; registered triggers are added dynamically. + const trigger: WorkflowYaml['triggers'][number] = { type: 'manual' }; + Reflect.set(trigger, 'type', 'inference.aroundCompletion'); + return trigger; +}; + +const createWorkflow = (steps: WorkflowYaml['steps']): WorkflowDetailDto => ({ + id: 'workflow-1', + name: 'Inference protection', + enabled: true, + createdAt: '2026-01-01T00:00:00.000Z', + createdBy: 'system', + lastUpdatedAt: '2026-01-01T00:00:00.000Z', + lastUpdatedBy: 'system', + yaml: 'name: Inference protection', + valid: true, + definition: { + version: '1', + name: 'Inference protection', + enabled: true, + triggers: [createAroundCompletionTrigger()], + steps, + }, +}); + +const proceedStep: WorkflowYaml['steps'][number] = { + name: 'proceed', + type: 'call_site.proceed', + with: {}, +}; +const configuredStep: WorkflowYaml['steps'][number] = { + name: 'configured', + type: 'test.step', + with: { type: 'call_site.proceed' }, +}; +const conditionalProceedStep: WorkflowYaml['steps'][number] = { + name: 'conditional', + type: 'if', + condition: 'true', + steps: [proceedStep], +}; + +const event: AroundCompletionEvent = { + messages: [{ role: MessageRole.User, content: 'hello' }], +}; +const pii: PiiTokenizationContext = { + detectEntities: jest.fn().mockResolvedValue([]), + tokenize: jest.fn(), +}; +const proceed: InferenceProceedCapability = { + invoke: jest.fn().mockResolvedValue({ rawContent: 'protected' }), +}; + +describe('createWorkflowAnonymizationProvider', () => { + it('rejects an empty space ID before workflow resolution', async () => { + const management = createManagement(); + const provider = createProvider(management); + + await expect( + provider.execute({ + event, + namespace: '', + request: httpServerMock.createKibanaRequest(), + pii, + proceed, + }) + ).rejects.toThrow('non-empty space ID'); + expect(management.resolveWorkflowTriggerMatches).not.toHaveBeenCalled(); + }); + + it('returns without execution when no workflow matches', async () => { + const management = createManagement(); + management.resolveWorkflowTriggerMatches.mockResolvedValue({ + matched: [], + invalidConditionWorkflows: [], + }); + const ensureManagedWorkflow = jest.fn().mockResolvedValue(undefined); + const provider = createProvider(management, ensureManagedWorkflow); + + await expect( + provider.execute({ + event, + namespace: 'default', + request: httpServerMock.createKibanaRequest(), + pii, + proceed, + }) + ).resolves.toEqual({ matched: false }); + expect(ensureManagedWorkflow).toHaveBeenCalledWith('default', expect.anything()); + expect(management.resolveWorkflowTriggerMatches).toHaveBeenCalledWith( + 'inference.aroundCompletion', + aroundCompletionEventSchema.parse(event), + 'default' + ); + expect(management.executeWorkflowSynchronously).not.toHaveBeenCalled(); + }); + + describe('trigger resolution cache', () => { + it('skips the ES lookup on the second call with the same (namespace, agentId)', async () => { + const management = createManagement(); + management.resolveWorkflowTriggerMatches.mockResolvedValue({ + matched: [], + invalidConditionWorkflows: [], + }); + const provider = createProvider(management); + const request = httpServerMock.createKibanaRequest(); + const eventWithAgent = { ...event, agentId: 'agent-a' }; + + await provider.execute({ + event: eventWithAgent, + namespace: 'default', + request, + pii, + proceed, + }); + await provider.execute({ + event: eventWithAgent, + namespace: 'default', + request, + pii, + proceed, + }); + + expect(management.resolveWorkflowTriggerMatches).toHaveBeenCalledTimes(1); + }); + + it('makes a new ES call when the agentId differs', async () => { + const management = createManagement(); + management.resolveWorkflowTriggerMatches.mockResolvedValue({ + matched: [], + invalidConditionWorkflows: [], + }); + const provider = createProvider(management); + const request = httpServerMock.createKibanaRequest(); + + await provider.execute({ + event: { ...event, agentId: 'agent-a' }, + namespace: 'default', + request, + pii, + proceed, + }); + await provider.execute({ + event: { ...event, agentId: 'agent-b' }, + namespace: 'default', + request, + pii, + proceed, + }); + + expect(management.resolveWorkflowTriggerMatches).toHaveBeenCalledTimes(2); + }); + + it('makes a new ES call when the namespace differs', async () => { + const management = createManagement(); + management.resolveWorkflowTriggerMatches.mockResolvedValue({ + matched: [], + invalidConditionWorkflows: [], + }); + const provider = createProvider(management); + const request = httpServerMock.createKibanaRequest(); + + await provider.execute({ event, namespace: 'space-a', request, pii, proceed }); + await provider.execute({ event, namespace: 'space-b', request, pii, proceed }); + + expect(management.resolveWorkflowTriggerMatches).toHaveBeenCalledTimes(2); + }); + + it('makes a new ES call on every request when triggerCacheTtlMs is 0', async () => { + const management = createManagement(); + management.resolveWorkflowTriggerMatches.mockResolvedValue({ + matched: [], + invalidConditionWorkflows: [], + }); + const provider = createProvider(management, jest.fn().mockResolvedValue(undefined), 0); + const request = httpServerMock.createKibanaRequest(); + + await provider.execute({ event, namespace: 'default', request, pii, proceed }); + await provider.execute({ event, namespace: 'default', request, pii, proceed }); + + expect(management.resolveWorkflowTriggerMatches).toHaveBeenCalledTimes(2); + }); + + it('does not cache a resolution with invalid conditions', async () => { + const management = createManagement(); + management.resolveWorkflowTriggerMatches.mockResolvedValue({ + matched: [], + invalidConditionWorkflows: [{ id: 'broken', name: 'Broken' }], + }); + const provider = createProvider(management); + const request = httpServerMock.createKibanaRequest(); + + await expect( + provider.execute({ event, namespace: 'default', request, pii, proceed }) + ).rejects.toThrow('Broken'); + await expect( + provider.execute({ event, namespace: 'default', request, pii, proceed }) + ).rejects.toThrow('Broken'); + + expect(management.resolveWorkflowTriggerMatches).toHaveBeenCalledTimes(2); + }); + }); + + it('executes one matching workflow synchronously with request-local capabilities', async () => { + const management = createManagement(); + management.resolveWorkflowTriggerMatches.mockResolvedValue({ + matched: [createWorkflow([configuredStep, conditionalProceedStep])], + invalidConditionWorkflows: [], + }); + management.executeWorkflowSynchronously.mockResolvedValue({ + workflowExecutionId: 'execution-1', + result: { status: ExecutionStatus.COMPLETED, output: { content: 'restored content' } }, + }); + const provider = createProvider(management); + const request = httpServerMock.createKibanaRequest(); + const abortSignal = new AbortController().signal; + const parsedEvent = aroundCompletionEventSchema.parse(event); + + await expect( + provider.execute({ event, namespace: 'space-a', request, pii, proceed, abortSignal }) + ).resolves.toEqual({ matched: true, content: 'restored content' }); + expect(management.executeWorkflowSynchronously).toHaveBeenCalledWith({ + workflowId: 'workflow-1', + workflow: expect.objectContaining({ id: 'workflow-1' }), + context: { + event: parsedEvent, + spaceId: 'space-a', + triggeredBy: 'inference.aroundCompletion', + }, + spaceId: 'space-a', + request, + capabilities: expect.any(Array), + abortSignal, + }); + const [{ capabilities, context }] = management.executeWorkflowSynchronously.mock.calls[0]; + expect(context.event).not.toBe(event); + expect(capabilities?.map(({ id }) => id)).toEqual([ + PII_TOKENIZATION_CAPABILITY_ID, + INFERENCE_PROCEED_CAPABILITY_ID, + ]); + expect(resolvePiiTokenizationCapabilityValue(capabilities?.[0].value ?? {})).toBe(pii); + expect(resolveInferenceProceedCapabilityValue(capabilities?.[1].value ?? {})).toBe(proceed); + }); + + it('rejects invalid trigger conditions and overlapping workflow matches before execution', async () => { + const management = createManagement(); + const provider = createProvider(management); + + management.resolveWorkflowTriggerMatches.mockResolvedValue({ + matched: [createWorkflow([proceedStep])], + invalidConditionWorkflows: [{ id: 'invalid-workflow', name: 'Broken policy' }], + }); + await expect( + provider.execute({ + event, + namespace: 'space-a', + request: httpServerMock.createKibanaRequest(), + pii, + proceed, + }) + ).rejects.toThrow('Broken policy (invalid-workflow)'); + expect(management.executeWorkflowSynchronously).not.toHaveBeenCalled(); + + management.resolveWorkflowTriggerMatches.mockResolvedValue({ + matched: [createWorkflow([proceedStep]), createWorkflow([proceedStep])], + invalidConditionWorkflows: [], + }); + await expect( + provider.execute({ + event, + namespace: 'space-a', + request: httpServerMock.createKibanaRequest(), + pii, + proceed, + }) + ).rejects.toThrow('Multiple workflows matched'); + expect(management.executeWorkflowSynchronously).not.toHaveBeenCalled(); + }); + + it('rejects when the matched workflow has no proceed step', async () => { + const management = createManagement(); + management.resolveWorkflowTriggerMatches.mockResolvedValue({ + matched: [createWorkflow([])], + invalidConditionWorkflows: [], + }); + + await expect( + createProvider(management).execute({ + event, + namespace: 'space-a', + request: httpServerMock.createKibanaRequest(), + pii, + proceed, + }) + ).rejects.toThrow('must contain exactly one'); + }); + + it('rejects when the workflow completes without string content', async () => { + const management = createManagement(); + management.resolveWorkflowTriggerMatches.mockResolvedValue({ + matched: [createWorkflow([proceedStep])], + invalidConditionWorkflows: [], + }); + management.executeWorkflowSynchronously.mockResolvedValue({ + workflowExecutionId: 'execution-1', + result: { status: ExecutionStatus.COMPLETED, output: {} }, + }); + + await expect( + createProvider(management).execute({ + event, + namespace: 'space-a', + request: httpServerMock.createKibanaRequest(), + pii, + proceed, + }) + ).rejects.toThrow('did not return string content'); + }); +}); diff --git a/x-pack/platform/plugins/shared/inference_workflows/server/workflow_anonymization/create_workflow_anonymization_provider.ts b/x-pack/platform/plugins/shared/inference_workflows/server/workflow_anonymization/create_workflow_anonymization_provider.ts new file mode 100644 index 0000000000000..ff317e3da4eb2 --- /dev/null +++ b/x-pack/platform/plugins/shared/inference_workflows/server/workflow_anonymization/create_workflow_anonymization_provider.ts @@ -0,0 +1,169 @@ +/* + * 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 { ExecutionStatus, visitNestedSteps, type WorkflowYaml } from '@kbn/workflows'; +import type { WorkflowAnonymizationProvider } from '@kbn/inference-plugin/server'; +import type { WorkflowsManagementApi } from '@kbn/workflows-management-plugin/server'; +import { + createInferenceProceedCapabilityValue, + createPiiTokenizationCapabilityValue, + INFERENCE_PROCEED_CAPABILITY_ID, + PII_TOKENIZATION_CAPABILITY_ID, +} from '@kbn/inference-plugin/server'; +import { + aroundCompletionEventSchema, + CALL_SITE_PROCEED_STEP_ID, + INFERENCE_AROUND_COMPLETION_TRIGGER_ID, +} from '../../common/workflow_anonymization'; +import { triggerEvaluationsCounter } from './anonymization_metrics'; + +type WorkflowAnonymizationManagement = Pick< + WorkflowsManagementApi, + 'resolveWorkflowTriggerMatches' | 'executeWorkflowSynchronously' +>; + +type TriggerResolution = Awaited< + ReturnType +>; + +const countWorkflowStepType = (steps: WorkflowYaml['steps'], stepType: string): number => { + let count = 0; + visitNestedSteps(steps, ({ step }) => { + if (step.type === stepType) { + count += 1; + } + }); + return count; +}; + +export const createWorkflowAnonymizationProvider = ({ + management, + ensureManagedWorkflow, + triggerCacheTtlMs, +}: { + management: WorkflowAnonymizationManagement; + ensureManagedWorkflow: ( + spaceId: string, + request: Parameters[0]['request'] + ) => Promise; + triggerCacheTtlMs: number; +}): WorkflowAnonymizationProvider => { + // Cache keyed on (spaceId, agentId) — the two stable identifiers that workflow trigger + // conditions are expected to vary on. Conditions on sessionId or message content are not + // supported with caching; those use-cases are outside the anonymization policy model. + // TTL is controlled by xpack.inference.anonymization.triggerCacheTtlSeconds (0 = no cache). + const triggerCache = new Map(); + + return { + supportsSynchronousExecution: true, + execute: async ({ event, namespace, request, pii, proceed, abortSignal }) => { + if (namespace.trim().length === 0) { + throw new Error('Workflow anonymization requires a non-empty space ID'); + } + const parsedEvent = aroundCompletionEventSchema.parse(event); + await ensureManagedWorkflow(namespace, request); + + const cacheKey = `${namespace}:${parsedEvent.agentId ?? ''}`; + const now = Date.now(); + const cached = triggerCacheTtlMs > 0 ? triggerCache.get(cacheKey) : undefined; + let resolution: TriggerResolution; + + if (cached && cached.expiresAt > now) { + resolution = cached.result; + } else { + resolution = await management.resolveWorkflowTriggerMatches( + INFERENCE_AROUND_COMPLETION_TRIGGER_ID, + parsedEvent, + namespace + ); + // Only cache valid resolutions — invalid conditions are likely transient config errors + // that should be re-evaluated on the next request rather than held for the full TTL. + if (triggerCacheTtlMs > 0 && resolution.invalidConditionWorkflows.length === 0) { + // Prune expired entries on every write to prevent unbounded accumulation over the + // plugin lifetime (e.g. deleted spaces whose keys are never looked up again). + for (const [key, entry] of triggerCache) { + if (entry.expiresAt <= now) { + triggerCache.delete(key); + } + } + triggerCache.set(cacheKey, { + result: resolution, + expiresAt: now + triggerCacheTtlMs, + }); + } + } + + if (resolution.invalidConditionWorkflows.length > 0) { + // Fail closed across the space. Ignoring a malformed subscribed policy could allow an + // unprotected connector call when the workflow author intended that policy to match. + triggerEvaluationsCounter.add(1, { outcome: 'invalid_condition' }); + throw new Error( + `Invalid around-completion trigger condition in workflow(s): ${resolution.invalidConditionWorkflows + .map(({ id, name }) => `${name} (${id})`) + .join(', ')}` + ); + } + if (resolution.matched.length === 0) { + triggerEvaluationsCounter.add(1, { outcome: 'no_match' }); + return { matched: false }; + } + if (resolution.matched.length > 1) { + triggerEvaluationsCounter.add(1, { outcome: 'overlap_conflict' }); + throw new Error('Multiple workflows matched the around-completion inference event'); + } + + const [workflow] = resolution.matched; + const proceedCount = countWorkflowStepType( + workflow.definition?.steps ?? [], + CALL_SITE_PROCEED_STEP_ID + ); + if (proceedCount !== 1) { + throw new Error( + `Workflow "${workflow.id}" must contain exactly one ${CALL_SITE_PROCEED_STEP_ID} step` + ); + } + + const response = await management.executeWorkflowSynchronously({ + workflowId: workflow.id, + // Pass the already-fetched DTO to avoid a second ES read on the inference hot path. + // executeWorkflowSynchronously validates the DTO (enabled/valid/definition) before executing. + workflow, + context: { + event: parsedEvent, + spaceId: namespace, + triggeredBy: INFERENCE_AROUND_COMPLETION_TRIGGER_ID, + }, + spaceId: namespace, + request, + capabilities: [ + { id: PII_TOKENIZATION_CAPABILITY_ID, value: createPiiTokenizationCapabilityValue(pii) }, + { + id: INFERENCE_PROCEED_CAPABILITY_ID, + value: createInferenceProceedCapabilityValue(proceed), + }, + ], + abortSignal, + }); + + if (response.result?.status !== ExecutionStatus.COMPLETED) { + const stepError = response.result?.error?.message; + throw new Error( + stepError + ? `Workflow "${workflow.id}" did not complete successfully: ${stepError}` + : `Workflow "${workflow.id}" did not complete successfully` + ); + } + const content = response.result.output?.content; + if (typeof content !== 'string') { + throw new Error(`Workflow "${workflow.id}" did not return string content`); + } + + triggerEvaluationsCounter.add(1, { outcome: 'matched' }); + return { matched: true, content }; + }, + }; +}; diff --git a/x-pack/platform/plugins/shared/inference_workflows/server/workflow_anonymization/message_records.test.ts b/x-pack/platform/plugins/shared/inference_workflows/server/workflow_anonymization/message_records.test.ts new file mode 100644 index 0000000000000..0ce21f595f9ff --- /dev/null +++ b/x-pack/platform/plugins/shared/inference_workflows/server/workflow_anonymization/message_records.test.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 { MessageRole } from '@kbn/inference-common'; +import { createCompletionTextRecords } from './message_records'; + +describe('createCompletionTextRecords', () => { + it('throws when tool-call arguments exceed the maximum nesting depth', () => { + // Build a structure that is 101 levels deep — one beyond MAX_STRUCTURED_DEPTH (100). + // Tool-call arguments are user-influenced and could be arbitrarily nested; the guard + // ensures we fail closed rather than recurse until OOM. + const deeplyNestedArgs: Record = {}; + let level: Record = deeplyNestedArgs; + for (let i = 0; i < 101; i++) { + const next: Record = {}; + level.nested = next; + level = next; + } + level.value = 'secret'; + + expect(() => + createCompletionTextRecords({ + messages: [ + { + role: MessageRole.Assistant, + content: null, + toolCalls: [ + { + toolCallId: 'call-1', + function: { name: 'test', arguments: deeplyNestedArgs }, + }, + ], + }, + ], + }) + ).toThrow('maximum nesting depth'); + }); +}); diff --git a/x-pack/platform/plugins/shared/inference_workflows/server/workflow_anonymization/message_records.ts b/x-pack/platform/plugins/shared/inference_workflows/server/workflow_anonymization/message_records.ts new file mode 100644 index 0000000000000..857a27f9a4d79 --- /dev/null +++ b/x-pack/platform/plugins/shared/inference_workflows/server/workflow_anonymization/message_records.ts @@ -0,0 +1,204 @@ +/* + * 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 { MessageRole, type Message } from '@kbn/inference-common'; +import type { PiiTextRecord } from '@kbn/inference-plugin/server'; + +type RecordValues = ReadonlyMap; + +// Fail-closed guard: tool-call arguments are user-influenced and could be arbitrarily nested. +// Exceeding this depth throws rather than silently truncating or OOM-ing. +const MAX_STRUCTURED_DEPTH = 100; + +const collectStructuredStrings = ( + value: unknown, + path: string, + records: PiiTextRecord[], + depth = 0 +): void => { + if (depth > MAX_STRUCTURED_DEPTH) { + throw new Error( + `PII scanner exceeded maximum nesting depth (${MAX_STRUCTURED_DEPTH}) at path "${path}"` + ); + } + + if (typeof value === 'string') { + records.push({ id: path, text: value }); + return; + } + + if (Array.isArray(value)) { + value.forEach((item, index) => + collectStructuredStrings(item, `${path}/${index}`, records, depth + 1) + ); + return; + } + + if (!value || typeof value !== 'object') { + return; + } + + const objectValue = value as Record; + if (objectValue.type === 'image') { + return; + } + + Object.entries(objectValue).forEach(([key, entry]) => + collectStructuredStrings(entry, `${path}/${key}`, records, depth + 1) + ); +}; + +const replaceStructuredStrings = ( + value: T, + path: string, + values: RecordValues, + depth = 0 +): T => { + if (depth > MAX_STRUCTURED_DEPTH) { + throw new Error( + `PII scanner exceeded maximum nesting depth (${MAX_STRUCTURED_DEPTH}) at path "${path}"` + ); + } + + if (typeof value === 'string') { + return (values.get(path) ?? value) as T; + } + + if (Array.isArray(value)) { + return value.map((item, index) => + replaceStructuredStrings(item, `${path}/${index}`, values, depth + 1) + ) as T; + } + + if (!value || typeof value !== 'object') { + return value; + } + + const objectValue = value as Record; + if (objectValue.type === 'image') { + return value; + } + + return Object.fromEntries( + Object.entries(objectValue).map(([key, entry]) => [ + key, + replaceStructuredStrings(entry, `${path}/${key}`, values, depth + 1), + ]) + ) as T; +}; + +const collectMessageStrings = (message: Message, index: number, records: PiiTextRecord[]): void => { + const path = `/messages/${index}`; + + if (message.role === MessageRole.User) { + if (typeof message.content === 'string') { + records.push({ id: `${path}/content`, text: message.content }); + return; + } + + message.content.forEach((content, contentIndex) => { + if (content.type === 'text') { + records.push({ id: `${path}/content/${contentIndex}/text`, text: content.text }); + } + }); + return; + } + + if (message.role === MessageRole.Assistant) { + if (typeof message.content === 'string') { + records.push({ id: `${path}/content`, text: message.content }); + } + message.toolCalls?.forEach((toolCall, toolCallIndex) => + collectStructuredStrings( + toolCall.function.arguments, + `${path}/toolCalls/${toolCallIndex}/function/arguments`, + records + ) + ); + return; + } + + collectStructuredStrings(message.response, `${path}/response`, records); + if (message.data !== undefined) { + collectStructuredStrings(message.data, `${path}/data`, records); + } +}; + +const replaceMessageStrings = (message: Message, index: number, values: RecordValues): Message => { + const path = `/messages/${index}`; + + if (message.role === MessageRole.User) { + return { + ...message, + content: + typeof message.content === 'string' + ? values.get(`${path}/content`) ?? message.content + : message.content.map((content, contentIndex) => + content.type === 'text' + ? { + ...content, + text: values.get(`${path}/content/${contentIndex}/text`) ?? content.text, + } + : content + ), + }; + } + + if (message.role === MessageRole.Assistant) { + return { + ...message, + content: + typeof message.content === 'string' + ? values.get(`${path}/content`) ?? message.content + : message.content, + toolCalls: message.toolCalls?.map((toolCall, toolCallIndex) => ({ + ...toolCall, + function: { + ...toolCall.function, + arguments: replaceStructuredStrings( + toolCall.function.arguments, + `${path}/toolCalls/${toolCallIndex}/function/arguments`, + values + ), + }, + })), + }; + } + + return { + ...message, + response: replaceStructuredStrings(message.response, `${path}/response`, values), + ...(message.data !== undefined + ? { data: replaceStructuredStrings(message.data, `${path}/data`, values) } + : {}), + }; +}; + +export const createCompletionTextRecords = ({ + system, + messages, +}: { + system?: string; + messages: readonly Message[]; +}): { + records: readonly PiiTextRecord[]; + replace(values: RecordValues): { system?: string; messages: Message[] }; +} => { + const records: PiiTextRecord[] = []; + if (system !== undefined) { + records.push({ id: '/system', text: system }); + } + messages.forEach((message, index) => collectMessageStrings(message, index, records)); + + return { + records, + replace: (values) => ({ + ...(system !== undefined ? { system: values.get('/system') ?? system } : {}), + messages: messages.map((message, index) => replaceMessageStrings(message, index, values)), + }), + }; +}; diff --git a/x-pack/platform/plugins/shared/inference_workflows/server/workflow_anonymization/pii_restore_step.test.ts b/x-pack/platform/plugins/shared/inference_workflows/server/workflow_anonymization/pii_restore_step.test.ts new file mode 100644 index 0000000000000..e6ed358808503 --- /dev/null +++ b/x-pack/platform/plugins/shared/inference_workflows/server/workflow_anonymization/pii_restore_step.test.ts @@ -0,0 +1,33 @@ +/* + * 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 { executePiiRestore, piiRestoreStepDefinition } from './pii_restore_step'; + +describe('piiRestoreStepDefinition', () => { + it('restores protected values through the step definition handler', async () => { + const result = await piiRestoreStepDefinition.handler({ + input: { + rawContent: 'Contact EMAIL_TOKEN', + tokenMap: { + EMAIL_TOKEN: { original: 'person@example.com', entityClass: 'EMAIL' }, + }, + }, + } as unknown as Parameters[0]); + expect(result).toEqual({ output: { content: 'Contact person@example.com' } }); + }); + + it('restores protected values in final response content', () => { + expect( + executePiiRestore({ + rawContent: 'Contact EMAIL_TOKEN', + tokenMap: { + EMAIL_TOKEN: { original: 'person@example.com', entityClass: 'EMAIL' }, + }, + }) + ).toEqual({ content: 'Contact person@example.com' }); + }); +}); diff --git a/x-pack/platform/plugins/shared/inference_workflows/server/workflow_anonymization/pii_restore_step.ts b/x-pack/platform/plugins/shared/inference_workflows/server/workflow_anonymization/pii_restore_step.ts new file mode 100644 index 0000000000000..94acacba12254 --- /dev/null +++ b/x-pack/platform/plugins/shared/inference_workflows/server/workflow_anonymization/pii_restore_step.ts @@ -0,0 +1,31 @@ +/* + * 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 { createServerStepDefinition } from '@kbn/workflows-extensions/server'; +import { piiRestoreCommonDefinition } from '../../common/workflow_anonymization'; +import { restoreTokens } from './token_map'; + +export const executePiiRestore = ({ + rawContent, + tokenMap, +}: { + rawContent: string; + tokenMap: Parameters[1]; +}): { content: string } => ({ content: restoreTokens(rawContent, tokenMap) }); + +export const piiRestoreStepHandler = async ({ + input, +}: { + input: Parameters[0]; +}) => ({ + output: executePiiRestore(input), +}); + +export const piiRestoreStepDefinition = createServerStepDefinition({ + ...piiRestoreCommonDefinition, + handler: piiRestoreStepHandler, +}); diff --git a/x-pack/platform/plugins/shared/inference_workflows/server/workflow_anonymization/token_map.test.ts b/x-pack/platform/plugins/shared/inference_workflows/server/workflow_anonymization/token_map.test.ts new file mode 100644 index 0000000000000..ff53ef460711e --- /dev/null +++ b/x-pack/platform/plugins/shared/inference_workflows/server/workflow_anonymization/token_map.test.ts @@ -0,0 +1,21 @@ +/* + * 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 { restoreTokens } from './token_map'; + +describe('token map transformations', () => { + const tokenMap = { + SHORT_TOKEN: { original: 'secret', entityClass: 'ENTITY_NAME' }, + LONG_TOKEN: { original: 'secret value', entityClass: 'ENTITY_NAME' }, + }; + + it('restores every occurrence of call-local tokens', () => { + expect(restoreTokens('LONG_TOKEN then SHORT_TOKEN and SHORT_TOKEN', tokenMap)).toBe( + 'secret value then secret and secret' + ); + }); +}); diff --git a/x-pack/platform/plugins/shared/inference_workflows/server/workflow_anonymization/token_map.ts b/x-pack/platform/plugins/shared/inference_workflows/server/workflow_anonymization/token_map.ts new file mode 100644 index 0000000000000..2c0b0b1290b2b --- /dev/null +++ b/x-pack/platform/plugins/shared/inference_workflows/server/workflow_anonymization/token_map.ts @@ -0,0 +1,25 @@ +/* + * 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 { TokenMap } from '../../common/workflow_anonymization'; + +const replaceStringValues = (text: string, replacements: Record): string => { + const keys = Object.keys(replacements).filter((k) => k.length > 0); + if (keys.length === 0) return text; + const sorted = [...keys].sort((a, b) => b.length - a.length); + let result = text; + for (const key of sorted) { + result = result.split(key).join(replacements[key]); + } + return result; +}; + +export const restoreTokens = (value: string, tokenMap: TokenMap): string => + replaceStringValues( + value, + Object.fromEntries(Object.entries(tokenMap).map(([token, entry]) => [token, entry.original])) + ); diff --git a/x-pack/platform/plugins/shared/inference_workflows/tsconfig.json b/x-pack/platform/plugins/shared/inference_workflows/tsconfig.json index 036b4a1dd242d..2b07bc26fe039 100644 --- a/x-pack/platform/plugins/shared/inference_workflows/tsconfig.json +++ b/x-pack/platform/plugins/shared/inference_workflows/tsconfig.json @@ -3,7 +3,7 @@ "compilerOptions": { "outDir": "target/types" }, - "include": ["../../../../../typings/**/*", "common/**/*", "public/**/*", "server/**/*"], + "include": ["../../../../../typings/**/*", "common/**/*", "public/**/*", "server/**/*", "integration_tests/**/*"], "exclude": ["target/**/*"], "kbn_references": [ "@kbn/core", @@ -11,8 +11,14 @@ "@kbn/zod", "@kbn/workflows", "@kbn/workflows-extensions", + "@kbn/workflows-management-plugin", "@kbn/inference-plugin", "@kbn/inference-common", - "@kbn/search-inference-endpoints" + "@kbn/search-inference-endpoints", + "@kbn/core-http-server-mocks", + "@kbn/core-logging-server-mocks", + "@kbn/core-spaces-common", + "@kbn/spaces-plugin", + "@kbn/config-schema" ] }