improve ModelStage utility type

packages/core/saved-objects/core-saved-objects-migration-server-internal/src/zdt/model/stages/init.ts

@@ -13,7 +13,7 @@ import { isTypeof } from '../../actions';
 import type { State } from '../../state';
 import type { ModelStage } from '../types';
 
-export const init: ModelStage<'INIT'> = (state, res, context): State => {
+export const init: ModelStage<'INIT', 'DONE' | 'FATAL'> = (state, res, context): State => {
   if (Either.isLeft(res)) {
     const left = res.left;
     if (isTypeof(left, 'incompatible_cluster_routing_allocation')) {

packages/core/saved-objects/core-saved-objects-migration-server-internal/src/zdt/model/types.ts

@@ -8,7 +8,12 @@
 
 import { ExcludeRetryableEsError } from '../../model/types';
 import type { MigratorContext } from '../context';
-import type { State, AllActionStates, StateFromActionState } from '../state';
+import type {
+  AllActionStates,
+  AllControlStates,
+  StateFromActionState,
+  StateFromControlState,
+} from '../state';
 import type { ResponseType } from '../next';
 
 /**
@@ -21,8 +26,8 @@ export type StateActionResponse<T extends AllActionStates> = ExcludeRetryableEsE
 /**
  * Defines a stage delegation function for the model
  */
-export type ModelStage<T extends AllActionStates> = (
+export type ModelStage<T extends AllActionStates, R extends AllControlStates> = (
   state: StateFromActionState<T>,
   res: StateActionResponse<T>,
   context: MigratorContext
-) => State;
+) => StateFromControlState<T | R>;

packages/core/saved-objects/core-saved-objects-migration-server-internal/src/zdt/state/index.ts

@@ -15,5 +15,6 @@ export type {
   AllActionStates,
   AllControlStates,
   StateFromActionState,
+  StateFromControlState,
 } from './types';
 export { createInitialState } from './create_initial_state';

packages/core/saved-objects/core-saved-objects-migration-server-internal/src/zdt/state/types.ts

@@ -49,4 +49,9 @@ export interface ControlStateMap {
 /**
  * Utility type to reverse lookup an `AllControlStates` to it's corresponding State subtype.
  */
-export type StateFromActionState<T extends AllActionStates> = ControlStateMap[T];
+export type StateFromControlState<T extends AllControlStates> = ControlStateMap[T];
+
+/**
+ * Utility type to reverse lookup an `AllActionStates` to it's corresponding State subtype.
+ */
+export type StateFromActionState<T extends AllActionStates> = StateFromControlState<T>;