diff --git a/framework/src/android/net/ConnectivityManager.java b/framework/src/android/net/ConnectivityManager.java
index 4f18fa2..f8a1293 100644
--- a/framework/src/android/net/ConnectivityManager.java
+++ b/framework/src/android/net/ConnectivityManager.java
@@ -1200,11 +1200,14 @@
 
     /** @hide */
     public static final long FEATURE_USE_DECLARED_METHODS_FOR_CALLBACKS = 1L;
+    /** @hide */
+    public static final long FEATURE_QUEUE_NETWORK_AGENT_EVENTS_IN_SYSTEM_SERVER = 1L << 1;
 
     /** @hide */
     @Retention(RetentionPolicy.SOURCE)
     @LongDef(flag = true, prefix = "FEATURE_", value = {
-            FEATURE_USE_DECLARED_METHODS_FOR_CALLBACKS
+            FEATURE_USE_DECLARED_METHODS_FOR_CALLBACKS,
+            FEATURE_QUEUE_NETWORK_AGENT_EVENTS_IN_SYSTEM_SERVER
     })
     public @interface ConnectivityManagerFeature {}
 
@@ -4881,7 +4884,8 @@
         return 0;
     }
 
-    private boolean isFeatureEnabled(@ConnectivityManagerFeature long connectivityManagerFeature) {
+    /** @hide */
+    public boolean isFeatureEnabled(@ConnectivityManagerFeature long connectivityManagerFeature) {
         synchronized (mEnabledConnectivityManagerFeaturesLock) {
             if (mEnabledConnectivityManagerFeatures == null) {
                 try {
diff --git a/framework/src/android/net/NetworkAgent.java b/framework/src/android/net/NetworkAgent.java
index 08f5ecd..95b45ce 100644
--- a/framework/src/android/net/NetworkAgent.java
+++ b/framework/src/android/net/NetworkAgent.java
@@ -37,6 +37,7 @@
 import android.telephony.data.NrQosSessionAttributes;
 import android.util.Log;
 
+import com.android.internal.annotations.GuardedBy;
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.net.module.util.FrameworkConnectivityStatsLog;
 
@@ -116,15 +117,33 @@
     private final ArrayList<RegistryAction> mPreConnectedQueue = new ArrayList<>();
     private volatile long mLastBwRefreshTime = 0;
     private static final long BW_REFRESH_MIN_WIN_MS = 500;
+
+    private final boolean mQueueRemoved;
+
     private boolean mBandwidthUpdateScheduled = false;
     private AtomicBoolean mBandwidthUpdatePending = new AtomicBoolean(false);
     @NonNull
     private NetworkInfo mNetworkInfo;
     @NonNull
     private final Object mRegisterLock = new Object();
-    // TODO : move the preconnected queue to the system server and remove this
+    // TODO : when ConnectivityFlags.QUEUE_NETWORK_AGENT_EVENTS_IN_SYSTEM_SERVER is
+    // not chickened out this is never read. Remove when retiring this flag.
     private boolean mConnected = false;
 
+    /** @hide */
+    @Retention(RetentionPolicy.SOURCE)
+    @IntDef(prefix = { "STATE_" }, value = {
+        STATE_CREATED,
+        STATE_REGISTERED,
+        STATE_UNREGISTERED
+    })
+    public @interface NetworkAgentState {}
+    private static final int STATE_CREATED = 0;
+    private static final int STATE_REGISTERED = 1;
+    private static final int STATE_UNREGISTERED = 2;
+    @GuardedBy("mRegisterLock")
+    private int mState = STATE_CREATED;
+
     /**
      * The ID of the {@link NetworkProvider} that created this object, or
      * {@link NetworkProvider#ID_NONE} if unknown.
@@ -506,6 +525,18 @@
         return ni;
     }
 
+    /**
+     * Returns whether a given ConnectivityManager feature is enabled.
+     *
+     * Tests can override this.
+     * @hide
+     */
+    @VisibleForTesting
+    public boolean isFeatureEnabled(@NonNull Context context,
+            @ConnectivityManager.ConnectivityManagerFeature long feature) {
+        return context.getSystemService(ConnectivityManager.class).isFeatureEnabled(feature);
+    }
+
     // Temporary backward compatibility constructor
     public NetworkAgent(@NonNull Context context, @NonNull Looper looper, @NonNull String logTag,
             @NonNull NetworkCapabilities nc, @NonNull LinkProperties lp, int score,
@@ -588,6 +619,10 @@
             @Nullable LocalNetworkConfig localNetworkConfig, @NonNull NetworkScore score,
             @NonNull NetworkAgentConfig config, int providerId, @NonNull NetworkInfo ni) {
         mHandler = new NetworkAgentHandler(looper);
+        // If the feature is enabled, then events are queued in the system
+        // server, and it's removed from this NetworkAgent.
+        mQueueRemoved = isFeatureEnabled(context,
+                ConnectivityManager.FEATURE_QUEUE_NETWORK_AGENT_EVENTS_IN_SYSTEM_SERVER);
         LOG_TAG = logTag;
         mNetworkInfo = new NetworkInfo(ni);
         this.providerId = providerId;
@@ -609,24 +644,31 @@
         public void handleMessage(Message msg) {
             switch (msg.what) {
                 case EVENT_AGENT_CONNECTED: {
-                    // TODO : move the pre-connected queue to the system server, and remove
-                    // handling this EVENT_AGENT_CONNECTED message.
-                    synchronized (mPreConnectedQueue) {
-                        if (mConnected) {
-                            log("Received new connection while already connected!");
-                        } else {
-                            if (VDBG) log("NetworkAgent fully connected");
-                            for (RegistryAction a : mPreConnectedQueue) {
-                                try {
-                                    a.execute(mRegistry);
-                                } catch (RemoteException e) {
-                                    Log.wtf(LOG_TAG, "Communication error with registry", e);
-                                    // Fall through
+                    if (mQueueRemoved) {
+                        // No handling. This message is legacy from a time where the
+                        // agent had to wait until the registry was sent to it, which
+                        // would only happen after the corresponding NetworkMonitor
+                        // was created.
+                        mConnected = true; // never read, but mConnected = false would be confusing
+                    } else {
+                        // Feature chickened out, keep the old queueing behavior
+                        synchronized (mRegisterLock) {
+                            if (mConnected) {
+                                log("Received new connection while already connected!");
+                            } else {
+                                if (VDBG) log("NetworkAgent fully connected");
+                                for (RegistryAction a : mPreConnectedQueue) {
+                                    try {
+                                        a.execute(mRegistry);
+                                    } catch (RemoteException e) {
+                                        Log.wtf(LOG_TAG, "Communication error with registry", e);
+                                        // Fall through
+                                    }
                                 }
+                                mPreConnectedQueue.clear();
                             }
-                            mPreConnectedQueue.clear();
+                            mConnected = true;
                         }
-                        mConnected = true;
                     }
                     break;
                 }
@@ -634,7 +676,8 @@
                     if (DBG) log("NetworkAgent channel lost");
                     // let the client know CS is done with us.
                     onNetworkUnwanted();
-                    synchronized (mPreConnectedQueue) {
+                    synchronized (mRegisterLock) {
+                        mState = STATE_UNREGISTERED;
                         mConnected = false;
                     }
                     break;
@@ -757,8 +800,19 @@
     public Network register() {
         if (VDBG) log("Registering NetworkAgent");
         synchronized (mRegisterLock) {
-            if (mNetwork != null) {
-                throw new IllegalStateException("Agent already registered");
+            if (mQueueRemoved) {
+                switch (mState) {
+                    case STATE_REGISTERED:
+                        throw new IllegalStateException("Agent already registered");
+                    case STATE_UNREGISTERED:
+                        throw new IllegalStateException("Agent already unregistered");
+                    default: // CREATED, this is the normal case
+                }
+            } else {
+                // Feature is chickened out, do the old processing
+                if (mNetwork != null) {
+                    throw new IllegalStateException("Agent already registered");
+                }
             }
             final ConnectivityManager cm = (ConnectivityManager) mInitialConfiguration.context
                     .getSystemService(Context.CONNECTIVITY_SERVICE);
@@ -787,6 +841,7 @@
             } else {
                 mNetwork = result.network;
                 mRegistry = result.registry;
+                mState = STATE_REGISTERED;
             }
             mInitialConfiguration = null; // All this memory can now be GC'd
         }
@@ -936,6 +991,7 @@
             mNetwork = network;
             mInitialConfiguration = null;
             mRegistry = registry;
+            mState = STATE_REGISTERED;
         }
         return new NetworkAgentBinder(mHandler);
     }
@@ -961,30 +1017,49 @@
         return mNetwork;
     }
 
-    private void queueOrSendMessage(@NonNull RegistryAction action) {
-        synchronized (mPreConnectedQueue) {
-            if (mNetwork == null && !Process.isApplicationUid(Process.myUid())) {
-                // Theoretically, it should not be valid to queue messages here before
-                // registering the NetworkAgent. However, practically, with the way
-                // queueing works right now, it ends up working out just fine.
-                // Log a statistic so that we know if this is happening in the
-                // wild. The check for isApplicationUid is to prevent logging the
-                // metric from test code.
+    private void logTerribleErrorMessageBeforeConnect() {
+        FrameworkConnectivityStatsLog.write(
+                FrameworkConnectivityStatsLog.CORE_NETWORKING_TERRIBLE_ERROR_OCCURRED,
+                FrameworkConnectivityStatsLog.CORE_NETWORKING_TERRIBLE_ERROR_OCCURRED__ERROR_TYPE__TYPE_MESSAGE_QUEUED_BEFORE_CONNECT
+        );
+    }
 
-                FrameworkConnectivityStatsLog.write(
-                        FrameworkConnectivityStatsLog.CORE_NETWORKING_TERRIBLE_ERROR_OCCURRED,
-                        FrameworkConnectivityStatsLog.CORE_NETWORKING_TERRIBLE_ERROR_OCCURRED__ERROR_TYPE__TYPE_MESSAGE_QUEUED_BEFORE_CONNECT
-                );
-            }
-            if (mConnected) {
-                try {
-                    action.execute(mRegistry);
-                } catch (RemoteException e) {
-                    Log.wtf(LOG_TAG, "Error executing registry action", e);
-                    // Fall through: the channel is asynchronous and does not report errors back
+    private void send(@NonNull RegistryAction action) {
+        synchronized (mRegisterLock) {
+            if (mQueueRemoved) {
+                if (mState <= STATE_CREATED) {
+                    // Log a terrible error. There is nothing to do with this message
+                    // so drop it.
+                    logTerribleErrorMessageBeforeConnect();
+                    Log.e(LOG_TAG, "Agent not yet registered, ignoring command");
+                    return;
+                }
+                if (mState >= STATE_UNREGISTERED) {
+                    // This should not crash for two reasons : first, the agent may
+                    // be disconnected by ConnectivityService at any time and the message
+                    // typically arrives on another thread, so it's not feasible for
+                    // apps to check before sending, they'd have to always catch. Second,
+                    // historically this hasn't thrown and some code may be relying on
+                    // the historical behavior.
+                    Log.e(LOG_TAG, "Agent already unregistered, ignoring command");
+                    return;
                 }
             } else {
-                mPreConnectedQueue.add(action);
+                if (null == mNetwork) {
+                    // Log a terrible error but still enqueue the message for backward
+                    // compatibility.
+                    logTerribleErrorMessageBeforeConnect();
+                }
+                if (!mConnected) {
+                    mPreConnectedQueue.add(action);
+                    return;
+                }
+            }
+            try {
+                action.execute(mRegistry);
+            } catch (RemoteException e) {
+                Log.wtf(LOG_TAG, "Error executing registry action", e);
+                // Fall through: the channel is asynchronous and does not report errors back
             }
         }
     }
@@ -995,8 +1070,9 @@
      */
     public void sendLinkProperties(@NonNull LinkProperties linkProperties) {
         Objects.requireNonNull(linkProperties);
-        final LinkProperties lp = new LinkProperties(linkProperties);
-        queueOrSendMessage(reg -> reg.sendLinkProperties(lp));
+        // Copy the object because if the agent is running in the system server
+        // then the same instance will be seen by the registry
+        send(reg -> reg.sendLinkProperties(new LinkProperties(linkProperties)));
     }
 
     /**
@@ -1022,7 +1098,7 @@
             @SuppressLint("NullableCollection") @Nullable List<Network> underlyingNetworks) {
         final ArrayList<Network> underlyingArray = (underlyingNetworks != null)
                 ? new ArrayList<>(underlyingNetworks) : null;
-        queueOrSendMessage(reg -> reg.sendUnderlyingNetworks(underlyingArray));
+        send(reg -> reg.sendUnderlyingNetworks(underlyingArray));
     }
 
     /**
@@ -1032,7 +1108,7 @@
     public void markConnected() {
         mNetworkInfo.setDetailedState(NetworkInfo.DetailedState.CONNECTED, null /* reason */,
                 mNetworkInfo.getExtraInfo());
-        queueOrSendNetworkInfo(mNetworkInfo);
+        sendNetworkInfo(mNetworkInfo);
     }
 
     /**
@@ -1045,7 +1121,12 @@
         // When unregistering an agent nobody should use the extrainfo (or reason) any more.
         mNetworkInfo.setDetailedState(NetworkInfo.DetailedState.DISCONNECTED, null /* reason */,
                 null /* extraInfo */);
-        queueOrSendNetworkInfo(mNetworkInfo);
+        synchronized (mRegisterLock) {
+            if (mState >= STATE_REGISTERED) {
+                sendNetworkInfo(mNetworkInfo);
+            }
+            mState = STATE_UNREGISTERED;
+        }
     }
 
     /**
@@ -1068,7 +1149,7 @@
      */
     public void setTeardownDelayMillis(
             @IntRange(from = 0, to = MAX_TEARDOWN_DELAY_MS) int teardownDelayMillis) {
-        queueOrSendMessage(reg -> reg.sendTeardownDelayMs(teardownDelayMillis));
+        send(reg -> reg.sendTeardownDelayMs(teardownDelayMillis));
     }
 
     /**
@@ -1107,7 +1188,7 @@
      */
     public void unregisterAfterReplacement(
             @IntRange(from = 0, to = MAX_TEARDOWN_DELAY_MS) int timeoutMillis) {
-        queueOrSendMessage(reg -> reg.sendUnregisterAfterReplacement(timeoutMillis));
+        send(reg -> reg.sendUnregisterAfterReplacement(timeoutMillis));
     }
 
     /**
@@ -1125,7 +1206,7 @@
     @SystemApi
     public void setLegacySubtype(final int legacySubtype, @NonNull final String legacySubtypeName) {
         mNetworkInfo.setSubtype(legacySubtype, legacySubtypeName);
-        queueOrSendNetworkInfo(mNetworkInfo);
+        sendNetworkInfo(mNetworkInfo);
     }
 
     /**
@@ -1147,7 +1228,7 @@
     @Deprecated
     public void setLegacyExtraInfo(@Nullable final String extraInfo) {
         mNetworkInfo.setExtraInfo(extraInfo);
-        queueOrSendNetworkInfo(mNetworkInfo);
+        sendNetworkInfo(mNetworkInfo);
     }
 
     /**
@@ -1155,13 +1236,9 @@
      * @hide TODO: expose something better.
      */
     @UnsupportedAppUsage(maxTargetSdk = Build.VERSION_CODES.P, trackingBug = 115609023)
-    public final void sendNetworkInfo(NetworkInfo networkInfo) {
-        queueOrSendNetworkInfo(networkInfo);
-    }
-
-    private void queueOrSendNetworkInfo(NetworkInfo networkInfo) {
+    private void sendNetworkInfo(final NetworkInfo networkInfo) {
         final NetworkInfo ni = new NetworkInfo(networkInfo);
-        queueOrSendMessage(reg -> reg.sendNetworkInfo(ni));
+        send(reg -> reg.sendNetworkInfo(ni));
     }
 
     /**
@@ -1174,7 +1251,7 @@
         mLastBwRefreshTime = System.currentTimeMillis();
         final NetworkCapabilities nc =
                 new NetworkCapabilities(networkCapabilities, NetworkCapabilities.REDACT_NONE);
-        queueOrSendMessage(reg -> reg.sendNetworkCapabilities(nc));
+        send(reg -> reg.sendNetworkCapabilities(nc));
     }
 
     /**
@@ -1186,7 +1263,7 @@
         Objects.requireNonNull(config);
         // If the agent doesn't have NET_CAPABILITY_LOCAL_NETWORK, this will be ignored by
         // ConnectivityService with a Log.wtf.
-        queueOrSendMessage(reg -> reg.sendLocalNetworkConfig(config));
+        send(reg -> reg.sendLocalNetworkConfig(config));
     }
 
     /**
@@ -1196,7 +1273,7 @@
      */
     public void sendNetworkScore(@NonNull NetworkScore score) {
         Objects.requireNonNull(score);
-        queueOrSendMessage(reg -> reg.sendScore(score));
+        send(reg -> reg.sendScore(score));
     }
 
     /**
@@ -1246,8 +1323,7 @@
      * @hide should move to NetworkAgentConfig.
      */
     public void explicitlySelected(boolean explicitlySelected, boolean acceptUnvalidated) {
-        queueOrSendMessage(reg -> reg.sendExplicitlySelected(
-                explicitlySelected, acceptUnvalidated));
+        send(reg -> reg.sendExplicitlySelected(explicitlySelected, acceptUnvalidated));
     }
 
     /**
@@ -1387,7 +1463,7 @@
      */
     public final void sendSocketKeepaliveEvent(int slot,
             @SocketKeepalive.KeepaliveEvent int event) {
-        queueOrSendMessage(reg -> reg.sendSocketKeepaliveEvent(slot, event));
+        send(reg -> reg.sendSocketKeepaliveEvent(slot, event));
     }
     /** @hide TODO delete once callers have moved to sendSocketKeepaliveEvent */
     public void onSocketKeepaliveEvent(int slot, int reason) {
@@ -1493,11 +1569,11 @@
             @NonNull final QosSessionAttributes attributes) {
         Objects.requireNonNull(attributes, "The attributes must be non-null");
         if (attributes instanceof EpsBearerQosSessionAttributes) {
-            queueOrSendMessage(ra -> ra.sendEpsQosSessionAvailable(qosCallbackId,
+            send(reg -> reg.sendEpsQosSessionAvailable(qosCallbackId,
                     new QosSession(sessionId, QosSession.TYPE_EPS_BEARER),
                     (EpsBearerQosSessionAttributes)attributes));
         } else if (attributes instanceof NrQosSessionAttributes) {
-            queueOrSendMessage(ra -> ra.sendNrQosSessionAvailable(qosCallbackId,
+            send(reg -> reg.sendNrQosSessionAvailable(qosCallbackId,
                     new QosSession(sessionId, QosSession.TYPE_NR_BEARER),
                     (NrQosSessionAttributes)attributes));
         }
@@ -1512,7 +1588,7 @@
      */
     public final void sendQosSessionLost(final int qosCallbackId,
             final int sessionId, final int qosSessionType) {
-        queueOrSendMessage(ra -> ra.sendQosSessionLost(qosCallbackId,
+        send(reg -> reg.sendQosSessionLost(qosCallbackId,
                 new QosSession(sessionId, qosSessionType)));
     }
 
@@ -1526,7 +1602,7 @@
      */
     public final void sendQosCallbackError(final int qosCallbackId,
             @QosCallbackException.ExceptionType final int exceptionType) {
-        queueOrSendMessage(ra -> ra.sendQosCallbackError(qosCallbackId, exceptionType));
+        send(reg -> reg.sendQosCallbackError(qosCallbackId, exceptionType));
     }
 
     /**
@@ -1543,7 +1619,7 @@
             throw new IllegalArgumentException("Duration must be within ["
                     + MIN_LINGER_TIMER_MS + "," + Integer.MAX_VALUE + "]ms");
         }
-        queueOrSendMessage(ra -> ra.sendLingerDuration((int) durationMs));
+        send(reg -> reg.sendLingerDuration((int) durationMs));
     }
 
     /**
@@ -1552,7 +1628,7 @@
      */
     public void sendAddDscpPolicy(@NonNull final DscpPolicy policy) {
         Objects.requireNonNull(policy);
-        queueOrSendMessage(ra -> ra.sendAddDscpPolicy(policy));
+        send(reg -> reg.sendAddDscpPolicy(policy));
     }
 
     /**
@@ -1560,14 +1636,14 @@
      * @param policyId the ID corresponding to a specific DSCP Policy.
      */
     public void sendRemoveDscpPolicy(final int policyId) {
-        queueOrSendMessage(ra -> ra.sendRemoveDscpPolicy(policyId));
+        send(reg -> reg.sendRemoveDscpPolicy(policyId));
     }
 
     /**
      * Remove all the DSCP policies on this network.
      */
     public void sendRemoveAllDscpPolicies() {
-        queueOrSendMessage(ra -> ra.sendRemoveAllDscpPolicies());
+        send(reg -> reg.sendRemoveAllDscpPolicies());
     }
 
     /** @hide */
diff --git a/service/src/com/android/server/ConnectivityService.java b/service/src/com/android/server/ConnectivityService.java
index 329b338..69d83e8 100644
--- a/service/src/com/android/server/ConnectivityService.java
+++ b/service/src/com/android/server/ConnectivityService.java
@@ -158,6 +158,7 @@
 import static com.android.server.connectivity.ConnectivityFlags.INGRESS_TO_VPN_ADDRESS_FILTERING;
 import static com.android.server.connectivity.ConnectivityFlags.NAMESPACE_TETHERING_BOOT;
 import static com.android.server.connectivity.ConnectivityFlags.QUEUE_CALLBACKS_FOR_FROZEN_APPS;
+import static com.android.server.connectivity.ConnectivityFlags.QUEUE_NETWORK_AGENT_EVENTS_IN_SYSTEM_SERVER;
 import static com.android.server.connectivity.ConnectivityFlags.REQUEST_RESTRICTED_WIFI;
 import static com.android.server.connectivity.ConnectivityFlags.WIFI_DATA_INACTIVITY_TIMEOUT;
 
@@ -9418,10 +9419,11 @@
         if (DBG) log("registerNetworkAgent " + nai);
         mDeps.getNetworkStack().makeNetworkMonitor(
                 nai.network, name, new NetworkMonitorCallbacks(nai));
-        // NetworkAgentInfo registration will finish when the NetworkMonitor is created.
-        // If the network disconnects or sends any other event before that, messages are deferred by
-        // NetworkAgent until nai.connect(), which will be called when finalizing the
-        // registration. TODO : have NetworkAgentInfo defer them instead.
+        // NetworkAgentInfo registration is done, but CS will only accept messages when the
+        // NetworkMonitor is created. If the network disconnects or sends any other event
+        // before that, messages are deferred by the Tracker Handler until it is (by asking
+        // NetworkAgentInfo to do it). The window is very small unless the NetworkStack
+        // doesn't reply immediately, which would mean a broken system anyway.
         final NetworkAndAgentRegistryParcelable result = new NetworkAndAgentRegistryParcelable();
         result.network = nai.network;
         result.registry = nai.getRegistry();
@@ -15048,6 +15050,9 @@
         if (mUseDeclaredMethodsForCallbacksEnabled) {
             features |= ConnectivityManager.FEATURE_USE_DECLARED_METHODS_FOR_CALLBACKS;
         }
+        if (mQueueNetworkAgentEventsInSystemServer) {
+            features |= ConnectivityManager.FEATURE_QUEUE_NETWORK_AGENT_EVENTS_IN_SYSTEM_SERVER;
+        }
         return features;
     }
 
@@ -15056,6 +15061,8 @@
         switch (featureFlag) {
             case INGRESS_TO_VPN_ADDRESS_FILTERING:
                 return mIngressToVpnAddressFiltering;
+            case QUEUE_NETWORK_AGENT_EVENTS_IN_SYSTEM_SERVER:
+                return mQueueNetworkAgentEventsInSystemServer;
             default:
                 throw new IllegalArgumentException("Unknown flag: " + featureFlag);
         }
diff --git a/staticlibs/testutils/devicetests/com/android/testutils/TestableNetworkAgent.kt b/staticlibs/testutils/devicetests/com/android/testutils/TestableNetworkAgent.kt
index bfbbc34..0413ed4 100644
--- a/staticlibs/testutils/devicetests/com/android/testutils/TestableNetworkAgent.kt
+++ b/staticlibs/testutils/devicetests/com/android/testutils/TestableNetworkAgent.kt
@@ -17,6 +17,7 @@
 package com.android.testutils
 
 import android.content.Context
+import android.net.ConnectivityManager.FEATURE_QUEUE_NETWORK_AGENT_EVENTS_IN_SYSTEM_SERVER
 import android.net.InetAddresses.parseNumericAddress
 import android.net.KeepalivePacketData
 import android.net.LinkAddress
@@ -28,6 +29,7 @@
 import android.net.NetworkCapabilities.TRANSPORT_TEST
 import android.net.NetworkProvider
 import android.net.NetworkRequest
+import android.net.NetworkScore
 import android.net.QosFilter
 import android.net.Uri
 import android.os.Looper
@@ -64,16 +66,21 @@
 import kotlin.test.assertNotNull
 import kotlin.test.assertNull
 import kotlin.test.assertTrue
+import kotlin.test.fail
 import org.junit.Assert.assertArrayEquals
 
 // Any legal score (0~99) for the test network would do, as it is going to be kept up by the
 // requests filed by the test and should never match normal internet requests. 70 is the default
 // score of Ethernet networks, it's as good a value as any other.
-private const val TEST_NETWORK_SCORE = 70
+private val TEST_NETWORK_SCORE = NetworkScore.Builder().setLegacyInt(70).build()
 
 private class Provider(context: Context, looper: Looper) :
             NetworkProvider(context, looper, "NetworkAgentTest NetworkProvider")
 
+private val enabledFeatures = mutableMapOf(
+    FEATURE_QUEUE_NETWORK_AGENT_EVENTS_IN_SYSTEM_SERVER to true
+)
+
 public open class TestableNetworkAgent(
     context: Context,
     looper: Looper,
@@ -81,8 +88,17 @@
     val lp: LinkProperties,
     conf: NetworkAgentConfig
 ) : NetworkAgent(context, looper, TestableNetworkAgent::class.java.simpleName /* tag */,
-        nc, lp, TEST_NETWORK_SCORE, conf, Provider(context, looper)) {
+    nc, lp, TEST_NETWORK_SCORE, conf, Provider(context, looper)) {
+
+    override fun isFeatureEnabled(context: Context, feature: Long): Boolean {
+        when (val it = enabledFeatures.get(feature)) {
+            null -> fail("Unmocked feature $feature, see TestableNetworkAgent.enabledFeatures")
+            else -> return it
+        }
+    }
+
     companion object {
+        fun setFeatureEnabled(flag: Long, enabled: Boolean) = enabledFeatures.set(flag, enabled)
 
         /**
          * Convenience method to create a [NetworkRequest] matching [TestableNetworkAgent]s from
diff --git a/tests/cts/net/src/android/net/cts/NetworkAgentTest.kt b/tests/cts/net/src/android/net/cts/NetworkAgentTest.kt
index 1fa9e3a..bd9bd2a 100644
--- a/tests/cts/net/src/android/net/cts/NetworkAgentTest.kt
+++ b/tests/cts/net/src/android/net/cts/NetworkAgentTest.kt
@@ -384,9 +384,9 @@
             initialLp = lp,
             initialNc = nc
         )
-        agent.setTeardownDelayMillis(0)
         // Connect the agent and verify initial status callbacks.
         agent.register()
+        agent.setTeardownDelayMillis(0)
         agent.markConnected()
         agent.expectCallback<OnNetworkCreated>()
         agent.expectPostConnectionCallbacks(expectedInitSignalStrengthThresholds)
@@ -1934,4 +1934,19 @@
         // VPN networks are always created as soon as the agent is registered.
         doTestNativeNetworkCreation(expectCreatedImmediately = true, intArrayOf(TRANSPORT_VPN))
     }
+
+    @Test(expected = IllegalStateException::class)
+    fun testRegisterAgain() {
+        val agent = createNetworkAgent()
+        agent.register()
+        agent.unregister()
+        agent.register()
+    }
+
+    @Test
+    fun testUnregisterBeforeRegister() {
+        // For backward compatibility, this shouldn't crash.
+        val agent = createNetworkAgent()
+        agent.unregister()
+    }
 }