diff --git a/include/psa/crypto_builtin_primitives.h b/include/psa/crypto_builtin_primitives.h
index c76bc78..2830b61 100644
--- a/include/psa/crypto_builtin_primitives.h
+++ b/include/psa/crypto_builtin_primitives.h
@@ -111,4 +111,42 @@
 
 #define MBEDTLS_PSA_CIPHER_OPERATION_INIT { 0, 0, 0, { 0 } }
 
+
+/* EC-JPAKE operation definitions */
+
+#include "mbedtls/ecjpake.h"
+
+#if defined(MBEDTLS_PSA_BUILTIN_ALG_JPAKE)
+#define MBEDTLS_PSA_BUILTIN_PAKE  1
+#endif
+
+/* Note: the format for mbedtls_ecjpake_read/write function has an extra
+ * length byte for each step, plus an extra 3 bytes for ECParameters in the
+ * server's 2nd round. */
+#define MBEDTLS_PSA_PAKE_BUFFER_SIZE ((3 + 1 + 65 + 1 + 65 + 1 + 32) * 2)
+
+typedef struct {
+    psa_algorithm_t MBEDTLS_PRIVATE(alg);
+    unsigned int MBEDTLS_PRIVATE(state);
+    unsigned int MBEDTLS_PRIVATE(sequence);
+#if defined(MBEDTLS_PSA_BUILTIN_PAKE)
+    unsigned int MBEDTLS_PRIVATE(input_step);
+    unsigned int MBEDTLS_PRIVATE(output_step);
+    uint8_t *MBEDTLS_PRIVATE(password);
+    size_t MBEDTLS_PRIVATE(password_len);
+    uint8_t MBEDTLS_PRIVATE(role);
+    uint8_t MBEDTLS_PRIVATE(buffer[MBEDTLS_PSA_PAKE_BUFFER_SIZE]);
+    size_t MBEDTLS_PRIVATE(buffer_length);
+    size_t MBEDTLS_PRIVATE(buffer_offset);
+#endif
+    /* Context structure for the Mbed TLS EC-JPAKE implementation. */
+    union {
+        unsigned int MBEDTLS_PRIVATE(dummy);
+        mbedtls_ecjpake_context MBEDTLS_PRIVATE(pake);
+    } MBEDTLS_PRIVATE(ctx);
+
+} mbedtls_psa_pake_operation_t;
+
+#define MBEDTLS_PSA_PAKE_OPERATION_INIT { { 0 } }
+
 #endif /* PSA_CRYPTO_BUILTIN_PRIMITIVES_H */
diff --git a/include/psa/crypto_driver_contexts_primitives.h b/include/psa/crypto_driver_contexts_primitives.h
index 620a4b3..7d09620 100644
--- a/include/psa/crypto_driver_contexts_primitives.h
+++ b/include/psa/crypto_driver_contexts_primitives.h
@@ -45,6 +45,8 @@
 #include <libtestdriver1/include/psa/crypto.h>
 #endif
 
+#include "mbedtls/ecjpake.h"
+
 #if defined(PSA_CRYPTO_DRIVER_TEST)
 
 #if defined(MBEDTLS_TEST_LIBTESTDRIVER1) && \
@@ -87,6 +89,32 @@
 #define MBEDTLS_OPAQUE_TEST_DRIVER_CIPHER_OPERATION_INIT \
     { 0, MBEDTLS_TRANSPARENT_TEST_DRIVER_CIPHER_OPERATION_INIT }
 
+#if defined(MBEDTLS_TEST_LIBTESTDRIVER1) && \
+    defined(LIBTESTDRIVER1_MBEDTLS_PSA_BUILTIN_PAKE)
+
+typedef libtestdriver1_mbedtls_psa_pake_operation_t
+    mbedtls_transparent_test_driver_pake_operation_t;
+typedef libtestdriver1_psa_pake_operation_t
+    mbedtls_opaque_test_driver_pake_operation_t;
+
+#define MBEDTLS_TRANSPARENT_TEST_DRIVER_PAKE_OPERATION_INIT \
+    LIBTESTDRIVER1_MBEDTLS_PSA_PAKE_OPERATION_INIT
+#define MBEDTLS_OPAQUE_TEST_DRIVER_PAKE_OPERATION_INIT \
+    LIBTESTDRIVER1_MBEDTLS_PSA_PAKE_OPERATION_INIT
+
+#else
+typedef mbedtls_psa_pake_operation_t
+    mbedtls_transparent_test_driver_pake_operation_t;
+typedef mbedtls_psa_pake_operation_t
+    mbedtls_opaque_test_driver_pake_operation_t;
+
+#define MBEDTLS_TRANSPARENT_TEST_DRIVER_PAKE_OPERATION_INIT \
+    MBEDTLS_PSA_PAKE_OPERATION_INIT
+#define MBEDTLS_OPAQUE_TEST_DRIVER_PAKE_OPERATION_INIT \
+    MBEDTLS_PSA_PAKE_OPERATION_INIT
+
+#endif /* MBEDTLS_TEST_LIBTESTDRIVER1 && LIBTESTDRIVER1_MBEDTLS_PSA_BUILTIN_PAKE */
+
 #endif /* PSA_CRYPTO_DRIVER_TEST */
 
 /* Define the context to be used for an operation that is executed through the
@@ -113,5 +141,14 @@
 #endif
 } psa_driver_cipher_context_t;
 
+typedef union {
+    unsigned dummy; /* Make sure this union is always non-empty */
+    mbedtls_psa_pake_operation_t mbedtls_ctx;
+#if defined(PSA_CRYPTO_DRIVER_TEST)
+    mbedtls_transparent_test_driver_pake_operation_t transparent_test_driver_ctx;
+    mbedtls_opaque_test_driver_pake_operation_t opaque_test_driver_ctx;
+#endif
+} psa_driver_pake_context_t;
+
 #endif /* PSA_CRYPTO_DRIVER_CONTEXTS_PRIMITIVES_H */
 /* End of automatically generated file. */
diff --git a/include/psa/crypto_extra.h b/include/psa/crypto_extra.h
index bd1b5af..da74bed 100644
--- a/include/psa/crypto_extra.h
+++ b/include/psa/crypto_extra.h
@@ -1826,14 +1826,7 @@
 /** Returns a suitable initializer for a PAKE operation object of type
  * psa_pake_operation_t.
  */
-#if defined(MBEDTLS_PSA_BUILTIN_PAKE)
-#define PSA_PAKE_OPERATION_INIT { PSA_ALG_NONE, 0, 0, 0, 0,              \
-                                  NULL, 0,               \
-                                  PSA_PAKE_ROLE_NONE, { 0 }, 0, 0,         \
-                                  { .dummy = 0 } }
-#else
-#define PSA_PAKE_OPERATION_INIT { PSA_ALG_NONE, 0, 0, { 0 } }
-#endif
+#define PSA_PAKE_OPERATION_INIT { 0, { .dummy = 0 } }
 
 struct psa_pake_cipher_suite_s {
     psa_algorithm_t algorithm;
@@ -1904,35 +1897,15 @@
     }
 }
 
-#if defined(MBEDTLS_PSA_BUILTIN_ALG_JPAKE)
-#include <mbedtls/ecjpake.h>
-/* Note: the format for mbedtls_ecjpake_read/write function has an extra
- * length byte for each step, plus an extra 3 bytes for ECParameters in the
- * server's 2nd round. */
-#define MBEDTLS_PSA_PAKE_BUFFER_SIZE ((3 + 1 + 65 + 1 + 65 + 1 + 32) * 2)
-#endif
-
 struct psa_pake_operation_s {
-    psa_algorithm_t MBEDTLS_PRIVATE(alg);
-    unsigned int MBEDTLS_PRIVATE(state);
-    unsigned int MBEDTLS_PRIVATE(sequence);
-#if defined(MBEDTLS_PSA_BUILTIN_PAKE)
-    unsigned int MBEDTLS_PRIVATE(input_step);
-    unsigned int MBEDTLS_PRIVATE(output_step);
-    uint8_t *MBEDTLS_PRIVATE(password);
-    size_t MBEDTLS_PRIVATE(password_len);
-    psa_pake_role_t MBEDTLS_PRIVATE(role);
-    uint8_t MBEDTLS_PRIVATE(buffer[MBEDTLS_PSA_PAKE_BUFFER_SIZE]);
-    size_t MBEDTLS_PRIVATE(buffer_length);
-    size_t MBEDTLS_PRIVATE(buffer_offset);
-#endif
-    union {
-#if defined(MBEDTLS_PSA_BUILTIN_ALG_JPAKE)
-        mbedtls_ecjpake_context ecjpake;
-#endif
-        /* Make the union non-empty even with no supported algorithms. */
-        uint8_t dummy;
-    } MBEDTLS_PRIVATE(ctx);
+    /** Unique ID indicating which driver got assigned to do the
+     * operation. Since driver contexts are driver-specific, swapping
+     * drivers halfway through the operation is not supported.
+     * ID values are auto-generated in psa_crypto_driver_wrappers.h
+     * ID value zero means the context is not valid or not assigned to
+     * any driver (i.e. none of the driver contexts are active). */
+    unsigned int MBEDTLS_PRIVATE(id);
+    psa_driver_pake_context_t MBEDTLS_PRIVATE(ctx);
 };
 
 static inline struct psa_pake_cipher_suite_s psa_pake_cipher_suite_init(void)