diff --git a/package.json b/package.json
index 4f4c670b7..29d488f5c 100644
--- a/package.json
+++ b/package.json
@@ -196,7 +196,7 @@
     "@babel/preset-typescript": "7.17.12",
     "@electron/fuses": "1.5.0",
     "@mixer/parallel-prettier": "2.0.1",
-    "@signalapp/mock-server": "2.9.1",
+    "@signalapp/mock-server": "2.10.0",
     "@storybook/addon-a11y": "6.5.6",
     "@storybook/addon-actions": "6.5.6",
     "@storybook/addon-controls": "6.5.6",
diff --git a/ts/textsecure/AccountManager.ts b/ts/textsecure/AccountManager.ts
index 0dfb2f0d5..c515c82d6 100644
--- a/ts/textsecure/AccountManager.ts
+++ b/ts/textsecure/AccountManager.ts
@@ -520,6 +520,7 @@ export default class AccountManager extends EventTarget {
     let password = Bytes.toBase64(getRandomBytes(16));
     password = password.substring(0, password.length - 2);
     const registrationId = generateRegistrationId();
+    const pniRegistrationId = generateRegistrationId();
 
     const previousNumber = storage.user.getNumber();
     const previousACI = storage.user.getUuid(UUIDKind.ACI)?.toString();
@@ -537,14 +538,15 @@ export default class AccountManager extends EventTarget {
       }`
     );
 
-    const response = await this.server.confirmCode(
+    const response = await this.server.confirmCode({
       number,
-      verificationCode,
-      password,
+      code: verificationCode,
+      newPassword: password,
       registrationId,
-      encryptedDeviceName,
-      { accessKey }
-    );
+      pniRegistrationId,
+      deviceName: encryptedDeviceName,
+      accessKey,
+    });
 
     const ourUuid = UUID.cast(response.uuid);
     const ourPni = UUID.cast(response.pni);
@@ -663,9 +665,7 @@ export default class AccountManager extends EventTarget {
     const registrationIdMap = {
       ...(storage.get('registrationIdMap') || {}),
       [ourUuid]: registrationId,
-
-      // TODO: DESKTOP-3318
-      [ourPni]: registrationId,
+      [ourPni]: pniRegistrationId,
     };
 
     await storage.put('identityKeyMap', identityKeyMap);
@@ -714,20 +714,7 @@ export default class AccountManager extends EventTarget {
         [pni.toString()]: identityKeyPair,
       };
 
-      const aci = storage.user.getCheckedUuid(UUIDKind.ACI);
-      const oldRegistrationIdMap = storage.get('registrationIdMap') || {};
-      const registrationIdMap = {
-        ...oldRegistrationIdMap,
-
-        // TODO: DESKTOP-3318
-        [pni.toString()]: oldRegistrationIdMap[aci.toString()],
-      };
-
-      await Promise.all([
-        storage.put('identityKeyMap', identityKeyMap),
-        storage.put('registrationIdMap', registrationIdMap),
-      ]);
-
+      await storage.put('identityKeyMap', identityKeyMap);
       await storage.protocol.hydrateCaches();
     });
 
@@ -743,6 +730,33 @@ export default class AccountManager extends EventTarget {
         );
         await this.server.registerKeys(keys, UUIDKind.PNI);
         await this.confirmKeys(keys, UUIDKind.PNI);
+
+        const pni = storage.user.getCheckedUuid(UUIDKind.PNI);
+
+        // Repair registration id
+        const deviceId = storage.user.getDeviceId();
+        const { devices } = await this.server.getKeysForIdentifier(
+          pni.toString(),
+          deviceId
+        );
+        const us = devices.find(
+          remoteDevice => remoteDevice.deviceId === deviceId
+        );
+        if (us) {
+          const oldRegistrationIdMap = storage.get('registrationIdMap') || {};
+          const oldRegistrationId = oldRegistrationIdMap[pni.toString()];
+          if (oldRegistrationId !== us.registrationId) {
+            log.warn(
+              'updatePNIIdentity: repairing PNI registration id from ' +
+                `${oldRegistrationId} to ${us.registrationId}`
+            );
+          }
+          const registrationIdMap = {
+            ...oldRegistrationIdMap,
+            [pni.toString()]: us.registrationId,
+          };
+          await storage.put('registrationIdMap', registrationIdMap);
+        }
       } catch (error) {
         log.error(
           'updatePNIIdentity: Failed to upload PNI prekeys. Moving on',
diff --git a/ts/textsecure/WebAPI.ts b/ts/textsecure/WebAPI.ts
index 2bf668600..3b930428e 100644
--- a/ts/textsecure/WebAPI.ts
+++ b/ts/textsecure/WebAPI.ts
@@ -798,17 +798,22 @@ const verifyAciResponse = z
 export type VerifyAciRequestType = Array<{ aci: string; fingerprint: string }>;
 export type VerifyAciResponseType = z.infer<typeof verifyAciResponse>;
 
+export type ConfirmCodeOptionsType = Readonly<{
+  number: string;
+  code: string;
+  newPassword: string;
+  registrationId: number;
+  pniRegistrationId: number;
+  deviceName?: string | null;
+  accessKey?: Uint8Array;
+}>;
+
 export type WebAPIType = {
   startRegistration(): unknown;
   finishRegistration(baton: unknown): void;
   cdsLookup: (options: CdsLookupOptionsType) => Promise<CDSResponseType>;
   confirmCode: (
-    number: string,
-    code: string,
-    newPassword: string,
-    registrationId: number,
-    deviceName?: string | null,
-    options?: { accessKey?: Uint8Array }
+    options: ConfirmCodeOptionsType
   ) => Promise<ConfirmCodeResultType>;
   createGroup: (
     group: Proto.IGroup,
@@ -1845,14 +1850,15 @@ export function initialize({
       current.resolve();
     }
 
-    async function confirmCode(
-      number: string,
-      code: string,
-      newPassword: string,
-      registrationId: number,
-      deviceName?: string | null,
-      options: { accessKey?: Uint8Array } = {}
-    ) {
+    async function confirmCode({
+      number,
+      code,
+      newPassword,
+      registrationId,
+      pniRegistrationId,
+      deviceName,
+      accessKey,
+    }: ConfirmCodeOptionsType) {
       const capabilities: CapabilitiesUploadType = {
         announcementGroup: true,
         giftBadges: true,
@@ -1863,12 +1869,12 @@ export function initialize({
         stories: true,
       };
 
-      const { accessKey } = options;
       const jsonData = {
         capabilities,
         fetchesMessages: true,
         name: deviceName || undefined,
         registrationId,
+        pniRegistrationId,
         supportsSms: false,
         unidentifiedAccessKey: accessKey
           ? Bytes.toBase64(accessKey)
diff --git a/yarn.lock b/yarn.lock
index 796a0815e..f90e6c6c4 100644
--- a/yarn.lock
+++ b/yarn.lock
@@ -1975,10 +1975,10 @@
     node-gyp-build "^4.2.3"
     uuid "^8.3.0"
 
-"@signalapp/mock-server@2.9.1":
-  version "2.9.1"
-  resolved "https://registry.yarnpkg.com/@signalapp/mock-server/-/mock-server-2.9.1.tgz#f6cde436d2f6665fbacda425c8b661dedd733831"
-  integrity sha512-KkMnIcBVmZwGLPzsri2PukKioTg50fsthvAmyjG3iDiuQ9KUaTPKpTf0DT/jFXzkQmwUCjs4MKIQqJYjXZWvvw==
+"@signalapp/mock-server@2.10.0":
+  version "2.10.0"
+  resolved "https://registry.yarnpkg.com/@signalapp/mock-server/-/mock-server-2.10.0.tgz#a27246e7b912caebc0bef628303e11689bf9b74c"
+  integrity sha512-kHos3n8lNBhivUecEFG4g1rvYpJ6oPgzKMOsaI+vN8R1R4Pc63WXxrLsxqAI2QmAngD+nmOgbjwAvKyH4MN0+w==
   dependencies:
     "@signalapp/libsignal-client" "^0.20.0"
     debug "^4.3.2"