Nexus CancelCommand in CHASM (#9288)

Ported command handler for Nexus "cancel" command from HSM to CHASM.

CHASM migration.

- [ ] built
- [ ] run locally and tested manually
- [x] covered by existing tests
- [x] added new unit test(s)
- [ ] added new functional test(s)

---------

Co-authored-by: Roey Berman <roey.berman@gmail.com>

chasm/lib/nexusoperation/cancellation.go

@@ -8,6 +8,7 @@ import (
 var _ chasm.Component = (*Cancellation)(nil)
 var _ chasm.StateMachine[nexusoperationpb.CancellationStatus] = (*Cancellation)(nil)
 
+// Cancellation is a CHASM component that represents a pending cancellation of a Nexus operation.
 type Cancellation struct {
 	chasm.UnimplementedComponent
 
@@ -15,10 +16,11 @@ type Cancellation struct {
 	*nexusoperationpb.CancellationState
 }
 
-func NewCancellation() *Cancellation {
-	return &Cancellation{}
+func newCancellation(state *nexusoperationpb.CancellationState) *Cancellation {
+	return &Cancellation{CancellationState: state}
 }
 
+// LifecycleState maps the cancellation's status to a CHASM lifecycle state.
 func (o *Cancellation) LifecycleState(_ chasm.Context) chasm.LifecycleState {
 	switch o.Status {
 	case nexusoperationpb.CANCELLATION_STATUS_SUCCEEDED:
@@ -31,10 +33,12 @@ func (o *Cancellation) LifecycleState(_ chasm.Context) chasm.LifecycleState {
 	}
 }
 
+// StateMachineState returns the current cancellation status.
 func (o *Cancellation) StateMachineState() nexusoperationpb.CancellationStatus {
 	return o.Status
 }
 
+// SetStateMachineState sets the cancellation status.
 func (o *Cancellation) SetStateMachineState(status nexusoperationpb.CancellationStatus) {
 	o.Status = status
 }

chasm/lib/nexusoperation/config.go

@@ -10,7 +10,7 @@ import (
 	"go.temporal.io/server/common/rpc/interceptor"
 )
 
-var ChasmNexusEnabled = dynamicconfig.NewGlobalBoolSetting(
+var ChasmNexusEnabled = dynamicconfig.NewNamespaceBoolSetting(
 	"nexusoperation.enableChasm",
 	false,
 	`Feature flag that controls whether the legacy HSM-based implementation (when flag is false; default) or the newer
@@ -150,7 +150,7 @@ Added for safety. Defaults to true. Likely to be removed in future server versio
 
 type Config struct {
 	ChasmEnabled                        dynamicconfig.BoolPropertyFnWithNamespaceFilter
-	ChasmNexusEnabled                   dynamicconfig.BoolPropertyFn
+	ChasmNexusEnabled                   dynamicconfig.BoolPropertyFnWithNamespaceFilter
 	RequestTimeout                      dynamicconfig.DurationPropertyFnWithDestinationFilter
 	MinRequestTimeout                   dynamicconfig.DurationPropertyFnWithNamespaceFilter
 	MaxConcurrentOperations             dynamicconfig.IntPropertyFnWithNamespaceFilter

chasm/lib/nexusoperation/operation.go

@@ -1,6 +1,7 @@
 package nexusoperation
 
 import (
+	"go.temporal.io/api/serviceerror"
 	"go.temporal.io/server/chasm"
 	nexusoperationpb "go.temporal.io/server/chasm/lib/nexusoperation/gen/nexusoperationpb/v1"
 )
@@ -8,8 +9,15 @@ import (
 var _ chasm.Component = (*Operation)(nil)
 var _ chasm.StateMachine[nexusoperationpb.OperationStatus] = (*Operation)(nil)
 
+// ErrCancellationAlreadyRequested is returned when a cancellation has already been requested for an operation.
+var ErrCancellationAlreadyRequested = serviceerror.NewFailedPrecondition("cancellation already requested")
+
+// ErrOperationAlreadyCompleted is returned when trying to cancel an operation that has already completed.
+var ErrOperationAlreadyCompleted = serviceerror.NewFailedPrecondition("operation already completed")
+
 type OperationStore any
 
+// Operation is a CHASM component that represents a Nexus operation.
 type Operation struct {
 	chasm.UnimplementedComponent
 
@@ -19,12 +27,17 @@ type Operation struct {
 	// Pointer to an implementation of the "store". For a workflow-based Nexus operation
 	// this is a parent pointer back to the workflow. For a standalone Nexus operation this is nil.
 	Store chasm.ParentPtr[OperationStore]
+
+	// Cancellation is a child component that manages sending the cancel request to the Nexus endpoint.
+	Cancellation chasm.Field[*Cancellation]
 }
 
+// NewOperation creates a new Operation component with the given persisted state.
 func NewOperation(state *nexusoperationpb.OperationState) *Operation {
 	return &Operation{OperationState: state}
 }
 
+// LifecycleState maps the operation's status to a CHASM lifecycle state.
 func (o *Operation) LifecycleState(_ chasm.Context) chasm.LifecycleState {
 	switch o.Status {
 	case nexusoperationpb.OPERATION_STATUS_SUCCEEDED:
@@ -38,10 +51,36 @@ func (o *Operation) LifecycleState(_ chasm.Context) chasm.LifecycleState {
 	}
 }
 
+// StateMachineState returns the current operation status.
 func (o *Operation) StateMachineState() nexusoperationpb.OperationStatus {
 	return o.Status
 }
 
+// SetStateMachineState sets the operation status.
 func (o *Operation) SetStateMachineState(status nexusoperationpb.OperationStatus) {
 	o.Status = status
 }
+
+// Cancel requests cancellation of the operation. It creates a Cancellation child component and, if the
+// operation has already started, schedules the cancellation request to be sent to the Nexus endpoint.
+func (o *Operation) Cancel(ctx chasm.MutableContext, requestedEventID int64) error {
+	if !TransitionCanceled.Possible(o) {
+		return ErrOperationAlreadyCompleted
+	}
+	if _, ok := o.Cancellation.TryGet(ctx); ok {
+		return ErrCancellationAlreadyRequested
+	}
+
+	cancellation := newCancellation(&nexusoperationpb.CancellationState{
+		RequestedEventId: requestedEventID,
+	})
+	o.Cancellation = chasm.NewComponentField(ctx, cancellation)
+
+	// Once started, the handler returns a token that can be used in the cancelation request.
+	// Until then, no need to schedule the cancelation.
+	if o.Status == nexusoperationpb.OPERATION_STATUS_STARTED {
+		return transitionCancellationScheduled.Apply(cancellation, ctx, EventCancellationScheduled{})
+	}
+
+	return nil
+}

chasm/lib/nexusoperation/operation_statemachine.go

@@ -161,8 +161,8 @@ var transitionSucceeded = chasm.NewTransition(
 	},
 	nexusoperationpb.OPERATION_STATUS_SUCCEEDED,
 	func(o *Operation, ctx chasm.MutableContext, event EventSucceeded) error {
-		// Terminal state - no tasks to emit
-		// The component will be deleted after this transition
+		// Terminal state - no tasks to emit.
+		// The component will be deleted after this transition.
 		return nil
 	},
 )
@@ -179,9 +179,8 @@ var transitionFailed = chasm.NewTransition(
 	},
 	nexusoperationpb.OPERATION_STATUS_FAILED,
 	func(o *Operation, ctx chasm.MutableContext, event EventFailed) error {
-		// Terminal state - no tasks to emit
-		// Not recording the last attempt information here since the state machine
-		// will be deleted immediately after this transition
+		// Terminal state - no tasks to emit.
+		// The component will be deleted after this transition.
 		return nil
 	},
 )
@@ -190,7 +189,7 @@ var transitionFailed = chasm.NewTransition(
 type EventCanceled struct {
 }
 
-var transitionCanceled = chasm.NewTransition(
+var TransitionCanceled = chasm.NewTransition(
 	[]nexusoperationpb.OperationStatus{
 		nexusoperationpb.OPERATION_STATUS_SCHEDULED,
 		nexusoperationpb.OPERATION_STATUS_STARTED,
@@ -198,8 +197,8 @@ var transitionCanceled = chasm.NewTransition(
 	},
 	nexusoperationpb.OPERATION_STATUS_CANCELED,
 	func(o *Operation, ctx chasm.MutableContext, event EventCanceled) error {
-		// Terminal state - no tasks to emit
-		// The component will be deleted after this transition
+		// Terminal state - no tasks to emit.
+		// The component will be deleted after this transition.
 		return nil
 	},
 )
@@ -216,8 +215,8 @@ var transitionTimedOut = chasm.NewTransition(
 	},
 	nexusoperationpb.OPERATION_STATUS_TIMED_OUT,
 	func(o *Operation, ctx chasm.MutableContext, event EventTimedOut) error {
-		// Terminal state - no tasks to emit
-		// The component will be deleted after this transition
+		// Terminal state - no tasks to emit.
+		// The component will be deleted after this transition.
 		return nil
 	},
 )

chasm/lib/nexusoperation/operation_statemachine_test.go

@@ -394,7 +394,7 @@ func TestTransitionCanceled(t *testing.T) {
 
 			event := EventCanceled{}
 
-			err := transitionCanceled.Apply(operation, ctx, event)
+			err := TransitionCanceled.Apply(operation, ctx, event)
 			require.NoError(t, err)
 
 			require.Equal(t, nexusoperationpb.OPERATION_STATUS_CANCELED, operation.Status)

chasm/lib/nexusoperation/workflow/commands.go

@@ -46,7 +46,7 @@ func registerCommandHandlers(
 	}
 	return registry.Register(
 		enumspb.COMMAND_TYPE_REQUEST_CANCEL_NEXUS_OPERATION,
-		handleCancelCommand,
+		h.handleCancelCommand,
 	)
 }
 
@@ -68,6 +68,10 @@ func (ch *commandHandler) handleScheduleCommand(
 		}
 	}
 
+	if !ch.config.ChasmNexusEnabled(nsName) {
+		return command.ErrNotSupported
+	}
+
 	attrs := cmd.GetScheduleNexusOperationCommandAttributes()
 	if attrs == nil {
 		return command.FailWorkflowTaskError{
@@ -286,13 +290,89 @@ func (ch *commandHandler) handleScheduleCommand(
 	return nil
 }
 
-func handleCancelCommand(
+func (ch *commandHandler) handleCancelCommand(
 	chasmCtx chasm.MutableContext,
 	wf *chasmworkflow.Workflow,
 	validator command.Validator,
 	cmd *commandpb.Command,
 	opts command.HandlerOptions,
 ) error {
-	// TODO: Implement CHASM nexus operation cancellation
-	return serviceerror.NewUnimplemented("CHASM nexus operation cancellation not yet implemented")
+	if !ch.config.Enabled() {
+		return command.FailWorkflowTaskError{
+			Cause:   enumspb.WORKFLOW_TASK_FAILED_CAUSE_FEATURE_DISABLED,
+			Message: "Nexus operations disabled",
+		}
+	}
+
+	nsName := chasmCtx.NamespaceEntry().Name().String()
+	if !ch.config.ChasmNexusEnabled(nsName) {
+		return command.ErrNotSupported
+	}
+
+	attrs := cmd.GetRequestCancelNexusOperationCommandAttributes()
+	if attrs == nil {
+		return command.FailWorkflowTaskError{
+			Cause:   enumspb.WORKFLOW_TASK_FAILED_CAUSE_BAD_REQUEST_CANCEL_NEXUS_OPERATION_ATTRIBUTES,
+			Message: "empty CancelNexusOperationCommandAttributes",
+		}
+	}
+
+	key := strconv.FormatInt(attrs.ScheduledEventId, 10)
+	operationField, operationFound := wf.Operations[key]
+	hasBufferedEvent := func() bool {
+		return wf.HasAnyBufferedEvent(makeNexusOperationTerminalEventFilter(attrs.ScheduledEventId))
+	}
+
+	if !operationFound && !hasBufferedEvent() {
+		return command.FailWorkflowTaskError{
+			Cause:   enumspb.WORKFLOW_TASK_FAILED_CAUSE_BAD_REQUEST_CANCEL_NEXUS_OPERATION_ATTRIBUTES,
+			Message: fmt.Sprintf("requested cancelation for a non-existing or already completed operation with scheduled event ID of %d", attrs.ScheduledEventId),
+		}
+	}
+
+	// Always create the event even if there's a buffered completion to avoid breaking replay in the SDK.
+	// The event will be applied before the completion since buffered events are reordered and put at the end of the
+	// batch, after command events from the workflow task.
+	event := wf.AddHistoryEvent(enumspb.EVENT_TYPE_NEXUS_OPERATION_CANCEL_REQUESTED, func(he *historypb.HistoryEvent) {
+		he.Attributes = &historypb.HistoryEvent_NexusOperationCancelRequestedEventAttributes{
+			NexusOperationCancelRequestedEventAttributes: &historypb.NexusOperationCancelRequestedEventAttributes{
+				ScheduledEventId:             attrs.ScheduledEventId,
+				WorkflowTaskCompletedEventId: opts.WorkflowTaskCompletedEventID,
+			},
+		}
+		he.UserMetadata = cmd.UserMetadata
+	})
+
+	if !operationFound {
+		// Operation not found but there's a buffered terminal event. The workflow couldn't know
+		// the operation completed while its task was in flight. Ignore.
+		return nil
+	}
+
+	op := operationField.Get(chasmCtx)
+	err := op.Cancel(chasmCtx, event.GetEventId())
+	if errors.Is(err, nexusoperation.ErrCancellationAlreadyRequested) {
+		return command.FailWorkflowTaskError{
+			Cause:   enumspb.WORKFLOW_TASK_FAILED_CAUSE_BAD_REQUEST_CANCEL_NEXUS_OPERATION_ATTRIBUTES,
+			Message: fmt.Sprintf("cancelation was already requested for an operation with scheduled event ID %d", attrs.ScheduledEventId),
+		}
+	}
+	return err
+}
+
+func makeNexusOperationTerminalEventFilter(scheduledEventID int64) func(event *historypb.HistoryEvent) bool {
+	return func(event *historypb.HistoryEvent) bool {
+		switch event.EventType {
+		case enumspb.EVENT_TYPE_NEXUS_OPERATION_COMPLETED:
+			return event.GetNexusOperationCompletedEventAttributes().GetScheduledEventId() == scheduledEventID
+		case enumspb.EVENT_TYPE_NEXUS_OPERATION_FAILED:
+			return event.GetNexusOperationFailedEventAttributes().GetScheduledEventId() == scheduledEventID
+		case enumspb.EVENT_TYPE_NEXUS_OPERATION_CANCELED:
+			return event.GetNexusOperationCanceledEventAttributes().GetScheduledEventId() == scheduledEventID
+		case enumspb.EVENT_TYPE_NEXUS_OPERATION_TIMED_OUT:
+			return event.GetNexusOperationTimedOutEventAttributes().GetScheduledEventId() == scheduledEventID
+		default:
+			return false
+		}
+	}
 }

chasm/lib/nexusoperation/workflow/commands_test.go

@@ -25,6 +25,7 @@ import (
 	"go.temporal.io/server/common/namespace"
 	commonnexus "go.temporal.io/server/common/nexus"
 	"go.temporal.io/server/common/nexus/nexustest"
+	"go.temporal.io/server/service/history/historybuilder"
 	"go.temporal.io/server/service/history/tests"
 	"google.golang.org/protobuf/proto"
 	"google.golang.org/protobuf/types/known/durationpb"
@@ -46,14 +47,22 @@ func operationKey(eventID int64) string {
 type testContext struct {
 	chasmCtx        *chasm.MockMutableContext
 	wf              *chasmworkflow.Workflow
+	backend         *chasm.MockNodeBackend
 	execInfo        *persistencespb.WorkflowExecutionInfo
 	scheduleHandler command.Handler
 	cancelHandler   command.Handler
 	history         *historypb.History
 }
 
+func (tcx *testContext) setHasAnyBufferedEvent(value bool) {
+	tcx.backend.HandleHasAnyBufferedEvent = func(filter historybuilder.BufferedEventFilter) bool {
+		return value
+	}
+}
+
 var defaultConfig = &nexusoperation.Config{
 	Enabled:                            dynamicconfig.GetBoolPropertyFn(true),
+	ChasmNexusEnabled:                  dynamicconfig.GetBoolPropertyFnFilteredByNamespace(true),
 	MaxServiceNameLength:               dynamicconfig.GetIntPropertyFnFilteredByNamespace(len("service")),
 	MaxOperationNameLength:             dynamicconfig.GetIntPropertyFnFilteredByNamespace(len("op")),
 	MaxConcurrentOperations:            dynamicconfig.GetIntPropertyFnFilteredByNamespace(2),
@@ -123,6 +132,7 @@ func newTestContext(t *testing.T, cfg *nexusoperation.Config) testContext {
 	return testContext{
 		chasmCtx:        chasmCtx,
 		wf:              wf,
+		backend:         backend,
 		execInfo:        execInfo,
 		history:         history,
 		scheduleHandler: scheduleHandler,
@@ -143,6 +153,16 @@ func TestHandleScheduleCommand(t *testing.T) {
 		require.Empty(t, tcx.history.Events)
 	})
 
+	t.Run("chasm nexus not enabled", func(t *testing.T) {
+		tcx := newTestContext(t, &nexusoperation.Config{
+			Enabled:           dynamicconfig.GetBoolPropertyFn(true),
+			ChasmNexusEnabled: dynamicconfig.GetBoolPropertyFnFilteredByNamespace(false),
+		})
+		err := tcx.scheduleHandler(tcx.chasmCtx, tcx.wf, commandValidator{maxPayloadSize: 1}, &commandpb.Command{}, command.HandlerOptions{WorkflowTaskCompletedEventID: 1})
+		require.ErrorIs(t, err, command.ErrNotSupported)
+		require.Empty(t, tcx.history.Events)
+	})
+
 	t.Run("empty attributes", func(t *testing.T) {
 		tcx := newTestContext(t, defaultConfig)
 		err := tcx.scheduleHandler(tcx.chasmCtx, tcx.wf, commandValidator{maxPayloadSize: 1}, &commandpb.Command{}, command.HandlerOptions{WorkflowTaskCompletedEventID: 1})
@@ -632,7 +652,6 @@ func TestHandleScheduleCommand(t *testing.T) {
 
 func TestHandleCancelCommand(t *testing.T) {
 	t.Run("feature disabled", func(t *testing.T) {
-		t.Skip("requires CHASM nexus operation cancellation implementation")
 		tcx := newTestContext(t, &nexusoperation.Config{
 			Enabled: dynamicconfig.GetBoolPropertyFn(false),
 		})
@@ -644,8 +663,17 @@ func TestHandleCancelCommand(t *testing.T) {
 		require.Empty(t, tcx.history.Events)
 	})
 
+	t.Run("chasm nexus not enabled", func(t *testing.T) {
+		tcx := newTestContext(t, &nexusoperation.Config{
+			Enabled:           dynamicconfig.GetBoolPropertyFn(true),
+			ChasmNexusEnabled: dynamicconfig.GetBoolPropertyFnFilteredByNamespace(false),
+		})
+		err := tcx.cancelHandler(tcx.chasmCtx, tcx.wf, commandValidator{maxPayloadSize: 1}, &commandpb.Command{}, command.HandlerOptions{WorkflowTaskCompletedEventID: 1})
+		require.ErrorIs(t, err, command.ErrNotSupported)
+		require.Empty(t, tcx.history.Events)
+	})
+
 	t.Run("empty attributes", func(t *testing.T) {
-		t.Skip("requires CHASM nexus operation cancellation implementation")
 		tcx := newTestContext(t, defaultConfig)
 		err := tcx.cancelHandler(tcx.chasmCtx, tcx.wf, commandValidator{maxPayloadSize: 1}, &commandpb.Command{}, command.HandlerOptions{WorkflowTaskCompletedEventID: 1})
 		var failWFTErr command.FailWorkflowTaskError
@@ -656,7 +684,6 @@ func TestHandleCancelCommand(t *testing.T) {
 	})
 
 	t.Run("operation not found", func(t *testing.T) {
-		t.Skip("requires CHASM nexus operation cancellation implementation")
 		tcx := newTestContext(t, defaultConfig)
 
 		err := tcx.cancelHandler(tcx.chasmCtx, tcx.wf, commandValidator{maxPayloadSize: 1}, &commandpb.Command{
@@ -674,7 +701,6 @@ func TestHandleCancelCommand(t *testing.T) {
 	})
 
 	t.Run("operation already completed", func(t *testing.T) {
-		t.Skip("requires CHASM nexus operation cancellation implementation")
 		tcx := newTestContext(t, defaultConfig)
 
 		err := tcx.scheduleHandler(tcx.chasmCtx, tcx.wf, commandValidator{maxPayloadSize: 1}, &commandpb.Command{
@@ -691,6 +717,7 @@ func TestHandleCancelCommand(t *testing.T) {
 		event := tcx.history.Events[0]
 
 		// TODO: Complete the operation using CHASM equivalent of CompletedEventDefinition.
+		tcx.wf.RemoveNexusOperation(operationKey(event.EventId))
 
 		// Try to cancel - should fail since operation is completed/deleted.
 		err = tcx.cancelHandler(tcx.chasmCtx, tcx.wf, commandValidator{maxPayloadSize: 1}, &commandpb.Command{
@@ -708,9 +735,8 @@ func TestHandleCancelCommand(t *testing.T) {
 	})
 
 	t.Run("operation already completed - completion buffered", func(t *testing.T) {
-		t.Skip("requires CHASM nexus operation cancellation and buffered event support")
 		tcx := newTestContext(t, defaultConfig)
-		// TODO: CHASM equivalent of HasAnyBufferedEvent setup.
+		tcx.setHasAnyBufferedEvent(true) // simulate buffered terminal event
 
 		err := tcx.scheduleHandler(tcx.chasmCtx, tcx.wf, commandValidator{maxPayloadSize: 1}, &commandpb.Command{
 			Attributes: &commandpb.Command_ScheduleNexusOperationCommandAttributes{
@@ -726,6 +752,7 @@ func TestHandleCancelCommand(t *testing.T) {
 		event := tcx.history.Events[0]
 
 		// TODO: Complete the operation using CHASM equivalent of CompletedEventDefinition.
+		tcx.wf.RemoveNexusOperation(operationKey(event.EventId))
 
 		// Try to cancel - should succeed because there's a buffered completion.
 		err = tcx.cancelHandler(tcx.chasmCtx, tcx.wf, commandValidator{maxPayloadSize: 1}, &commandpb.Command{
@@ -741,8 +768,7 @@ func TestHandleCancelCommand(t *testing.T) {
 		require.Equal(t, event.EventId, crAttrs.ScheduledEventId)
 	})
 
-	t.Run("sets event attributes with UserMetadata and spawns cancelation child machine", func(t *testing.T) {
-		t.Skip("requires CHASM nexus operation cancellation implementation")
+	t.Run("sets event attributes with UserMetadata and spawns cancelation child", func(t *testing.T) {
 		tcx := newTestContext(t, defaultConfig)
 		err := tcx.scheduleHandler(tcx.chasmCtx, tcx.wf, commandValidator{maxPayloadSize: 1}, &commandpb.Command{
 			Attributes: &commandpb.Command_ScheduleNexusOperationCommandAttributes{
@@ -753,6 +779,14 @@ func TestHandleCancelCommand(t *testing.T) {
 				},
 			},
 		}, command.HandlerOptions{WorkflowTaskCompletedEventID: 1})
+		require.NoError(t, err)
+
+		// TODO: Replace with CHASM equivalent of ScheduledEventDefinition.Apply().
+		event := tcx.history.Events[0]
+		key := operationKey(event.EventId)
+		op := tcx.wf.Operations[key].Get(tcx.chasmCtx)
+		op.SetStateMachineState(nexusoperationpb.OPERATION_STATUS_SCHEDULED)
+
 		userMetadata := &sdkpb.UserMetadata{
 			Summary: &commonpb.Payload{
 				Metadata: map[string][]byte{"test_key": []byte(`test_val`)},
@@ -763,9 +797,7 @@ func TestHandleCancelCommand(t *testing.T) {
 				Data:     []byte(`Test Details Data`),
 			},
 		}
-		require.NoError(t, err)
 		require.Len(t, tcx.history.Events, 1)
-		event := tcx.history.Events[0]
 
 		err = tcx.cancelHandler(tcx.chasmCtx, tcx.wf, commandValidator{maxPayloadSize: 1}, &commandpb.Command{
 			Attributes: &commandpb.Command_RequestCancelNexusOperationCommandAttributes{
@@ -777,9 +809,8 @@ func TestHandleCancelCommand(t *testing.T) {
 		}, command.HandlerOptions{WorkflowTaskCompletedEventID: 1})
 		require.NoError(t, err)
 
-		key := operationKey(event.EventId)
-		_, ok := tcx.wf.Operations[key]
-		require.True(t, ok)
+		opField, operationFound := tcx.wf.Operations[key]
+		require.True(t, operationFound)
 
 		require.Len(t, tcx.history.Events, 2)
 		crAttrs := tcx.history.Events[1].GetNexusOperationCancelRequestedEventAttributes()
@@ -788,7 +819,10 @@ func TestHandleCancelCommand(t *testing.T) {
 		savedUserMetadata := tcx.history.Events[1].GetUserMetadata()
 		require.EqualExportedValues(t, userMetadata, savedUserMetadata)
 
-		// TODO: Verify cancelation child component exists (CHASM equivalent of HSM CancelationMachineKey check).
+		// Verify cancelation child component exists.
+		op = opField.Get(tcx.chasmCtx)
+		_, hasCancellation := op.Cancellation.TryGet(tcx.chasmCtx)
+		require.True(t, hasCancellation)
 	})
 }
 

chasm/lib/workflow/command/handler.go

@@ -1,12 +1,18 @@
 package command
 
 import (
+	"errors"
+
 	commandpb "go.temporal.io/api/command/v1"
 	enumspb "go.temporal.io/api/enums/v1"
 	"go.temporal.io/server/chasm"
 	chasmworkflow "go.temporal.io/server/chasm/lib/workflow"
 )
 
+// ErrNotSupported is returned by a [Handler] when the command type is registered but not supported;
+// for example, because of a disabled feature flag.
+var ErrNotSupported = errors.New("command not supported")
+
 type HandlerOptions struct {
 	WorkflowTaskCompletedEventID int64
 }

chasm/lib/workflow/workflow.go

@@ -12,6 +12,7 @@ import (
 	callbackspb "go.temporal.io/server/chasm/lib/callback/gen/callbackpb/v1"
 	"go.temporal.io/server/chasm/lib/nexusoperation"
 	"go.temporal.io/server/common/nexus/nexusrpc"
+	"go.temporal.io/server/service/history/historybuilder"
 	"google.golang.org/protobuf/types/known/emptypb"
 	"google.golang.org/protobuf/types/known/timestamppb"
 )
@@ -29,7 +30,7 @@ type Workflow struct {
 	// Callbacks map is used to store the callbacks for the workflow.
 	Callbacks chasm.Map[string, *callback.Callback]
 
-	// Operations map is used to store the nexus operations for the workflow.
+	// Operations map is used to store the Nexus operations for the workflow.
 	Operations chasm.Map[string, *nexusoperation.Operation]
 }
 
@@ -132,6 +133,7 @@ func (w *Workflow) AddCompletionCallbacks(
 	return nil
 }
 
+// AddNexusOperation adds a Nexus operation component to the workflow.
 func (w *Workflow) AddNexusOperation(
 	ctx chasm.MutableContext,
 	key string,
@@ -143,10 +145,22 @@ func (w *Workflow) AddNexusOperation(
 	w.Operations[key] = chasm.NewComponentField(ctx, op)
 }
 
+// AddHistoryEvent adds a history event to the workflow.
 func (w *Workflow) AddHistoryEvent(t enumspb.EventType, setAttributes func(*historypb.HistoryEvent)) *historypb.HistoryEvent {
 	return w.MSPointer.AddHistoryEvent(t, setAttributes)
 }
 
+// HasAnyBufferedEvent returns true if the workflow has any buffered event matching the given filter.
+func (w *Workflow) HasAnyBufferedEvent(filter historybuilder.BufferedEventFilter) bool {
+	return w.MSPointer.HasAnyBufferedEvent(filter)
+}
+
+// RemoveNexusOperation removes a Nexus operation from the workflow.
+func (w *Workflow) RemoveNexusOperation(key string) {
+	delete(w.Operations, key)
+}
+
+// PendingNexusOperationCount returns the number of pending Nexus operations in the workflow.
 func (w *Workflow) PendingNexusOperationCount() int {
 	return len(w.Operations)
 }

chasm/ms_pointer.go

@@ -33,6 +33,11 @@ func (m MSPointer) AddHistoryEvent(t enumspb.EventType, setAttributes func(*hist
 	return m.backend.AddHistoryEvent(t, setAttributes)
 }
 
+// HasAnyBufferedEvent returns true if there is at least one buffered event that matches the provided filter.
+func (m MSPointer) HasAnyBufferedEvent(filter func(*historypb.HistoryEvent) bool) bool {
+	return m.backend.HasAnyBufferedEvent(filter)
+}
+
 // GetNexusCompletion retrieves the Nexus operation completion data for the given request ID from the underlying mutable state.
 func (m MSPointer) GetNexusCompletion(ctx Context, requestID string) (nexusrpc.CompleteOperationOptions, error) {
 	return m.backend.GetNexusCompletion(ctx.goContext(), requestID)

chasm/node_backend_mock.go

@@ -22,15 +22,16 @@ type MockNodeBackend struct {
 	// Optional function overrides. If nil, methods return zero-values.
 	HandleGetExecutionState           func() *persistencespb.WorkflowExecutionState
 	HandleGetExecutionInfo            func() *persistencespb.WorkflowExecutionInfo
-	HandleGetApproximatePersistedSize func() int
 	HandleGetCurrentVersion           func() int64
 	HandleNextTransitionCount         func() int64
+	HandleGetApproximatePersistedSize func() int
 	HandleCurrentVersionedTransition  func() *persistencespb.VersionedTransition
 	HandleGetWorkflowKey              func() definition.WorkflowKey
 	HandleUpdateWorkflowStateStatus   func(state enumsspb.WorkflowExecutionState, status enumspb.WorkflowExecutionStatus) (bool, error)
 	HandleIsWorkflow                  func() bool
 	HandleGetNexusCompletion          func(ctx context.Context, requestID string) (nexusrpc.CompleteOperationOptions, error)
 	HandleAddHistoryEvent             func(t enumspb.EventType, setAttributes func(*historypb.HistoryEvent)) *historypb.HistoryEvent
+	HandleHasAnyBufferedEvent         func(filter func(*historypb.HistoryEvent) bool) bool
 	HandleGetNamespaceEntry           func() *namespace.Namespace
 	HandleEndpointRegistry            func() EndpointRegistry
 
@@ -191,6 +192,13 @@ func (m *MockNodeBackend) AddHistoryEvent(t enumspb.EventType, setAttributes fun
 	return nil
 }
 
+func (m *MockNodeBackend) HasAnyBufferedEvent(filter func(*historypb.HistoryEvent) bool) bool {
+	if m.HandleHasAnyBufferedEvent != nil {
+		return m.HandleHasAnyBufferedEvent(filter)
+	}
+	return false
+}
+
 func (m *MockNodeBackend) GetNamespaceEntry() *namespace.Namespace {
 	if m.HandleGetNamespaceEntry != nil {
 		return m.HandleGetNamespaceEntry()

chasm/tree.go

@@ -207,6 +207,7 @@ type (
 		GetWorkflowKey() definition.WorkflowKey
 		AddTasks(...tasks.Task)
 		AddHistoryEvent(t enumspb.EventType, setAttributes func(*historypb.HistoryEvent)) *historypb.HistoryEvent
+		HasAnyBufferedEvent(filter func(*historypb.HistoryEvent) bool) bool
 		DeleteCHASMPureTasks(maxScheduledTime time.Time)
 		UpdateWorkflowStateStatus(
 			state enumsspb.WorkflowExecutionState,

service/history/api/respondworkflowtaskcompleted/workflow_task_completed_handler.go

@@ -20,8 +20,6 @@ import (
 	"go.temporal.io/api/workflowservice/v1"
 	"go.temporal.io/server/api/historyservice/v1"
 	"go.temporal.io/server/api/matchingservice/v1"
-	"go.temporal.io/server/chasm"
-	chasmworkflow "go.temporal.io/server/chasm/lib/workflow"
 	chasmcommand "go.temporal.io/server/chasm/lib/workflow/command"
 	"go.temporal.io/server/common"
 	"go.temporal.io/server/common/backoff"
@@ -326,43 +324,33 @@ func (handler *workflowTaskCompletedHandler) handleCommand(
 		return nil, handler.handleCommandProtocolMessage(ctx, command.GetProtocolMessageCommandAttributes(), msgs)
 
 	default:
-		var commandHandler chasmcommand.Handler
-		var chasmCtx chasm.MutableContext
-		var chasmWorkflow *chasmworkflow.Workflow
-
 		// TODO: need to handle migration between HSM and CHASM
 
 		handlerOpts := chasmcommand.HandlerOptions{
 			WorkflowTaskCompletedEventID: handler.workflowTaskCompletedID,
 		}
-		if handler.mutableState.ChasmEnabled() {
-			// Use CHASM command handler.
-			chasmWorkflow, chasmCtx, err = handler.mutableState.ChasmWorkflowComponent(ctx)
-			if err != nil {
-				return nil, err
-			}
+		validator := commandValidator{sizeChecker: handler.sizeLimitChecker, commandType: command.GetCommandType()}
 
-			var ok bool
-			commandHandler, ok = handler.chasmCommandRegistry.Handler(command.GetCommandType())
-			if !ok {
-				return nil, serviceerror.NewInvalidArgumentf("Unknown command type: %v", command.GetCommandType())
-			}
-		} else {
-			// Use HSM command handler.
-			legacyHandler, ok := handler.commandHandlerRegistry.Handler(command.GetCommandType())
-			if !ok {
-				return nil, serviceerror.NewInvalidArgumentf("Unknown command type: %v", command.GetCommandType())
-			}
-
-			// Wrap HSM command handler to match CHASM command handler signature.
-			commandHandler = func(_ chasm.MutableContext, _ *chasmworkflow.Workflow, v chasmcommand.Validator, cmd *commandpb.Command, opts chasmcommand.HandlerOptions) error {
-				return legacyHandler(ctx, handler.mutableState, v, opts.WorkflowTaskCompletedEventID, cmd)
+		// Try CHASM command handler first, fall back to HSM if not supported.
+		handledByCHASM := false
+		if handler.config.ChasmEnabled(handler.mutableState.GetNamespaceEntry().Name().String()) {
+			if chasmHandler, ok := handler.chasmCommandRegistry.Handler(command.GetCommandType()); ok {
+				chasmWorkflow, chasmCtx, chasmErr := handler.mutableState.ChasmWorkflowComponent(ctx)
+				if chasmErr != nil {
+					return nil, chasmErr
+				}
+				err = chasmHandler(chasmCtx, chasmWorkflow, validator, command, handlerOpts)
+				handledByCHASM = !errors.Is(err, chasmcommand.ErrNotSupported)
 			}
 		}
+		if !handledByCHASM {
+			hsmHandler, ok := handler.commandHandlerRegistry.Handler(command.GetCommandType())
+			if !ok {
+				return nil, serviceerror.NewInvalidArgumentf("Unknown command type: %v", command.GetCommandType())
+			}
+			err = hsmHandler(ctx, handler.mutableState, validator, handlerOpts.WorkflowTaskCompletedEventID, command)
+		}
 
-		// Invoke command handler.
-		validator := commandValidator{sizeChecker: handler.sizeLimitChecker, commandType: command.GetCommandType()}
-		err = commandHandler(chasmCtx, chasmWorkflow, validator, command, handlerOpts)
 		var failWFTErr chasmcommand.FailWorkflowTaskError
 		if errors.As(err, &failWFTErr) {
 			if failWFTErr.TerminateWorkflow {

service/history/api/respondworkflowtaskcompleted/workflow_task_completed_handler_test.go

@@ -79,17 +79,20 @@ func TestCommandProtocolMessage(t *testing.T) {
 		out.ms.EXPECT().GetNamespaceEntry().Return(tests.LocalNamespaceEntry).AnyTimes()
 		out.ms.EXPECT().GetCurrentVersion().Return(tests.LocalNamespaceEntry.FailoverVersion(tests.WorkflowID)).AnyTimes()
 
+		dcClient := dynamicconfig.StaticClient(nil)
 		if opts.chasmEnabled {
-			out.ms.EXPECT().ChasmEnabled().Return(true)
 			out.chasmCommandRegistry = chasmcommand.NewRegistry()
 			mockCtx := &chasm.MockMutableContext{}
 			wf := chasmworkflow.NewWorkflow(mockCtx, chasm.MSPointer{})
 			out.ms.EXPECT().ChasmWorkflowComponent(gomock.Any()).Return(wf, mockCtx, nil)
+			dcClient = dynamicconfig.StaticClient(map[dynamicconfig.Key]any{
+				dynamicconfig.EnableChasm.Key(): true,
+			})
 		}
 
 		out.updates = update.NewRegistry(out.ms)
 		var effects effect.Buffer
-		col := dynamicconfig.NewCollection(dynamicconfig.StaticClient(nil), logger)
+		col := dynamicconfig.NewCollection(dcClient, logger)
 		config := configs.NewConfig(col, 1)
 		mockMeta := persistence.NewMockMetadataManager(ctrl)
 		nsReg := nsregistry.NewRegistry(

service/history/configs/config.go

@@ -59,6 +59,7 @@ type Config struct {
 	HistoryCacheTTL                       dynamicconfig.DurationPropertyFn
 	HistoryCacheNonUserContextLockTimeout dynamicconfig.DurationPropertyFn
 	HistoryCacheBackgroundEvict           dynamicconfig.TypedPropertyFn[dynamicconfig.CacheBackgroundEvictSettings]
+	ChasmEnabled                          dynamicconfig.BoolPropertyFnWithNamespaceFilter
 	EnableWorkflowExecutionTimeoutTimer   dynamicconfig.BoolPropertyFn
 	EnableUpdateWorkflowModeIgnoreCurrent dynamicconfig.BoolPropertyFn
 	EnableTransitionHistory               dynamicconfig.BoolPropertyFnWithNamespaceFilter
@@ -474,6 +475,7 @@ func NewConfig(
 		HistoryCacheTTL:                       dynamicconfig.HistoryCacheTTL.Get(dc),
 		HistoryCacheNonUserContextLockTimeout: dynamicconfig.HistoryCacheNonUserContextLockTimeout.Get(dc),
 		HistoryCacheBackgroundEvict:           dynamicconfig.HistoryCacheBackgroundEvict.Get(dc),
+		ChasmEnabled:                          dynamicconfig.EnableChasm.Get(dc),
 		EnableWorkflowExecutionTimeoutTimer:   dynamicconfig.EnableWorkflowExecutionTimeoutTimer.Get(dc),
 		EnableUpdateWorkflowModeIgnoreCurrent: dynamicconfig.EnableUpdateWorkflowModeIgnoreCurrent.Get(dc),
 		EnableTransitionHistory:               dynamicconfig.EnableTransitionHistory.Get(dc),

service/history/historybuilder/history_builder.go

@@ -53,7 +53,7 @@ type (
 
 	TaskIDGenerator func(number int) ([]int64, error)
 
-	BufferedEventFilter func(*historypb.HistoryEvent) bool
+	BufferedEventFilter = func(*historypb.HistoryEvent) bool
 )
 
 func New(