diff --git a/ChangeLog.d/support-ecdh-kdf-with-ecdh-key.txt b/ChangeLog.d/support-ecdh-kdf-with-ecdh-key.txt
index 842d5ffd3..6660dc3d5 100644
--- a/ChangeLog.d/support-ecdh-kdf-with-ecdh-key.txt
+++ b/ChangeLog.d/support-ecdh-kdf-with-ecdh-key.txt
@@ -2,3 +2,5 @@ Features
    * In PSA, allow using a key declared with a base key agreement algorithm
      in combined key agreement and derivation operations, as long as the key
      agreement algorithm in use matches the algorithm the key was declared with.
+     This is currently non-standard behaviour, but expected to make it into a
+     future revision of the PSA Crypto standard.
diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index 69852fa56..ad6f052af 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -1085,7 +1085,8 @@ static int psa_key_algorithm_permits( psa_algorithm_t policy_alg,
                 ( requested_alg & ~PSA_ALG_HASH_MASK ) );
     }
     /* If policy_alg is a generic key agreement operation, then using it for
-     * a key derivation with that key agreement is also compliant. */
+     * a key derivation with that key agreement should also be allowed. This
+     * behaviour is expected to be defined in a future specification version. */
     if( PSA_ALG_IS_RAW_KEY_AGREEMENT( policy_alg ) &&
         PSA_ALG_IS_KEY_AGREEMENT( requested_alg ) )
     {
@@ -5876,12 +5877,6 @@ static psa_status_t psa_key_agreement_internal( psa_key_derivation_operation_t *
                                                 PSA_KEY_TYPE_DERIVE,
                                                 shared_secret,
                                                 shared_secret_length );
-
-    /* If a private key has been added as SECRET, we allow the derived
-     * key material to be used as a key in PSA Crypto. */
-    if( step == PSA_KEY_DERIVATION_INPUT_SECRET )
-        operation->can_output_key = 1;
-
 exit:
     mbedtls_platform_zeroize( shared_secret, shared_secret_length );
     return( status );
@@ -5906,6 +5901,13 @@ psa_status_t psa_key_derivation_key_agreement( psa_key_derivation_operation_t *o
                                          peer_key, peer_key_length );
     if( status != PSA_SUCCESS )
         psa_key_derivation_abort( operation );
+    else
+    {
+        /* If a private key has been added as SECRET, we allow the derived
+         * key material to be used as a key in PSA Crypto. */
+        if( step == PSA_KEY_DERIVATION_INPUT_SECRET )
+            operation->can_output_key = 1;
+    }
     return( status );
 }
 
diff --git a/tests/suites/test_suite_psa_crypto.function b/tests/suites/test_suite_psa_crypto.function
index 328508dd2..fc563cb15 100644
--- a/tests/suites/test_suite_psa_crypto.function
+++ b/tests/suites/test_suite_psa_crypto.function
@@ -5233,7 +5233,7 @@ void key_agreement_setup( int alg_arg,
 {
     psa_key_handle_t our_key = 0;
     psa_algorithm_t alg = alg_arg;
-    psa_algorithm_t key_alg = our_key_alg_arg;
+    psa_algorithm_t our_key_alg = our_key_alg_arg;
     psa_key_type_t our_key_type = our_key_type_arg;
     psa_key_derivation_operation_t operation = PSA_KEY_DERIVATION_OPERATION_INIT;
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
@@ -5243,7 +5243,7 @@ void key_agreement_setup( int alg_arg,
     PSA_ASSERT( psa_crypto_init( ) );
 
     psa_set_key_usage_flags( &attributes, PSA_KEY_USAGE_DERIVE );
-    psa_set_key_algorithm( &attributes, key_alg );
+    psa_set_key_algorithm( &attributes, our_key_alg );
     psa_set_key_type( &attributes, our_key_type );
     PSA_ASSERT( psa_import_key( &attributes,
                                 our_key_data->x, our_key_data->len,