diff --git a/attestation-agent/attestation-agent/src/lib.rs b/attestation-agent/attestation-agent/src/lib.rs
index 550c9fd1c..b8a453728 100644
--- a/attestation-agent/attestation-agent/src/lib.rs
+++ b/attestation-agent/attestation-agent/src/lib.rs
@@ -180,7 +180,7 @@ impl AttestationAPIs for AttestationAgent {
         Ok(evidence.into_bytes())
     }
 
-    async fn get_derived_key(&self, key_id: &[u8]) -> Result<Vec<u8>> {
+    async fn get_derived_key(&self, key_id: &[u8], context: Vec<u8>) -> Result<Vec<u8>> {
         let context = Vec::new(); // Empty context vector as per SNP implementation
         self.attester.get_derived_key(key_id, context).await
     }
diff --git a/attestation-agent/kbs_protocol/src/evidence_provider/mock.rs b/attestation-agent/kbs_protocol/src/evidence_provider/mock.rs
index 0d7179408..79b2f3e08 100644
--- a/attestation-agent/kbs_protocol/src/evidence_provider/mock.rs
+++ b/attestation-agent/kbs_protocol/src/evidence_provider/mock.rs
@@ -19,7 +19,7 @@ impl EvidenceProvider for MockedEvidenceProvider {
         Ok("test evidence".into())
     }
 
-    async fn get_derived_key(&self, _key_id: &[u8], context: Vec<u8>) -> Result<Vec<u8>> {
+    async fn get_derived_key(&self, _key_id: &[u8], _context: Vec<u8>) -> Result<Vec<u8>> {
         Ok(vec![0u8; 32]) // Return a mock 32-byte key filled with zeros
     }
 
diff --git a/attestation-agent/kbs_protocol/src/evidence_provider/native.rs b/attestation-agent/kbs_protocol/src/evidence_provider/native.rs
index 02ab06616..47c02911c 100644
--- a/attestation-agent/kbs_protocol/src/evidence_provider/native.rs
+++ b/attestation-agent/kbs_protocol/src/evidence_provider/native.rs
@@ -35,9 +35,9 @@ impl EvidenceProvider for NativeEvidenceProvider {
         Ok(detect_tee_type())
     }
 
-    async fn get_derived_key(&self, key_id: &[u8]) -> Result<Vec<u8>> {
+    async fn get_derived_key(&self, key_id: &[u8], context: Vec<u8>) -> Result<Vec<u8>> {
         self.0
-            .get_derived_key(key_id)
+            .get_derived_key(key_id, context)
             .await
             .map_err(|e| Error::GetDerivedKey(e.to_string()))
     }