diff --git a/core/java/android/os/WorkSource.java b/core/java/android/os/WorkSource.java
index 401b4a36a7438..bb043b0109d5f 100644
--- a/core/java/android/os/WorkSource.java
+++ b/core/java/android/os/WorkSource.java
@@ -407,11 +407,11 @@ public class WorkSource implements Parcelable {
     }
 
     public boolean remove(WorkSource other) {
-        if (mNum <= 0 || other.mNum <= 0) {
+        if (isEmpty() || other.isEmpty()) {
             return false;
         }
 
-        boolean uidRemoved = false;
+        boolean uidRemoved;
         if (mNames == null && other.mNames == null) {
             uidRemoved = removeUids(other);
         } else {
@@ -427,13 +427,8 @@ public class WorkSource implements Parcelable {
         }
 
         boolean chainRemoved = false;
-        if (other.mChains != null) {
-            if (mChains != null) {
-                chainRemoved = mChains.removeAll(other.mChains);
-            }
-        } else if (mChains != null) {
-            mChains.clear();
-            chainRemoved = true;
+        if (other.mChains != null && mChains != null) {
+            chainRemoved = mChains.removeAll(other.mChains);
         }
 
         return uidRemoved || chainRemoved;
diff --git a/core/tests/coretests/src/android/os/WorkSourceTest.java b/core/tests/coretests/src/android/os/WorkSourceTest.java
index 90b457561180e..566ac4daf9502 100644
--- a/core/tests/coretests/src/android/os/WorkSourceTest.java
+++ b/core/tests/coretests/src/android/os/WorkSourceTest.java
@@ -331,4 +331,24 @@ public class WorkSourceTest extends TestCase {
         wc.addNode(200, "tag2");
         assertEquals(100, wc.getAttributionUid());
     }
+
+    public void testRemove_fromChainedWorkSource() {
+        WorkSource ws1 = new WorkSource();
+        ws1.createWorkChain().addNode(50, "foo");
+        ws1.createWorkChain().addNode(75, "bar");
+        ws1.add(100);
+
+        WorkSource ws2 = new WorkSource();
+        ws2.add(100);
+
+        assertTrue(ws1.remove(ws2));
+        assertEquals(2, ws1.getWorkChains().size());
+        assertEquals(50, ws1.getWorkChains().get(0).getAttributionUid());
+        assertEquals(75, ws1.getWorkChains().get(1).getAttributionUid());
+
+        ws2.createWorkChain().addNode(50, "foo");
+        assertTrue(ws1.remove(ws2));
+        assertEquals(1, ws1.getWorkChains().size());
+        assertEquals(75, ws1.getWorkChains().get(0).getAttributionUid());
+    }
 }