feat: add agent context provider

Signed-off-by: Timo Glastra <timo@animo.id>

packages/core/src/agent/Agent.ts

@@ -3,10 +3,10 @@ import type { InboundTransport } from '../transport/InboundTransport'
 import type { OutboundTransport } from '../transport/OutboundTransport'
 import type { InitConfig } from '../types'
 import type { Wallet } from '../wallet/Wallet'
-import type { AgentContext } from './AgentContext'
 import type { AgentDependencies } from './AgentDependencies'
 import type { AgentMessageReceivedEvent } from './Events'
 import type { TransportSession } from './TransportService'
+import type { AgentContext } from './context'
 import type { Subscription } from 'rxjs'
 import type { DependencyContainer } from 'tsyringe'
 
@@ -38,12 +38,12 @@ import { WalletModule } from '../wallet/WalletModule'
 import { WalletError } from '../wallet/error'
 
 import { AgentConfig } from './AgentConfig'
-import { DefaultAgentContext } from './AgentContext'
 import { EventEmitter } from './EventEmitter'
 import { AgentEventTypes } from './Events'
 import { MessageReceiver } from './MessageReceiver'
 import { MessageSender } from './MessageSender'
 import { TransportService } from './TransportService'
+import { DefaultAgentContextProvider, DefaultAgentContext } from './context'
 
 export class Agent {
   protected agentConfig: AgentConfig
@@ -115,16 +115,23 @@ export class Agent {
       )
     }
 
+    this.walletService = this.container.resolve(IndyWallet)
+
+    // Bind the default agent context to the container for use in modules etc.
+    this.agentContext = new DefaultAgentContext(this.walletService, this.agentConfig, 'default')
+    this.container.registerInstance(InjectionSymbols.AgentContext, this.agentContext)
+
+    // If no agent context has been registered we use the default agent context.
+    if (!this.container.isRegistered(InjectionSymbols.AgentContextProvider)) {
+      const agentContextProvider = new DefaultAgentContextProvider(this.agentContext)
+      this.container.registerInstance(InjectionSymbols.AgentContextProvider, agentContextProvider)
+    }
+
     // Resolve instances after everything is registered
     this.eventEmitter = this.container.resolve(EventEmitter)
     this.messageSender = this.container.resolve(MessageSender)
     this.messageReceiver = this.container.resolve(MessageReceiver)
     this.transportService = this.container.resolve(TransportService)
-    this.walletService = this.container.resolve(IndyWallet)
-
-    // Bind the default agent context to the container for use in modules etc.
-    this.agentContext = new DefaultAgentContext(this.walletService, this.agentConfig)
-    this.container.registerInstance(InjectionSymbols.AgentContext, this.agentContext)
 
     // We set the modules in the constructor because that allows to set them as read-only
     this.connections = this.container.resolve(ConnectionsModule)
@@ -147,8 +154,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,
           })
         )
       )
@@ -278,8 +286,18 @@ export class Agent {
     return this.walletService.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() {

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 { inject, scoped, Lifecycle } from 'tsyringe'
 

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,10 +1,10 @@
 import type { ConnectionRecord } from '../modules/connections'
 import type { InboundTransport } from '../transport'
 import type { PlaintextMessage, EncryptedMessage } 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 { inject, Lifecycle, scoped } from 'tsyringe'
 
@@ -21,6 +21,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'
 
@@ -32,6 +33,7 @@ export class MessageReceiver {
   private dispatcher: Dispatcher
   private logger: Logger
   private connectionService: ConnectionService
+  private agentContextProvider: AgentContextProvider
   public readonly inboundTransports: InboundTransport[] = []
 
   public constructor(
@@ -40,13 +42,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
   }
 
@@ -55,17 +59,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

@@ -3,10 +3,10 @@ import type { DidDocument, Key } 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 { inject, Lifecycle, scoped } from 'tsyringe'
 

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

@@ -0,0 +1,14 @@
+import type { Wallet } from '../../wallet'
+import type { AgentConfig } from '../AgentConfig'
+
+export interface AgentContext {
+  readonly wallet: Wallet
+  readonly config: AgentConfig
+
+  /**
+   * An identifier that allows to correlate this context across usages. An example of the contextCorrelationId could be
+   * the id of the `TenantRecord` that is associated with this context. The AgentContextProvider can use this identifier to
+   * correlate an inbound message to a specific context (if the message is not encrypted, it's impossible to correlate it to a tenant)
+   */
+  readonly contextCorrelationId: string
+}

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/DefaultAgentContext.ts

@@ -0,0 +1,15 @@
+import type { Wallet } from '../../wallet'
+import type { AgentConfig } from '../AgentConfig'
+import type { AgentContext } from './AgentContext'
+
+export class DefaultAgentContext implements AgentContext {
+  public readonly wallet: Wallet
+  public readonly config: AgentConfig
+  public readonly contextCorrelationId: string
+
+  public constructor(wallet: Wallet, config: AgentConfig, contextCorrelationId: string) {
+    this.wallet = wallet
+    this.config = config
+    this.contextCorrelationId = contextCorrelationId
+  }
+}

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

@@ -0,0 +1,20 @@
+import type { AgentContext } from './AgentContext'
+import type { AgentContextProvider } from './AgentContextProvider'
+
+/**
+ * Default implementation of AgentContextProvider.
+ *
+ * Holds a single `AgentContext` instance that will be used for all messages, i.e. a
+ * a single tenant agent.
+ */
+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,20 @@
+import type { AgentContextProvider } from '../AgentContextProvider'
+
+import { getAgentConfig } from '../../../../tests/helpers'
+import { MockAgentContext } from '../../../../tests/mocks'
+import { DefaultAgentContextProvider } from '../DefaultAgentContextProvider'
+
+const agentConfig = getAgentConfig('DefaultAgentContextProvider')
+
+describe('DefaultAgentContextProvider', () => {
+  describe('getContextForInboundMessage()', () => {
+    test('returns the agent context provided in the constructor', async () => {
+      const agentContext = new MockAgentContext(agentConfig)
+      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,4 @@
+export * from './AgentContext'
+export * from './DefaultAgentContext'
+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 { ConnectionRecord } from '../../modules/connections'
 import type { Key } from '../../modules/dids'
-import type { AgentContext } from '../AgentContext'
 import type { AgentMessage } from '../AgentMessage'
+import type { AgentContext } from '../context'
 
 import { AriesFrameworkError } from '../../error'
 

packages/core/src/constants.ts

@@ -3,6 +3,7 @@ export const InjectionSymbols = {
   StorageService: Symbol('StorageService'),
   Logger: Symbol('Logger'),
   AgentContext: Symbol('AgentContext'),
+  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

@@ -604,6 +604,7 @@ export class OutOfBandModule {
       payload: {
         message: plaintextMessage,
         connection: connectionRecord,
+        contextCorrelationId: this.agentContext.contextCorrelationId,
       },
     })
   }
@@ -644,6 +645,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 { scoped, Lifecycle } from 'tsyringe'

packages/core/src/modules/routing/__tests__/mediationRecipient.test.ts

@@ -234,12 +234,14 @@ describe('MediationRecipientService', () => {
         type: AgentEventTypes.AgentMessageReceived,
         payload: {
           message: { first: 'value' },
+          contextCorrelationId: agentContext.contextCorrelationId,
         },
       })
       expect(eventEmitter.emit).toHaveBeenNthCalledWith(2, agentContext, {
         type: AgentEventTypes.AgentMessageReceived,
         payload: {
           message: { second: 'value' },
+          contextCorrelationId: agentContext.contextCorrelationId,
         },
       })
     })

packages/core/src/modules/routing/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

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