Merge pull request #455 from xmtp/nmolnar/strongly-typed-client

bench/encode.ts

@@ -17,7 +17,7 @@ const encodeV1 = () => {
       const alice = await Client.create(newWallet(), { env: 'local' })
       const bobKeys = (await newPrivateKeyBundle()).getPublicKeyBundle()
 
-      const message = randomBytes(size)
+      const message = randomBytes(size).toString()
       const timestamp = new Date()
 
       // The returned function is the actual benchmark. Everything above is setup

src/Client.ts

@@ -11,8 +11,8 @@ import { utils } from 'ethers'
 import { Signer } from './types/Signer'
 import { Conversations } from './conversations'
 import { ContentTypeText, TextCodec } from './codecs/Text'
-import { ContentTypeId, ContentCodec } from './MessageContent'
-import { compress } from './Compression'
+import { ContentTypeId, ContentCodec, EncodedContent } from './MessageContent'
+import { compress, decompress } from './Compression'
 import { content as proto, messageApi } from '@xmtp/proto'
 import { decodeContactBundle, encodeContactBundle } from './ContactBundle'
 import HttpApiClient, {
@@ -41,6 +41,7 @@ import {
 } from './keystore/persistence'
 import { hasMetamaskWithSnaps } from './keystore/snapHelpers'
 import { version as snapVersion, package as snapPackage } from './snapInfo.json'
+import { ExtractDecodedType } from './types/client'
 const { Compression } = proto
 
 // eslint-disable @typescript-eslint/explicit-module-boundary-types
@@ -244,7 +245,8 @@ export function defaultOptions(opts?: Partial<ClientOptions>): ClientOptions {
  * Client class initiates connection to the XMTP network.
  * Should be created with `await Client.create(options)`
  */
-export default class Client {
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+export default class Client<ContentTypes = any> {
   address: string
   keystore: Keystore
   apiClient: ApiClient
@@ -256,7 +258,7 @@ export default class Client {
   > // addresses and key bundles that we have witnessed
 
   private _backupClient: BackupClient
-  private readonly _conversations: Conversations
+  private readonly _conversations: Conversations<ContentTypes>
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
   private _codecs: Map<string, ContentCodec<any>>
   private _maxContentSize: number
@@ -286,7 +288,7 @@ export default class Client {
   /**
    * @type {Conversations}
    */
-  get conversations(): Conversations {
+  get conversations(): Conversations<ContentTypes> {
     return this._conversations
   }
 
@@ -304,10 +306,16 @@ export default class Client {
    * @param wallet the wallet as a Signer instance
    * @param opts specify how to to connect to the network
    */
-  static async create(
+
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  static async create<ContentCodecs extends ContentCodec<any>[] = []>(
     wallet: Signer | null,
-    opts?: Partial<ClientOptions>
-  ): Promise<Client> {
+    opts?: Partial<ClientOptions> & { codecs?: ContentCodecs }
+  ): Promise<
+    Client<
+      ExtractDecodedType<[...ContentCodecs, TextCodec][number]> | undefined
+    >
+  > {
     const options = defaultOptions(opts)
     const apiClient = options.apiClientFactory(options)
     const keystore = await bootstrapKeystore(options, apiClient, wallet)
@@ -317,12 +325,9 @@ export default class Client {
     const address = publicKeyBundle.walletSignatureAddress()
     apiClient.setAuthenticator(new KeystoreAuthenticator(keystore))
     const backupClient = await Client.setupBackupClient(address, options.env)
-    const client = new Client(
-      publicKeyBundle,
-      apiClient,
-      backupClient,
-      keystore
-    )
+    const client = new Client<
+      ExtractDecodedType<[...ContentCodecs, TextCodec][number]> | undefined
+    >(publicKeyBundle, apiClient, backupClient, keystore)
     await client.init(options)
     return client
   }
@@ -337,9 +342,9 @@ export default class Client {
    * impersonate a user on the XMTP network and read the user's
    * messages.
    */
-  static async getKeys(
+  static async getKeys<U>(
     wallet: Signer | null,
-    opts?: Partial<ClientOptions>
+    opts?: Partial<ClientOptions> & { codecs?: U }
   ): Promise<Uint8Array> {
     const client = await Client.create(wallet, opts)
     const keys = await client.keystore.getPrivateKeyBundle()
@@ -596,10 +601,13 @@ export default class Client {
    * messages of the given Content Type
    */
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
-  registerCodec(codec: ContentCodec<any>): void {
+  registerCodec<Codec extends ContentCodec<any>>(
+    codec: Codec
+  ): Client<ContentTypes | ExtractDecodedType<Codec>> {
     const id = codec.contentType
     const key = `${id.authorityId}/${id.typeId}`
     this._codecs.set(key, codec)
+    return this
   }
 
   /**
@@ -624,8 +632,7 @@ export default class Client {
    * with the given options
    */
   async encodeContent(
-    // eslint-disable-next-line @typescript-eslint/no-explicit-any
-    content: any,
+    content: ContentTypes,
     options?: SendOptions
   ): Promise<Uint8Array> {
     const contentType = options?.contentType || ContentTypeText
@@ -646,6 +653,39 @@ export default class Client {
     return proto.EncodedContent.encode(encoded).finish()
   }
 
+  async decodeContent(contentBytes: Uint8Array): Promise<{
+    content: ContentTypes
+    contentType: ContentTypeId
+    error?: Error
+    contentFallback?: string
+  }> {
+    const encodedContent = proto.EncodedContent.decode(contentBytes)
+
+    if (!encodedContent.type) {
+      throw new Error('missing content type')
+    }
+
+    let content: any // eslint-disable-line @typescript-eslint/no-explicit-any
+    const contentType = new ContentTypeId(encodedContent.type)
+    let error: Error | undefined
+
+    await decompress(encodedContent, 1000)
+
+    const codec = this.codecFor(contentType)
+    if (codec) {
+      content = codec.decode(encodedContent as EncodedContent, this)
+    } else {
+      error = new Error('unknown content type ' + contentType)
+    }
+
+    return {
+      content,
+      contentType,
+      error,
+      contentFallback: encodedContent.fallback,
+    }
+  }
+
   listInvitations(opts?: ListMessagesOptions): Promise<messageApi.Envelope[]> {
     return this.listEnvelopes(
       buildUserInviteTopic(this.address),

src/Message.ts

@@ -4,23 +4,14 @@ import {
   ConversationV2,
 } from './conversations/Conversation'
 import type Client from './Client'
-import {
-  message as proto,
-  content as protoContent,
-  conversationReference,
-} from '@xmtp/proto'
+import { message as proto, conversationReference } from '@xmtp/proto'
 import Long from 'long'
 import Ciphertext from './crypto/Ciphertext'
 import { PublicKeyBundle, PublicKey } from './crypto'
 import { bytesToHex } from './crypto/utils'
 import { sha256 } from './crypto/encryption'
-import {
-  ContentTypeFallback,
-  ContentTypeId,
-  EncodedContent,
-} from './MessageContent'
+import { ContentTypeId } from './MessageContent'
 import { dateToNs, nsToDate } from './utils'
-import { decompress } from './Compression'
 import { Keystore } from './keystore'
 import { buildDecryptV1Request, getResultOrThrow } from './utils/keystore'
 
@@ -228,16 +219,17 @@ export class MessageV2 extends MessageBase implements proto.MessageV2 {
 
 export type Message = MessageV1 | MessageV2
 
-export class DecodedMessage {
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+export class DecodedMessage<ContentTypes = any> {
   id: string
   messageVersion: 'v1' | 'v2'
   senderAddress: string
   recipientAddress?: string
   sent: Date
   contentTopic: string
-  conversation: Conversation
+  conversation: Conversation<ContentTypes>
   contentType: ContentTypeId
-  content: any // eslint-disable-line @typescript-eslint/no-explicit-any
+  content: ContentTypes
   error?: Error
   contentBytes: Uint8Array
   contentFallback?: string
@@ -255,7 +247,7 @@ export class DecodedMessage {
     sent,
     error,
     contentFallback,
-  }: Omit<DecodedMessage, 'toBytes'>) {
+  }: Omit<DecodedMessage<ContentTypes>, 'toBytes'>) {
     this.id = id
     this.messageVersion = messageVersion
     this.senderAddress = senderAddress
@@ -283,10 +275,10 @@ export class DecodedMessage {
     }).finish()
   }
 
-  static async fromBytes(
+  static async fromBytes<ContentTypes>(
     data: Uint8Array,
-    client: Client
-  ): Promise<DecodedMessage> {
+    client: Client<ContentTypes>
+  ): Promise<DecodedMessage<ContentTypes>> {
     const protoVal = proto.DecodedMessage.decode(data)
     const messageVersion = protoVal.messageVersion
 
@@ -299,7 +291,7 @@ export class DecodedMessage {
     }
 
     const { content, contentType, error, contentFallback } =
-      await decodeContent(protoVal.contentBytes, client)
+      await client.decodeContent(protoVal.contentBytes)
 
     return new DecodedMessage({
       ...protoVal,
@@ -317,16 +309,16 @@ export class DecodedMessage {
     })
   }
 
-  static fromV1Message(
+  static fromV1Message<ContentTypes>(
     message: MessageV1,
-    content: any, // eslint-disable-line @typescript-eslint/no-explicit-any
+    content: ContentTypes,
     contentType: ContentTypeId,
     contentBytes: Uint8Array,
     contentTopic: string,
-    conversation: Conversation,
+    conversation: Conversation<ContentTypes>,
     error?: Error,
     contentFallback?: string
-  ): DecodedMessage {
+  ): DecodedMessage<ContentTypes> {
     const { id, senderAddress, recipientAddress, sent } = message
     if (!senderAddress) {
       throw new Error('Sender address is required')
@@ -347,17 +339,17 @@ export class DecodedMessage {
     })
   }
 
-  static fromV2Message(
+  static fromV2Message<ContentTypes>(
     message: MessageV2,
-    content: any, // eslint-disable-line @typescript-eslint/no-explicit-any
+    content: ContentTypes,
     contentType: ContentTypeId,
     contentTopic: string,
     contentBytes: Uint8Array,
-    conversation: Conversation,
+    conversation: Conversation<ContentTypes>,
     senderAddress: string,
     error?: Error,
     contentFallback?: string
-  ): DecodedMessage {
+  ): DecodedMessage<ContentTypes> {
     const { id, sent } = message
 
     return new DecodedMessage({
@@ -376,39 +368,11 @@ export class DecodedMessage {
   }
 }
 
-export async function decodeContent(contentBytes: Uint8Array, client: Client) {
-  const encodedContent = protoContent.EncodedContent.decode(contentBytes)
-
-  if (!encodedContent.type) {
-    throw new Error('missing content type')
-  }
-
-  let content: any // eslint-disable-line @typescript-eslint/no-explicit-any
-  const contentType = new ContentTypeId(encodedContent.type)
-  let error: Error | undefined
-
-  await decompress(encodedContent, 1000)
-
-  const codec = client.codecFor(contentType)
-  if (codec) {
-    content = codec.decode(encodedContent as EncodedContent, client)
-  } else {
-    error = new Error('unknown content type ' + contentType)
-  }
-
-  return {
-    content,
-    contentType,
-    error,
-    contentFallback: encodedContent.fallback,
-  }
-}
-
-function conversationReferenceToConversation(
+function conversationReferenceToConversation<ContentTypes>(
   reference: conversationReference.ConversationReference,
-  client: Client,
+  client: Client<ContentTypes>,
   version: DecodedMessage['messageVersion']
-): Conversation {
+): Conversation<ContentTypes> {
   if (version === 'v1') {
     return new ConversationV1(
       client,
@@ -427,3 +391,10 @@ function conversationReferenceToConversation(
   }
   throw new Error(`Unknown conversation version ${version}`)
 }
+
+export function decodeContent<ContentTypes>(
+  contentBytes: Uint8Array,
+  client: Client<ContentTypes>
+) {
+  return client.decodeContent(contentBytes)
+}

src/Stream.ts

@@ -16,9 +16,10 @@ export type ContentTopicUpdater<M> = (msg: M) => string[] | undefined
  * Stream implements an Asynchronous Iterable over messages received from a topic.
  * As such can be used with constructs like for-await-of, yield*, array destructing, etc.
  */
-export default class Stream<T> {
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+export default class Stream<T, ClientType = any> {
   topics: string[]
-  client: Client
+  client: Client<ClientType>
   // queue of incoming Waku messages
   messages: T[]
   // queue of already pending Promises
@@ -32,7 +33,7 @@ export default class Stream<T> {
   onConnectionLost?: OnConnectionLostCallback
 
   constructor(
-    client: Client,
+    client: Client<ClientType>,
     topics: string[],
     decoder: MessageDecoder<T>,
     contentTopicUpdater?: ContentTopicUpdater<T>,
@@ -100,13 +101,13 @@ export default class Stream<T> {
     )
   }
 
-  static async create<T>(
-    client: Client,
+  static async create<T, ClientType = string>(
+    client: Client<ClientType>,
     topics: string[],
     decoder: MessageDecoder<T>,
     contentTopicUpdater?: ContentTopicUpdater<T>,
     onConnectionLost?: OnConnectionLostCallback
-  ): Promise<Stream<T>> {
+  ): Promise<Stream<T, ClientType>> {
     const stream = new Stream(
       client,
       topics,