diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index 69852fa..ad6f052 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -1085,7 +1085,8 @@
                 ( 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 @@
                                                 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 @@
                                          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 );
 }