diff --git a/core/java/android/net/NetworkProvider.java b/core/java/android/net/NetworkProvider.java
index d31218d..a17a498 100644
--- a/core/java/android/net/NetworkProvider.java
+++ b/core/java/android/net/NetworkProvider.java
@@ -51,13 +51,6 @@
     public static final int ID_NONE = -1;
 
     /**
-     * A hardcoded ID for NetworkAgents representing VPNs. These agents are not created by any
-     * provider, so they use this constant for clarity instead of NONE.
-     * @hide only used by ConnectivityService.
-     */
-    public static final int ID_VPN = -2;
-
-    /**
      * The first providerId value that will be allocated.
      * @hide only used by ConnectivityService.
      */
diff --git a/services/core/java/com/android/server/ConnectivityService.java b/services/core/java/com/android/server/ConnectivityService.java
index 5420ee2..f056117 100644
--- a/services/core/java/com/android/server/ConnectivityService.java
+++ b/services/core/java/com/android/server/ConnectivityService.java
@@ -4821,15 +4821,6 @@
         }
     }
 
-    private void updateVpnCapabilities(Vpn vpn, @Nullable NetworkCapabilities nc) {
-        ensureRunningOnConnectivityServiceThread();
-        NetworkAgentInfo vpnNai = getNetworkAgentInfoForNetId(vpn.getNetId());
-        if (vpnNai == null || nc == null) {
-            return;
-        }
-        updateCapabilities(vpnNai.getCurrentScore(), vpnNai, nc);
-    }
-
     @Override
     public boolean updateLockdownVpn() {
         if (Binder.getCallingUid() != Process.SYSTEM_UID) {
@@ -5140,7 +5131,7 @@
         }
     }
 
-    private void onUserStart(int userId) {
+    private void onUserStarted(int userId) {
         synchronized (mVpns) {
             Vpn userVpn = mVpns.get(userId);
             if (userVpn != null) {
@@ -5155,7 +5146,7 @@
         }
     }
 
-    private void onUserStop(int userId) {
+    private void onUserStopped(int userId) {
         synchronized (mVpns) {
             Vpn userVpn = mVpns.get(userId);
             if (userVpn == null) {
@@ -5169,28 +5160,22 @@
 
     private void onUserAdded(int userId) {
         mPermissionMonitor.onUserAdded(userId);
-        Network defaultNetwork = getNetwork(getDefaultNetwork());
         synchronized (mVpns) {
             final int vpnsSize = mVpns.size();
             for (int i = 0; i < vpnsSize; i++) {
                 Vpn vpn = mVpns.valueAt(i);
                 vpn.onUserAdded(userId);
-                NetworkCapabilities nc = vpn.updateCapabilities(defaultNetwork);
-                updateVpnCapabilities(vpn, nc);
             }
         }
     }
 
     private void onUserRemoved(int userId) {
         mPermissionMonitor.onUserRemoved(userId);
-        Network defaultNetwork = getNetwork(getDefaultNetwork());
         synchronized (mVpns) {
             final int vpnsSize = mVpns.size();
             for (int i = 0; i < vpnsSize; i++) {
                 Vpn vpn = mVpns.valueAt(i);
                 vpn.onUserRemoved(userId);
-                NetworkCapabilities nc = vpn.updateCapabilities(defaultNetwork);
-                updateVpnCapabilities(vpn, nc);
             }
         }
     }
@@ -5272,9 +5257,9 @@
             if (userId == UserHandle.USER_NULL) return;
 
             if (Intent.ACTION_USER_STARTED.equals(action)) {
-                onUserStart(userId);
+                onUserStarted(userId);
             } else if (Intent.ACTION_USER_STOPPED.equals(action)) {
-                onUserStop(userId);
+                onUserStopped(userId);
             } else if (Intent.ACTION_USER_ADDED.equals(action)) {
                 onUserAdded(userId);
             } else if (Intent.ACTION_USER_REMOVED.equals(action)) {
@@ -8276,13 +8261,12 @@
             return false;
         }
 
-        final Network[] underlyingNetworks;
-        synchronized (mVpns) {
-            final Vpn vpn = getVpnIfOwner(callbackUid);
-            underlyingNetworks = (vpn == null) ? null : vpn.getUnderlyingNetworks();
-        }
-        if (underlyingNetworks != null) {
-            if (Arrays.asList(underlyingNetworks).contains(nai.network)) return true;
+        for (NetworkAgentInfo virtual : mNetworkAgentInfos.values()) {
+            if (virtual.supportsUnderlyingNetworks()
+                    && virtual.networkCapabilities.getOwnerUid() == callbackUid
+                    && ArrayUtils.contains(virtual.declaredUnderlyingNetworks, nai.network)) {
+                return true;
+            }
         }
 
         // Administrator UIDs also contains the Owner UID
diff --git a/tests/net/java/com/android/server/ConnectivityServiceTest.java b/tests/net/java/com/android/server/ConnectivityServiceTest.java
index 561c6ba..c917e66 100644
--- a/tests/net/java/com/android/server/ConnectivityServiceTest.java
+++ b/tests/net/java/com/android/server/ConnectivityServiceTest.java
@@ -1058,7 +1058,9 @@
 
         public void setUids(Set<UidRange> uids) {
             mNetworkCapabilities.setUids(uids);
-            updateCapabilitiesInternal(null /* defaultNetwork */, true);
+            if (mAgentRegistered) {
+                mMockNetworkAgent.setNetworkCapabilities(mNetworkCapabilities, true);
+            }
         }
 
         public void setVpnType(int vpnType) {
@@ -1089,6 +1091,10 @@
             mMockNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_VPN, lp,
                     mNetworkCapabilities);
             mMockNetworkAgent.waitForIdle(TIMEOUT_MS);
+            verify(mNetworkManagementService, times(1))
+                    .addVpnUidRanges(eq(mMockVpn.getNetId()), eq(uids.toArray(new UidRange[0])));
+            verify(mNetworkManagementService, never())
+                    .removeVpnUidRanges(eq(mMockVpn.getNetId()), any());
             mAgentRegistered = true;
             mNetworkCapabilities.set(mMockNetworkAgent.getNetworkCapabilities());
             mNetworkAgent = mMockNetworkAgent.getNetworkAgent();
@@ -1143,28 +1149,6 @@
             mMockNetworkAgent.sendLinkProperties(lp);
         }
 
-        private NetworkCapabilities updateCapabilitiesInternal(Network defaultNetwork,
-                boolean sendToConnectivityService) {
-            if (!mAgentRegistered) return null;
-            super.updateCapabilities(defaultNetwork);
-            // Because super.updateCapabilities will update the capabilities of the agent but
-            // not the mock agent, the mock agent needs to know about them.
-            copyCapabilitiesToNetworkAgent(sendToConnectivityService);
-            return new NetworkCapabilities(mNetworkCapabilities);
-        }
-
-        private void copyCapabilitiesToNetworkAgent(boolean sendToConnectivityService) {
-            if (null != mMockNetworkAgent) {
-                mMockNetworkAgent.setNetworkCapabilities(mNetworkCapabilities,
-                        sendToConnectivityService);
-            }
-        }
-
-        @Override
-        public NetworkCapabilities updateCapabilities(Network defaultNetwork) {
-            return updateCapabilitiesInternal(defaultNetwork, false);
-        }
-
         public void disconnect() {
             if (mMockNetworkAgent != null) mMockNetworkAgent.disconnect();
             mAgentRegistered = false;
@@ -6922,8 +6906,8 @@
         final Set<UidRange> vpnRange = Collections.singleton(UidRange.createForUser(VPN_USER));
         mMockVpn.establish(lp, VPN_UID, vpnRange);
 
-        // Connected VPN should have interface rules set up. There are two expected invocations,
-        // one during VPN uid update, one during VPN LinkProperties update
+        // A connected VPN should have interface rules set up. There are two expected invocations,
+        // one during the VPN initial connection, one during the VPN LinkProperties update.
         ArgumentCaptor<int[]> uidCaptor = ArgumentCaptor.forClass(int[].class);
         verify(mMockNetd, times(2)).firewallAddUidInterfaceRules(eq("tun0"), uidCaptor.capture());
         assertContainsExactly(uidCaptor.getAllValues().get(0), APP1_UID, APP2_UID);
@@ -7438,20 +7422,14 @@
         setupLocationPermissions(Build.VERSION_CODES.Q, true, AppOpsManager.OPSTR_FINE_LOCATION,
                 Manifest.permission.ACCESS_FINE_LOCATION);
 
-        // setUp() calls mockVpn() which adds a VPN with the Test Runner's uid. Configure it to be
-        // active
-        final VpnInfo info = new VpnInfo();
-        info.ownerUid = Process.myUid();
-        info.vpnIface = VPN_IFNAME;
-        mMockVpn.setVpnInfo(info);
-
         mMockVpn.establishForMyUid();
-        waitForIdle();
 
+        // Wait for networks to connect and broadcasts to be sent before removing permissions.
+        waitForIdle();
         mServiceContext.setPermission(android.Manifest.permission.NETWORK_STACK, PERMISSION_DENIED);
 
-
         assertTrue(mService.setUnderlyingNetworksForVpn(new Network[] {network}));
+        waitForIdle();
         assertTrue(
                 "Active VPN permission not applied",
                 mService.checkConnectivityDiagnosticsPermissions(
@@ -7459,6 +7437,7 @@
                         mContext.getOpPackageName()));
 
         assertTrue(mService.setUnderlyingNetworksForVpn(null));
+        waitForIdle();
         assertFalse(
                 "VPN shouldn't receive callback on non-underlying network",
                 mService.checkConnectivityDiagnosticsPermissions(
diff --git a/tests/net/java/com/android/server/connectivity/VpnTest.java b/tests/net/java/com/android/server/connectivity/VpnTest.java
index a553b58..e1e0efa 100644
--- a/tests/net/java/com/android/server/connectivity/VpnTest.java
+++ b/tests/net/java/com/android/server/connectivity/VpnTest.java
@@ -41,6 +41,7 @@
 import static org.mockito.ArgumentMatchers.anyBoolean;
 import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.ArgumentMatchers.argThat;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.atLeastOnce;
 import static org.mockito.Mockito.doAnswer;
@@ -86,10 +87,10 @@
 import android.os.Bundle;
 import android.os.ConditionVariable;
 import android.os.INetworkManagementService;
-import android.os.Looper;
 import android.os.Process;
 import android.os.UserHandle;
 import android.os.UserManager;
+import android.os.test.TestLooper;
 import android.provider.Settings;
 import android.security.Credentials;
 import android.security.KeyStore;
@@ -100,6 +101,7 @@
 import androidx.test.runner.AndroidJUnit4;
 
 import com.android.internal.R;
+import com.android.internal.net.LegacyVpnInfo;
 import com.android.internal.net.VpnConfig;
 import com.android.internal.net.VpnProfile;
 import com.android.server.IpSecService;
@@ -223,6 +225,8 @@
                 .thenReturn(mNotificationManager);
         when(mContext.getSystemService(eq(Context.CONNECTIVITY_SERVICE)))
                 .thenReturn(mConnectivityManager);
+        when(mContext.getSystemServiceName(eq(ConnectivityManager.class)))
+                .thenReturn(Context.CONNECTIVITY_SERVICE);
         when(mContext.getSystemService(eq(Context.IPSEC_SERVICE))).thenReturn(mIpSecManager);
         when(mContext.getString(R.string.config_customVpnAlwaysOnDisconnectedDialogComponent))
                 .thenReturn(Resources.getSystem().getString(
@@ -589,7 +593,7 @@
     }
 
     @Test
-    public void testNotificationShownForAlwaysOnApp() {
+    public void testNotificationShownForAlwaysOnApp() throws Exception {
         final UserHandle userHandle = UserHandle.of(primaryUser.id);
         final Vpn vpn = createVpn(primaryUser.id);
         setMockedUsers(primaryUser);
@@ -619,7 +623,6 @@
 
     @Test
     public void testCapabilities() {
-        final Vpn vpn = createVpn(primaryUser.id);
         setMockedUsers(primaryUser);
 
         final Network mobile = new Network(1);
@@ -1037,7 +1040,7 @@
         when(exception.getErrorType())
                 .thenReturn(IkeProtocolException.ERROR_TYPE_AUTHENTICATION_FAILED);
 
-        final Vpn vpn = startLegacyVpn(mVpnProfile);
+        final Vpn vpn = startLegacyVpn(createVpn(primaryUser.id), (mVpnProfile));
         final NetworkCallback cb = triggerOnAvailableAndGetCallback();
 
         // Wait for createIkeSession() to be called before proceeding in order to ensure consistent
@@ -1048,20 +1051,20 @@
         ikeCb.onClosedExceptionally(exception);
 
         verify(mConnectivityManager, timeout(TEST_TIMEOUT_MS)).unregisterNetworkCallback(eq(cb));
-        assertEquals(DetailedState.FAILED, vpn.getNetworkInfo().getDetailedState());
+        assertEquals(LegacyVpnInfo.STATE_FAILED, vpn.getLegacyVpnInfo().state);
     }
 
     @Test
     public void testStartPlatformVpnIllegalArgumentExceptionInSetup() throws Exception {
         when(mIkev2SessionCreator.createIkeSession(any(), any(), any(), any(), any(), any()))
                 .thenThrow(new IllegalArgumentException());
-        final Vpn vpn = startLegacyVpn(mVpnProfile);
+        final Vpn vpn = startLegacyVpn(createVpn(primaryUser.id), mVpnProfile);
         final NetworkCallback cb = triggerOnAvailableAndGetCallback();
 
         // Wait for createIkeSession() to be called before proceeding in order to ensure consistent
         // state
         verify(mConnectivityManager, timeout(TEST_TIMEOUT_MS)).unregisterNetworkCallback(eq(cb));
-        assertEquals(DetailedState.FAILED, vpn.getNetworkInfo().getDetailedState());
+        assertEquals(LegacyVpnInfo.STATE_FAILED, vpn.getLegacyVpnInfo().state);
     }
 
     private void setAndVerifyAlwaysOnPackage(Vpn vpn, int uid, boolean lockdownEnabled) {
@@ -1100,8 +1103,7 @@
         // a subsequent CL.
     }
 
-    public Vpn startLegacyVpn(final VpnProfile vpnProfile) throws Exception {
-        final Vpn vpn = createVpn(primaryUser.id);
+    private Vpn startLegacyVpn(final Vpn vpn, final VpnProfile vpnProfile) throws Exception {
         setMockedUsers(primaryUser);
 
         // Dummy egress interface
@@ -1118,7 +1120,7 @@
 
     @Test
     public void testStartPlatformVpn() throws Exception {
-        startLegacyVpn(mVpnProfile);
+        startLegacyVpn(createVpn(primaryUser.id), mVpnProfile);
         // TODO: Test the Ikev2VpnRunner started up properly. Relies on utility methods added in
         // a subsequent patch.
     }
@@ -1153,7 +1155,7 @@
                     legacyRunnerReady.open();
                     return new Network(102);
                 });
-        final Vpn vpn = startLegacyVpn(profile);
+        final Vpn vpn = startLegacyVpn(createVpn(primaryUser.id), profile);
         final TestDeps deps = (TestDeps) vpn.mDeps;
         try {
             // udppsk and 1701 are the values for TYPE_L2TP_IPSEC_PSK
@@ -1287,8 +1289,13 @@
         doReturn(UserHandle.of(userId)).when(asUserContext).getUser();
         when(mContext.createContextAsUser(eq(UserHandle.of(userId)), anyInt()))
                 .thenReturn(asUserContext);
-        return new Vpn(Looper.myLooper(), mContext, new TestDeps(), mNetService,
+        final TestLooper testLooper = new TestLooper();
+        final Vpn vpn = new Vpn(testLooper.getLooper(), mContext, new TestDeps(), mNetService,
                 userId, mKeyStore, mSystemServices, mIkev2SessionCreator);
+        verify(mConnectivityManager, times(1)).registerNetworkProvider(argThat(
+                provider -> provider.getName().contains("VpnNetworkProvider")
+        ));
+        return vpn;
     }
 
     private static void assertBlocked(Vpn vpn, int... uids) {