diff --git a/core/java/android/app/ContextImpl.java b/core/java/android/app/ContextImpl.java
index bb64c3492c4f0..a4256a9d8ab3e 100644
--- a/core/java/android/app/ContextImpl.java
+++ b/core/java/android/app/ContextImpl.java
@@ -250,6 +250,10 @@ class ContextImpl extends Context {
 
     private final Object mSync = new Object();
 
+    /**
+     * Whether this is created from {@link #createSystemContext(ActivityThread)} or
+     * {@link #createSystemUiContext(ContextImpl, int)} or any {@link Context} that system UI uses.
+     */
     private boolean mIsSystemOrSystemUiContext;
     private boolean mIsUiContext;
     private boolean mIsAssociatedWithDisplay;
@@ -1922,16 +1926,18 @@ class ContextImpl extends Context {
     /** @hide */
     @Override
     public boolean isUiContext() {
-        return mIsSystemOrSystemUiContext || mIsUiContext || isSystemOrSystemUI();
+        return mIsSystemOrSystemUiContext || mIsUiContext;
     }
 
     /**
      * Temporary workaround to permit incorrect usages of Context by SystemUI.
-     * TODO(b/149790106): Fix usages and remove.
+     * TODO(b/147647877): Fix usages and remove.
      */
-    private boolean isSystemOrSystemUI() {
-        return ActivityThread.isSystem() || checkPermission("android.permission.STATUS_BAR_SERVICE",
-                Binder.getCallingPid(), Binder.getCallingUid()) == PERMISSION_GRANTED;
+    private static boolean isSystemOrSystemUI(Context context) {
+        return ActivityThread.isSystem() || context.checkPermission(
+                "android.permission.STATUS_BAR_SERVICE",
+                Binder.getCallingPid(),
+                Binder.getCallingUid()) == PERMISSION_GRANTED;
     }
 
     private static boolean isUiComponent(String name) {
@@ -2467,7 +2473,7 @@ class ContextImpl extends Context {
 
     @Override
     public Display getDisplay() {
-        if (!mIsSystemOrSystemUiContext && !mIsAssociatedWithDisplay && !isSystemOrSystemUI()) {
+        if (!mIsSystemOrSystemUiContext && !mIsAssociatedWithDisplay) {
             throw new UnsupportedOperationException("Tried to obtain display from a Context not "
                     + "associated with  one. Only visual Contexts (such as Activity or one created "
                     + "with Context#createWindowContext) or ones created with "
@@ -2643,6 +2649,7 @@ class ContextImpl extends Context {
         ContextImpl context = new ContextImpl(null, mainThread, packageInfo, null, null, null, null,
                 0, null, opPackageName);
         context.setResources(packageInfo.getResources());
+        context.mIsSystemOrSystemUiContext = isSystemOrSystemUI(context);
         return context;
     }
 
@@ -2672,6 +2679,7 @@ class ContextImpl extends Context {
                 activityInfo.splitName, activityToken, null, 0, classLoader, null);
         context.mIsUiContext = true;
         context.mIsAssociatedWithDisplay = true;
+        context.mIsSystemOrSystemUiContext = isSystemOrSystemUI(context);
 
         // Clamp display ID to DEFAULT_DISPLAY if it is INVALID_DISPLAY.
         displayId = (displayId != Display.INVALID_DISPLAY) ? displayId : Display.DEFAULT_DISPLAY;
diff --git a/core/tests/coretests/src/android/content/ContextTest.java b/core/tests/coretests/src/android/content/ContextTest.java
index aab39bf611634..2057a8181147e 100644
--- a/core/tests/coretests/src/android/content/ContextTest.java
+++ b/core/tests/coretests/src/android/content/ContextTest.java
@@ -18,12 +18,15 @@ package android.content;
 
 import static android.view.Display.DEFAULT_DISPLAY;
 
+import static com.google.common.truth.Truth.assertThat;
+
 import static org.junit.Assert.assertEquals;
 
 import android.app.ActivityThread;
 import android.hardware.display.DisplayManager;
 import android.os.UserHandle;
 
+import androidx.test.core.app.ApplicationProvider;
 import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
 import androidx.test.platform.app.InstrumentationRegistry;
@@ -32,7 +35,6 @@ import org.junit.Test;
 import org.junit.runner.RunWith;
 
 /**
- *
  *  Build/Install/Run:
  *   atest FrameworksCoreTests:ContextTest
  */
@@ -47,6 +49,14 @@ public class ContextTest {
         assertEquals(systemContext.getDisplay().getDisplayId(), systemContext.getDisplayId());
     }
 
+    @Test
+    public void testDisplayIdForSystemUiContext() {
+        final Context systemUiContext =
+                ActivityThread.currentActivityThread().getSystemUiContext();
+
+        assertEquals(systemUiContext.getDisplay().getDisplayId(), systemUiContext.getDisplayId());
+    }
+
     @Test
     public void testDisplayIdForTestContext() {
         final Context testContext =
@@ -94,4 +104,27 @@ public class ContextTest {
                 InstrumentationRegistry.getInstrumentation().getTargetContext();
         testContext.startActivityAsUser(new Intent(), new UserHandle(UserHandle.USER_ALL));
     }
+
+    @Test
+    public void testIsUiContext_appContext_returnsFalse() {
+        final Context appContext = ApplicationProvider.getApplicationContext();
+
+        assertThat(appContext.isUiContext()).isFalse();
+    }
+
+    @Test
+    public void testIsUiContext_systemContext_returnsTrue() {
+        final Context systemContext =
+                ActivityThread.currentActivityThread().getSystemContext();
+
+        assertThat(systemContext.isUiContext()).isTrue();
+    }
+
+    @Test
+    public void testIsUiContext_systemUiContext_returnsTrue() {
+        final Context systemUiContext =
+                ActivityThread.currentActivityThread().getSystemUiContext();
+
+        assertThat(systemUiContext.isUiContext()).isTrue();
+    }
 }