feat: add agent context provider (#921)

Signed-off-by: Timo Glastra <[email protected]>

packages/core/src/agent/Agent.ts

@@ -42,14 +42,14 @@ import { WalletModule } from '../wallet/WalletModule'
 import { WalletError } from '../wallet/error'
 
 import { AgentConfig } from './AgentConfig'
-import { AgentContext } from './AgentContext'
 import { Dispatcher } from './Dispatcher'
 import { EnvelopeService } from './EnvelopeService'
 import { EventEmitter } from './EventEmitter'
 import { AgentEventTypes } from './Events'
 import { MessageReceiver } from './MessageReceiver'
 import { MessageSender } from './MessageSender'
 import { TransportService } from './TransportService'
+import { AgentContext, DefaultAgentContextProvider } from './context'
 
 export class Agent {
   protected agentConfig: AgentConfig
@@ -138,8 +138,9 @@ export class Agent {
       .pipe(
         takeUntil(this.stop$),
         concatMap((e) =>
-          this.messageReceiver.receiveMessage(this.agentContext, e.payload.message, {
+          this.messageReceiver.receiveMessage(e.payload.message, {
             connection: e.payload.connection,
+            contextCorrelationId: e.payload.contextCorrelationId,
           })
         )
       )
@@ -269,8 +270,18 @@ export class Agent {
     return this.agentContext.wallet.publicDid
   }
 
+  /**
+   * Receive a message. This should mainly be used for receiving connection-less messages.
+   *
+   * If you want to receive messages that originated from e.g. a transport make sure to use the {@link MessageReceiver}
+   * for this. The `receiveMessage` method on the `Agent` class will associate the current context to the message, which
+   * may not be what should happen (e.g. in case of multi tenancy).
+   */
   public async receiveMessage(inboundMessage: unknown, session?: TransportSession) {
-    return await this.messageReceiver.receiveMessage(this.agentContext, inboundMessage, { session })
+    return await this.messageReceiver.receiveMessage(inboundMessage, {
+      session,
+      contextCorrelationId: this.agentContext.contextCorrelationId,
+    })
   }
 
   public get injectionContainer() {
@@ -368,6 +379,16 @@ export class Agent {
       W3cVcModule
     )
 
-    dependencyManager.registerInstance(AgentContext, new AgentContext({ dependencyManager }))
+    // TODO: contextCorrelationId for base wallet
+    // Bind the default agent context to the container for use in modules etc.
+    dependencyManager.registerInstance(
+      AgentContext,
+      new AgentContext({ dependencyManager, contextCorrelationId: 'default' })
+    )
+
+    // If no agent context provider has been registered we use the default agent context provider.
+    if (!this.dependencyManager.isRegistered(InjectionSymbols.AgentContextProvider)) {
+      this.dependencyManager.registerSingleton(InjectionSymbols.AgentContextProvider, DefaultAgentContextProvider)
+    }
   }
 }

packages/core/src/agent/EnvelopeService.ts

@@ -1,6 +1,6 @@
 import type { EncryptedMessage, PlaintextMessage } from '../types'
-import type { AgentContext } from './AgentContext'
 import type { AgentMessage } from './AgentMessage'
+import type { AgentContext } from './context'
 
 import { InjectionSymbols } from '../constants'
 import { Key, KeyType } from '../crypto'

packages/core/src/agent/EventEmitter.ts

@@ -1,5 +1,5 @@
-import type { AgentContext } from './AgentContext'
 import type { BaseEvent } from './Events'
+import type { AgentContext } from './context'
 import type { EventEmitter as NativeEventEmitter } from 'events'
 
 import { fromEventPattern, Subject } from 'rxjs'

packages/core/src/agent/Events.ts

@@ -16,6 +16,7 @@ export interface AgentMessageReceivedEvent extends BaseEvent {
   payload: {
     message: unknown
     connection?: ConnectionRecord
+    contextCorrelationId?: string
   }
 }
 

packages/core/src/agent/MessageReceiver.ts

@@ -1,17 +1,17 @@
 import type { ConnectionRecord } from '../modules/connections'
 import type { InboundTransport } from '../transport'
 import type { EncryptedMessage, PlaintextMessage } from '../types'
-import type { AgentContext } from './AgentContext'
 import type { AgentMessage } from './AgentMessage'
 import type { DecryptedMessageContext } from './EnvelopeService'
 import type { TransportSession } from './TransportService'
+import type { AgentContext } from './context'
 
 import { InjectionSymbols } from '../constants'
 import { AriesFrameworkError } from '../error'
 import { Logger } from '../logger'
 import { ConnectionService } from '../modules/connections'
 import { ProblemReportError, ProblemReportMessage, ProblemReportReason } from '../modules/problem-reports'
-import { injectable, inject } from '../plugins'
+import { inject, injectable } from '../plugins'
 import { isValidJweStructure } from '../utils/JWE'
 import { JsonTransformer } from '../utils/JsonTransformer'
 import { canHandleMessageType, parseMessageType, replaceLegacyDidSovPrefixOnMessage } from '../utils/messageType'
@@ -20,6 +20,7 @@ import { Dispatcher } from './Dispatcher'
 import { EnvelopeService } from './EnvelopeService'
 import { MessageSender } from './MessageSender'
 import { TransportService } from './TransportService'
+import { AgentContextProvider } from './context'
 import { createOutboundMessage } from './helpers'
 import { InboundMessageContext } from './models/InboundMessageContext'
 
@@ -31,6 +32,7 @@ export class MessageReceiver {
   private dispatcher: Dispatcher
   private logger: Logger
   private connectionService: ConnectionService
+  private agentContextProvider: AgentContextProvider
   public readonly inboundTransports: InboundTransport[] = []
 
   public constructor(
@@ -39,13 +41,15 @@ export class MessageReceiver {
     messageSender: MessageSender,
     connectionService: ConnectionService,
     dispatcher: Dispatcher,
+    @inject(InjectionSymbols.AgentContextProvider) agentContextProvider: AgentContextProvider,
     @inject(InjectionSymbols.Logger) logger: Logger
   ) {
     this.envelopeService = envelopeService
     this.transportService = transportService
     this.messageSender = messageSender
     this.connectionService = connectionService
     this.dispatcher = dispatcher
+    this.agentContextProvider = agentContextProvider
     this.logger = logger
   }
 
@@ -54,17 +58,26 @@ export class MessageReceiver {
   }
 
   /**
-   * Receive and handle an inbound DIDComm message. It will decrypt the message, transform it
+   * Receive and handle an inbound DIDComm message. It will determine the agent context, decrypt the message, transform it
    * to it's corresponding message class and finally dispatch it to the dispatcher.
    *
    * @param inboundMessage the message to receive and handle
    */
   public async receiveMessage(
-    agentContext: AgentContext,
     inboundMessage: unknown,
-    { session, connection }: { session?: TransportSession; connection?: ConnectionRecord }
+    {
+      session,
+      connection,
+      contextCorrelationId,
+    }: { session?: TransportSession; connection?: ConnectionRecord; contextCorrelationId?: string } = {}
   ) {
-    this.logger.debug(`Agent ${agentContext.config.label} received message`)
+    this.logger.debug(`Agent received message`)
+
+    // Find agent context for the inbound message
+    const agentContext = await this.agentContextProvider.getContextForInboundMessage(inboundMessage, {
+      contextCorrelationId,
+    })
+
     if (this.isEncryptedMessage(inboundMessage)) {
       await this.receiveEncryptedMessage(agentContext, inboundMessage as EncryptedMessage, session)
     } else if (this.isPlaintextMessage(inboundMessage)) {

packages/core/src/agent/MessageSender.ts

@@ -4,10 +4,10 @@ import type { DidDocument } from '../modules/dids'
 import type { OutOfBandRecord } from '../modules/oob/repository'
 import type { OutboundTransport } from '../transport/OutboundTransport'
 import type { OutboundMessage, OutboundPackage, EncryptedMessage } from '../types'
-import type { AgentContext } from './AgentContext'
 import type { AgentMessage } from './AgentMessage'
 import type { EnvelopeKeys } from './EnvelopeService'
 import type { TransportSession } from './TransportService'
+import type { AgentContext } from './context'
 
 import { DID_COMM_TRANSPORT_QUEUE, InjectionSymbols } from '../constants'
 import { ReturnRouteTypes } from '../decorators/transport/TransportDecorator'

packages/core/src/agent/context/AgentContext.ts

@@ -0,0 +1,49 @@
+import type { DependencyManager } from '../../plugins'
+import type { Wallet } from '../../wallet'
+
+import { InjectionSymbols } from '../../constants'
+import { AgentConfig } from '../AgentConfig'
+
+export class AgentContext {
+  /**
+   * Dependency manager holds all dependencies for the current context. Possibly a child of a parent dependency manager,
+   * in which case all singleton dependencies from the parent context are also available to this context.
+   */
+  public readonly dependencyManager: DependencyManager
+
+  /**
+   * An identifier that allows to correlate this context across sessions. This identifier is created by the `AgentContextProvider`
+   * and should only be meaningful to the `AgentContextProvider`. The `contextCorrelationId` MUST uniquely identity the context and
+   * should be enough to start a new session.
+   *
+   * An example of the `contextCorrelationId` is for example the id of the `TenantRecord` that is associated with this context when using the tenant module.
+   * The `TenantAgentContextProvider` will set the `contextCorrelationId` to the `TenantRecord` id when creating the context, and will be able to create a context
+   * for a specific tenant using the `contextCorrelationId`.
+   */
+  public readonly contextCorrelationId: string
+
+  public constructor({
+    dependencyManager,
+    contextCorrelationId,
+  }: {
+    dependencyManager: DependencyManager
+    contextCorrelationId: string
+  }) {
+    this.dependencyManager = dependencyManager
+    this.contextCorrelationId = contextCorrelationId
+  }
+
+  /**
+   * Convenience method to access the agent config for the current context.
+   */
+  public get config() {
+    return this.dependencyManager.resolve(AgentConfig)
+  }
+
+  /**
+   * Convenience method to access the wallet for the current context.
+   */
+  public get wallet() {
+    return this.dependencyManager.resolve<Wallet>(InjectionSymbols.Wallet)
+  }
+}

packages/core/src/agent/context/AgentContextProvider.ts

@@ -0,0 +1,17 @@
+import type { AgentContext } from './AgentContext'
+
+export interface AgentContextProvider {
+  /**
+   * Find the agent context based for an inbound message. It's possible to provide a contextCorrelationId to make it
+   * easier for the context provider implementation to correlate inbound messages to the correct context. This can be useful if
+   * a plaintext message is passed and the context provider can't determine the context based on the recipient public keys
+   * of the inbound message.
+   *
+   * The implementation of this method could range from a very simple one that always returns the same context to
+   * a complex one that manages the context for a multi-tenant agent.
+   */
+  getContextForInboundMessage(
+    inboundMessage: unknown,
+    options?: { contextCorrelationId?: string }
+  ): Promise<AgentContext>
+}

packages/core/src/agent/context/DefaultAgentContextProvider.ts

@@ -0,0 +1,24 @@
+import type { AgentContextProvider } from './AgentContextProvider'
+
+import { injectable } from '../../plugins'
+
+import { AgentContext } from './AgentContext'
+
+/**
+ * Default implementation of AgentContextProvider.
+ *
+ * Holds a single `AgentContext` instance that will be used for all messages, i.e. a
+ * a single tenant agent.
+ */
+@injectable()
+export class DefaultAgentContextProvider implements AgentContextProvider {
+  private agentContext: AgentContext
+
+  public constructor(agentContext: AgentContext) {
+    this.agentContext = agentContext
+  }
+
+  public async getContextForInboundMessage(): Promise<AgentContext> {
+    return this.agentContext
+  }
+}

packages/core/src/agent/context/__tests__/DefaultAgentContextProvider.test.ts

@@ -0,0 +1,18 @@
+import type { AgentContextProvider } from '../AgentContextProvider'
+
+import { getAgentContext } from '../../../../tests/helpers'
+import { DefaultAgentContextProvider } from '../DefaultAgentContextProvider'
+
+const agentContext = getAgentContext()
+
+describe('DefaultAgentContextProvider', () => {
+  describe('getContextForInboundMessage()', () => {
+    test('returns the agent context provided in the constructor', async () => {
+      const agentContextProvider: AgentContextProvider = new DefaultAgentContextProvider(agentContext)
+
+      const message = {}
+
+      await expect(agentContextProvider.getContextForInboundMessage(message)).resolves.toBe(agentContext)
+    })
+  })
+})

packages/core/src/agent/context/index.ts

@@ -0,0 +1,3 @@
+export * from './AgentContext'
+export * from './AgentContextProvider'
+export * from './DefaultAgentContextProvider'

packages/core/src/agent/index.ts

@@ -1 +1 @@
-export * from './AgentContext'
+export * from './context'

packages/core/src/agent/models/InboundMessageContext.ts

@@ -1,7 +1,7 @@
 import type { Key } from '../../crypto'
 import type { ConnectionRecord } from '../../modules/connections'
-import type { AgentContext } from '../AgentContext'
 import type { AgentMessage } from '../AgentMessage'
+import type { AgentContext } from '../context'
 
 import { AriesFrameworkError } from '../../error'
 

packages/core/src/constants.ts

@@ -2,6 +2,7 @@ export const InjectionSymbols = {
   MessageRepository: Symbol('MessageRepository'),
   StorageService: Symbol('StorageService'),
   Logger: Symbol('Logger'),
+  AgentContextProvider: Symbol('AgentContextProvider'),
   AgentDependencies: Symbol('AgentDependencies'),
   Stop$: Symbol('Stop$'),
   FileSystem: Symbol('FileSystem'),

packages/core/src/index.ts

@@ -1,7 +1,8 @@
 // reflect-metadata used for class-transformer + class-validator
 import 'reflect-metadata'
 
-export { AgentContext } from './agent/AgentContext'
+export { AgentContext } from './agent'
+export { MessageReceiver } from './agent/MessageReceiver'
 export { Agent } from './agent/Agent'
 export { EventEmitter } from './agent/EventEmitter'
 export { Handler, HandlerInboundMessage } from './agent/Handler'

packages/core/src/modules/oob/OutOfBandModule.ts

@@ -625,6 +625,7 @@ export class OutOfBandModule {
       payload: {
         message: plaintextMessage,
         connection: connectionRecord,
+        contextCorrelationId: this.agentContext.contextCorrelationId,
       },
     })
   }
@@ -666,6 +667,7 @@ export class OutOfBandModule {
       type: AgentEventTypes.AgentMessageReceived,
       payload: {
         message: plaintextMessage,
+        contextCorrelationId: this.agentContext.contextCorrelationId,
       },
     })
   }

packages/core/src/modules/proofs/ProofResponseCoordinator.ts

@@ -1,4 +1,4 @@
-import type { AgentContext } from '../../agent/AgentContext'
+import type { AgentContext } from '../../agent/context'
 import type { ProofRecord } from './repository'
 
 import { injectable } from '../../plugins'

packages/core/src/modules/routing/protocol/pickup/v1/handlers/BatchHandler.ts

@@ -24,6 +24,7 @@ export class BatchHandler implements Handler {
         type: AgentEventTypes.AgentMessageReceived,
         payload: {
           message: message.message,
+          contextCorrelationId: messageContext.agentContext.contextCorrelationId,
         },
       })
     })

packages/core/src/modules/routing/services/MediationRecipientService.ts

@@ -300,6 +300,7 @@ export class MediationRecipientService {
         type: AgentEventTypes.AgentMessageReceived,
         payload: {
           message: attachment.getDataAsJson<EncryptedMessage>(),
+          contextCorrelationId: messageContext.agentContext.contextCorrelationId,
         },
       })
     }