diff --git a/include/psa/crypto_builtin_composites.h b/include/psa/crypto_builtin_composites.h
index 1d11b003e..b65922b9c 100644
--- a/include/psa/crypto_builtin_composites.h
+++ b/include/psa/crypto_builtin_composites.h
@@ -76,6 +76,58 @@ typedef struct
 
 #define MBEDTLS_PSA_MAC_OPERATION_INIT {0, {0}}
 
+#if defined(MBEDTLS_PSA_BUILTIN_ALG_GCM) || \
+    defined(MBEDTLS_PSA_BUILTIN_ALG_CCM) || \
+    defined(MBEDTLS_PSA_BUILTIN_ALG_CHACHA20_POLY1305)
+#define MBEDTLS_PSA_BUILTIN_AEAD  1
+#endif
+
+/* Context structure for the Mbed TLS AEAD implementation. */
+typedef struct
+{
+    psa_algorithm_t alg;
+    psa_key_type_t key_type;
+
+    unsigned int is_encrypt : 1;
+    unsigned int ad_started : 1;
+    unsigned int body_started : 1;
+
+    uint8_t tag_length;
+
+    /* Buffers for AD/data - only required until CCM gets proper multipart
+       support. */
+    uint8_t *ad_buffer;
+    size_t ad_length;
+
+    uint8_t *body_buffer;
+    size_t body_length;
+
+    uint8_t *tag_buffer;
+
+    /* buffer to store Nonce - only required until CCM and GCM get proper
+       multipart support. */
+    uint8_t *nonce;
+    size_t nonce_length;
+
+    union
+    {
+        unsigned dummy; /* Enable easier initializing of the union. */
+#if defined(MBEDTLS_PSA_BUILTIN_ALG_CCM)
+        mbedtls_ccm_context ccm;
+#endif /* MBEDTLS_PSA_BUILTIN_ALG_CCM */
+#if defined(MBEDTLS_PSA_BUILTIN_ALG_GCM)
+        mbedtls_gcm_context gcm;
+#endif /* MBEDTLS_PSA_BUILTIN_ALG_GCM */
+#if defined(MBEDTLS_PSA_BUILTIN_ALG_CHACHA20_POLY1305)
+        mbedtls_chachapoly_context chachapoly;
+#endif /* MBEDTLS_PSA_BUILTIN_ALG_CHACHA20_POLY1305 */
+
+    } ctx;
+
+} mbedtls_psa_aead_operation_t;
+
+#define MBEDTLS_PSA_AEAD_OPERATION_INIT {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, {0}}
+
 /*
  * BEYOND THIS POINT, TEST DRIVER DECLARATIONS ONLY.
  */
@@ -87,6 +139,10 @@ typedef mbedtls_psa_mac_operation_t mbedtls_opaque_test_driver_mac_operation_t;
 #define MBEDTLS_TRANSPARENT_TEST_DRIVER_MAC_OPERATION_INIT MBEDTLS_PSA_MAC_OPERATION_INIT
 #define MBEDTLS_OPAQUE_TEST_DRIVER_MAC_OPERATION_INIT MBEDTLS_PSA_MAC_OPERATION_INIT
 
+typedef mbedtls_psa_aead_operation_t mbedtls_transparent_test_driver_aead_operation_t;
+
+#define MBEDTLS_TRANSPARENT_TEST_DRIVER_AEAD_OPERATION_INIT MBEDTLS_PSA_AEAD_OPERATION_INIT
+
 #endif /* PSA_CRYPTO_DRIVER_TEST */
 
 #endif /* PSA_CRYPTO_BUILTIN_COMPOSITES_H */
diff --git a/include/psa/crypto_builtin_primitives.h b/include/psa/crypto_builtin_primitives.h
index b67b23ff1..75801a178 100644
--- a/include/psa/crypto_builtin_primitives.h
+++ b/include/psa/crypto_builtin_primitives.h
@@ -118,58 +118,6 @@ typedef struct {
 
 #define MBEDTLS_PSA_CIPHER_OPERATION_INIT {0, 0, 0, {0}}
 
-#if defined(MBEDTLS_PSA_BUILTIN_ALG_GCM) || \
-    defined(MBEDTLS_PSA_BUILTIN_ALG_CCM) || \
-    defined(MBEDTLS_PSA_BUILTIN_ALG_CHACHA20_POLY1305)
-#define MBEDTLS_PSA_BUILTIN_AEAD  1
-#endif
-
-/* Context structure for the Mbed TLS cipher implementation. */
-typedef struct
-{
-    psa_algorithm_t alg;
-    psa_key_type_t key_type;
-
-    unsigned int is_encrypt : 1;
-    unsigned int ad_started : 1;
-    unsigned int body_started : 1;
-
-    uint8_t tag_length;
-
-    /* Buffers for AD/data - only required until CCM gets proper multipart
-       support. */
-    uint8_t *ad_buffer;
-    size_t ad_length;
-
-    uint8_t *body_buffer;
-    size_t body_length;
-
-    uint8_t *tag_buffer;
-
-    /* buffer to store Nonce - only required until CCM and GCM get proper
-       multipart support. */
-    uint8_t *nonce;
-    size_t nonce_length;
-
-    union
-    {
-        unsigned dummy; /* Enable easier initializing of the union. */
-#if defined(MBEDTLS_PSA_BUILTIN_ALG_CCM)
-        mbedtls_ccm_context ccm;
-#endif /* MBEDTLS_PSA_BUILTIN_ALG_CCM */
-#if defined(MBEDTLS_PSA_BUILTIN_ALG_GCM)
-        mbedtls_gcm_context gcm;
-#endif /* MBEDTLS_PSA_BUILTIN_ALG_GCM */
-#if defined(MBEDTLS_PSA_BUILTIN_ALG_CHACHA20_POLY1305)
-        mbedtls_chachapoly_context chachapoly;
-#endif /* MBEDTLS_PSA_BUILTIN_ALG_CHACHA20_POLY1305 */
-
-    } ctx;
-
-} mbedtls_psa_aead_operation_t;
-
-#define MBEDTLS_PSA_AEAD_OPERATION_INIT {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, {0}}
-
 /*
  * BEYOND THIS POINT, TEST DRIVER DECLARATIONS ONLY.
  */
@@ -182,9 +130,6 @@ typedef mbedtls_psa_hash_operation_t mbedtls_transparent_test_driver_hash_operat
 typedef mbedtls_psa_cipher_operation_t
         mbedtls_transparent_test_driver_cipher_operation_t;
 
-typedef mbedtls_psa_aead_operation_t
-        mbedtls_transparent_test_driver_aead_operation_t;
-
 typedef struct {
     unsigned int initialised : 1;
     mbedtls_transparent_test_driver_cipher_operation_t ctx;
diff --git a/include/psa/crypto_driver_contexts_composites.h b/include/psa/crypto_driver_contexts_composites.h
index 239fdcb33..957986c22 100644
--- a/include/psa/crypto_driver_contexts_composites.h
+++ b/include/psa/crypto_driver_contexts_composites.h
@@ -58,5 +58,13 @@ typedef union {
 #endif
 } psa_driver_mac_context_t;
 
+typedef union {
+    unsigned dummy; /* Make sure this union is always non-empty */
+    mbedtls_psa_aead_operation_t mbedtls_ctx;
+#if defined(PSA_CRYPTO_DRIVER_TEST)
+    mbedtls_transparent_test_driver_aead_operation_t transparent_test_driver_ctx;
+#endif
+} psa_driver_aead_context_t;
+
 #endif /* PSA_CRYPTO_DRIVER_CONTEXTS_COMPOSITES_H */
 /* End of automatically generated file. */
diff --git a/include/psa/crypto_driver_contexts_primitives.h b/include/psa/crypto_driver_contexts_primitives.h
index 4fba9eb03..104d4bdb6 100644
--- a/include/psa/crypto_driver_contexts_primitives.h
+++ b/include/psa/crypto_driver_contexts_primitives.h
@@ -65,13 +65,5 @@ typedef union {
 #endif
 } psa_driver_cipher_context_t;
 
-typedef union {
-    unsigned dummy; /* Make sure this union is always non-empty */
-    mbedtls_psa_aead_operation_t mbedtls_ctx;
-#if defined(PSA_CRYPTO_DRIVER_TEST)
-    mbedtls_transparent_test_driver_aead_operation_t transparent_test_driver_ctx;
-#endif
-} psa_driver_aead_context_t;
-
 #endif /* PSA_CRYPTO_DRIVER_CONTEXTS_PRIMITIVES_H */
 /* End of automatically generated file. */