diff --git a/packages/dispatcher/src/dispatcher.ts b/packages/dispatcher/src/dispatcher.ts index 6b3406f1..db74a379 100644 --- a/packages/dispatcher/src/dispatcher.ts +++ b/packages/dispatcher/src/dispatcher.ts @@ -27,7 +27,7 @@ import { UserTypingHandler, } from './handler' // end-codegen-imports -import { UpdateInfoForError } from './handler' +import { UpdateInfo } from './handler' import { filters, UpdateFilter } from './filters' import { handlers } from './builders' import { ChatMemberUpdate } from './updates' @@ -165,14 +165,23 @@ export class Dispatcher { private _handlersCount: Record = {} - private _errorHandler?: < - T extends Exclude - >( + private _errorHandler?: ( err: Error, - update: UpdateInfoForError, + update: UpdateInfo & T, state?: UpdateState ) => MaybeAsync + private _preUpdateHandler?: ( + update: UpdateInfo & T, + state?: UpdateState + ) => MaybeAsync + + private _postUpdateHandler?: ( + handled: boolean, + update: UpdateInfo & T, + state?: UpdateState + ) => MaybeAsync + /** * Create a new dispatcher, that will be used as a child, * optionally providing a custom key delegate @@ -284,12 +293,13 @@ export class Dispatcher { * @param update Update to process * @param users Map of users * @param chats Map of chats + * @returns Whether the update was handled */ async dispatchUpdateNow( update: tl.TypeUpdate | tl.TypeMessage, users: UsersIndex, chats: ChatsIndex - ): Promise { + ): Promise { return this._dispatchUpdateNowImpl(update, users, chats) } @@ -303,8 +313,8 @@ export class Dispatcher { parsedState?: UpdateState | null, parsedScene?: string | null, forceScene?: true - ): Promise { - if (!this._client) return + ): Promise { + if (!this._client) return false const isRawMessage = update && tl.isAnyMessage(update) @@ -342,11 +352,11 @@ export class Dispatcher { if (this._scene) { if (this._scene !== parsedScene) // should not happen, but just in case - return + return false } else { if (!this._scenes || !(parsedScene in this._scenes)) // not registered scene - return + return false return this._scenes[parsedScene]._dispatchUpdateNowImpl( update, @@ -392,64 +402,50 @@ export class Dispatcher { } } - outer: for (const grp of this._groupsOrder) { - const group = this._groups[grp] + let shouldDispatch = true + let shouldDispatchChildren = true + let wasHandled = false - if (update && !isRawMessage && 'raw' in group) { - const handlers = group['raw'] as RawUpdateHandler[] + const updateInfo = { type: parsedType, data: parsed } + switch ( + await this._preUpdateHandler?.( + updateInfo as any, + parsedState as any + ) + ) { + case 'stop': + shouldDispatch = false + break + case 'stop-children': + return false + } - for (const h of handlers) { - let result: void | PropagationAction + if (shouldDispatch) { + outer: for (const grp of this._groupsOrder) { + const group = this._groups[grp] - if ( - !h.check || - (await h.check( - this._client, - update as any, - users!, - chats! - )) - ) { - result = await h.callback( - this._client, - update as any, - users!, - chats! - ) - } else continue + if (update && !isRawMessage && 'raw' in group) { + const handlers = group['raw'] as RawUpdateHandler[] - switch (result) { - case 'continue': - continue - case 'stop': - break outer - case 'stop-children': - return - } - - break - } - } - - if (parsedType && parsedType in group) { - // raw is not handled here, so we can safely assume this - const handlers = group[parsedType] as Exclude< - UpdateHandler, - RawUpdateHandler - >[] - - try { for (const h of handlers) { let result: void | PropagationAction if ( !h.check || - (await h.check(parsed, parsedState as never)) + (await h.check( + this._client, + update as any, + users!, + chats! + )) ) { result = await h.callback( - parsed, - parsedState as never + this._client, + update as any, + users!, + chats! ) + wasHandled = true } else continue switch (result) { @@ -458,61 +454,109 @@ export class Dispatcher { case 'stop': break outer case 'stop-children': - return - case 'scene': { - if (!parsedState) - throw new MtCuteArgumentError( - 'Cannot use ToScene without state' - ) - - const scene = parsedState['_scene'] - - if (!scene) - throw new MtCuteArgumentError( - 'Cannot use ToScene without entering a scene' - ) - - return this._scenes[ - scene - ]._dispatchUpdateNowImpl( - update, - users, - chats, - parsed, - parsedType, - undefined, - scene, - true - ) - } + shouldDispatchChildren = false + break outer } break } - } catch (e) { - if (this._errorHandler) { - const handled = await this._errorHandler( - e, - { type: parsedType, data: parsed }, - parsedState as never - ) - if (!handled) throw e - } else { - throw e + } + + if (parsedType && parsedType in group) { + // raw is not handled here, so we can safely assume this + const handlers = group[parsedType] as Exclude< + UpdateHandler, + RawUpdateHandler + >[] + + try { + for (const h of handlers) { + let result: void | PropagationAction + + if ( + !h.check || + (await h.check(parsed, parsedState as never)) + ) { + result = await h.callback( + parsed, + parsedState as never + ) + wasHandled = true + } else continue + + switch (result) { + case 'continue': + continue + case 'stop': + break outer + case 'stop-children': + shouldDispatchChildren = false + break outer + case 'scene': { + if (!parsedState) + throw new MtCuteArgumentError( + 'Cannot use ToScene without state' + ) + + const scene = parsedState['_scene'] + + if (!scene) + throw new MtCuteArgumentError( + 'Cannot use ToScene without entering a scene' + ) + + return this._scenes[ + scene + ]._dispatchUpdateNowImpl( + update, + users, + chats, + parsed, + parsedType, + undefined, + scene, + true + ) + } + } + + break + } + } catch (e) { + if (this._errorHandler) { + const handled = await this._errorHandler( + e, + updateInfo as any, + parsedState as never + ) + if (!handled) throw e + } else { + throw e + } } } } } - for (const child of this._children) { - await child._dispatchUpdateNowImpl( - update, - users, - chats, - parsed, - parsedType - ) + if (shouldDispatchChildren) { + for (const child of this._children) { + wasHandled ||= await child._dispatchUpdateNowImpl( + update, + users, + chats, + parsed, + parsedType + ) + } } + + this._postUpdateHandler?.( + wasHandled, + updateInfo as any, + parsedState as any + ) + + return wasHandled } /** @@ -611,11 +655,11 @@ export class Dispatcher { * * @param handler Error handler */ - onError( + onError( handler: | (( err: Error, - update: UpdateInfoForError, + update: UpdateInfo & T, state?: UpdateState ) => MaybeAsync) | null @@ -624,13 +668,62 @@ export class Dispatcher { else this._errorHandler = undefined } + /** + * Register pre-update middleware. + * + * This is used locally within this dispatcher + * (does not affect children/parent) before processing + * an update, and can be used to skip this update. + * + * There can be at most one pre-update middleware. + * Pass `null` to remove it. + * + * @param handler Pre-update middleware + */ + onPreUpdate( + handler: + | (( + update: UpdateInfo & T, + state?: UpdateState + ) => MaybeAsync) + | null + ): void { + if (handler) this._preUpdateHandler = handler + else this._preUpdateHandler = undefined + } + + /** + * Register post-update middleware. + * + * This is used locally within this dispatcher + * (does not affect children/parent) after successfully + * processing an update, and can be used for stats. + * + * There can be at most one post-update middleware. + * Pass `null` to remove it. + * + * @param handler Pre-update middleware + */ + onPostUpdate( + handler: + | (( + handled: boolean, + update: UpdateInfo & T, + state?: UpdateState + ) => MaybeAsync) + | null + ): void { + if (handler) this._postUpdateHandler = handler + else this._postUpdateHandler = undefined + } + /** * Set error handler that will propagate * the error to the parent dispatcher */ propagateErrorToParent>( err: Error, - update: UpdateInfoForError, + update: UpdateInfo, state?: UpdateState ): MaybeAsync { if (!this.parent) diff --git a/packages/dispatcher/src/handler.ts b/packages/dispatcher/src/handler.ts index 595864a7..40677531 100644 --- a/packages/dispatcher/src/handler.ts +++ b/packages/dispatcher/src/handler.ts @@ -30,13 +30,13 @@ type ParsedUpdateHandler = BaseUpdateHandler< (update: Update, state: State) => MaybeAsync > -export type UpdateInfoForError = T extends ParsedUpdateHandler< +export type UpdateInfo = T extends ParsedUpdateHandler< infer K, infer Q > ? { - type: K - data: Q + readonly type: K + readonly data: Q } : never diff --git a/packages/dispatcher/tests/dispatcher.spec.ts b/packages/dispatcher/tests/dispatcher.spec.ts index cf6b85fc..6f0b7bfe 100644 --- a/packages/dispatcher/tests/dispatcher.spec.ts +++ b/packages/dispatcher/tests/dispatcher.spec.ts @@ -1,11 +1,6 @@ import { describe, it } from 'mocha' import { expect } from 'chai' -import { - ContinuePropagation, - Dispatcher, - handlers, - StopPropagation, -} from '../src' +import { Dispatcher, handlers, PropagationAction } from '../src' import { TelegramClient } from '@mtcute/client' describe('Dispatcher', () => { @@ -21,19 +16,19 @@ describe('Dispatcher', () => { dp.onRawUpdate((cl, upd) => { log.push('(wrap) received ' + upd._) - return ContinuePropagation + return PropagationAction.Continue }) dp.addUpdateHandler( handlers.rawUpdate((cl, upd) => { log.push('(factory) received ' + upd._) - return ContinuePropagation + return PropagationAction.Continue }) ) dp.addUpdateHandler({ type: 'raw', callback: (cl, upd) => { log.push('(raw) received ' + upd._) - return ContinuePropagation + return PropagationAction.Continue }, }) @@ -58,7 +53,7 @@ describe('Dispatcher', () => { dp.onRawUpdate((cl, upd) => { log.push('(no) received ' + upd._) - return ContinuePropagation + return PropagationAction.Continue }) dp.onRawUpdate( @@ -66,7 +61,7 @@ describe('Dispatcher', () => { (cl, upd) => { log.push('(true) received ' + upd._) - return ContinuePropagation + return PropagationAction.Continue } ) @@ -75,7 +70,7 @@ describe('Dispatcher', () => { (cl, upd) => { log.push('(false) received ' + upd._) - return ContinuePropagation + return PropagationAction.Continue } ) @@ -125,7 +120,7 @@ describe('Dispatcher', () => { dp.onRawUpdate((cl, upd) => { log.push('(grp0) received ' + upd._) - return ContinuePropagation + return PropagationAction.Continue }, 0) dp.onRawUpdate((cl, upd) => { log.push('(grp0 2) received ' + upd._) @@ -157,7 +152,7 @@ describe('Dispatcher', () => { dp.onRawUpdate((cl, upd) => { log.push('(grp0) received ' + upd._) - return ContinuePropagation + return PropagationAction.Continue }, 0) dp.onRawUpdate((cl, upd) => { log.push('(grp0 2) received ' + upd._) @@ -165,7 +160,7 @@ describe('Dispatcher', () => { dp.onRawUpdate((cl, upd) => { log.push('(grp1) received ' + upd._) - return StopPropagation + return PropagationAction.Continue }, 1) dp.onRawUpdate((cl, upd) => { log.push('(grp1 2) received ' + upd._) @@ -217,12 +212,12 @@ describe('Dispatcher', () => { dp.onRawUpdate((cl, upd) => { log.push('(parent 0) received ' + upd._) - return ContinuePropagation + return PropagationAction.Continue }, 0) dp.onRawUpdate((cl, upd) => { log.push('(parent 1) received ' + upd._) - return StopPropagation + return PropagationAction.Stop }, 1) dp.onRawUpdate((cl, upd) => { @@ -231,12 +226,12 @@ describe('Dispatcher', () => { child.onRawUpdate((cl, upd) => { log.push('(child 0) received ' + upd._) - return ContinuePropagation + return PropagationAction.Continue }, 0) child.onRawUpdate((cl, upd) => { log.push('(child 1) received ' + upd._) - return StopPropagation + return PropagationAction.Stop }, 1) child.onRawUpdate((cl, upd) => {