diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index 4ab0c63..7190aa4 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -3214,6 +3214,25 @@
     return( status );
 }
 
+/* Helper function to get the base algorithm from its variants. */
+static psa_algorithm_t psa_aead_get_base_algorithm(psa_algorithm_t alg)
+{
+    switch( PSA_ALG_AEAD_WITH_SHORTENED_TAG( alg, 0 ) )
+    {
+        case PSA_ALG_AEAD_WITH_SHORTENED_TAG( PSA_ALG_CCM, 0 ):
+            return( PSA_ALG_CCM );
+
+        case PSA_ALG_AEAD_WITH_SHORTENED_TAG( PSA_ALG_GCM, 0 ):
+            return( PSA_ALG_GCM );
+
+        case PSA_ALG_AEAD_WITH_SHORTENED_TAG( PSA_ALG_CHACHA20_POLY1305, 0 ):
+            return( PSA_ALG_CHACHA20_POLY1305 );
+
+        default:
+            return( PSA_ERROR_NOT_SUPPORTED );
+    }
+}
+
 /* Set the key for a multipart authenticated encryption operation. */
 psa_status_t psa_aead_encrypt_setup( psa_aead_operation_t *operation,
                                      mbedtls_svc_key_id_t key,
@@ -3226,6 +3245,12 @@
     if( !PSA_ALG_IS_AEAD( alg ) || PSA_ALG_IS_WILDCARD( alg ) )
         return( PSA_ERROR_NOT_SUPPORTED );
 
+    if( operation->key_set || operation->nonce_set ||
+        operation->ad_started || operation->body_started )
+    {
+        return( PSA_ERROR_BAD_STATE );
+    }
+
     status = psa_get_and_lock_key_slot_with_policy(
                  key, &slot, PSA_KEY_USAGE_ENCRYPT, alg );
 
@@ -3242,6 +3267,7 @@
                                                     &attributes, slot->key.data,
                                                     slot->key.bytes, alg );
 
+    operation->key_type = psa_get_key_type( &attributes );
 
     unlock_status = psa_unlock_key_slot( slot );
 
@@ -3250,6 +3276,12 @@
         return( unlock_status );
     }
 
+    if( status == PSA_SUCCESS )
+    {
+        operation->alg = psa_aead_get_base_algorithm( alg );
+        operation->key_set = 1;
+    }
+
     return( status );
 }
 
@@ -3265,6 +3297,12 @@
     if( !PSA_ALG_IS_AEAD( alg ) || PSA_ALG_IS_WILDCARD( alg ) )
         return( PSA_ERROR_NOT_SUPPORTED );
 
+    if( operation->key_set || operation->nonce_set ||
+        operation->ad_started || operation->body_started )
+    {
+        return( PSA_ERROR_BAD_STATE );
+    }
+
     status = psa_get_and_lock_key_slot_with_policy(
                  key, &slot, PSA_KEY_USAGE_DECRYPT, alg );
 
@@ -3281,6 +3319,7 @@
                                                     &attributes, slot->key.data,
                                                     slot->key.bytes, alg );
 
+    operation->key_type = psa_get_key_type( &attributes );
 
     unlock_status = psa_unlock_key_slot( slot );
 
@@ -3289,6 +3328,12 @@
         return( unlock_status );
     }
 
+    if( status == PSA_SUCCESS )
+    {
+        operation->alg = psa_aead_get_base_algorithm( alg );
+        operation->key_set = 1;
+    }
+
     return( status );
 }
 
@@ -3341,14 +3386,23 @@
                                  const uint8_t *nonce,
                                  size_t nonce_length )
 {
+    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
+
     if( !operation->key_set || operation->nonce_set ||
         operation->ad_started || operation->body_started )
     {
         return( PSA_ERROR_BAD_STATE );
     }
 
-    return( psa_driver_wrapper_aead_set_nonce( operation, nonce,
-                                               nonce_length ) );
+    status = psa_driver_wrapper_aead_set_nonce( operation, nonce,
+                                                nonce_length );
+
+    if( status == PSA_SUCCESS )
+    {
+        operation->nonce_set = 1;
+    }
+
+    return( status );
 }
 
 /* Declare the lengths of the message and additional data for multipart AEAD. */
@@ -3356,26 +3410,44 @@
                                    size_t ad_length,
                                    size_t plaintext_length )
 {
+    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
+
     if( !operation->key_set || operation->lengths_set )
     {
         return( PSA_ERROR_BAD_STATE );
     }
 
-    return( psa_driver_wrapper_aead_set_lengths( operation, ad_length,
-                                                 plaintext_length ) );
+    status = psa_driver_wrapper_aead_set_lengths( operation, ad_length,
+                                                  plaintext_length );
+
+    if( status == PSA_SUCCESS )
+    {
+        operation->lengths_set = 1;
+    }
+
+    return status;
 }
  /* Pass additional data to an active multipart AEAD operation. */
 psa_status_t psa_aead_update_ad( psa_aead_operation_t *operation,
                                  const uint8_t *input,
                                  size_t input_length )
 {
+    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
+
     if( !operation->nonce_set || !operation->key_set )
     {
         return( PSA_ERROR_BAD_STATE );
     }
 
-    return( psa_driver_wrapper_aead_update_ad( operation, input,
-                                               input_length ) );
+    status = psa_driver_wrapper_aead_update_ad( operation, input,
+                                                input_length );
+
+    if( status == PSA_SUCCESS )
+    {
+        operation->ad_started = 1;
+    }
+
+    return status;
 }
 
 /* Encrypt or decrypt a message fragment in an active multipart AEAD
@@ -3387,6 +3459,7 @@
                               size_t output_size,
                               size_t *output_length )
 {
+    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
 
     *output_length = 0;
 
@@ -3395,9 +3468,16 @@
         return( PSA_ERROR_BAD_STATE );
     }
 
-    return( psa_driver_wrapper_aead_update( operation, input, input_length,
-                                            output, output_size,
-                                            output_length ) );
+    status = psa_driver_wrapper_aead_update( operation, input, input_length,
+                                             output, output_size,
+                                             output_length );
+
+    if( status == PSA_SUCCESS )
+    {
+        operation->body_started = 1;
+    }
+
+    return status;
 }
 
 /* Finish encrypting a message in a multipart AEAD operation. */
@@ -3422,6 +3502,7 @@
                                             ciphertext_size,
                                             ciphertext_length,
                                             tag, tag_size, tag_length ) );
+
 }
 
 /* Finish authenticating and decrypting a message in a multipart AEAD
@@ -3466,7 +3547,6 @@
     operation->key_set = 0;
     operation->nonce_set = 0;
     operation->lengths_set = 0;
-    operation->is_encrypt = 0;
     operation->ad_started = 0;
     operation->body_started = 0;
 
diff --git a/library/psa_crypto_aead.c b/library/psa_crypto_aead.c
index f5b4dc5..8f8b74e 100644
--- a/library/psa_crypto_aead.c
+++ b/library/psa_crypto_aead.c
@@ -20,7 +20,6 @@
 
 #include "common.h"
 
-
 #if defined(MBEDTLS_PSA_CRYPTO_C)
 
 #include "psa_crypto_aead.h"
@@ -55,7 +54,7 @@
 
 
 static psa_status_t psa_aead_setup(
-    psa_aead_operation_t *operation,
+    mbedtls_psa_aead_operation_t *operation,
     const psa_key_attributes_t *attributes,
     const uint8_t *key_buffer,
     psa_algorithm_t alg )
@@ -66,12 +65,6 @@
     mbedtls_cipher_id_t cipher_id;
     size_t full_tag_length = 0;
 
-    if( operation->key_set || operation->nonce_set ||
-        operation->ad_started || operation->body_started )
-    {
-        return( PSA_ERROR_BAD_STATE );
-    }
-
     key_bits = attributes->core.bits;
 
     cipher_info = mbedtls_cipher_info_from_psa( alg,
@@ -146,12 +139,12 @@
         > full_tag_length )
         return( PSA_ERROR_INVALID_ARGUMENT );
 
-    operation->tag_length = PSA_AEAD_TAG_LENGTH( attributes->core.type,
+    operation->key_type = psa_get_key_type( attributes );
+
+    operation->tag_length = PSA_AEAD_TAG_LENGTH( operation->key_type,
                                                  key_bits,
                                                  alg );
 
-    operation->key_set = 1;
-
     return( PSA_SUCCESS );
 }
 
@@ -165,7 +158,7 @@
     uint8_t *ciphertext, size_t ciphertext_size, size_t *ciphertext_length )
 {
     psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
-    psa_aead_operation_t operation = PSA_AEAD_OPERATION_INIT;
+    mbedtls_psa_aead_operation_t operation = MBEDTLS_PSA_AEAD_OPERATION_INIT;
     uint8_t *tag;
     (void) key_buffer_size;
 
@@ -275,7 +268,7 @@
     uint8_t *plaintext, size_t plaintext_size, size_t *plaintext_length )
 {
     psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
-    psa_aead_operation_t operation = PSA_AEAD_OPERATION_INIT;
+    mbedtls_psa_aead_operation_t operation = MBEDTLS_PSA_AEAD_OPERATION_INIT;
     const uint8_t *tag = NULL;
     (void) key_buffer_size;
 
@@ -354,7 +347,8 @@
 
 /* Set the key and algorithm for a multipart authenticated encryption
  * operation. */
-psa_status_t mbedtls_psa_aead_encrypt_setup( psa_aead_operation_t *operation,
+psa_status_t mbedtls_psa_aead_encrypt_setup( mbedtls_psa_aead_operation_t
+                                                                    *operation,
                                              const psa_key_attributes_t
                                                                     *attributes,
                                              const uint8_t *key_buffer,
@@ -377,7 +371,8 @@
 
 /* Set the key and algorithm for a multipart authenticated decryption
  * operation. */
-psa_status_t mbedtls_psa_aead_decrypt_setup( psa_aead_operation_t *operation,
+psa_status_t mbedtls_psa_aead_decrypt_setup( mbedtls_psa_aead_operation_t
+                                                                    *operation,
                                              const psa_key_attributes_t
                                                                     *attributes,
                                              const uint8_t *key_buffer,
@@ -399,7 +394,8 @@
 }
 
 /* Set a nonce for the multipart AEAD operation*/
-psa_status_t mbedtls_psa_aead_set_nonce( psa_aead_operation_t *operation,
+psa_status_t mbedtls_psa_aead_set_nonce( mbedtls_psa_aead_operation_t
+                                                                    *operation,
                                          const uint8_t *nonce,
                                          size_t nonce_length )
 {
@@ -454,15 +450,11 @@
         return ( PSA_ERROR_NOT_SUPPORTED );
     }
 
-    if( status == PSA_SUCCESS )
-    {
-        operation->nonce_set = 1;
-    }
-
     return( status );
 }
  /* Declare the lengths of the message and additional data for AEAD. */
-psa_status_t mbedtls_psa_aead_set_lengths( psa_aead_operation_t *operation,
+psa_status_t mbedtls_psa_aead_set_lengths( mbedtls_psa_aead_operation_t
+                                                                    *operation,
                                            size_t ad_length,
                                            size_t plaintext_length )
 {
@@ -512,7 +504,8 @@
 }
 
 /* Pass additional data to an active multipart AEAD operation. */
-psa_status_t mbedtls_psa_aead_update_ad( psa_aead_operation_t *operation,
+psa_status_t mbedtls_psa_aead_update_ad( mbedtls_psa_aead_operation_t
+                                                                    *operation,
                                          const uint8_t *input,
                                          size_t input_length )
 {
@@ -611,7 +604,7 @@
 
 /* Encrypt or decrypt a message fragment in an active multipart AEAD
  * operation.*/
-psa_status_t mbedtls_psa_aead_update( psa_aead_operation_t *operation,
+psa_status_t mbedtls_psa_aead_update( mbedtls_psa_aead_operation_t *operation,
                                       const uint8_t *input,
                                       size_t input_length,
                                       uint8_t *output,
@@ -786,7 +779,7 @@
 
 /* Common checks for both mbedtls_psa_aead_finish() and
    mbedtls_psa_aead_verify() */
-static psa_status_t mbedtls_psa_aead_finish_checks( psa_aead_operation_t
+static psa_status_t mbedtls_psa_aead_finish_checks( mbedtls_psa_aead_operation_t
                                                                     *operation,
                                                     size_t output_size,
                                                     size_t tag_size )
@@ -828,7 +821,7 @@
 }
 
 /* Finish encrypting a message in a multipart AEAD operation. */
-psa_status_t mbedtls_psa_aead_finish( psa_aead_operation_t *operation,
+psa_status_t mbedtls_psa_aead_finish( mbedtls_psa_aead_operation_t *operation,
                                       uint8_t *ciphertext,
                                       size_t ciphertext_size,
                                       size_t *ciphertext_length,
@@ -903,7 +896,7 @@
 
 /* Finish authenticating and decrypting a message in a multipart AEAD
  * operation.*/
-psa_status_t mbedtls_psa_aead_verify( psa_aead_operation_t *operation,
+psa_status_t mbedtls_psa_aead_verify( mbedtls_psa_aead_operation_t *operation,
                                       uint8_t *plaintext,
                                       size_t plaintext_size,
                                       size_t *plaintext_length,
@@ -1033,7 +1026,7 @@
 }
 
 /* Abort an AEAD operation */
-psa_status_t mbedtls_psa_aead_abort( psa_aead_operation_t *operation )
+psa_status_t mbedtls_psa_aead_abort( mbedtls_psa_aead_operation_t *operation )
 {
     switch( operation->alg )
     {
@@ -1054,6 +1047,11 @@
 #endif /* MBEDTLS_PSA_BUILTIN_ALG_CHACHA20_POLY1305 */
     }
 
+    operation->lengths_set = 0;
+    operation->is_encrypt = 0;
+    operation->ad_started = 0;
+    operation->body_started = 0;
+
     mbedtls_free(operation->ad_buffer);
     operation->ad_buffer = NULL;
     operation->ad_length = 0;
diff --git a/library/psa_crypto_aead.h b/library/psa_crypto_aead.h
index a9d2687..4b6d6cd 100644
--- a/library/psa_crypto_aead.h
+++ b/library/psa_crypto_aead.h
@@ -160,37 +160,39 @@
  * -# Allocate an operation object which will be passed to all the functions
  *    listed here.
  * -# Initialize the operation object with one of the methods described in the
- *    documentation for #psa_aead_operation_t, e.g.
- *    #PSA_AEAD_OPERATION_INIT.
+ *    documentation for #mbedtls_psa_aead_operation_t, e.g.
+ *    #MBEDTLS_PSA_AEAD_OPERATION_INIT.
  * -# Call mbedtls_psa_aead_encrypt_setup() to specify the algorithm and key.
  * -# If needed, call mbedtls_psa_aead_set_lengths() to specify the length of
  *    the inputs to the subsequent calls to mbedtls_psa_aead_update_ad() and
- *    mbedtls_psa_aead_update(). See the documentation of mbedtls_psa_aead_set_lengths()
- *    for details.
+ *    mbedtls_psa_aead_update(). See the documentation of
+ *    mbedtls_psa_aead_set_lengths() for details.
  * -# Call either psa_aead_generate_nonce() or
  *    mbedtls_psa_aead_set_nonce() to generate or set the nonce. You should use
  *    psa_aead_generate_nonce() unless the protocol you are implementing
  *    requires a specific nonce value.
- * -# Call mbedtls_psa_aead_update_ad() zero, one or more times, passing a fragment
- *    of the non-encrypted additional authenticated data each time.
+ * -# Call mbedtls_psa_aead_update_ad() zero, one or more times, passing
+ *    a fragment of the non-encrypted additional authenticated data each time.
  * -# Call mbedtls_psa_aead_update() zero, one or more times, passing a fragment
  *    of the message to encrypt each time.
  * -# Call mbedtls_psa_aead_finish().
  *
- * If an error occurs at any step after a call to mbedtls_psa_aead_encrypt_setup(),
- * the operation will need to be reset by a call to mbedtls_psa_aead_abort(). The
- * application may call mbedtls_psa_aead_abort() at any time after the operation
- * has been initialized.
+ * If an error occurs at any step after a call to
+ * mbedtls_psa_aead_encrypt_setup(), the operation will need to be reset by a
+ * call to mbedtls_psa_aead_abort(). The application may call
+ * mbedtls_psa_aead_abort() at any time after the operation has been
+ * initialized.
  *
- * After a successful call to mbedtls_psa_aead_encrypt_setup(), the application must
- * eventually terminate the operation. The following events terminate an
+ * After a successful call to mbedtls_psa_aead_encrypt_setup(), the application
+ * must eventually terminate the operation. The following events terminate an
  * operation:
  * - A successful call to mbedtls_psa_aead_finish().
  * - A call to mbedtls_psa_aead_abort().
  *
  * \param[in,out] operation     The operation object to set up. It must have
  *                              been initialized as per the documentation for
- *                              #mbedtls_psa_aead_operation_t and not yet in use.
+ *                              #mbedtls_psa_aead_operation_t and not yet in
+ *                              use.
  * \param[in]  attributes       The attributes of the key to use for the
  *                              operation.
  * \param[in]  key_buffer       The buffer containing the key context.
@@ -219,9 +221,12 @@
  *         It is implementation-dependent whether a failure to initialize
  *         results in this error code.
  */
-psa_status_t mbedtls_psa_aead_encrypt_setup(psa_aead_operation_t *operation,
-                                            const psa_key_attributes_t *attributes,
-                                            const uint8_t *key_buffer, size_t key_buffer_size,
+psa_status_t mbedtls_psa_aead_encrypt_setup(mbedtls_psa_aead_operation_t
+                                                                     *operation,
+                                            const psa_key_attributes_t
+                                                                    *attributes,
+                                            const uint8_t *key_buffer,
+                                            size_t key_buffer_size,
                                             psa_algorithm_t alg);
 
 /** Set the key for a multipart authenticated decryption operation.
@@ -236,34 +241,36 @@
  * -# Allocate an operation object which will be passed to all the functions
  *    listed here.
  * -# Initialize the operation object with one of the methods described in the
- *    documentation for #psa_aead_operation_t, e.g.
+ *    documentation for #mbedtls_psa_aead_operation_t, e.g.
  *    #PSA_AEAD_OPERATION_INIT.
  * -# Call mbedtls_psa_aead_decrypt_setup() to specify the algorithm and key.
- * -# If needed, call mbedtls_psa_aead_set_lengths() to specify the length of the
- *    inputs to the subsequent calls to mbedtls_psa_aead_update_ad() and
- *    mbedtls_psa_aead_update(). See the documentation of mbedtls_psa_aead_set_lengths()
- *    for details.
+ * -# If needed, call mbedtls_psa_aead_set_lengths() to specify the length of
+ *    the inputs to the subsequent calls to mbedtls_psa_aead_update_ad() and
+ *    mbedtls_psa_aead_update(). See the documentation of
+ *    mbedtls_psa_aead_set_lengths() for details.
  * -# Call mbedtls_psa_aead_set_nonce() with the nonce for the decryption.
- * -# Call mbedtls_psa_aead_update_ad() zero, one or more times, passing a fragment
- *    of the non-encrypted additional authenticated data each time.
+ * -# Call mbedtls_psa_aead_update_ad() zero, one or more times, passing a
+ *    fragment of the non-encrypted additional authenticated data each time.
  * -# Call mbedtls_psa_aead_update() zero, one or more times, passing a fragment
  *    of the ciphertext to decrypt each time.
  * -# Call mbedtls_psa_aead_verify().
  *
- * If an error occurs at any step after a call to mbedtls_psa_aead_decrypt_setup(),
- * the operation will need to be reset by a call to mbedtls_psa_aead_abort(). The
- * application may call mbedtls_psa_aead_abort() at any time after the operation
- * has been initialized.
+ * If an error occurs at any step after a call to
+ * mbedtls_psa_aead_decrypt_setup(), the operation will need to be reset by a
+ * call to mbedtls_psa_aead_abort(). The application may call
+ * mbedtls_psa_aead_abort() at any time after the operation has been
+ * initialized.
  *
- * After a successful call to mbedtls_psa_aead_decrypt_setup(), the application must
- * eventually terminate the operation. The following events terminate an
+ * After a successful call to mbedtls_psa_aead_decrypt_setup(), the application
+ * must eventually terminate the operation. The following events terminate an
  * operation:
  * - A successful call to mbedtls_psa_aead_verify().
  * - A call to mbedtls_psa_aead_abort().
  *
  * \param[in,out] operation     The operation object to set up. It must have
  *                              been initialized as per the documentation for
- *                              #psa_aead_operation_t and not yet in use.
+ *                              #mbedtls_psa_aead_operation_t and not yet in
+ *                              use.
  * \param[in]  attributes       The attributes of the key to use for the
  *                              operation.
  * \param[in]  key_buffer       The buffer containing the key context.
@@ -292,9 +299,12 @@
  *         It is implementation-dependent whether a failure to initialize
  *         results in this error code.
  */
-psa_status_t mbedtls_psa_aead_decrypt_setup(psa_aead_operation_t *operation,
-                                            const psa_key_attributes_t *attributes,
-                                            const uint8_t *key_buffer, size_t key_buffer_size,
+psa_status_t mbedtls_psa_aead_decrypt_setup(mbedtls_psa_aead_operation_t
+                                                                     *operation,
+                                            const psa_key_attributes_t
+                                                                    *attributes,
+                                            const uint8_t *key_buffer,
+                                            size_t key_buffer_size,
                                             psa_algorithm_t alg);
 
 /** Set the nonce for an authenticated encryption or decryption operation.
@@ -313,9 +323,9 @@
  * If this function returns an error status, the operation enters an error
  * state and must be aborted by calling mbedtls_psa_aead_abort().
  *
- * \note When encrypting, applications should use mbedtls_psa_aead_generate_nonce()
- * instead of this function, unless implementing a protocol that requires
- * a non-random IV.
+ * \note When encrypting, applications should use
+ * mbedtls_psa_aead_generate_nonce() instead of this function, unless
+ * implementing a protocol that requires a non-random IV.
  *
  * \param[in,out] operation     Active AEAD operation.
  * \param[in] nonce             Buffer containing the nonce to use.
@@ -338,7 +348,7 @@
  *         It is implementation-dependent whether a failure to initialize
  *         results in this error code.
  */
-psa_status_t mbedtls_psa_aead_set_nonce(psa_aead_operation_t *operation,
+psa_status_t mbedtls_psa_aead_set_nonce(mbedtls_psa_aead_operation_t *operation,
                                         const uint8_t *nonce,
                                         size_t nonce_length);
 
@@ -350,10 +360,10 @@
  *       specification for transparent drivers.
  *
  * The application must call this function before calling
- * mbedtls_psa_aead_update_ad() or mbedtls_psa_aead_update() if the algorithm for
- * the operation requires it. If the algorithm does not require it,
- * calling this function is optional, but if this function is called
- * then the implementation must enforce the lengths.
+ * mbedtls_psa_aead_update_ad() or mbedtls_psa_aead_update() if the algorithm
+ * for the operation requires it. If the algorithm does not require it, calling
+ * this function is optional, but if this function is called then the
+ * implementation must enforce the lengths.
  *
  * You may call this function before or after setting the nonce with
  * mbedtls_psa_aead_set_nonce() or psa_aead_generate_nonce().
@@ -375,8 +385,8 @@
  *         Success.
  * \retval #PSA_ERROR_BAD_STATE
  *         The operation state is not valid (it must be active, and
- *         mbedtls_psa_aead_update_ad() and mbedtls_psa_aead_update() must not have been
- *         called yet).
+ *         mbedtls_psa_aead_update_ad() and mbedtls_psa_aead_update() must not
+ *         have been called yet).
  * \retval #PSA_ERROR_INVALID_ARGUMENT
  *         At least one of the lengths is not acceptable for the chosen
  *         algorithm.
@@ -389,7 +399,8 @@
  *         It is implementation-dependent whether a failure to initialize
  *         results in this error code.
  */
-psa_status_t mbedtls_psa_aead_set_lengths(psa_aead_operation_t *operation,
+psa_status_t mbedtls_psa_aead_set_lengths(mbedtls_psa_aead_operation_t
+                                                                     *operation,
                                           size_t ad_length,
                                           size_t plaintext_length);
 
@@ -407,18 +418,19 @@
  * data to encrypt or decrypt with mbedtls_psa_aead_update().
  *
  * Before calling this function, you must:
- * 1. Call either mbedtls_psa_aead_encrypt_setup() or mbedtls_psa_aead_decrypt_setup().
- * 2. Set the nonce with psa_aead_generate_nonce() or
- *    mbedtls_psa_aead_set_nonce().
+ * 1. Call either mbedtls_psa_aead_encrypt_setup() or
+ *    mbedtls_psa_aead_decrypt_setup(). 2. Set the nonce with
+ *    psa_aead_generate_nonce() or mbedtls_psa_aead_set_nonce().
  *
  * If this function returns an error status, the operation enters an error
  * state and must be aborted by calling mbedtls_psa_aead_abort().
  *
- * \warning When decrypting, until mbedtls_psa_aead_verify() has returned #PSA_SUCCESS,
- *          there is no guarantee that the input is valid. Therefore, until
- *          you have called mbedtls_psa_aead_verify() and it has returned #PSA_SUCCESS,
- *          treat the input as untrusted and prepare to undo any action that
- *          depends on the input if mbedtls_psa_aead_verify() returns an error status.
+ * \warning When decrypting, until mbedtls_psa_aead_verify() has returned
+ *          #PSA_SUCCESS, there is no guarantee that the input is valid.
+ *          Therefore, until you have called mbedtls_psa_aead_verify() and it
+ *          has returned #PSA_SUCCESS, treat the input as untrusted and prepare
+ *          to undo any action that depends on the input if
+ *          mbedtls_psa_aead_verify() returns an error status.
  *
  * \note    For the time being #PSA_ALG_CCM and #PSA_ALG_GCM require the entire
  *          additional data to be passed in in one go, i.e. only call
@@ -448,7 +460,7 @@
  *         It is implementation-dependent whether a failure to initialize
  *         results in this error code.
  */
-psa_status_t mbedtls_psa_aead_update_ad(psa_aead_operation_t *operation,
+psa_status_t mbedtls_psa_aead_update_ad(mbedtls_psa_aead_operation_t *operation,
                                         const uint8_t *input,
                                         size_t input_length);
 
@@ -460,9 +472,9 @@
  *       transparent drivers.
  *
  * Before calling this function, you must:
- * 1. Call either mbedtls_psa_aead_encrypt_setup() or mbedtls_psa_aead_decrypt_setup().
- *    The choice of setup function determines whether this function
- *    encrypts or decrypts its input.
+ * 1. Call either mbedtls_psa_aead_encrypt_setup() or
+ *    mbedtls_psa_aead_decrypt_setup(). The choice of setup function determines
+ *    whether this function encrypts or decrypts its input.
  * 2. Set the nonce with psa_aead_generate_nonce() or
  * mbedtls_psa_aead_set_nonce(). 3. Call mbedtls_psa_aead_update_ad() to pass
  * all the additional data.
@@ -537,7 +549,7 @@
  *         It is implementation-dependent whether a failure to initialize
  *         results in this error code.
  */
-psa_status_t mbedtls_psa_aead_update(psa_aead_operation_t *operation,
+psa_status_t mbedtls_psa_aead_update(mbedtls_psa_aead_operation_t *operation,
                                      const uint8_t *input,
                                      size_t input_length,
                                      uint8_t *output,
@@ -618,7 +630,7 @@
  *         It is implementation-dependent whether a failure to initialize
  *         results in this error code.
  */
-psa_status_t mbedtls_psa_aead_finish(psa_aead_operation_t *operation,
+psa_status_t mbedtls_psa_aead_finish(mbedtls_psa_aead_operation_t *operation,
                                      uint8_t *ciphertext,
                                      size_t ciphertext_size,
                                      size_t *ciphertext_length,
@@ -703,7 +715,7 @@
  *         It is implementation-dependent whether a failure to initialize
  *         results in this error code.
  */
-psa_status_t mbedtls_psa_aead_verify(psa_aead_operation_t *operation,
+psa_status_t mbedtls_psa_aead_verify(mbedtls_psa_aead_operation_t *operation,
                                      uint8_t *plaintext,
                                      size_t plaintext_size,
                                      size_t *plaintext_length,
@@ -723,11 +735,11 @@
  * mbedtls_psa_aead_encrypt_setup() or mbedtls_psa_aead_decrypt_setup() again.
  *
  * You may call this function any time after the operation object has
- * been initialized as described in #psa_aead_operation_t.
+ * been initialized as described in #mbedtls_psa_aead_operation_t.
  *
  * In particular, calling mbedtls_psa_aead_abort() after the operation has been
- * terminated by a call to mbedtls_psa_aead_abort(), mbedtls_psa_aead_finish() or
- * mbedtls_psa_aead_verify() is safe and has no effect.
+ * terminated by a call to mbedtls_psa_aead_abort(), mbedtls_psa_aead_finish()
+ * or mbedtls_psa_aead_verify() is safe and has no effect.
  *
  * \param[in,out] operation     Initialized AEAD operation.
  *
@@ -740,7 +752,7 @@
  *         It is implementation-dependent whether a failure to initialize
  *         results in this error code.
  */
-psa_status_t mbedtls_psa_aead_abort(psa_aead_operation_t *operation);
+psa_status_t mbedtls_psa_aead_abort(mbedtls_psa_aead_operation_t *operation);
 
 
 #endif /* PSA_CRYPTO_AEAD */
diff --git a/library/psa_crypto_driver_wrappers.c b/library/psa_crypto_driver_wrappers.c
index 7faedb3..7590800 100644
--- a/library/psa_crypto_driver_wrappers.c
+++ b/library/psa_crypto_driver_wrappers.c
@@ -1310,10 +1310,9 @@
 
 #if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
 #if defined(PSA_CRYPTO_DRIVER_TEST)
-            status = PSA_ERROR_NOT_SUPPORTED;
             status = mbedtls_test_transparent_aead_encrypt_setup(
-                        operation, attributes,
-                        key_buffer, key_buffer_size,
+                        &operation->ctx.transparent_test_driver_ctx,
+                        attributes, key_buffer, key_buffer_size,
                         alg );
             /* Declared with fallback == true */
             operation->id = PSA_CRYPTO_TRANSPARENT_TEST_DRIVER_ID;
@@ -1325,7 +1324,7 @@
 
             /* Fell through, meaning no accelerator supports this operation */
             status = mbedtls_psa_aead_encrypt_setup(
-                        operation, attributes,
+                        &operation->ctx.mbedtls_ctx, attributes,
                         key_buffer, key_buffer_size,
                         alg );
 
@@ -1360,9 +1359,9 @@
 
 #if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
 #if defined(PSA_CRYPTO_DRIVER_TEST)
-            status = PSA_ERROR_NOT_SUPPORTED;
             status = mbedtls_test_transparent_aead_decrypt_setup(
-                        operation, attributes,
+                        &operation->ctx.transparent_test_driver_ctx,
+                         attributes,
                         key_buffer, key_buffer_size,
                         alg );
             /* Declared with fallback == true */
@@ -1375,7 +1374,8 @@
 
             /* Fell through, meaning no accelerator supports this operation */
             status = mbedtls_psa_aead_decrypt_setup(
-                        operation, attributes,
+                        &operation->ctx.mbedtls_ctx,
+                         attributes,
                         key_buffer, key_buffer_size,
                         alg );
 
@@ -1401,16 +1401,18 @@
     {
 #if defined(MBEDTLS_PSA_BUILTIN_CIPHER)
         case PSA_CRYPTO_MBED_TLS_DRIVER_ID:
-            return( mbedtls_psa_aead_set_nonce( operation, nonce,
+            return( mbedtls_psa_aead_set_nonce( &operation->ctx.mbedtls_ctx,
+                                                nonce,
                                                 nonce_length ) );
 
-#endif /* MBEDTLS_PSA_BUILTIN_CIPHER */
+#endif /* MBEDTLS_PSA_BUILTIN_AEAD */
 
 #if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
 #if defined(PSA_CRYPTO_DRIVER_TEST)
         case PSA_CRYPTO_TRANSPARENT_TEST_DRIVER_ID:
             return( mbedtls_test_transparent_aead_set_nonce(
-                        operation, nonce, nonce_length ) );
+                         &operation->ctx.transparent_test_driver_ctx,
+                         nonce, nonce_length ) );
 
         /* Add cases for opaque driver here */
 
@@ -1431,18 +1433,20 @@
 {
     switch( operation->id )
     {
-#if defined(MBEDTLS_PSA_BUILTIN_CIPHER)
+#if defined(MBEDTLS_PSA_BUILTIN_AEAD)
         case PSA_CRYPTO_MBED_TLS_DRIVER_ID:
-            return( mbedtls_psa_aead_set_lengths( operation, ad_length,
+            return( mbedtls_psa_aead_set_lengths( &operation->ctx.mbedtls_ctx,
+                                                  ad_length,
                                                   plaintext_length ) );
 
-#endif /* MBEDTLS_PSA_BUILTIN_CIPHER */
+#endif /* MBEDTLS_PSA_BUILTIN_AEAD */
 
 #if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
 #if defined(PSA_CRYPTO_DRIVER_TEST)
         case PSA_CRYPTO_TRANSPARENT_TEST_DRIVER_ID:
             return( mbedtls_test_transparent_aead_set_lengths(
-                        operation, ad_length, plaintext_length ) );
+                        &operation->ctx.transparent_test_driver_ctx,
+                        ad_length, plaintext_length ) );
 
         /* Add cases for opaque driver here */
 
@@ -1463,18 +1467,20 @@
 {
     switch( operation->id )
     {
-#if defined(MBEDTLS_PSA_BUILTIN_CIPHER)
+#if defined(MBEDTLS_PSA_BUILTIN_AEAD)
         case PSA_CRYPTO_MBED_TLS_DRIVER_ID:
-            return( mbedtls_psa_aead_update_ad( operation, input,
+            return( mbedtls_psa_aead_update_ad( &operation->ctx.mbedtls_ctx,
+                                                input,
                                                 input_length ) );
 
-#endif /* MBEDTLS_PSA_BUILTIN_CIPHER */
+#endif /* MBEDTLS_PSA_BUILTIN_AEAD */
 
 #if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
 #if defined(PSA_CRYPTO_DRIVER_TEST)
         case PSA_CRYPTO_TRANSPARENT_TEST_DRIVER_ID:
             return( mbedtls_test_transparent_aead_update_ad(
-                        operation, input, input_length ) );
+                        &operation->ctx.transparent_test_driver_ctx,
+                        input, input_length ) );
 
         /* Add cases for opaque driver here */
 
@@ -1498,19 +1504,21 @@
 {
     switch( operation->id )
     {
-#if defined(MBEDTLS_PSA_BUILTIN_CIPHER)
+#if defined(MBEDTLS_PSA_BUILTIN_AEAD)
         case PSA_CRYPTO_MBED_TLS_DRIVER_ID:
-            return( mbedtls_psa_aead_update( operation, input, input_length,
+            return( mbedtls_psa_aead_update( &operation->ctx.mbedtls_ctx,
+                                             input, input_length,
                                              output, output_size,
                                              output_length ) );
 
-#endif /* MBEDTLS_PSA_BUILTIN_CIPHER */
+#endif /* MBEDTLS_PSA_BUILTIN_AEAD */
 
 #if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
 #if defined(PSA_CRYPTO_DRIVER_TEST)
         case PSA_CRYPTO_TRANSPARENT_TEST_DRIVER_ID:
             return( mbedtls_test_transparent_aead_update(
-                        operation, input, input_length, output, output_size,
+                        &operation->ctx.transparent_test_driver_ctx,
+                        input, input_length, output, output_size,
                         output_length ) );
 
         /* Add cases for opaque driver here */
@@ -1539,20 +1547,22 @@
 {
     switch( operation->id )
     {
-#if defined(MBEDTLS_PSA_BUILTIN_CIPHER)
+#if defined(MBEDTLS_PSA_BUILTIN_AEAD)
         case PSA_CRYPTO_MBED_TLS_DRIVER_ID:
-            return( mbedtls_psa_aead_finish( operation, ciphertext,
+            return( mbedtls_psa_aead_finish( &operation->ctx.mbedtls_ctx,
+                                             ciphertext,
                                              ciphertext_size,
                                              ciphertext_length, tag,
                                              tag_size, tag_length ) );
 
-#endif /* MBEDTLS_PSA_BUILTIN_CIPHER */
+#endif /* MBEDTLS_PSA_BUILTIN_AEAD */
 
 #if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
 #if defined(PSA_CRYPTO_DRIVER_TEST)
         case PSA_CRYPTO_TRANSPARENT_TEST_DRIVER_ID:
             return( mbedtls_test_transparent_aead_finish(
-                        operation, ciphertext, ciphertext_size,
+                        &operation->ctx.transparent_test_driver_ctx,
+                        ciphertext, ciphertext_size,
                         ciphertext_length, tag, tag_size, tag_length ) );
 
         /* Add cases for opaque driver here */
@@ -1581,19 +1591,22 @@
 {
     switch( operation->id )
     {
-#if defined(MBEDTLS_PSA_BUILTIN_CIPHER)
+#if defined(MBEDTLS_PSA_BUILTIN_AEAD)
         case PSA_CRYPTO_MBED_TLS_DRIVER_ID:
-            return( mbedtls_psa_aead_verify( operation, plaintext,
-                                             plaintext_size, plaintext_length,
+            return( mbedtls_psa_aead_verify( &operation->ctx.mbedtls_ctx,
+                                             plaintext,
+                                             plaintext_size,
+                                             plaintext_length,
                                              tag, tag_length ) );
 
-#endif /* MBEDTLS_PSA_BUILTIN_CIPHER */
+#endif /* MBEDTLS_PSA_BUILTIN_AEAD */
 
 #if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
 #if defined(PSA_CRYPTO_DRIVER_TEST)
         case PSA_CRYPTO_TRANSPARENT_TEST_DRIVER_ID:
             return( mbedtls_test_transparent_aead_verify(
-                        operation, plaintext, plaintext_size,
+                        &operation->ctx.transparent_test_driver_ctx,
+                        plaintext, plaintext_size,
                         plaintext_length, tag, tag_length ) );
 
         /* Add cases for opaque driver here */
@@ -1616,16 +1629,17 @@
 {
     switch( operation->id )
     {
-#if defined(MBEDTLS_PSA_BUILTIN_CIPHER)
+#if defined(MBEDTLS_PSA_BUILTIN_AEAD)
         case PSA_CRYPTO_MBED_TLS_DRIVER_ID:
-            return( mbedtls_psa_aead_abort( operation ) );
+            return( mbedtls_psa_aead_abort( &operation->ctx.mbedtls_ctx ) );
 
-#endif /* MBEDTLS_PSA_BUILTIN_CIPHER */
+#endif /* MBEDTLS_PSA_BUILTIN_AEAD */
 
 #if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
 #if defined(PSA_CRYPTO_DRIVER_TEST)
         case PSA_CRYPTO_TRANSPARENT_TEST_DRIVER_ID:
-            return( mbedtls_test_transparent_aead_abort( operation ) );
+            return( mbedtls_test_transparent_aead_abort(
+               &operation->ctx.transparent_test_driver_ctx ) );
 
         /* Add cases for opaque driver here */