feat: add message handler middleware and fallback

packages/core/src/agent/Dispatcher.ts

@@ -1,16 +1,19 @@
 import type { AgentMessage } from './AgentMessage'
 import type { AgentMessageProcessedEvent } from './Events'
+import type { MessageHandlerMiddleware } from './MessageHandlerMiddleware'
 import type { InboundMessageContext } from './models/InboundMessageContext'
 
 import { InjectionSymbols } from '../constants'
-import { CredoError } from '../error/CredoError'
+import { CredoError } from '../error'
 import { Logger } from '../logger'
+import { ProblemReportError, ProblemReportReason } from '../modules/problem-reports'
 import { injectable, inject } from '../plugins'
-import { parseMessageType } from '../utils/messageType'
+import { canHandleMessageType, parseMessageType } from '../utils/messageType'
 
 import { ProblemReportMessage } from './../modules/problem-reports/messages/ProblemReportMessage'
 import { EventEmitter } from './EventEmitter'
 import { AgentEventTypes } from './Events'
+import { MessageHandlerMiddlewareRunner } from './MessageHandlerMiddleware'
 import { MessageHandlerRegistry } from './MessageHandlerRegistry'
 import { MessageSender } from './MessageSender'
 import { OutboundMessageContext } from './models'
@@ -34,22 +37,58 @@ class Dispatcher {
     this.logger = logger
   }
 
+  private defaultHandlerMiddleware: MessageHandlerMiddleware = async (inboundMessageContext, next) => {
+    let messageHandler = inboundMessageContext.messageHandler
+
+    if (!messageHandler && inboundMessageContext.agentContext.dependencyManager.fallbackMessageHandler) {
+      messageHandler = {
+        supportedMessages: [],
+        handle: inboundMessageContext.agentContext.dependencyManager.fallbackMessageHandler,
+      }
+    }
+
+    if (!messageHandler) {
+      throw new ProblemReportError(
+        `Error handling message ${inboundMessageContext.message.id} with type ${inboundMessageContext.message.type}. The message type is not supported`,
+        {
+          problemCode: ProblemReportReason.MessageParseFailure,
+        }
+      )
+    }
+
+    const outboundMessage = await messageHandler.handle(inboundMessageContext)
+    if (outboundMessage) {
+      inboundMessageContext.setResponseMessage(outboundMessage)
+    }
+
+    await next()
+  }
+
   public async dispatch(messageContext: InboundMessageContext): Promise<void> {
     const { agentContext, connection, senderKey, recipientKey, message } = messageContext
-    const messageHandler = this.messageHandlerRegistry.getHandlerForMessageType(message.type)
 
-    if (!messageHandler) {
-      throw new CredoError(`No handler for message type "${message.type}" found`)
+    // Set default handler if available, middleware can still override the message handler
+    const messageHandler = this.messageHandlerRegistry.getHandlerForMessageType(message.type)
+    if (messageHandler) {
+      messageContext.setMessageHandler(messageHandler)
     }
 
-    let outboundMessage: OutboundMessageContext<AgentMessage> | void
+    let outboundMessage: OutboundMessageContext<AgentMessage> | undefined
 
     try {
-      outboundMessage = await messageHandler.handle(messageContext)
+      const middlewares = [...agentContext.dependencyManager.messageHandlerMiddlewares, this.defaultHandlerMiddleware]
+      await MessageHandlerMiddlewareRunner.run(middlewares, messageContext)
+
+      outboundMessage = messageContext.responseMessage
     } catch (error) {
       const problemReportMessage = error.problemReport
 
       if (problemReportMessage instanceof ProblemReportMessage && messageContext.connection) {
+        const messageType = parseMessageType(messageContext.message.type)
+        if (canHandleMessageType(ProblemReportMessage, messageType)) {
+          throw new CredoError(`Not sending problem report in response to problem report: ${message}`)
+        }
+
         const { protocolUri: problemReportProtocolUri } = parseMessageType(problemReportMessage.type)
         const { protocolUri: inboundProtocolUri } = parseMessageType(messageContext.message.type)
 
@@ -91,6 +130,7 @@ class Dispatcher {
 
       await this.messageSender.sendMessage(outboundMessage)
     }
+
     // Emit event that allows to hook into received messages
     this.eventEmitter.emit<AgentMessageProcessedEvent>(agentContext, {
       type: AgentEventTypes.AgentMessageProcessed,

packages/core/src/agent/MessageHandlerMiddleware.ts

@@ -0,0 +1,26 @@
+import type { InboundMessageContext } from './models/InboundMessageContext'
+
+export interface MessageHandlerMiddleware {
+  (inboundMessageContext: InboundMessageContext, next: () => Promise<void>): Promise<void>
+}
+
+export class MessageHandlerMiddlewareRunner {
+  public static async run(middlewares: MessageHandlerMiddleware[], inboundMessageContext: InboundMessageContext) {
+    const compose = (middlewares: MessageHandlerMiddleware[]) => {
+      return async function (inboundMessageContext: InboundMessageContext) {
+        let index = -1
+        async function dispatch(i: number): Promise<void> {
+          if (i <= index) throw new Error('next() called multiple times')
+          index = i
+          const fn = middlewares[i]
+          if (!fn) return
+          await fn(inboundMessageContext, () => dispatch(i + 1))
+        }
+        await dispatch(0)
+      }
+    }
+
+    const composed = compose(middlewares)
+    await composed(inboundMessageContext)
+  }
+}

packages/core/src/agent/MessageReceiver.ts

@@ -1,4 +1,3 @@
-import type { AgentMessage } from './AgentMessage'
 import type { DecryptedMessageContext } from './EnvelopeService'
 import type { TransportSession } from './TransportService'
 import type { AgentContext } from './context'
@@ -16,6 +15,7 @@ import { isValidJweStructure } from '../utils/JWE'
 import { JsonTransformer } from '../utils/JsonTransformer'
 import { canHandleMessageType, parseMessageType, replaceLegacyDidSovPrefixOnMessage } from '../utils/messageType'
 
+import { AgentMessage } from './AgentMessage'
 import { Dispatcher } from './Dispatcher'
 import { EnvelopeService } from './EnvelopeService'
 import { MessageHandlerRegistry } from './MessageHandlerRegistry'
@@ -250,13 +250,7 @@ export class MessageReceiver {
     replaceLegacyDidSovPrefixOnMessage(message)
 
     const messageType = message['@type']
-    const MessageClass = this.messageHandlerRegistry.getMessageClassForMessageType(messageType)
-
-    if (!MessageClass) {
-      throw new ProblemReportError(`No message class found for message type "${messageType}"`, {
-        problemCode: ProblemReportReason.MessageParseFailure,
-      })
-    }
+    const MessageClass = this.messageHandlerRegistry.getMessageClassForMessageType(messageType) ?? AgentMessage
 
     // Cast the plain JSON object to specific instance of Message extended from AgentMessage
     let messageTransformed: AgentMessage