diff --git a/src/agent.jazz b/src/agent.jazz
index 2a8aa93..5e29e9e 100644
--- a/src/agent.jazz
+++ b/src/agent.jazz
@@ -27,13 +27,11 @@ inline fn generate_key(reg u64 shared_mem sync_mem private_mem, stack u64 key_id
     key[31] &= 127;
     key[31] |= 64;
 
-    key_id += 1; 
-
     addr = key_id * 32;
     addr += private_mem;
 
     for i=0 to 4 {
-        (u64)[private_mem + i * 8] = key[u64 i];
+        (u64)[addr + i * 8] = key[u64 i];
     }
 
     idptr = (u64)[sync_mem + 16];
@@ -44,7 +42,9 @@ inline fn generate_key(reg u64 shared_mem sync_mem private_mem, stack u64 key_id
     pkptr = (u64)[sync_mem + 24];
     pkptr += shared_mem;
 
-    jade_scalarmult_curve25519_amd64_mulx_base(pkptr, private_mem);
+    jade_scalarmult_curve25519_amd64_mulx_base(pkptr, addr);
+
+    key_id += 1; 
 
     return key_id;
 }