diff --git a/include/psa/crypto.h b/include/psa/crypto.h
index 6005269c3..f1731f694 100644
--- a/include/psa/crypto.h
+++ b/include/psa/crypto.h
@@ -2147,29 +2147,21 @@ psa_status_t psa_generator_abort(psa_crypto_generator_t *generator);
 psa_status_t psa_key_derivation_setup(psa_crypto_generator_t *generator,
                                       psa_algorithm_t alg);
 
-/** Provide an input for key derivation.
+/** Provide an input for key derivation or key agreement.
  *
- * Which inputs are required and in what order depends on the type of
- * key derivation algorithm.
+ * Which inputs are required and in what order depends on the algorithm.
+ * Refer to the documentation of each key derivation or key agreement
+ * algorithm for information.
  *
- * - For HKDF (#PSA_ALG_HKDF), the following inputs are supported:
- *     - #PSA_KDF_STEP_SALT is the salt used in the "extract" step.
- *       It is optional; if omitted, the derivation uses an empty salt.
- *     - #PSA_KDF_STEP_SECRET is the secret key used in the "extract" step.
- *       It may be a key of type #PSA_KEY_TYPE_DERIVE with the
- *       usage flag #PSA_KEY_USAGE_DERIVE.
- *     - #PSA_KDF_STEP_INFO is the info string used in the "expand" step.
- *   You must pass #PSA_KDF_STEP_SALT before #PSA_KDF_STEP_SECRET.
- *   #PSA_KDF_STEP_INFO may be passed at any time before starting to
- *   generate output.
+ * This function passes direct inputs. Some inputs must be passed as keys
+ * using psa_key_derivation_input_key() instead of this function. Refer to
+ * the documentation of individual step types for information.
  *
  * \param[in,out] generator       The generator object to use. It must
  *                                have been set up with
  *                                psa_key_derivation_setup() and must not
  *                                have produced any output yet.
  * \param step                    Which step the input data is for.
- *                                See above for the permitted values
- *                                depending on the algorithm.
  * \param[in] data                Input data to use.
  * \param data_length             Size of the \p data buffer in bytes.
  *
@@ -2177,6 +2169,8 @@ psa_status_t psa_key_derivation_setup(psa_crypto_generator_t *generator,
  *         Success.
  * \retval #PSA_ERROR_INVALID_ARGUMENT
  *         \c step is not compatible with the generator's algorithm.
+ * \retval #PSA_ERROR_INVALID_ARGUMENT
+ *         \c step does not allow direct inputs.
  * \retval #PSA_ERROR_INSUFFICIENT_MEMORY
  * \retval #PSA_ERROR_COMMUNICATION_FAILURE
  * \retval #PSA_ERROR_HARDWARE_FAILURE
@@ -2195,17 +2189,23 @@ psa_status_t psa_key_derivation_input_bytes(psa_crypto_generator_t *generator,
 
 /** Provide an input for key derivation in the form of a key.
  *
- * See the descrition of psa_key_derivation_input_bytes() regarding
- * what inputs are supported and in what order. An input step may only be
- * a key if the descrition of psa_key_derivation_input_bytes() explicitly
- * allows it.
+ * Which inputs are required and in what order depends on the algorithm.
+ * Refer to the documentation of each key derivation or key agreement
+ * algorithm for information.
+ *
+ * This function passes key inputs. Some inputs must be passed as keys
+ * of the appropriate type using this function, while others must be
+ * passed as direct inputs using psa_key_derivation_input_bytes(). Refer to
+ * the documentation of individual step types for information.
  *
  * \param[in,out] generator       The generator object to use. It must
  *                                have been set up with
  *                                psa_key_derivation_setup() and must not
  *                                have produced any output yet.
  * \param step                    Which step the input data is for.
- * \param handle                  Handle to the secret key.
+ * \param handle                  Handle to the key. It must have an
+ *                                appropriate type for \p step and must
+ *                                allow the usage #PSA_KEY_USAGE_DERIVE.
  *
  * \retval #PSA_SUCCESS
  *         Success.
@@ -2214,6 +2214,8 @@ psa_status_t psa_key_derivation_input_bytes(psa_crypto_generator_t *generator,
  * \retval #PSA_ERROR_NOT_PERMITTED
  * \retval #PSA_ERROR_INVALID_ARGUMENT
  *         \c step is not compatible with the generator's algorithm.
+ * \retval #PSA_ERROR_INVALID_ARGUMENT
+ *         \c step does not allow key inputs.
  * \retval #PSA_ERROR_INSUFFICIENT_MEMORY
  * \retval #PSA_ERROR_COMMUNICATION_FAILURE
  * \retval #PSA_ERROR_HARDWARE_FAILURE
diff --git a/include/psa/crypto_values.h b/include/psa/crypto_values.h
index 5c81acdbe..fedd35c5c 100644
--- a/include/psa/crypto_values.h
+++ b/include/psa/crypto_values.h
@@ -1104,6 +1104,15 @@
  *
  * For example, `PSA_ALG_HKDF(PSA_ALG_SHA256)` is HKDF using HMAC-SHA-256.
  *
+ * This key derivation algorithm uses the following inputs:
+ * - #PSA_KDF_STEP_SALT is the salt used in the "extract" step.
+ *   It is optional; if omitted, the derivation uses an empty salt.
+ * - #PSA_KDF_STEP_SECRET is the secret key used in the "extract" step.
+ * - #PSA_KDF_STEP_INFO is the info string used in the "expand" step.
+ * You must pass #PSA_KDF_STEP_SALT before #PSA_KDF_STEP_SECRET.
+ * You may pass #PSA_KDF_STEP_INFO at any time after steup and before
+ * starting to generate output.
+ *
  * \param hash_alg      A hash algorithm (\c PSA_ALG_XXX value such that
  *                      #PSA_ALG_IS_HASH(\p hash_alg) is true).
  *
@@ -1421,11 +1430,44 @@
  * @{
  */
 
+/** A secret input for key derivation.
+ *
+ * This must be a key of type #PSA_KEY_TYPE_DERIVE.
+ */
 #define PSA_KDF_STEP_SECRET              ((psa_key_derivation_step_t)0x0101)
+
+/** A label for key derivation.
+ *
+ * This must be a direct input.
+ */
 #define PSA_KDF_STEP_LABEL               ((psa_key_derivation_step_t)0x0201)
+
+/** A salt for key derivation.
+ *
+ * This must be a direct input.
+ */
 #define PSA_KDF_STEP_SALT                ((psa_key_derivation_step_t)0x0202)
+
+/** An information string for key derivation.
+ *
+ * This must be a direct input.
+ */
 #define PSA_KDF_STEP_INFO                ((psa_key_derivation_step_t)0x0203)
-#define PSA_KDF_STEP_PEER_KEY            ((psa_key_derivation_step_t)0x0301)
+
+/** The private key in a key agreement.
+ *
+ * This must be a key pair of the appropriate type for the key agreement
+ * algorithm.
+ */
+#define PSA_KDF_STEP_OUR_KEY             ((psa_key_derivation_step_t)0x0301)
+
+/** A label for key derivation.
+ *
+ * This may be a key pair of the appropriate type for the key agreement
+ * algorithm, or a direct input which is parsed as a public key in the
+ * same format as psa_import_key().
+ */
+#define PSA_KDF_STEP_PEER_KEY            ((psa_key_derivation_step_t)0x0302)
 
 /**@}*/
 
diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index 916c52fc6..6269fba71 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -4128,10 +4128,11 @@ static psa_status_t psa_hkdf_input( psa_hkdf_generator_t *hkdf,
 }
 #endif /* MBEDTLS_MD_C */
 
-psa_status_t psa_key_derivation_input_bytes( psa_crypto_generator_t *generator,
-                                             psa_key_derivation_step_t step,
-                                             const uint8_t *data,
-                                             size_t data_length )
+static psa_status_t psa_key_derivation_input_raw(
+    psa_crypto_generator_t *generator,
+    psa_key_derivation_step_t step,
+    const uint8_t *data,
+    size_t data_length )
 {
     psa_status_t status;
 
@@ -4165,6 +4166,23 @@ psa_status_t psa_key_derivation_input_bytes( psa_crypto_generator_t *generator,
     return( status );
 }
 
+psa_status_t psa_key_derivation_input_bytes( psa_crypto_generator_t *generator,
+                                             psa_key_derivation_step_t step,
+                                             const uint8_t *data,
+                                             size_t data_length )
+{
+    switch( step )
+    {
+        case PSA_KDF_STEP_LABEL:
+        case PSA_KDF_STEP_SALT:
+        case PSA_KDF_STEP_INFO:
+            return( psa_key_derivation_input_raw( generator, step,
+                                                  data, data_length ) );
+        default:
+            return( PSA_ERROR_INVALID_ARGUMENT );
+    }
+}
+
 psa_status_t psa_key_derivation_input_key( psa_crypto_generator_t *generator,
                                            psa_key_derivation_step_t step,
                                            psa_key_handle_t handle )
@@ -4176,18 +4194,21 @@ psa_status_t psa_key_derivation_input_key( psa_crypto_generator_t *generator,
                                     generator->alg );
     if( status != PSA_SUCCESS )
         return( status );
+    // TODO: for a key agreement algorithm, allow the corresponding key type and step
     if( slot->type != PSA_KEY_TYPE_DERIVE )
         return( PSA_ERROR_INVALID_ARGUMENT );
     /* Don't allow a key to be used as an input that is usually public.
      * This is debatable. It's ok from a cryptographic perspective to
      * use secret material as an input that is usually public. However
-     * this is usually not intended, so be conservative at least for now. */
+     * the material should be dedicated to a particular input step,
+     * otherwise this may allow the key to be used in an unintended way
+     * and leak values derived from the key. So be conservative. */
     if( step != PSA_KDF_STEP_SECRET )
         return( PSA_ERROR_INVALID_ARGUMENT );
-    return( psa_key_derivation_input_bytes( generator,
-                                            step,
-                                            slot->data.raw.data,
-                                            slot->data.raw.bytes ) );
+    return( psa_key_derivation_input_raw( generator,
+                                          step,
+                                          slot->data.raw.data,
+                                          slot->data.raw.bytes ) );
 }