feat(wrapped-keys): add versioned key update support and client APIs

.changeset/young-pianos-tickle.md

@@ -0,0 +1,6 @@
+---
+'@lit-protocol/wrapped-keys': minor
+'@lit-protocol/e2e': minor
+---
+
+Wrapped-keys now supports updating ciphertext/ACCs via a new PUT endpoint, returns version history when requested.

docs/docs.json

@@ -133,6 +133,7 @@
                   "sdk/sdk-reference/wrapped-keys/functions/exportPrivateKey",
                   "sdk/sdk-reference/wrapped-keys/functions/importPrivateKey",
                   "sdk/sdk-reference/wrapped-keys/functions/getEncryptedKey",
+                  "sdk/sdk-reference/wrapped-keys/functions/updateEncryptedKey",
                   "sdk/sdk-reference/wrapped-keys/functions/listEncryptedKeyMetadata",
                   "sdk/sdk-reference/wrapped-keys/functions/storeEncryptedKey",
                   "sdk/sdk-reference/wrapped-keys/functions/storeEncryptedKeyBatch",
@@ -246,4 +247,4 @@
       "discord": "https://litgateway.com/discord"
     }
   }
-}
+}

docs/sdk/sdk-reference/wrapped-keys/functions/getEncryptedKey.mdx

@@ -22,14 +22,18 @@ Fetches the encrypted ciphertext and metadata for a stored wrapped key without d
   Identifier of the wrapped key to retrieve.
 </ParamField>
 
+<ParamField path="params.includeVersions" type="boolean">
+  Optional flag to return the `versions` array of prior states.
+</ParamField>
+
 <ParamField path="params.userMaxPrice" type="bigint">
   Optional price ceiling for the Lit Action.
 </ParamField>
 
 ## Returns
 
 <ResponseField name="result" type="StoredKeyData">
-  Includes ciphertext, `dataToEncryptHash`, memo, key type, public key, and PKP address.
+  Includes ciphertext, `dataToEncryptHash`, memo, key type, public key, PKP address, and (when requested) `versions` + `updatedAt`.
 </ResponseField>
 
 ## Example
@@ -39,5 +43,6 @@ const storedKey = await wrappedKeysApi.getEncryptedKey({
   pkpSessionSigs,
   litClient,
   id,
+  includeVersions: true, // optional
 });
 ```

docs/sdk/sdk-reference/wrapped-keys/functions/updateEncryptedKey.mdx

@@ -0,0 +1,65 @@
+---
+title: updateEncryptedKey
+---
+
+# Function
+
+> **updateEncryptedKey**(`params`)
+
+Re-encrypts an existing wrapped key and appends the previous state to the `versions` array.
+
+## Parameters
+
+<ParamField path="params.pkpSessionSigs" type="SessionSigsMap" required>
+  Session signatures identifying the PKP that owns the wrapped key.
+</ParamField>
+
+<ParamField path="params.litClient" type="LitClient" required>
+  Lit client instance used to talk to the wrapped-keys service.
+</ParamField>
+
+<ParamField path="params.id" type="string" required>
+  Identifier of the wrapped key to update.
+</ParamField>
+
+<ParamField path="params.ciphertext" type="string" required>
+  New base64-encrypted ciphertext to store.
+</ParamField>
+
+<ParamField path="params.memo" type="string">
+  Optional updated memo for the wrapped key.
+</ParamField>
+
+<ParamField path="params.evmContractConditions" type="string">
+  Optional updated ACCs for EVM contract conditions.
+</ParamField>
+
+<ParamField path="params.userMaxPrice" type="bigint">
+  Optional price ceiling for the Lit Action.
+</ParamField>
+
+## Returns
+
+<ResponseField name="result" type="UpdateEncryptedKeyResult">
+  Returns `id`, `pkpAddress`, and `updatedAt` of the updated wrapped key.
+</ResponseField>
+
+## Example
+
+```ts
+const res = await wrappedKeysApi.updateEncryptedKey({
+  pkpSessionSigs,
+  litClient,
+  id,
+  ciphertext: 'base64-new',
+  memo: 'rotated memo',
+});
+
+const withHistory = await wrappedKeysApi.getEncryptedKey({
+  pkpSessionSigs,
+  litClient,
+  id,
+  includeVersions: true,
+});
+console.log(withHistory.versions); // [{ ciphertext: 'old', ... }]
+```

packages/e2e/src/test-helpers/executeJs/wrappedKeys.ts

@@ -57,6 +57,21 @@ namespace TestHelper {
   export const randomHash = (input: string) =>
     createHash('sha256').update(input).digest('hex');
 
+  export const createEvmContractConditions = () =>
+    JSON.stringify([
+      {
+        contractAddress: ZERO_ADDRESS,
+        standardContractType: 'ERC20',
+        chain: EVM_CHAIN,
+        method: 'balanceOf',
+        parameters: [':userAddress'],
+        returnValueTest: {
+          comparator: '>=',
+          value: '0',
+        },
+      },
+    ]);
+
   export const createStorePayload = (memo = randomMemo('store')) => {
     const ciphertext = randomCiphertext();
     return {
@@ -341,6 +356,109 @@ export const registerWrappedKeysTests = () => {
         expect(storedKey.dataToEncryptHash).toBeTruthy();
       });
 
+      test('updateEncryptedKey rotates ciphertext and returns version history', async () => {
+        const pkpSessionSigs = await TestHelper.createPkpSessionSigs({
+          testEnv,
+          alice,
+          delegationAuthSig: aliceDelegationAuthSig,
+        });
+
+        const initialPayload = TestHelper.createStorePayload(
+          TestHelper.randomMemo('update-before')
+        );
+
+        const { id } = await wrappedKeysApi.storeEncryptedKey({
+          pkpSessionSigs,
+          litClient: testEnv.litClient,
+          ...initialPayload,
+        });
+
+        const newCiphertext = TestHelper.randomCiphertext();
+        const newMemo = TestHelper.randomMemo('update-after');
+
+        const updateResult = await wrappedKeysApi.updateEncryptedKey({
+          pkpSessionSigs,
+          litClient: testEnv.litClient,
+          id,
+          ciphertext: newCiphertext,
+          memo: newMemo,
+        });
+
+        expect(updateResult.id).toBe(id);
+        expect(updateResult.pkpAddress).toBe(alice.pkp!.ethAddress);
+        expect(updateResult.updatedAt).toBeTruthy();
+
+        const fetched = await wrappedKeysApi.getEncryptedKey({
+          pkpSessionSigs,
+          litClient: testEnv.litClient,
+          id,
+          includeVersions: true,
+        });
+
+        expect(fetched.ciphertext).toBe(newCiphertext);
+        expect(fetched.memo).toBe(newMemo);
+        expect(fetched.updatedAt).toBeTruthy();
+        expect(fetched.versions).toBeDefined();
+        expect(fetched.versions?.length).toBe(1);
+        expect(fetched.versions?.[0].ciphertext).toBe(
+          initialPayload.ciphertext
+        );
+        expect(fetched.versions?.[0].memo).toBe(initialPayload.memo);
+      });
+
+      test('updateEncryptedKey with evmContractConditions stores conditions in version history', async () => {
+        const pkpSessionSigs = await TestHelper.createPkpSessionSigs({
+          testEnv,
+          alice,
+          delegationAuthSig: aliceDelegationAuthSig,
+        });
+
+        const initialPayload = TestHelper.createStorePayload(
+          TestHelper.randomMemo('update-evm-before')
+        );
+
+        const { id } = await wrappedKeysApi.storeEncryptedKey({
+          pkpSessionSigs,
+          litClient: testEnv.litClient,
+          ...initialPayload,
+        });
+
+        const newCiphertext = TestHelper.randomCiphertext();
+        const newMemo = TestHelper.randomMemo('update-evm-after');
+        const evmContractConditions = TestHelper.createEvmContractConditions();
+
+        const updateResult = await wrappedKeysApi.updateEncryptedKey({
+          pkpSessionSigs,
+          litClient: testEnv.litClient,
+          id,
+          ciphertext: newCiphertext,
+          memo: newMemo,
+          evmContractConditions,
+        });
+
+        expect(updateResult.id).toBe(id);
+        expect(updateResult.pkpAddress).toBe(alice.pkp!.ethAddress);
+        expect(updateResult.updatedAt).toBeTruthy();
+
+        const fetched = await wrappedKeysApi.getEncryptedKey({
+          pkpSessionSigs,
+          litClient: testEnv.litClient,
+          id,
+          includeVersions: true,
+        });
+
+        expect(fetched.ciphertext).toBe(newCiphertext);
+        expect(fetched.memo).toBe(newMemo);
+        expect(fetched.updatedAt).toBeTruthy();
+        expect(fetched.versions).toBeDefined();
+        expect(fetched.versions?.length).toBe(1);
+        expect(fetched.versions?.[0].ciphertext).toBe(
+          initialPayload.ciphertext
+        );
+        expect(fetched.versions?.[0].memo).toBe(initialPayload.memo);
+        expect(fetched.versions?.[0].evmContractConditions).toBeUndefined();
+      });
+
       test('importPrivateKey persists an externally generated key', async () => {
         const pkpSessionSigs = await TestHelper.createPkpSessionSigs({
           testEnv,

packages/wrapped-keys/src/index.ts

@@ -9,6 +9,7 @@ import {
   signTransactionWithEncryptedKey,
   storeEncryptedKey,
   storeEncryptedKeyBatch,
+  updateEncryptedKey,
 } from './lib/api';
 import {
   CHAIN_ETHEREUM,
@@ -43,6 +44,7 @@ export const api = {
   storeEncryptedKey,
   storeEncryptedKeyBatch,
   batchGeneratePrivateKeys,
+  updateEncryptedKey,
 };
 
 export const config = {
@@ -76,6 +78,9 @@ export type {
   StoreEncryptedKeyResult,
   StoredKeyData,
   StoredKeyMetadata,
+  UpdateEncryptedKeyParams,
+  UpdateEncryptedKeyResult,
+  WrappedKeyVersion,
 } from './lib/types';
 
 export type {

packages/wrapped-keys/src/lib/api/get-encrypted-key.ts

@@ -15,7 +15,7 @@ import { GetEncryptedKeyDataParams, StoredKeyData } from '../types';
 export async function getEncryptedKey(
   params: GetEncryptedKeyDataParams
 ): Promise<StoredKeyData> {
-  const { pkpSessionSigs, litClient, id } = params;
+  const { pkpSessionSigs, litClient, id, includeVersions } = params;
 
   const sessionSig = getFirstSessionSig(pkpSessionSigs);
   const pkpAddress = getPkpAddressFromSessionSig(sessionSig);
@@ -26,5 +26,6 @@ export async function getEncryptedKey(
     id,
     sessionSig,
     litNetwork,
+    includeVersions,
   });
 }

packages/wrapped-keys/src/lib/api/index.ts

@@ -8,6 +8,7 @@ import { signMessageWithEncryptedKey } from './sign-message-with-encrypted-key';
 import { signTransactionWithEncryptedKey } from './sign-transaction-with-encrypted-key';
 import { storeEncryptedKey } from './store-encrypted-key';
 import { storeEncryptedKeyBatch } from './store-encrypted-key-batch';
+import { updateEncryptedKey } from './update-encrypted-key';
 
 export {
   listEncryptedKeyMetadata,
@@ -20,4 +21,5 @@ export {
   storeEncryptedKeyBatch,
   getEncryptedKey,
   batchGeneratePrivateKeys,
+  updateEncryptedKey,
 };

packages/wrapped-keys/src/lib/api/update-encrypted-key.ts

@@ -0,0 +1,40 @@
+import { updatePrivateKey } from '../service-client';
+import { UpdateEncryptedKeyParams, UpdateEncryptedKeyResult } from '../types';
+import {
+  getFirstSessionSig,
+  getLitNetworkFromClient,
+  getPkpAddressFromSessionSig,
+} from './utils';
+
+/**
+ * Updates an existing wrapped key and appends the previous state to versions.
+ *
+ * @param { UpdateEncryptedKeyParams } params Parameters required to update the encrypted private key
+ * @returns { Promise<UpdateEncryptedKeyResult> } An object containing the id, pkpAddress, and updatedAt timestamp of the updated key
+ */
+export async function updateEncryptedKey(
+  params: UpdateEncryptedKeyParams
+): Promise<UpdateEncryptedKeyResult> {
+  const {
+    pkpSessionSigs,
+    litClient,
+    id,
+    ciphertext,
+    evmContractConditions,
+    memo,
+  } = params;
+
+  const sessionSig = getFirstSessionSig(pkpSessionSigs);
+  const pkpAddress = getPkpAddressFromSessionSig(sessionSig);
+  const litNetwork = getLitNetworkFromClient(litClient);
+
+  return updatePrivateKey({
+    pkpAddress,
+    id,
+    sessionSig,
+    ciphertext,
+    evmContractConditions,
+    memo,
+    litNetwork,
+  });
+}