Add PSA interuptable key agreement APIs

Signed-off-by: Waleed Elmelegy <waleed.elmelegy@arm.com>
diff --git a/tf-psa-crypto/include/psa/crypto.h b/tf-psa-crypto/include/psa/crypto.h
index 58b6887..62f1dca 100644
--- a/tf-psa-crypto/include/psa/crypto.h
+++ b/tf-psa-crypto/include/psa/crypto.h
@@ -5123,7 +5123,7 @@
 
 psa_status_t psa_key_agreement_iop_setup(
     psa_key_agreement_iop_t *operation,
-    psa_key_id_t private_key,
+    mbedtls_svc_key_id_t private_key,
     const uint8_t *peer_key,
     size_t peer_key_length,
     psa_algorithm_t alg,
@@ -5208,7 +5208,7 @@
  */
 psa_status_t psa_key_agreement_iop_complete(
     psa_key_agreement_iop_t *operation,
-    psa_key_id_t *key);
+    mbedtls_svc_key_id_t *key);
 
 /**
  * \brief                       Abort a key agreement operation.
diff --git a/tf-psa-crypto/include/psa/crypto_builtin_composites.h b/tf-psa-crypto/include/psa/crypto_builtin_composites.h
index c9c0c6b..c5dde0d 100644
--- a/tf-psa-crypto/include/psa/crypto_builtin_composites.h
+++ b/tf-psa-crypto/include/psa/crypto_builtin_composites.h
@@ -33,6 +33,10 @@
 #endif
 #include "mbedtls/chachapoly.h"
 
+#if defined(MBEDTLS_PSA_BUILTIN_ALG_ECDH)
+#include "mbedtls/ecdh.h"
+#endif
+
 /*
  * MAC multi-part operation definitions.
  */
@@ -227,4 +231,22 @@
 #define MBEDTLS_PSA_GENERATE_KEY_IOP_INIT { 0 }
 #endif
 
+/* Context structure for the Mbed TLS interruptible key agreement implementation. */
+typedef struct {
+#if defined(MBEDTLS_PSA_BUILTIN_ALG_ECDH) && defined(MBEDTLS_ECP_RESTARTABLE)
+    mbedtls_ecdh_context MBEDTLS_PRIVATE(ctx);
+    const psa_key_attributes_t *MBEDTLS_PRIVATE(attributes);
+    uint32_t MBEDTLS_PRIVATE(num_ops);
+#else
+    /* Make the struct non-empty if algs not supported. */
+    unsigned MBEDTLS_PRIVATE(dummy);
+#endif
+} mbedtls_psa_key_agreement_interruptible_operation_t;
+
+#if defined(MBEDTLS_PSA_BUILTIN_ALG_ECDH) && defined(MBEDTLS_ECP_RESTARTABLE)
+#define MBEDTLS_PSA_KEY_AGREEMENT_INTERRUPTIBLE_OPERATION_INIT { { 0 }, { 0 }, 0 }
+#else
+#define MBEDTLS_PSA_KEY_AGREEMENT_INTERRUPTIBLE_OPERATION_INIT { 0 }
+#endif
+
 #endif /* PSA_CRYPTO_BUILTIN_COMPOSITES_H */
diff --git a/tf-psa-crypto/include/psa/crypto_driver_contexts_composites.h b/tf-psa-crypto/include/psa/crypto_driver_contexts_composites.h
index 5a484fc..086c0c2 100644
--- a/tf-psa-crypto/include/psa/crypto_driver_contexts_composites.h
+++ b/tf-psa-crypto/include/psa/crypto_driver_contexts_composites.h
@@ -147,5 +147,10 @@
 #endif
 } psa_driver_pake_context_t;
 
+typedef union {
+    unsigned dummy; /* Make sure this union is always non-empty */
+    mbedtls_psa_key_agreement_interruptible_operation_t mbedtls_ctx;
+} psa_driver_key_agreement_interruptible_context_t;
+
 #endif /* PSA_CRYPTO_DRIVER_CONTEXTS_COMPOSITES_H */
 /* End of automatically generated file. */
diff --git a/tf-psa-crypto/include/psa/crypto_struct.h b/tf-psa-crypto/include/psa/crypto_struct.h
index 76ef5c4..b300557 100644
--- a/tf-psa-crypto/include/psa/crypto_struct.h
+++ b/tf-psa-crypto/include/psa/crypto_struct.h
@@ -508,14 +508,15 @@
      * any driver (i.e. none of the driver contexts are active).
      */
     unsigned int MBEDTLS_PRIVATE(id);
-
+    psa_driver_key_agreement_interruptible_context_t MBEDTLS_PRIVATE(ctx);
+    uint32_t MBEDTLS_PRIVATE(num_ops);
 #endif
 };
 
 #if defined(MBEDTLS_PSA_CRYPTO_CLIENT) && !defined(MBEDTLS_PSA_CRYPTO_C)
 #define PSA_KEY_AGREEMENT_IOP_INIT { 0 }
 #else
-#define PSA_KEY_AGREEMENT_IOP_INIT { 0 }
+#define PSA_KEY_AGREEMENT_IOP_INIT { 0, { 0 }, 0 }
 #endif
 
 static inline struct psa_key_agreement_iop_s