diff --git a/kernel/ethosu_mailbox.h b/kernel/ethosu_mailbox.h
index edf922b..c192b54 100644
--- a/kernel/ethosu_mailbox.h
+++ b/kernel/ethosu_mailbox.h
@@ -51,8 +51,9 @@
 };
 
 struct ethosu_mailbox_msg {
-	int  id;
-	void (*fail)(struct ethosu_mailbox_msg *msg);
+	int      id;
+	uint32_t type;
+	void     (*fail)(struct ethosu_mailbox_msg *msg);
 };
 
 /****************************************************************************
@@ -93,7 +94,8 @@
  * Return: a valid pointer on success, otherwise an error ptr.
  */
 struct ethosu_mailbox_msg *ethosu_mailbox_find(struct ethosu_mailbox *mbox,
-					       int msg_id);
+					       int msg_id,
+					       uint32_t msg_type);
 
 /**
  * ethosu_mailbox_fail() - Fail mailbox messages