Merge branch 'feat/custom-messages-type'

.github/workflows/release.yml

@@ -34,7 +34,7 @@ jobs:
       - run: |
           dfx identity use action
           mops import-identity --no-encrypt -- "$(dfx identity export action)"
-          mops publish --no-docs
+          mops publish
           echo "version=$(cat mops.toml | grep "version =" | cut -d\" -f2)" >> "$GITHUB_OUTPUT"
         env:
           CARGO_REGISTRY_TOKEN: ${{ secrets.CARGO_REGISTRY_TOKEN }}

README.md

@@ -22,17 +22,23 @@ In order for the frontend clients and the Gateway to work properly, the canister
 ```
 import "./ws_types.did";
 
+// define here your message type
+type MyMessageType = {
+  some_field : text;
+};
+
 service : {
-  "ws_register" : (CanisterWsRegisterArguments) -> (CanisterWsRegisterResult);
   "ws_open" : (CanisterWsOpenArguments) -> (CanisterWsOpenResult);
   "ws_close" : (CanisterWsCloseArguments) -> (CanisterWsCloseResult);
-  "ws_message" : (CanisterWsMessageArguments) -> (CanisterWsMessageResult);
+  "ws_message" : (CanisterWsMessageArguments, opt MyMessageType) -> (CanisterWsMessageResult);
   "ws_get_messages" : (CanisterWsGetMessagesArguments) -> (CanisterWsGetMessagesResult) query;
 };
 ```
 This snipped is copied from the [service.example.did](./did/service.example.did) file and the types imported are defined in the [ws_types.did](./did/ws_types.did) file.
 
-**Note**: `dfx` should already generate the Candid interface for you, so you don't need to write it yourself.
+To define your message type, you can use the [Candid reference docs](https://internetcomputer.org/docs/current/references/candid-ref). We suggest you to define your message type using a [variant](https://internetcomputer.org/docs/current/references/candid-ref#type-variant--n--t--), so that you can support different messages over the same websocket instance and make it safe for future updates.
+
+**Note**: `dfx` should already generate the Candid interface for you, so you don't need to write any `.did` file yourself.
 
 ## Development
 

did/service.example.did

@@ -1,8 +1,13 @@
 import "./ws_types.did";
 
+// define your message type here
+type MyMessageType = {
+  some_field : text;
+};
+
 service : {
   "ws_open" : (CanisterWsOpenArguments) -> (CanisterWsOpenResult);
   "ws_close" : (CanisterWsCloseArguments) -> (CanisterWsCloseResult);
-  "ws_message" : (CanisterWsMessageArguments) -> (CanisterWsMessageResult);
+  "ws_message" : (CanisterWsMessageArguments, opt MyMessageType) -> (CanisterWsMessageResult);
   "ws_get_messages" : (CanisterWsGetMessagesArguments) -> (CanisterWsGetMessagesResult) query;
 };

src/lib.mo

@@ -512,8 +512,42 @@ module {
 	public type OnOpenCallback = (OnOpenCallbackArgs) -> async ();
 
 	/// Arguments passed to the `on_message` handler.
+	/// The `message` argument is the message received from the client, serialized in Candid.
+	/// To deserialize the message, use [from_candid].
+	///
+	/// # Example
+	/// This example is the deserialize equivalent of the [ws_send]'s example serialize one.
+	/// ```motoko
+	/// import IcWebSocketCdk "mo:ic-websocket-cdk";
+	///
+	/// actor MyCanister {
+	///   // ...
+	///
+	///   type MyMessage = {
+	///     some_field: Text;
+	///   };
+	///
+	///   // initialize the CDK
+	///
+	///   func on_message(args : IcWebSocketCdk.OnMessageCallbackArgs) : async () {
+	///     let received_message: ?MyMessage = from_candid(args.message);
+	///     switch (received_message) {
+	///       case (?received_message) {
+	///         Debug.print("Received message: some_field: " # received_message.some_field);
+	///       };
+	///       case (invalid_arg) {
+	///         return #Err("invalid argument: " # debug_show (invalid_arg));
+	///       };
+	///     };
+	///   };
+	///
+	///   // ...
+	/// }
+	/// ```
 	public type OnMessageCallbackArgs = {
+		/// The principal of the client sending the message to the canister.
 		client_principal : ClientPrincipal;
+		/// The message received from the client, serialized in Candid. See [OnMessageCallbackArgs] for an example on how to deserialize the message.
 		message : Blob;
 	};
 	/// Handler initialized by the canister and triggered by the CDK once a message is received by
@@ -1104,12 +1138,34 @@ module {
 		#Ok;
 	};
 
-	/// Sends a message to the client.
+	/// Sends a message to the client. The message must already be serialized **using Candid**.
+	/// Use [to_candid] to serialize the message.
 	///
-	/// Under the hood, the message is certified, and then it is added to the queue of messages
+	/// Under the hood, the message is certified and added to the queue of messages
 	/// that the WS Gateway will poll in the next iteration.
-	/// **Note**: you have to serialize the message to a `Blob` before calling this method.
-	/// Use the `to_candid` function.
+	///
+	/// # Example
+	/// This example is the serialize equivalent of the [OnMessageCallbackArgs]'s example deserialize one.
+	/// ```motoko
+	/// import IcWebSocketCdk "mo:ic-websocket-cdk";
+	///
+	/// actor MyCanister {
+	///   // ...
+	///
+	///   type MyMessage = {
+	///     some_field: Text;
+	///   };
+	///
+	///   // initialize the CDK
+	///
+	///   // at some point in your code
+	///   let msg : MyMessage = {
+	///     some_field: "Hello, World!";
+	///   };
+	///
+	///   IcWebSocketCdk.ws_send(ws_state, client_principal, to_candid(msg));
+	/// }
+	/// ```
 	public func ws_send(ws_state : IcWebSocketState, client_principal : ClientPrincipal, msg_bytes : Blob) : async CanisterWsSendResult {
 		_ws_send(ws_state, client_principal, msg_bytes, false);
 	};
@@ -1280,7 +1336,31 @@ module {
 		};
 
 		/// Handles the WS messages received either directly from the client or relayed by the WS Gateway.
-		public func ws_message(caller : Principal, args : CanisterWsMessageArguments) : async CanisterWsMessageResult {
+		///
+		/// The second argument is only needed to expose the type of the message on the canister Candid interface and get automatic types generation on the client side.
+		/// This way, on the client you have the same types and you don't have to care about serializing and deserializing the messages sent through IC WebSocket.
+		///
+		/// # Example
+		/// ```motoko
+		/// import IcWebSocketCdk "mo:ic-websocket-cdk";
+		///
+		/// actor MyCanister {
+		///   // ...
+		///
+		///   type MyMessage = {
+		///     some_field: Text;
+		///   };
+		///
+		///   // declare also the other methods: ws_open, ws_close, ws_get_messages
+		///
+		///   public shared ({ caller }) func ws_message(args : IcWebSocketCdk.CanisterWsMessageArguments, msg_type : ?MyMessage) : async IcWebSocketCdk.CanisterWsMessageResult {
+		///     await ws.ws_message(caller, args, msg_type);
+		///   };
+		///
+		///   // ...
+		/// }
+		/// ```
+		public func ws_message(caller : Principal, args : CanisterWsMessageArguments, _msg_type : ?Any) : async CanisterWsMessageResult {
 			// check if client registered its principal by calling ws_open
 			let registered_client_key = switch (WS_STATE.get_client_key_from_principal(caller)) {
 				case (#Err(err)) {
@@ -1359,7 +1439,7 @@ module {
 			};
 		};
 
-		/// Sends a message to the client. See [ws_send] function for reference.
+		/// Sends a message to the client. See [IcWebSocketCdk.ws_send] function for reference.
 		public func send(client_principal : ClientPrincipal, msg_bytes : Blob) : async CanisterWsSendResult {
 			await ws_send(WS_STATE, client_principal, msg_bytes);
 		};

tests/integration/utils/api.ts

@@ -1,9 +1,10 @@
 // helpers for functions that are called frequently in tests
 
 import { ActorSubclass } from "@dfinity/agent";
-import { IDL } from "@dfinity/candid";
 import { anonymousClient, gateway1Data } from "./actors";
-import type { CanisterOutputCertifiedMessages, ClientKey, ClientPrincipal, WebsocketMessage, _SERVICE } from "../../src/declarations/test_canister/test_canister.did";
+import { IDL } from "@dfinity/candid";
+import { extractApplicationMessageIdlFromActor } from "./idl";
+import type { AppMessage, CanisterOutputCertifiedMessages, ClientKey, ClientPrincipal, WebsocketMessage, _SERVICE } from "../../src/declarations/test_canister/test_canister.did";
 
 type GenericResult<T> = {
   Ok: T,
@@ -53,7 +54,7 @@ type WsMessageArgs = {
 export const wsMessage = async (args: WsMessageArgs, throwIfError = false) => {
   const res = await args.actor.ws_message({
     msg: args.message,
-  });
+  }, []);
 
   return resolveResult(res, throwIfError);
 };
@@ -118,16 +119,12 @@ export const initializeCdk = async (args: InitializeCdkArgs) => {
 type WsSendArgs = {
   clientPrincipal: ClientPrincipal,
   actor: ActorSubclass<_SERVICE>,
-  messages: Array<{
-    text: string,
-  }>,
+  messages: Array<AppMessage>,
 };
 
 export const wsSend = async (args: WsSendArgs, throwIfError = false) => {
-  const serializedMessages = args.messages.map((msg) => {
-    return new Uint8Array(IDL.encode([IDL.Record({ 'text': IDL.Text })], [msg]));
-  });
-  const res = await args.actor.ws_send(args.clientPrincipal, serializedMessages);
+  const messagesBytes = args.messages.map((msg) => IDL.encode([extractApplicationMessageIdlFromActor<AppMessage>(args.actor)], [msg])).map((m) => new Uint8Array(m));
+  const res = await args.actor.ws_send(args.clientPrincipal, messagesBytes);
 
   return resolveResult(res, throwIfError);
 };

tests/integration/utils/idl.ts

@@ -1,6 +1,6 @@
 import { IDL } from "@dfinity/candid";
-import { Cbor } from "@dfinity/agent";
-import type { CanisterOutputMessage, ClientKey, WebsocketMessage } from "../../src/declarations/test_canister/test_canister.did";
+import { Actor, ActorSubclass, Cbor } from "@dfinity/agent";
+import type { CanisterOutputMessage, ClientKey, WebsocketMessage, _SERVICE } from "../../src/declarations/test_canister/test_canister.did";
 
 export const ClientPrincipalIdl = IDL.Principal;
 export const ClientKeyIdl = IDL.Record({
@@ -65,3 +65,30 @@ export const getWebsocketMessageFromCanisterMessage = (msg: CanisterOutputMessag
   const websocketMessage: WebsocketMessage = Cbor.decode(msg.content as Uint8Array);
   return websocketMessage;
 }
+
+/**
+ * Extracts the message type from the canister service definition.
+ * 
+ * @throws {Error} if the canister does not implement the ws_message method
+ * @throws {Error} if the application message type is not optional
+ * 
+ * COPIED from IC WebSocket JS SDK.
+ */
+export const extractApplicationMessageIdlFromActor = <T>(actor: ActorSubclass<_SERVICE>): IDL.Type<T> => {
+  const wsMessageMethod = Actor.interfaceOf(actor)._fields.find((f) => f[0] === "ws_message");
+
+  if (!wsMessageMethod) {
+    throw new Error("Canister does not implement ws_message method");
+  }
+
+  if (wsMessageMethod[1].argTypes.length !== 2) {
+    throw new Error("ws_message method must have 2 arguments");
+  }
+
+  const applicationMessageArg = wsMessageMethod[1].argTypes[1] as IDL.OptClass<T>;
+  if (!(applicationMessageArg instanceof IDL.OptClass)) {
+    throw new Error("Application message type must be optional in the ws_message arguments");
+  }
+
+  return applicationMessageArg["_type"]; // extract the underlying option type
+};

tests/src/test_canister/main.mo

@@ -9,6 +9,10 @@ actor class TestCanister(
   init_keep_alive_timeout_ms : Nat64,
 ) {
 
+  type AppMessage = {
+    text : Text;
+  };
+
   var ws_state = IcWebSocketCdk.IcWebSocketState(init_gateway_principal);
 
   func on_open(args : IcWebSocketCdk.OnOpenCallbackArgs) : async () {
@@ -49,8 +53,8 @@ actor class TestCanister(
   };
 
   // method called by the WS Gateway to send a message of type GatewayMessage to the canister
-  public shared ({ caller }) func ws_message(args : IcWebSocketCdk.CanisterWsMessageArguments) : async IcWebSocketCdk.CanisterWsMessageResult {
-    await ws.ws_message(caller, args);
+  public shared ({ caller }) func ws_message(args : IcWebSocketCdk.CanisterWsMessageArguments, msg_type : ?AppMessage) : async IcWebSocketCdk.CanisterWsMessageResult {
+    await ws.ws_message(caller, args, msg_type);
   };
 
   // method called by the WS Gateway to get messages for all the clients it serves