fix(dispatcher): state fixes + better wizard interface

packages/dispatcher/src/state/key.ts

@@ -1,4 +1,6 @@
-import { assertNever, CallbackQuery, MaybeAsync, Message } from '@mtcute/client'
+import { assertNever, MaybeAsync } from '@mtcute/client'
+
+import { CallbackQueryContext, MessageContext } from '../context/index.js'
 
 /**
  * Function that determines how the state key is derived.
@@ -8,7 +10,7 @@ import { assertNever, CallbackQuery, MaybeAsync, Message } from '@mtcute/client'
  * @param msg  Message or callback from which to derive the key
  * @param scene  Current scene UID, or `null` if none
  */
-export type StateKeyDelegate = (upd: Message | CallbackQuery) => MaybeAsync<string | null>
+export type StateKeyDelegate = (upd: MessageContext | CallbackQueryContext) => MaybeAsync<string | null>
 
 /**
  * Default state key delegate.
@@ -22,7 +24,7 @@ export type StateKeyDelegate = (upd: Message | CallbackQuery) => MaybeAsync<stri
  *    - If in group/channel/supergroup (i.e. `upd.chatType !== 'user'`), `upd.chatId + '_' + upd.user.id`
  */
 export const defaultStateKeyDelegate: StateKeyDelegate = (upd): string | null => {
-    if (upd.constructor === Message) {
+    if (upd._name === 'new_message') {
         switch (upd.chat.chatType) {
             case 'private':
             case 'bot':
@@ -37,7 +39,7 @@ export const defaultStateKeyDelegate: StateKeyDelegate = (upd): string | null =>
         }
     }
 
-    if (upd.constructor === CallbackQuery) {
+    if (upd._name === 'callback_query') {
         if (upd.isInline) return null
         if (upd.chatType === 'user') return `${upd.user.id}`
 

packages/dispatcher/src/wizard.ts

@@ -24,7 +24,7 @@ export enum WizardSceneAction {
 }
 
 interface WizardInternalState {
-    $step: number
+    $step?: number
 }
 
 /**
@@ -51,34 +51,68 @@ export class WizardScene<State, SceneName extends string = string> extends Dispa
         return this._steps
     }
 
+    /**
+     * Go to the Nth step
+     */
+    async goToStep(state: UpdateState<WizardInternalState, SceneName>, step: number) {
+        if (step >= this._steps) {
+            await state.exit()
+        } else {
+            await state.merge({ $step: step }, this._defaultState)
+        }
+    }
+
+    /**
+     * Skip N steps
+     */
+    async skip(state: UpdateState<WizardInternalState, SceneName>, count = 1) {
+        const { $step } = (await state.get()) || {}
+        if ($step === undefined) throw new Error('Wizard state is not initialized')
+
+        return this.goToStep(state, $step + count)
+    }
+
+    /**
+     * Filter that will only pass if the current step is `step`
+     */
+    static onNthStep(step: number) {
+        const filter = filters.state<WizardInternalState>((it) => it.$step === step)
+
+        if (step === 0) return filters.or(filters.stateEmpty, filter)
+
+        return filter
+    }
+
+    /**
+     * Filter that will only pass if the current step is the one after last one added
+     */
+    onCurrentStep() {
+        return WizardScene.onNthStep(this._steps)
+    }
+
     /**
      * Add a step to the wizard
      */
     addStep(
-        handler: (msg: MessageContext, state: UpdateState<State, SceneName>) => MaybeAsync<WizardSceneAction | number>,
+        handler: (
+            msg: MessageContext,
+            state: UpdateState<State & WizardInternalState, SceneName>,
+        ) => MaybeAsync<WizardSceneAction | number>,
     ): void {
         const step = this._steps++
 
-        const filter = filters.state<WizardInternalState>((it) => it.$step === step)
+        this.onNewMessage(WizardScene.onNthStep(step), async (msg, state) => {
-
-        this.onNewMessage(step === 0 ? filters.or(filters.stateEmpty, filter) : filter, async (msg, state) => {
             const result = await handler(msg, state)
 
             if (typeof result === 'number') {
-                await state.merge({ $step: result }, this._defaultState)
+                await this.goToStep(state, result)
 
                 return
             }
 
             switch (result) {
                 case 'next': {
-                    const next = step + 1
+                    await this.goToStep(state, step + 1)
-
-                    if (next === this._steps) {
-                        await state.exit()
-                    } else {
-                        await state.merge({ $step: next }, this._defaultState)
-                    }
                     break
                 }
                 case 'exit':