diff --git a/components/TARGET_PSA/TARGET_MBED_SPM/COMPONENT_SPE/spm_server.c b/components/TARGET_PSA/TARGET_MBED_SPM/COMPONENT_SPE/spm_server.c
index 46be590e47..c0ab6a01de 100644
--- a/components/TARGET_PSA/TARGET_MBED_SPM/COMPONENT_SPE/spm_server.c
+++ b/components/TARGET_PSA/TARGET_MBED_SPM/COMPONENT_SPE/spm_server.c
@@ -165,6 +165,12 @@ static void copy_message_to_spm(spm_ipc_channel_t *channel, psa_msg_t *user_msg)
     user_msg->type = channel->msg_type;
     user_msg->rhandle = channel->rhandle;
     user_msg->handle = handle;
+
+    if (channel->src_partition == NULL) {
+        user_msg->client_id = PSA_NSPE_IDENTIFIER;
+    } else {
+        user_msg->client_id = channel->src_partition->partition_id;
+    }
 }
 
 static spm_ipc_channel_t *spm_rot_service_queue_dequeue(spm_rot_service_t *rot_service)
@@ -199,31 +205,30 @@ static spm_ipc_channel_t *spm_rot_service_queue_dequeue(spm_rot_service_t *rot_s
     return ret;
 }
 
-
-static uint32_t psa_wait(bool wait_any, uint32_t bitmask, uint32_t timeout)
+psa_signal_t psa_wait(psa_signal_t signal_mask, uint32_t timeout)
 {
     spm_partition_t *curr_partition = get_active_partition();
     SPM_ASSERT(NULL != curr_partition); // active thread in SPM must be in partition DB
 
-    uint32_t flags_interrupts = curr_partition->flags_interrupts | PSA_DOORBELL;
-    uint32_t flags_all = curr_partition->flags_rot_srv | flags_interrupts;
+    psa_signal_t flags_interrupts = curr_partition->flags_interrupts | PSA_DOORBELL;
+    psa_signal_t flags_all = curr_partition->flags_rot_srv | flags_interrupts;
 
     // In case we're waiting for any signal the bitmask must contain all the flags, otherwise
     // we should be waiting for a subset of interrupt signals.
-    if (wait_any) {
-        bitmask = flags_all;
+
+    if (signal_mask == PSA_WAIT_ANY) {
+        signal_mask = flags_all;
     } else {
-        // Make sure the interrupt mask contains only a subset of interrupt signal mask.
-        if (bitmask != (flags_interrupts & bitmask)) {
-            SPM_PANIC("interrupt mask 0x%x must have only bits from 0x%x!\n",
-                      bitmask, flags_interrupts);
+        if ((~flags_all) & signal_mask)  {
+            SPM_PANIC("signal mask 0x%x must have only bits from 0x%x!\n",
+                      signal_mask, flags_all);
         }
     }
 
-    uint32_t asserted_signals = osThreadFlagsWait(
-                                    bitmask,
+    psa_signal_t asserted_signals = osThreadFlagsWait(
+                                    signal_mask,
                                     osFlagsWaitAny | osFlagsNoClear,
-                                    (PSA_BLOCK == timeout) ? osWaitForever : timeout
+                                    (PSA_BLOCK & timeout) ? osWaitForever : 0
                                 );
 
     // Asserted_signals must be a subset of the supported ROT_SRV and interrupt signals.
@@ -233,16 +238,6 @@ static uint32_t psa_wait(bool wait_any, uint32_t bitmask, uint32_t timeout)
     return (osFlagsErrorTimeout == asserted_signals) ? 0 : asserted_signals;
 }
 
-uint32_t psa_wait_any(uint32_t timeout)
-{
-    return psa_wait(true, 0, timeout);
-}
-
-uint32_t psa_wait_interrupt(uint32_t interrupt_mask, uint32_t timeout)
-{
-    return psa_wait(false, interrupt_mask, timeout);
-}
-
 void psa_get(psa_signal_t signum, psa_msg_t *msg)
 {
     spm_partition_t *curr_partition = get_active_partition();
@@ -536,17 +531,6 @@ void psa_clear(void)
     }
 }
 
-int32_t psa_identity(psa_handle_t msg_handle)
-{
-    spm_active_msg_t *active_msg = get_msg_from_handle(msg_handle);
-    SPM_ASSERT(active_msg->channel != NULL);
-    if (active_msg->channel->src_partition == NULL) {
-        return PSA_NSPE_IDENTIFIER;
-    }
-
-    return active_msg->channel->src_partition->partition_id;
-}
-
 void psa_set_rhandle(psa_handle_t msg_handle, void *rhandle)
 {
     spm_active_msg_t *active_msg = get_msg_from_handle(msg_handle);
diff --git a/components/TARGET_PSA/TARGET_MBED_SPM/COMPONENT_SPE/spm_server.h b/components/TARGET_PSA/TARGET_MBED_SPM/COMPONENT_SPE/spm_server.h
index abd3990ad2..f3126ab3f4 100644
--- a/components/TARGET_PSA/TARGET_MBED_SPM/COMPONENT_SPE/spm_server.h
+++ b/components/TARGET_PSA/TARGET_MBED_SPM/COMPONENT_SPE/spm_server.h
@@ -42,35 +42,13 @@ extern "C" {
 /**
  * Return the signals that have been asserted.@n
  *
+ * @param[in] signal_mask A set of signals to query.
  * @param[in] timeout timeout value:@n
  *            @a PSA_BLOCK block the caller until any signal is asserted.@n
  *            @a PSA_POLL Returns immediately.
- * @return 32-bit value of asserted signals.
+ * @return asserted signals.
  */
-uint32_t psa_wait_any(uint32_t timeout);
-
-/**
- * Return interrupt and doorbell signals that have been asserted based on the bitmask provided.@n
- * The mask contains a set of signals the caller is interested in handling and must be a subset
- * of combined interrupt and doorbell mask for the calling partition.
- *
- * @param[in] interrupt_mask mask of signals.
- * @param[in] timeout timeout value:@n
- *            @a PSA_BLOCK block the caller until one of the requested signals is asserted.@n
- *            @a PSA_POLL Returns immediately.
- * @return 32-bit value of asserted signals.
- */
-uint32_t psa_wait_interrupt(uint32_t interrupt_mask, uint32_t timeout);
-
-/**
- * Return the partition ID of the caller.
- *
- * @note Bit[31] is set if the caller is from the NSPE.
- *
- * @param[in] msg_handle Handle for the caller's message.
- * @return Caller's partition ID
- */
-int32_t psa_identity(psa_handle_t msg_handle);
+psa_signal_t psa_wait(psa_signal_t signal_mask, uint32_t timeout);
 
 /**
  * Get the message that corresponds to a given signal.
diff --git a/components/TARGET_PSA/TARGET_MBED_SPM/psa_defs.h b/components/TARGET_PSA/TARGET_MBED_SPM/psa_defs.h
index cfe3da5b40..4e6b54e8bf 100644
--- a/components/TARGET_PSA/TARGET_MBED_SPM/psa_defs.h
+++ b/components/TARGET_PSA/TARGET_MBED_SPM/psa_defs.h
@@ -55,6 +55,7 @@ extern "C" {
 
 #define PSA_POLL  (0x00000000UL) /**< Returns immediately even if none of the requested signals is asserted.*/
 #define PSA_BLOCK (0x80000000UL) /**< Block the caller until one of the requested signals is asserted.*/
+#define PSA_WAIT_ANY (0xFFFFFFFFUL) /**< A mask value that includes all Secure Partition signals.*/
 
 #define PSA_MINOR_VERSION_POLICY_RELAXED (0UL) /**< Don't perform minor version check during psa_connect().*/
 #define PSA_MINOR_VERSION_POLICY_STRICT (1UL)  /**< Force minor version check during psa_connect().*/
@@ -86,6 +87,7 @@ typedef psa_status_t error_t;
 typedef struct psa_msg {
     uint32_t type;                   /**< The message type.*/
     psa_handle_t handle;             /**< Handle for the internal message structure.*/
+    int32_t client_id;               /**< Message origin.*/
     void *rhandle;                   /**< Reverse handle.*/
     size_t in_size[PSA_MAX_IOVEC];   /**< Array of sizes in bytes of the message payloads.*/
     size_t out_size[PSA_MAX_IOVEC];  /**< Array of sizes in bytes of the response buffers.*/
@@ -111,7 +113,6 @@ typedef struct psa_outvec {
 }
 #endif
 
-
 /** @}*/ // end of SPM group
 
 #endif  /* __MBED_PSA_DEFS_H__ */