From 2548731089e4949ab79d7955840071dd6b1a6337 Mon Sep 17 00:00:00 2001
From: Brent Bovenzi <brent@astronomer.io>
Date: Sat, 25 Jan 2025 10:39:21 -0500
Subject: [PATCH] Add icons to task states (#46028)

* Add icons to task states

* Move stateColor to theme

* Fix graph
---
 airflow/ui/src/components/DagRunInfo.tsx      |   2 +-
 airflow/ui/src/components/MetricsBadge.tsx    |   9 +-
 airflow/ui/src/components/StateIcon.tsx       |  70 +++++++
 .../ui/src/components/TaskInstanceTooltip.tsx |   1 +
 airflow/ui/src/components/ui/Status.tsx       |  26 ++-
 .../ui/src/layouts/Details/Graph/Graph.tsx    |   3 +-
 .../ui/src/layouts/Details/Graph/TaskNode.tsx |   9 +-
 airflow/ui/src/main.tsx                       |   5 +-
 .../ui/src/pages/Dag/Overview/Overview.tsx    |   5 +-
 airflow/ui/src/pages/Dag/Tasks/TaskCard.tsx   |   4 +-
 .../ui/src/pages/Dag/Tasks/TaskRecentRuns.tsx |   3 +-
 airflow/ui/src/pages/DagsList/RecentRuns.tsx  |   3 +-
 .../src/pages/Dashboard/Health/HealthTag.tsx  |  12 +-
 .../HistoricalMetrics/AssetEvents.tsx         |   2 +-
 .../HistoricalMetrics/DagRunMetrics.tsx       |   2 +-
 .../HistoricalMetrics/MetricSection.tsx       |   5 +-
 .../HistoricalMetrics/TaskInstanceMetrics.tsx |   2 +-
 .../pages/Dashboard/Stats/DAGImportErrors.tsx |   3 +-
 .../pages/Dashboard/Stats/DagFilterButton.tsx |   9 +-
 .../ui/src/pages/Dashboard/Stats/Stats.tsx    |  14 +-
 airflow/ui/src/pages/Pools/PoolBar.tsx        |  23 ++-
 airflow/ui/src/theme.ts                       | 176 ++++++++++++++++++
 airflow/ui/src/utils/stateColor.ts            |  37 ----
 23 files changed, 319 insertions(+), 106 deletions(-)
 create mode 100644 airflow/ui/src/components/StateIcon.tsx
 create mode 100644 airflow/ui/src/theme.ts
 delete mode 100644 airflow/ui/src/utils/stateColor.ts

diff --git a/airflow/ui/src/components/DagRunInfo.tsx b/airflow/ui/src/components/DagRunInfo.tsx
index dd5569dd45a2a..dca0c3e2f3fbe 100644
--- a/airflow/ui/src/components/DagRunInfo.tsx
+++ b/airflow/ui/src/components/DagRunInfo.tsx
@@ -74,7 +74,7 @@ const DagRunInfo = ({
     >
       <Box>
         <Time datetime={dataIntervalStart} mr={2} showTooltip={false} />
-        {state === undefined ? undefined : <Status state={state}>{state}</Status>}
+        {state === undefined ? undefined : <Status state={state} />}
       </Box>
     </Tooltip>
   ) : undefined;
diff --git a/airflow/ui/src/components/MetricsBadge.tsx b/airflow/ui/src/components/MetricsBadge.tsx
index d744d2f6eb2e7..f763a9510b2a8 100644
--- a/airflow/ui/src/components/MetricsBadge.tsx
+++ b/airflow/ui/src/components/MetricsBadge.tsx
@@ -17,15 +17,16 @@
  * under the License.
  */
 import { Badge, type BadgeProps } from "@chakra-ui/react";
+import type { ReactNode } from "react";
 
 type MetricBadgeProps = {
-  readonly backgroundColor: string;
-  readonly color?: string;
+  readonly icon?: ReactNode;
   readonly runs?: number;
 } & BadgeProps;
 
-export const MetricsBadge = ({ backgroundColor, color = "fg.inverted", runs }: MetricBadgeProps) => (
-  <Badge bg={backgroundColor} borderRadius={15} color={color} minWidth={10} mr={1} px={4} py={1} size="md">
+export const MetricsBadge = ({ icon, runs, ...rest }: MetricBadgeProps) => (
+  <Badge {...rest} borderRadius={15} minWidth={10} mr={1} px={4} py={1} size="md" variant="solid">
+    {icon}
     {runs}
   </Badge>
 );
diff --git a/airflow/ui/src/components/StateIcon.tsx b/airflow/ui/src/components/StateIcon.tsx
new file mode 100644
index 0000000000000..4102cae75df21
--- /dev/null
+++ b/airflow/ui/src/components/StateIcon.tsx
@@ -0,0 +1,70 @@
+/*!
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+import type { IconBaseProps } from "react-icons";
+import {
+  FiActivity,
+  FiAlertCircle,
+  FiAlertOctagon,
+  FiCalendar,
+  FiCheckCircle,
+  FiCircle,
+  FiRepeat,
+  FiSkipForward,
+  FiSlash,
+  FiWatch,
+} from "react-icons/fi";
+import { LuCalendarSync, LuRedo2 } from "react-icons/lu";
+import { PiQueue } from "react-icons/pi";
+
+import type { TaskInstanceState } from "openapi/requests/types.gen";
+
+type Props = {
+  readonly state?: TaskInstanceState | null;
+} & IconBaseProps;
+
+export const StateIcon = ({ state, ...rest }: Props) => {
+  switch (state) {
+    case "deferred":
+      return <FiWatch {...rest} />;
+    case "failed":
+      return <FiAlertOctagon {...rest} />;
+    case "queued":
+      return <PiQueue {...rest} />;
+    case "removed":
+      return <FiSlash {...rest} />;
+    case "restarting":
+      return <FiRepeat {...rest} />;
+    case "running":
+      return <FiActivity {...rest} />;
+    case "scheduled":
+      return <FiCalendar {...rest} />;
+    case "skipped":
+      return <FiSkipForward {...rest} />;
+    case "success":
+      return <FiCheckCircle {...rest} />;
+    case "up_for_reschedule":
+      return <LuCalendarSync {...rest} />;
+    case "up_for_retry":
+      return <LuRedo2 {...rest} />;
+    case "upstream_failed":
+      return <FiAlertCircle {...rest} />;
+    default:
+      return <FiCircle {...rest} />;
+  }
+};
diff --git a/airflow/ui/src/components/TaskInstanceTooltip.tsx b/airflow/ui/src/components/TaskInstanceTooltip.tsx
index bc44a23856ffd..ed09f9a9f93d3 100644
--- a/airflow/ui/src/components/TaskInstanceTooltip.tsx
+++ b/airflow/ui/src/components/TaskInstanceTooltip.tsx
@@ -38,6 +38,7 @@ const TaskInstanceTooltip = ({ children, positioning, taskInstance, ...rest }: P
       {...rest}
       content={
         <Box>
+          <Text>State: {taskInstance.state}</Text>
           {"dag_run_id" in taskInstance ? <Text>Run ID: {taskInstance.dag_run_id}</Text> : undefined}
           <Text>
             Start Date: <Time datetime={taskInstance.start_date} />
diff --git a/airflow/ui/src/components/ui/Status.tsx b/airflow/ui/src/components/ui/Status.tsx
index 91ef7e153096c..bd887ba7c5939 100644
--- a/airflow/ui/src/components/ui/Status.tsx
+++ b/airflow/ui/src/components/ui/Status.tsx
@@ -16,26 +16,24 @@
  * specific language governing permissions and limitations
  * under the License.
  */
-import { Status as ChakraStatus } from "@chakra-ui/react";
+import { Status as ChakraStatus, chakra } from "@chakra-ui/react";
 import * as React from "react";
 
 import type { DagRunState, TaskInstanceState } from "openapi/requests/types.gen";
-import { stateColor } from "src/utils/stateColor";
+
+import { StateIcon } from "../StateIcon";
 
 type StatusValue = DagRunState | TaskInstanceState | null;
 
 export type StatusProps = {
-  state: StatusValue;
+  state?: StatusValue;
 } & ChakraStatus.RootProps;
 
-export const Status = React.forwardRef<HTMLDivElement, StatusProps>(({ children, state, ...rest }, ref) => {
-  // "null" is actually a string on stateColor
-  const colorPalette = stateColor[state ?? "null"];
-
-  return (
-    <ChakraStatus.Root ref={ref} {...rest}>
-      <ChakraStatus.Indicator bg={colorPalette} />
-      {children}
-    </ChakraStatus.Root>
-  );
-});
+export const Status = React.forwardRef<HTMLDivElement, StatusProps>(({ children, state, ...rest }, ref) => (
+  <ChakraStatus.Root ref={ref} {...rest}>
+    <chakra.span color={`${state}.solid`}>
+      <StateIcon state={state} />
+    </chakra.span>
+    {children}
+  </ChakraStatus.Root>
+));
diff --git a/airflow/ui/src/layouts/Details/Graph/Graph.tsx b/airflow/ui/src/layouts/Details/Graph/Graph.tsx
index f6972255f61c4..ee2c270dd6ed9 100644
--- a/airflow/ui/src/layouts/Details/Graph/Graph.tsx
+++ b/airflow/ui/src/layouts/Details/Graph/Graph.tsx
@@ -24,7 +24,6 @@ import { useParams } from "react-router-dom";
 import { useGridServiceGridData, useStructureServiceStructureData } from "openapi/queries";
 import { useColorMode } from "src/context/colorMode";
 import { useOpenGroups } from "src/context/openGroups";
-import { stateColor } from "src/utils/stateColor";
 
 import Edge from "./Edge";
 import { JoinNode } from "./JoinNode";
@@ -42,7 +41,7 @@ const nodeColor = (
   }
 
   if (taskInstance?.state !== undefined && !isOpen) {
-    return stateColor[taskInstance.state ?? "null"];
+    return `var(--chakra-colors-${taskInstance.state}-solid)`;
   }
 
   if (isOpen && depth !== undefined && depth % 2 === 0) {
diff --git a/airflow/ui/src/layouts/Details/Graph/TaskNode.tsx b/airflow/ui/src/layouts/Details/Graph/TaskNode.tsx
index 1372e824a0c35..f6ba9b69bad71 100644
--- a/airflow/ui/src/layouts/Details/Graph/TaskNode.tsx
+++ b/airflow/ui/src/layouts/Details/Graph/TaskNode.tsx
@@ -24,7 +24,6 @@ import TaskInstanceTooltip from "src/components/TaskInstanceTooltip";
 import { Status } from "src/components/ui";
 import { useOpenGroups } from "src/context/openGroups";
 import { pluralize } from "src/utils";
-import { stateColor } from "src/utils/stateColor";
 
 import { NodeWrapper } from "./NodeWrapper";
 import { TaskName } from "./TaskName";
@@ -67,9 +66,7 @@ export const TaskNode = ({
           <Flex
             // Alternate background color for nested open groups
             bg={isOpen && depth !== undefined && depth % 2 === 0 ? "bg.muted" : "bg"}
-            borderColor={
-              taskInstance?.state === undefined ? "border" : stateColor[taskInstance.state ?? "null"]
-            }
+            borderColor={taskInstance?.state ? `${taskInstance.state}.solid` : "border"}
             borderRadius={5}
             borderWidth={isSelected ? 6 : 2}
             height={`${height}px`}
@@ -87,12 +84,12 @@ export const TaskNode = ({
                 label={label}
                 setupTeardownType={setupTeardownType}
               />
-              <Text color="fg.muted" fontSize="xs" mb={-1} mt={2} textTransform="capitalize">
+              <Text color="fg.muted" fontSize="sm" mt={1} textTransform="capitalize">
                 {isGroup ? "Task Group" : operator}
               </Text>
               {taskInstance === undefined ? undefined : (
                 <HStack>
-                  <Status fontSize="xs" state={taskInstance.state}>
+                  <Status fontSize="sm" state={taskInstance.state}>
                     {taskInstance.state}
                   </Status>
                   {taskInstance.try_number > 1 ? <MdRefresh /> : undefined}
diff --git a/airflow/ui/src/main.tsx b/airflow/ui/src/main.tsx
index a89f0f11cbbd3..a5f6a35cfc155 100644
--- a/airflow/ui/src/main.tsx
+++ b/airflow/ui/src/main.tsx
@@ -16,7 +16,7 @@
  * specific language governing permissions and limitations
  * under the License.
  */
-import { ChakraProvider, defaultSystem } from "@chakra-ui/react";
+import { ChakraProvider } from "@chakra-ui/react";
 import { QueryClientProvider } from "@tanstack/react-query";
 import axios, { type AxiosError } from "axios";
 import { StrictMode } from "react";
@@ -28,6 +28,7 @@ import { TimezoneProvider } from "src/context/timezone";
 import { router } from "src/router";
 
 import { queryClient } from "./queryClient";
+import { system } from "./theme";
 
 // redirect to login page if the API responds with unauthorized or forbidden errors
 axios.interceptors.response.use(
@@ -46,7 +47,7 @@ axios.interceptors.response.use(
 
 createRoot(document.querySelector("#root") as HTMLDivElement).render(
   <StrictMode>
-    <ChakraProvider value={defaultSystem}>
+    <ChakraProvider value={system}>
       <ColorModeProvider>
         <QueryClientProvider client={queryClient}>
           <TimezoneProvider>
diff --git a/airflow/ui/src/pages/Dag/Overview/Overview.tsx b/airflow/ui/src/pages/Dag/Overview/Overview.tsx
index a79ebfb4019b6..d89db680f746f 100644
--- a/airflow/ui/src/pages/Dag/Overview/Overview.tsx
+++ b/airflow/ui/src/pages/Dag/Overview/Overview.tsx
@@ -24,7 +24,6 @@ import { useParams } from "react-router-dom";
 import { useDagRunServiceGetDagRuns, useTaskInstanceServiceGetTaskInstances } from "openapi/queries";
 import TimeRangeSelector from "src/components/TimeRangeSelector";
 import { TrendCountButton } from "src/components/TrendCountButton";
-import { stateColor } from "src/utils/stateColor";
 
 const defaultHour = "168";
 
@@ -64,7 +63,7 @@ export const Overview = () => {
       </Box>
       <HStack>
         <TrendCountButton
-          colorPalette={stateColor.failed}
+          colorPalette="failed"
           count={failedTasks?.total_entries ?? 0}
           endDate={endDate}
           events={(failedTasks?.task_instances ?? []).map((ti) => ({
@@ -79,7 +78,7 @@ export const Overview = () => {
           startDate={startDate}
         />
         <TrendCountButton
-          colorPalette={stateColor.failed}
+          colorPalette="failed"
           count={failedRuns?.total_entries ?? 0}
           endDate={endDate}
           events={(failedRuns?.dag_runs ?? []).map((dr) => ({
diff --git a/airflow/ui/src/pages/Dag/Tasks/TaskCard.tsx b/airflow/ui/src/pages/Dag/Tasks/TaskCard.tsx
index 2d517dd4e5798..db8c01c1d77c4 100644
--- a/airflow/ui/src/pages/Dag/Tasks/TaskCard.tsx
+++ b/airflow/ui/src/pages/Dag/Tasks/TaskCard.tsx
@@ -63,9 +63,7 @@ export const TaskCard = ({ dagId, task, taskInstances }: Props) => (
             <Link asChild color="fg.info" fontSize="sm">
               <RouterLink to={getTaskInstanceLink(taskInstances[0])}>
                 <Time datetime={taskInstances[0].start_date} />
-                {taskInstances[0].state === null ? undefined : (
-                  <Status state={taskInstances[0].state}>{taskInstances[0].state}</Status>
-                )}
+                {taskInstances[0].state === null ? undefined : <Status state={taskInstances[0].state} />}
               </RouterLink>
             </Link>
           </TaskInstanceTooltip>
diff --git a/airflow/ui/src/pages/Dag/Tasks/TaskRecentRuns.tsx b/airflow/ui/src/pages/Dag/Tasks/TaskRecentRuns.tsx
index fee2b943582d0..3a7fc13d44949 100644
--- a/airflow/ui/src/pages/Dag/Tasks/TaskRecentRuns.tsx
+++ b/airflow/ui/src/pages/Dag/Tasks/TaskRecentRuns.tsx
@@ -24,7 +24,6 @@ import { Link } from "react-router-dom";
 import type { TaskInstanceResponse } from "openapi/requests/types.gen";
 import TaskInstanceTooltip from "src/components/TaskInstanceTooltip";
 import { getTaskInstanceLink } from "src/utils/links";
-import { stateColor } from "src/utils/stateColor";
 
 dayjs.extend(duration);
 
@@ -58,7 +57,7 @@ export const TaskRecentRuns = ({
             <Link to={getTaskInstanceLink(taskInstance)}>
               <Box p={1}>
                 <Box
-                  bg={stateColor[taskInstance.state]}
+                  bg={`${taskInstance.state}.solid`}
                   borderRadius="4px"
                   height={`${(taskInstance.duration / max) * BAR_HEIGHT}px`}
                   minHeight={1}
diff --git a/airflow/ui/src/pages/DagsList/RecentRuns.tsx b/airflow/ui/src/pages/DagsList/RecentRuns.tsx
index 99f33c9bd0204..4b80355123024 100644
--- a/airflow/ui/src/pages/DagsList/RecentRuns.tsx
+++ b/airflow/ui/src/pages/DagsList/RecentRuns.tsx
@@ -24,7 +24,6 @@ import { Link } from "react-router-dom";
 import type { DAGWithLatestDagRunsResponse } from "openapi/requests/types.gen";
 import Time from "src/components/Time";
 import { Tooltip } from "src/components/ui";
-import { stateColor } from "src/utils/stateColor";
 
 dayjs.extend(duration);
 
@@ -74,7 +73,7 @@ export const RecentRuns = ({
           <Link to={`/dags/${run.dag_id}/runs/${run.dag_run_id}/`}>
             <Box px={1}>
               <Box
-                bg={stateColor[run.state]}
+                bg={`${run.state}.solid`}
                 borderRadius="4px"
                 height={`${(run.duration / max) * BAR_HEIGHT}px`}
                 minHeight={1}
diff --git a/airflow/ui/src/pages/Dashboard/Health/HealthTag.tsx b/airflow/ui/src/pages/Dashboard/Health/HealthTag.tsx
index 31920db3133a9..35b5f75b6c274 100644
--- a/airflow/ui/src/pages/Dashboard/Health/HealthTag.tsx
+++ b/airflow/ui/src/pages/Dashboard/Health/HealthTag.tsx
@@ -16,8 +16,9 @@
  * specific language governing permissions and limitations
  * under the License.
  */
-import { Skeleton, TagLabel, Text } from "@chakra-ui/react";
+import { Skeleton, TagLabel, Text, HStack } from "@chakra-ui/react";
 
+import { StateIcon } from "src/components/StateIcon";
 import Time from "src/components/Time";
 import { Tag, Tooltip } from "src/components/ui";
 import { capitalize } from "src/utils";
@@ -37,6 +38,8 @@ export const HealthTag = ({
     return <Skeleton borderRadius="full" height={8} width={24} />;
   }
 
+  const state = status === "healthy" ? "success" : "failed";
+
   return (
     <Tooltip
       content={
@@ -49,8 +52,11 @@ export const HealthTag = ({
       }
       disabled={!Boolean(latestHeartbeat)}
     >
-      <Tag borderRadius="full" colorPalette={status === "healthy" ? "green" : "red"} size="lg">
-        <TagLabel>{title}</TagLabel>
+      <Tag borderRadius="full" colorPalette={state} size="lg">
+        <HStack>
+          <StateIcon state={state} />
+          <TagLabel>{title}</TagLabel>
+        </HStack>
       </Tag>
     </Tooltip>
   );
diff --git a/airflow/ui/src/pages/Dashboard/HistoricalMetrics/AssetEvents.tsx b/airflow/ui/src/pages/Dashboard/HistoricalMetrics/AssetEvents.tsx
index 21c425f0c5a76..2021a6add8fc7 100644
--- a/airflow/ui/src/pages/Dashboard/HistoricalMetrics/AssetEvents.tsx
+++ b/airflow/ui/src/pages/Dashboard/HistoricalMetrics/AssetEvents.tsx
@@ -51,7 +51,7 @@ export const AssetEvents = ({ assetSortBy, endDate, setAssetSortBy, startDate }:
     <Box borderRadius={5} borderWidth={1} ml={2} pb={2}>
       <Flex justify="space-between" mr={1} mt={0} pl={3} pt={1}>
         <HStack>
-          <MetricsBadge backgroundColor="blue.solid" runs={isLoading ? 0 : data?.total_entries} />
+          <MetricsBadge colorPalette="blue" runs={isLoading ? 0 : data?.total_entries} />
           <Heading marginEnd="auto" size="md">
             Asset Events
           </Heading>
diff --git a/airflow/ui/src/pages/Dashboard/HistoricalMetrics/DagRunMetrics.tsx b/airflow/ui/src/pages/Dashboard/HistoricalMetrics/DagRunMetrics.tsx
index 75b18c85fe64a..c298ee92aa2fe 100644
--- a/airflow/ui/src/pages/Dashboard/HistoricalMetrics/DagRunMetrics.tsx
+++ b/airflow/ui/src/pages/Dashboard/HistoricalMetrics/DagRunMetrics.tsx
@@ -33,7 +33,7 @@ const DAGRUN_STATES: Array<keyof DAGRunStates> = ["queued", "running", "success"
 export const DagRunMetrics = ({ dagRunStates, total }: DagRunMetricsProps) => (
   <Box borderRadius={5} borderWidth={1} p={2}>
     <HStack mb={4}>
-      <MetricsBadge backgroundColor="blue.solid" runs={total} />
+      <MetricsBadge colorPalette="blue" runs={total} />
       <Heading size="md">Dag Runs</Heading>
     </HStack>
     {DAGRUN_STATES.map((state) => (
diff --git a/airflow/ui/src/pages/Dashboard/HistoricalMetrics/MetricSection.tsx b/airflow/ui/src/pages/Dashboard/HistoricalMetrics/MetricSection.tsx
index c2455a8c96161..736943afd6771 100644
--- a/airflow/ui/src/pages/Dashboard/HistoricalMetrics/MetricSection.tsx
+++ b/airflow/ui/src/pages/Dashboard/HistoricalMetrics/MetricSection.tsx
@@ -20,7 +20,6 @@ import { Box, Flex, HStack, VStack, Text } from "@chakra-ui/react";
 
 import { MetricsBadge } from "src/components/MetricsBadge";
 import { capitalize } from "src/utils";
-import { stateColor } from "src/utils/stateColor";
 
 const BAR_WIDTH = 100;
 const BAR_HEIGHT = 5;
@@ -42,14 +41,14 @@ export const MetricSection = ({ runs, state, total }: MetricSectionProps) => {
     <VStack align="left" gap={1} mb={4} ml={0} pl={0}>
       <Flex justify="space-between">
         <HStack>
-          <MetricsBadge backgroundColor={stateColor[state as keyof typeof stateColor]} runs={runs} />
+          <MetricsBadge colorPalette={state} runs={runs} />
           <Text> {capitalize(state)} </Text>
         </HStack>
         <Text color="fg.muted"> {statePercent}% </Text>
       </Flex>
       <HStack gap={0} mt={2}>
         <Box
-          bg={stateColor[state as keyof typeof stateColor]}
+          bg={`${state}.solid`}
           borderLeftRadius={5}
           height={`${BAR_HEIGHT}px`}
           minHeight={2}
diff --git a/airflow/ui/src/pages/Dashboard/HistoricalMetrics/TaskInstanceMetrics.tsx b/airflow/ui/src/pages/Dashboard/HistoricalMetrics/TaskInstanceMetrics.tsx
index 1e92f42dfd373..143b0f2c737b5 100644
--- a/airflow/ui/src/pages/Dashboard/HistoricalMetrics/TaskInstanceMetrics.tsx
+++ b/airflow/ui/src/pages/Dashboard/HistoricalMetrics/TaskInstanceMetrics.tsx
@@ -39,7 +39,7 @@ const TASK_STATES: Array<keyof TaskInstanceStateCount> = [
 export const TaskInstanceMetrics = ({ taskInstanceStates, total }: TaskInstanceMetricsProps) => (
   <Box borderRadius={5} borderWidth={1} mt={2} p={2}>
     <HStack mb={4}>
-      <MetricsBadge backgroundColor="blue.solid" runs={total} />
+      <MetricsBadge colorPalette="blue" runs={total} />
       <Heading size="md">Task Instances</Heading>
     </HStack>
 
diff --git a/airflow/ui/src/pages/Dashboard/Stats/DAGImportErrors.tsx b/airflow/ui/src/pages/Dashboard/Stats/DAGImportErrors.tsx
index e28efa0f9e05a..9b6723ae849e5 100644
--- a/airflow/ui/src/pages/Dashboard/Stats/DAGImportErrors.tsx
+++ b/airflow/ui/src/pages/Dashboard/Stats/DAGImportErrors.tsx
@@ -22,7 +22,6 @@ import { FiChevronRight } from "react-icons/fi";
 import { useImportErrorServiceGetImportErrors } from "openapi/queries";
 import { ErrorAlert } from "src/components/ErrorAlert";
 import { MetricsBadge } from "src/components/MetricsBadge";
-import { stateColor } from "src/utils/stateColor";
 
 import { DAGImportErrorsModal } from "./DAGImportErrorsModal";
 
@@ -49,7 +48,7 @@ export const DAGImportErrors = () => {
           onClick={onOpen}
           variant="outline"
         >
-          <MetricsBadge backgroundColor={stateColor.failed} runs={importErrorsCount} />
+          <MetricsBadge colorPalette="failed" runs={importErrorsCount} />
           <Box alignItems="center" display="flex" gap={1}>
             <Text fontWeight="bold">Dag Import Errors</Text>
             <FiChevronRight />
diff --git a/airflow/ui/src/pages/Dashboard/Stats/DagFilterButton.tsx b/airflow/ui/src/pages/Dashboard/Stats/DagFilterButton.tsx
index b9498566fe85a..f5cdaf49f6014 100644
--- a/airflow/ui/src/pages/Dashboard/Stats/DagFilterButton.tsx
+++ b/airflow/ui/src/pages/Dashboard/Stats/DagFilterButton.tsx
@@ -17,6 +17,7 @@
  * under the License.
  */
 import { Box, Text, Button } from "@chakra-ui/react";
+import type { ReactNode } from "react";
 import { FiChevronRight } from "react-icons/fi";
 import { Link as RouterLink } from "react-router-dom";
 
@@ -26,18 +27,20 @@ import { capitalize } from "src/utils";
 // TODO: Add badge count once API is available
 
 export const DagFilterButton = ({
-  badgeColor,
+  colorPalette,
   filter,
+  icon,
   link,
 }: {
-  readonly badgeColor: string;
+  readonly colorPalette: string;
   readonly filter: string;
+  readonly icon?: ReactNode;
   readonly link: string;
 }) => (
   <RouterLink to={link}>
     <Button alignItems="center" borderRadius="md" display="flex" gap={2} variant="outline">
       <Box alignItems="center" display="flex" gap={1}>
-        <MetricsBadge backgroundColor={badgeColor} />
+        <MetricsBadge colorPalette={colorPalette} icon={icon} />
         <Text fontWeight="bold">{capitalize(filter)} Dags</Text>
         <FiChevronRight />
       </Box>
diff --git a/airflow/ui/src/pages/Dashboard/Stats/Stats.tsx b/airflow/ui/src/pages/Dashboard/Stats/Stats.tsx
index b71b54d42453f..87371c40edad2 100644
--- a/airflow/ui/src/pages/Dashboard/Stats/Stats.tsx
+++ b/airflow/ui/src/pages/Dashboard/Stats/Stats.tsx
@@ -19,7 +19,7 @@
 import { Box, Flex, Heading, HStack } from "@chakra-ui/react";
 import { FiClipboard } from "react-icons/fi";
 
-import { stateColor } from "src/utils/stateColor";
+import { StateIcon } from "src/components/StateIcon";
 
 import { DAGImportErrors } from "./DAGImportErrors";
 import { DagFilterButton } from "./DagFilterButton";
@@ -33,14 +33,20 @@ export const Stats = () => (
       </Heading>
     </Flex>
     <HStack>
-      <DagFilterButton badgeColor={stateColor.failed} filter="failed" link="dags?last_dag_run_state=failed" />
+      <DagFilterButton
+        colorPalette="failed"
+        filter="failed"
+        icon={<StateIcon state="failed" />}
+        link="dags?last_dag_run_state=failed"
+      />
       <DAGImportErrors />
       <DagFilterButton
-        badgeColor={stateColor.running}
+        colorPalette="running"
         filter="running"
+        icon={<StateIcon state="running" />}
         link="dags?last_dag_run_state=running"
       />
-      <DagFilterButton badgeColor="blue.solid" filter="active" link="dags?paused=false" />
+      <DagFilterButton colorPalette="blue" filter="active" link="dags?paused=false" />
     </HStack>
   </Box>
 );
diff --git a/airflow/ui/src/pages/Pools/PoolBar.tsx b/airflow/ui/src/pages/Pools/PoolBar.tsx
index 1fead035016c9..0cc8a1a1c47c7 100644
--- a/airflow/ui/src/pages/Pools/PoolBar.tsx
+++ b/airflow/ui/src/pages/Pools/PoolBar.tsx
@@ -19,21 +19,20 @@
  * under the License.
  */
 import { Box, Flex, HStack, Text, VStack } from "@chakra-ui/react";
-import React from "react";
-import { FiActivity, FiCalendar, FiCheckCircle, FiClock, FiList, FiXCircle } from "react-icons/fi";
+import { FiXCircle } from "react-icons/fi";
 
 import type { PoolResponse } from "openapi/requests/types.gen";
+import { StateIcon } from "src/components/StateIcon";
 import { Tooltip } from "src/components/ui";
 import { capitalize } from "src/utils";
-import { stateColor } from "src/utils/stateColor";
 
 const slots = {
-  open_slots: { color: stateColor.success, icon: FiCheckCircle },
-  occupied_slots: { color: stateColor.up_for_retry, icon: FiXCircle },
-  running_slots: { color: stateColor.running, icon: FiActivity },
-  queued_slots: { color: stateColor.queued, icon: FiList },
-  scheduled_slots: { color: stateColor.scheduled, icon: FiCalendar },
-  deferred_slots: { color: stateColor.deferred, icon: FiClock },
+  open_slots: { color: "success", icon: <StateIcon state="success" /> },
+  occupied_slots: { color: "up_for_retry", icon: <FiXCircle /> },
+  running_slots: { color: "running", icon: <StateIcon state="running" /> },
+  queued_slots: { color: "queued", icon: <StateIcon state="queued" /> },
+  scheduled_slots: { color: "scheduled", icon: <StateIcon state="scheduled" /> },
+  deferred_slots: { color: "deferred", icon: <StateIcon state="deferred" /> },
 };
 
 type PoolBarProps = {
@@ -50,7 +49,7 @@ const PoolBar = ({ pool }: PoolBarProps) => (
           </Text>
           {pool.include_deferred ? (
             <Tooltip content="Deferred Slots Included">
-              <FiClock size={18} />
+              <StateIcon size={18} state="deferred" />
             </Tooltip>
           ) : undefined}
         </HStack>
@@ -77,13 +76,13 @@ const PoolBar = ({ pool }: PoolBarProps) => (
             <Tooltip content={`${capitalize(slotKey.replace("_", " "))}: ${slotValue}`} key={slotKey}>
               <Flex
                 alignItems="center"
-                bg={color}
+                colorPalette={color}
                 flex={flexValue}
                 h="100%"
                 justifyContent="center"
                 position="relative"
               >
-                {React.createElement(icon, { size: 16, color: "white" })}
+                {icon}
               </Flex>
             </Tooltip>
           );
diff --git a/airflow/ui/src/theme.ts b/airflow/ui/src/theme.ts
new file mode 100644
index 0000000000000..706ae0ed57e41
--- /dev/null
+++ b/airflow/ui/src/theme.ts
@@ -0,0 +1,176 @@
+/*!
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/* eslint-disable perfectionist/sort-objects */
+import { createSystem, defaultConfig, defineConfig } from "@chakra-ui/react";
+
+const generateSemanticTokens = (color: string, darkContrast: string = "white") => ({
+  solid: { value: `{colors.${color}.600}` },
+  contrast: { value: { _light: "white", _dark: darkContrast } },
+  fg: { value: { _light: `{colors.${color}.800}`, _dark: `{colors.${color}.200}` } },
+  muted: { value: { _light: `{colors.${color}.200}`, _dark: `{colors.${color}.800}` } },
+  subtle: { value: { _light: `{colors.${color}.100}`, _dark: `{colors.${color}.900}` } },
+  emphasized: { value: { _light: `{colors.${color}.300}`, _dark: `{colors.${color}.700}` } },
+  focusRing: { value: { _light: `{colors.${color}.800}`, _dark: `{colors.${color}.200}` } },
+});
+
+const customConfig = defineConfig({
+  theme: {
+    tokens: {
+      colors: {
+        // Default green was too light
+        success: {
+          "50": { value: "#E0FFE0" },
+          "100": { value: "#C2FFC2" },
+          "200": { value: "#80FF80" },
+          "300": { value: "#42FF42" },
+          "400": { value: "#00FF00" },
+          "500": { value: "#00C200" },
+          "600": { value: "#008000" },
+          "700": { value: "#006100" },
+          "800": { value: "#004200" },
+          "900": { value: "#001F00" },
+          "950": { value: "#000F00" },
+        },
+        failed: defaultConfig.theme?.tokens?.colors?.red ?? {},
+        // Default gray was too dark
+        queued: {
+          "50": { value: "#F5F5F5" },
+          "100": { value: "#EBEBEB" },
+          "200": { value: "#D4D4D4" },
+          "300": { value: "#BFBFBF" },
+          "400": { value: "#ABABAB" },
+          "500": { value: "#969696" },
+          "600": { value: "#808080" },
+          "700": { value: "#616161" },
+          "800": { value: "#404040" },
+          "900": { value: "#212121" },
+          "950": { value: "#0F0F0F" },
+        },
+        skipped: defaultConfig.theme?.tokens?.colors?.pink ?? {},
+        up_for_reschedule: defaultConfig.theme?.tokens?.colors?.cyan ?? {},
+        up_for_retry: defaultConfig.theme?.tokens?.colors?.yellow ?? {},
+        upstream_failed: defaultConfig.theme?.tokens?.colors?.orange ?? {},
+        // lime
+        running: {
+          "50": { value: "#EFFBEF" },
+          "100": { value: "#DEF7DE" },
+          "200": { value: "#B9EEB9" },
+          "300": { value: "#98E698" },
+          "400": { value: "#78DE78" },
+          "500": { value: "#53D553" },
+          "600": { value: "#32CD32" },
+          "700": { value: "#269C26" },
+          "800": { value: "#196719" },
+          "900": { value: "#0D350D" },
+          "950": { value: "#061906" },
+        },
+        // violet
+        restarting: {
+          "50": { value: "#F6EBFF" },
+          "100": { value: "#EDD6FF" },
+          "200": { value: "#D9A8FF" },
+          "300": { value: "#C880FF" },
+          "400": { value: "#B657FF" },
+          "500": { value: "#A229FF" },
+          "600": { value: "#8F00FF" },
+          "700": { value: "#6E00C2" },
+          "800": { value: "#480080" },
+          "900": { value: "#260042" },
+          "950": { value: "#11001F" },
+        },
+        // mediumpurple
+        deferred: {
+          "50": { value: "#F6F3FC" },
+          "100": { value: "#EDE7F9" },
+          "200": { value: "#DACEF3" },
+          "300": { value: "#C8B6ED" },
+          "400": { value: "#B9A1E7" },
+          "500": { value: "#A689E1" },
+          "600": { value: "#9370DB" },
+          "700": { value: "#6432C8" },
+          "800": { value: "#412182" },
+          "900": { value: "#211041" },
+          "950": { value: "#100821" },
+        },
+        // tan
+        scheduled: {
+          "50": { value: "#FBF8F4" },
+          "100": { value: "#F8F3ED" },
+          "200": { value: "#F1E7DA" },
+          "300": { value: "#E8D9C4" },
+          "400": { value: "#E1CDB2" },
+          "500": { value: "#DAC1A0" },
+          "600": { value: "#D2B48C" },
+          "700": { value: "#B9894B" },
+          "800": { value: "#7D5C31" },
+          "900": { value: "#3E2E18" },
+          "950": { value: "#21180D" },
+        },
+        // lightblue
+        none: {
+          "50": { value: "#F7FBFD" },
+          "100": { value: "#F3F9FB" },
+          "200": { value: "#E4F2F7" },
+          "300": { value: "#D8ECF3" },
+          "400": { value: "#C8E5EE" },
+          "500": { value: "#BDDFEB" },
+          "600": { value: "#ADD8E6" },
+          "700": { value: "#5FB2CE" },
+          "800": { value: "#30819C" },
+          "900": { value: "#18414E" },
+          "950": { value: "#0C2027" },
+        },
+        // lightgrey
+        removed: {
+          "50": { value: "#FCFCFC" },
+          "100": { value: "#F7F7F7" },
+          "200": { value: "#F0F0F0" },
+          "300": { value: "#E8E8E8" },
+          "400": { value: "#E0E0E0" },
+          "500": { value: "#DBDBDB" },
+          "600": { value: "#D3D3D3" },
+          "700": { value: "#9E9E9E" },
+          "800": { value: "#696969" },
+          "900": { value: "#363636" },
+          "950": { value: "#1A1A1A" },
+        },
+      },
+    },
+    semanticTokens: {
+      colors: {
+        success: generateSemanticTokens("success"),
+        failed: defaultConfig.theme?.semanticTokens?.colors?.red ?? {},
+        queued: generateSemanticTokens("queued"),
+        skipped: defaultConfig.theme?.semanticTokens?.colors?.pink ?? {},
+        up_for_reschedule: defaultConfig.theme?.semanticTokens?.colors?.cyan ?? {},
+        up_for_retry: defaultConfig.theme?.semanticTokens?.colors?.yellow ?? {},
+        upstream_failed: defaultConfig.theme?.semanticTokens?.colors?.orange ?? {},
+        running: generateSemanticTokens("running"),
+        restarting: generateSemanticTokens("restarting"),
+        deferred: generateSemanticTokens("deferred"),
+        scheduled: generateSemanticTokens("scheduled"),
+        none: generateSemanticTokens("none", "black"),
+        removed: generateSemanticTokens("removed", "black"),
+      },
+    },
+  },
+});
+
+export const system = createSystem(defaultConfig, customConfig);
diff --git a/airflow/ui/src/utils/stateColor.ts b/airflow/ui/src/utils/stateColor.ts
deleted file mode 100644
index 555cd0fbdfdda..0000000000000
--- a/airflow/ui/src/utils/stateColor.ts
+++ /dev/null
@@ -1,37 +0,0 @@
-/*!
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements.  See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership.  The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License.  You may obtain a copy of the License at
- *
- *   http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing,
- * software distributed under the License is distributed on an
- * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
- * KIND, either express or implied.  See the License for the
- * specific language governing permissions and limitations
- * under the License.
- */
-
-// TODO: replace this with Airflow config values
-
-export const stateColor = {
-  deferred: "mediumpurple",
-  failed: "red",
-  none: "lightblue",
-  null: "lightblue",
-  queued: "gray",
-  removed: "lightgrey",
-  restarting: "violet",
-  running: "lime",
-  scheduled: "tan",
-  skipped: "hotpink",
-  success: "green",
-  up_for_reschedule: "turquoise",
-  up_for_retry: "gold",
-  upstream_failed: "orange",
-};