Fix service insight live update bugs (#22822)

* Fix live update bugs * support automator status * fix query search index * Update the logic to display the AutoPilot status and add playwright test * Fix tests --------- Co-authored-by: ulixius9 <mayursingal9@gmail.com>
2025-12-29 08:29:09 +00:00 · 2025-08-08 07:55:55 +05:30 · 2025-08-08 07:55:55 +05:30 · bce4bcd32f
commit bce4bcd32f
parent ec9b0ef030
40 changed files with 908 additions and 121 deletions
--- a/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/DataInsightSystemChartRepository.java
+++ b/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/DataInsightSystemChartRepository.java
@ -25,6 +25,7 @@ import org.openmetadata.schema.entity.app.AppType;
 import org.openmetadata.schema.entity.services.DatabaseService;
 import org.openmetadata.schema.entity.services.ingestionPipelines.IngestionPipeline;
 import org.openmetadata.schema.entity.services.ingestionPipelines.PipelineStatus;
+import org.openmetadata.schema.governance.workflows.WorkflowInstance;
 import org.openmetadata.schema.type.EntityReference;
 import org.openmetadata.schema.type.Include;
 import org.openmetadata.schema.utils.JsonUtils;
@ -33,6 +34,7 @@ import org.openmetadata.service.search.SearchClient;
 import org.openmetadata.service.socket.WebSocketManager;
 import org.openmetadata.service.socket.messages.ChartDataStreamMessage;
 import org.openmetadata.service.util.EntityUtil;
+import org.openmetadata.service.util.FullyQualifiedName;
 import org.openmetadata.service.util.ResultList;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@ -238,6 +240,67 @@ public class DataInsightSystemChartRepository extends EntityRepository<DataInsig
    return appStatusList;
  }

+  /**
+   * Get workflow instances for a specific entity link
+   * @param entityLink Entity link to filter workflow instances
+   * @param startTime Start timestamp for data range
+   * @param endTime End timestamp for data range
+   * @return List of workflow instances
+   */
+  private List<Map> getWorkflowInstances(String entityLink, long startTime, long endTime) {
+    List<Map> workflowInstances = new ArrayList<>();
+
+    try {
+      if (entityLink == null || entityLink.trim().isEmpty()) {
+        return workflowInstances;
+      }
+
+      // Get the workflow instance repository
+      WorkflowInstanceRepository workflowInstanceRepository =
+          (WorkflowInstanceRepository)
+              Entity.getEntityTimeSeriesRepository(Entity.WORKFLOW_INSTANCE);
+
+      if (workflowInstanceRepository == null) {
+        LOG.warn("WorkflowInstanceRepository not available");
+        return workflowInstances;
+      }
+
+      // Create filter for workflow instances
+      ListFilter filter = new ListFilter(null);
+      filter.addQueryParam("entityFQNHash", FullyQualifiedName.buildHash("AutoPilotWorkflow"));
+      filter.addQueryParam("entityLink", entityLink);
+
+      // Fetch workflow instances
+      ResultList<WorkflowInstance> instances =
+          workflowInstanceRepository.list(null, startTime, endTime, 100, filter, false);
+
+      if (instances != null && instances.getData() != null) {
+        for (WorkflowInstance instance : instances.getData()) {
+          Map<String, Object> instanceData = new HashMap<>();
+          instanceData.put("id", instance.getId().toString());
+          instanceData.put("workflowDefinitionId", instance.getWorkflowDefinitionId().toString());
+          instanceData.put("startedAt", instance.getStartedAt());
+          instanceData.put("endedAt", instance.getEndedAt());
+          instanceData.put("status", instance.getStatus().toString());
+          instanceData.put("timestamp", instance.getTimestamp());
+          instanceData.put("variables", instance.getVariables());
+          instanceData.put("entityLink", entityLink);
+          instanceData.put("type", "workflow");
+
+          workflowInstances.add(instanceData);
+        }
+      }
+
+      LOG.info(
+          "Found {} workflow instances for entity link: {}", workflowInstances.size(), entityLink);
+
+    } catch (Exception e) {
+      LOG.error("Error fetching workflow instances for entity link: {}", entityLink, e);
+    }
+
+    return workflowInstances;
+  }
+
  /**
   * Get the app repository
   */
@ -661,9 +724,11 @@ public class DataInsightSystemChartRepository extends EntityRepository<DataInsig
   * Check if there's already an active streaming session for the same criteria
   * @param chartNames Chart names being requested
   * @param serviceName Service name (can be null)
+   * @param entityLink Entity link (can be null)
   * @return Active session if exists, null otherwise
   */
-  private StreamingSession findActiveSession(String chartNames, String serviceName) {
+  private StreamingSession findActiveSession(
+      String chartNames, String serviceName, String entityLink) {
    return activeSessions.values().stream()
        .filter(session -> session.getChartNames().equals(chartNames))
        .filter(
@ -677,6 +742,17 @@ public class DataInsightSystemChartRepository extends EntityRepository<DataInsig
              }
              return false;
            })
+        .filter(
+            session -> {
+              // Handle null entityLink comparison
+              if (entityLink == null && session.getEntityLink() == null) {
+                return true;
+              }
+              if (entityLink != null && session.getEntityLink() != null) {
+                return entityLink.equals(session.getEntityLink());
+              }
+              return false;
+            })
        .findFirst()
        .orElse(null);
  }
@ -695,6 +771,7 @@ public class DataInsightSystemChartRepository extends EntityRepository<DataInsig
      String chartNames,
      String serviceName,
      String filter,
+      String entityLink,
      UUID userId,
      Long startTime,
      Long endTime) {
@ -707,7 +784,7 @@ public class DataInsightSystemChartRepository extends EntityRepository<DataInsig
    }

    // Check if there's already an active streaming session for the same criteria
-    StreamingSession existingSession = findActiveSession(chartNames, serviceName);
+    StreamingSession existingSession = findActiveSession(chartNames, serviceName, entityLink);
    if (existingSession != null) {
      // Add this user to the existing session
      existingSession.addUser(userId);
@ -731,7 +808,11 @@ public class DataInsightSystemChartRepository extends EntityRepository<DataInsig
          existingSession.getRemainingTime(),
          UPDATE_INTERVAL_MS,
          getIngestionPipelineStatus(serviceName),
-          getCollateAppStatus(serviceName));
+          getCollateAppStatus(serviceName),
+          getWorkflowInstances(
+              existingSession.getEntityLink(),
+              existingSession.getDataStartTime(),
+              existingSession.getDataEndTime()));

      // Calculate remaining time for existing session
      long remainingTime = existingSession.getRemainingTime();
@ -771,7 +852,7 @@ public class DataInsightSystemChartRepository extends EntityRepository<DataInsig

    try {
      String sessionId =
-          startStreaming(chartNames, serviceName, filter, userId, startTime, endTime);
+          startStreaming(chartNames, serviceName, filter, entityLink, userId, startTime, endTime);

      Map<String, Object> response = new HashMap<>();
      response.put("sessionId", sessionId);
@ -803,22 +884,24 @@ public class DataInsightSystemChartRepository extends EntityRepository<DataInsig
      String chartNames,
      String serviceName,
      String filter,
+      String entityLink,
      UUID userId,
      Long startTime,
      Long endTime) {
    String sessionId = UUID.randomUUID().toString();

    LOG.info(
-        "Starting chart data streaming session {} for user {} with charts: {} (time range: {} to {})",
+        "Starting chart data streaming session {} for user {} with charts: {} and entityLink: {} (time range: {} to {})",
        sessionId,
        userId,
        chartNames,
+        entityLink,
        startTime,
        endTime);

    StreamingSession session =
        new StreamingSession(
-            sessionId, chartNames, serviceName, filter, userId, startTime, endTime);
+            sessionId, chartNames, serviceName, filter, entityLink, userId, startTime, endTime);
    activeSessions.put(sessionId, session);

    // Send initial status message to all users in the session
@ -830,7 +913,8 @@ public class DataInsightSystemChartRepository extends EntityRepository<DataInsig
        STREAM_DURATION_MS,
        UPDATE_INTERVAL_MS,
        getIngestionPipelineStatus(serviceName),
-        getCollateAppStatus(serviceName));
+        getCollateAppStatus(serviceName),
+        getWorkflowInstances(entityLink, startTime, endTime));

    // Schedule the streaming task
    ScheduledFuture<?> future =
@ -862,7 +946,8 @@ public class DataInsightSystemChartRepository extends EntityRepository<DataInsig
        session.getFuture().cancel(true);
      }

-      sendMessageToAllUsers(session, "COMPLETED", null, null, 0L, 0L, List.of(), List.of());
+      sendMessageToAllUsers(
+          session, "COMPLETED", null, null, 0L, 0L, List.of(), List.of(), List.of());
      activeSessions.remove(sessionId);
    }
  }
@ -903,7 +988,8 @@ public class DataInsightSystemChartRepository extends EntityRepository<DataInsig
      if (session.getFuture() != null) {
        session.getFuture().cancel(true);
      }
-      sendMessageToAllUsers(session, "COMPLETED", null, null, 0L, 0L, List.of(), List.of());
+      sendMessageToAllUsers(
+          session, "COMPLETED", null, null, 0L, 0L, List.of(), List.of(), List.of());
      activeSessions.remove(sessionId);

      response.put("status", "stopped");
@ -919,7 +1005,9 @@ public class DataInsightSystemChartRepository extends EntityRepository<DataInsig
          session.getRemainingTime(),
          UPDATE_INTERVAL_MS,
          getIngestionPipelineStatus(session.getServiceName()),
-          getCollateAppStatus(session.getServiceName()));
+          getCollateAppStatus(session.getServiceName()),
+          getWorkflowInstances(
+              session.getEntityLink(), session.getDataStartTime(), session.getDataEndTime()));

      response.put("status", "user_removed");
      response.put("message", "User removed from streaming session");
@ -960,6 +1048,10 @@ public class DataInsightSystemChartRepository extends EntityRepository<DataInsig
      // Fetch ingestion pipeline status for the service
      List<Map> ingestionPipelineStatus = getIngestionPipelineStatus(session.getServiceName());

+      // Fetch workflow instances for the entity link
+      List<Map> workflowInstances =
+          getWorkflowInstances(session.getEntityLink(), startTime, endTime);
+
      // Send the data to all users in the session
      sendMessageToAllUsers(
          session,
@ -969,7 +1061,8 @@ public class DataInsightSystemChartRepository extends EntityRepository<DataInsig
          remainingTime,
          UPDATE_INTERVAL_MS,
          ingestionPipelineStatus,
-          getCollateAppStatus(session.getServiceName()));
+          getCollateAppStatus(session.getServiceName()),
+          workflowInstances);

    } catch (IOException e) {
      LOG.error("Error streaming chart data for session {}", session.getSessionId(), e);
@ -981,7 +1074,8 @@ public class DataInsightSystemChartRepository extends EntityRepository<DataInsig
          0L,
          0L,
          List.of(),
-          getCollateAppStatus(session.getServiceName()));
+          getCollateAppStatus(session.getServiceName()),
+          List.of());
      stopStreaming(session.getSessionId());
    } catch (Exception e) {
      LOG.error("Unexpected error in streaming session {}", session.getSessionId(), e);
@ -993,7 +1087,8 @@ public class DataInsightSystemChartRepository extends EntityRepository<DataInsig
          0L,
          0L,
          List.of(),
-          getCollateAppStatus(session.getServiceName()));
+          getCollateAppStatus(session.getServiceName()),
+          List.of());
      stopStreaming(session.getSessionId());
    }
  }
@ -1011,7 +1106,8 @@ public class DataInsightSystemChartRepository extends EntityRepository<DataInsig
      Long remainingTime,
      Long nextUpdate,
      List<Map> ingestionPipelineStatus,
-      List<Map> appStatus) {
+      List<Map> appStatus,
+      List<Map> workflowInstances) {
    ChartDataStreamMessage message =
        new ChartDataStreamMessage(
            sessionId,
@ -1023,7 +1119,8 @@ public class DataInsightSystemChartRepository extends EntityRepository<DataInsig
            remainingTime,
            nextUpdate,
            ingestionPipelineStatus,
-            appStatus);
+            appStatus,
+            workflowInstances);

    String messageJson = JsonUtils.pojoToJson(message);

@ -1043,7 +1140,8 @@ public class DataInsightSystemChartRepository extends EntityRepository<DataInsig
      Long remainingTime,
      Long nextUpdate,
      List<Map> ingestionPipelineStatus,
-      List<Map> appStatus) {
+      List<Map> appStatus,
+      List<Map> workflowInstances) {
    for (UUID userId : session.getUserIds()) {
      sendMessageToUser(
          userId,
@ -1055,7 +1153,8 @@ public class DataInsightSystemChartRepository extends EntityRepository<DataInsig
          remainingTime,
          nextUpdate,
          ingestionPipelineStatus,
-          appStatus);
+          appStatus,
+          workflowInstances);
    }
  }

@ -1086,6 +1185,7 @@ public class DataInsightSystemChartRepository extends EntityRepository<DataInsig
    private final String chartNames;
    private final String serviceName;
    private final String filter;
+    private final String entityLink;
    private final Set<UUID> userIds; // Multiple users can share the same session
    private final long startTime; // Session start time
    private final long dataStartTime; // Data range start time
@ -1097,6 +1197,7 @@ public class DataInsightSystemChartRepository extends EntityRepository<DataInsig
        String chartNames,
        String serviceName,
        String filter,
+        String entityLink,
        UUID userId,
        Long dataStartTime,
        Long dataEndTime) {
@ -1104,6 +1205,7 @@ public class DataInsightSystemChartRepository extends EntityRepository<DataInsig
      this.chartNames = chartNames;
      this.serviceName = serviceName;
      this.filter = filter;
+      this.entityLink = entityLink;
      this.userIds = new ConcurrentHashMap().newKeySet(); // Thread-safe set
      this.userIds.add(userId);
      this.startTime = System.currentTimeMillis(); // Session start time
@ -1142,6 +1244,10 @@ public class DataInsightSystemChartRepository extends EntityRepository<DataInsig
      return filter;
    }

+    public String getEntityLink() {
+      return entityLink;
+    }
+
    public Set<UUID> getUserIds() {
      return userIds;
    }
--- a/openmetadata-service/src/main/java/org/openmetadata/service/resources/datainsight/system/DataInsightSystemChartResource.java
+++ b/openmetadata-service/src/main/java/org/openmetadata/service/resources/datainsight/system/DataInsightSystemChartResource.java
@ -209,6 +209,11 @@ public class DataInsightSystemChartResource
              schema = @Schema(type = "string", example = "{\"query\":{...}}"))
          @QueryParam("filter")
          String filter,
+      @Parameter(
+              description = "Entity link for workflow instances filtering",
+              schema = @Schema(type = "String", example = "<#E::databaseService::sample_data>"))
+          @QueryParam("entityLink")
+          String entityLink,
      @Parameter(
              description = "Start time for data fetching (unix timestamp in milliseconds)",
              schema = @Schema(type = "long", example = "1426349294842"))
@ -229,7 +234,7 @@ public class DataInsightSystemChartResource
      // Call repository method to handle streaming
      Map<String, Object> response =
          repository.startChartDataStreaming(
-              chartNames, serviceName, filter, user.getId(), startTime, endTime);
+              chartNames, serviceName, filter, entityLink, user.getId(), startTime, endTime);

      // Check if there's an error in the response
      if (response.containsKey("error")) {
--- a/openmetadata-service/src/main/java/org/openmetadata/service/socket/messages/ChartDataStreamMessage.java
+++ b/openmetadata-service/src/main/java/org/openmetadata/service/socket/messages/ChartDataStreamMessage.java
@ -21,7 +21,8 @@ public class ChartDataStreamMessage {
      Long remainingTime,
      Long nextUpdate,
      List<Map> ingestionPipelineStatus,
-      List<Map> appStatus) {
+      List<Map> appStatus,
+      List<Map> workflowInstances) {
    this.sessionId = sessionId;
    this.status = status;
    this.serviceName = serviceName;
@ -32,6 +33,7 @@ public class ChartDataStreamMessage {
    this.nextUpdate = nextUpdate;
    this.ingestionPipelineStatus = ingestionPipelineStatus;
    this.appStatus = appStatus;
+    this.workflowInstances = workflowInstances;
  }

  @JsonProperty("sessionId")
@ -63,4 +65,7 @@ public class ChartDataStreamMessage {

  @JsonProperty("appStatus")
  private List<Map> appStatus;
+
+  @JsonProperty("workflowInstances")
+  private List<Map> workflowInstances;
 }
--- a/openmetadata-spec/src/main/resources/elasticsearch/en/query_index_mapping.json
+++ b/openmetadata-spec/src/main/resources/elasticsearch/en/query_index_mapping.json
@ -76,7 +76,13 @@
  "mappings": {
    "properties": {
      "id": {
-        "type": "keyword"
+        "type": "text",
+        "fields": {
+          "keyword": {
+            "type": "keyword",
+            "ignore_above": 256
+          }
+        }
      },
      "checksum": {
        "type": "text",
--- a/openmetadata-spec/src/main/resources/elasticsearch/jp/query_index_mapping.json
+++ b/openmetadata-spec/src/main/resources/elasticsearch/jp/query_index_mapping.json
@ -53,7 +53,13 @@
  "mappings": {
    "properties": {
      "id": {
-        "type": "keyword"
+        "type": "text",
+        "fields": {
+          "keyword": {
+            "type": "keyword",
+            "ignore_above": 256
+          }
+        }
      },
      "name": {
        "type": "text",
--- a/openmetadata-spec/src/main/resources/elasticsearch/zh/query_index_mapping.json
+++ b/openmetadata-spec/src/main/resources/elasticsearch/zh/query_index_mapping.json
@ -58,7 +58,13 @@
  "mappings": {
    "properties": {
      "id": {
-        "type": "keyword"
+        "type": "text",
+        "fields": {
+          "keyword": {
+            "type": "keyword",
+            "ignore_above": 256
+          }
+        }
      },
      "name": {
        "type": "text",
--- a/openmetadata-ui/src/main/resources/ui/playwright/e2e/Features/AutoPilot.spec.ts
+++ b/openmetadata-ui/src/main/resources/ui/playwright/e2e/Features/AutoPilot.spec.ts
@ -26,7 +26,6 @@ import {
  createNewPage,
  getApiContext,
  redirectToHomePage,
-  reloadAndWaitForNetworkIdle,
 } from '../../utils/common';
 import { getServiceCategoryFromService } from '../../utils/serviceIngestion';
 import { settingClick, SettingOptionsType } from '../../utils/sidebar';
@ -115,19 +114,67 @@ services.forEach((ServiceClass) => {
          state: 'detached',
        });

-        // Reload the page and wait for the network to be idle
-        await reloadAndWaitForNetworkIdle(page);
-
        // Check the auto pilot status
        await checkAutoPilotStatus(page, service);

-        // Reload the page and wait for the network to be idle
-        await reloadAndWaitForNetworkIdle(page);
-
        // Wait for the auto pilot status banner to be visible
        await expect(
          page.getByText('AutoPilot agents run completed successfully.')
        ).toBeVisible();
+
+        if (service.serviceType === 'Mysql') {
+          await page.getByTestId('agent-status-widget-view-more').click();
+
+          await page.waitForSelector(
+            '[data-testid="agent-status-card-Metadata"]',
+            {
+              state: 'visible',
+            }
+          );
+
+          // Check the agents statuses
+          await expect(
+            page.getByTestId('agent-status-card-Lineage')
+          ).toBeVisible();
+          await expect(
+            page.getByTestId('agent-status-card-Usage')
+          ).toBeVisible();
+          await expect(
+            page.getByTestId('agent-status-card-Auto Classification')
+          ).toBeVisible();
+          await expect(
+            page.getByTestId('agent-status-card-Profiler')
+          ).toBeVisible();
+
+          // Check the agents summary
+          await expect(
+            page
+              .getByTestId('agent-status-summary-item-Successful')
+              .getByTestId('pipeline-count')
+          ).toHaveText('3');
+          await expect(
+            page
+              .getByTestId('agent-status-summary-item-Pending')
+              .getByTestId('pipeline-count')
+          ).toHaveText('2');
+
+          // Check the total data assets count
+          await expect(
+            page
+              .getByTestId('total-data-assets-widget')
+              .getByTestId('Database-count')
+          ).toHaveText('1');
+          await expect(
+            page
+              .getByTestId('total-data-assets-widget')
+              .getByTestId('Database Schema-count')
+          ).toHaveText('2');
+          await expect(
+            page
+              .getByTestId('total-data-assets-widget')
+              .getByTestId('Table-count')
+          ).toHaveText('3');
+        }
      });

      test('Agents created by AutoPilot should be deleted', async ({
--- a/openmetadata-ui/src/main/resources/ui/playwright/support/entity/ingestion/ServiceBaseClass.ts
+++ b/openmetadata-ui/src/main/resources/ui/playwright/support/entity/ingestion/ServiceBaseClass.ts
@ -46,7 +46,7 @@ class ServiceBaseClass {
  protected entityName: string;
  protected shouldTestConnection: boolean;
  protected shouldAddIngestion: boolean;
-  protected shouldAddDefaultFilters: boolean;
+  public shouldAddDefaultFilters: boolean;
  protected entityFQN: string | null;
  public serviceResponseData: ResponseDataType = {} as ResponseDataType;

--- a/openmetadata-ui/src/main/resources/ui/src/components/MyData/Widgets/CuratedAssetsWidget/CuratedAssetsWidget.test.tsx
+++ b/openmetadata-ui/src/main/resources/ui/src/components/MyData/Widgets/CuratedAssetsWidget/CuratedAssetsWidget.test.tsx
@ -78,6 +78,10 @@ jest.mock('../../../../utils/ServiceUtilClassBase', () => ({
  getServiceTypeLogo: jest.fn().mockReturnValue('test-logo.png'),
 }));

+jest.mock('../../../common/RichTextEditor/RichTextEditorPreviewerV1', () =>
+  jest.fn().mockImplementation(({ markdown }) => <div>{markdown}</div>)
+);
+
 jest.mock('../../../common/ErrorWithPlaceholder/ErrorPlaceHolder', () =>
  jest.fn().mockImplementation(({ children, icon, type, className }) => (
    <div className={className} data-testid="error-placeholder" data-type={type}>
--- a/openmetadata-ui/src/main/resources/ui/src/components/ServiceInsights/AgentsStatusWidget/AgentsStatusWidget.interface.ts
+++ b/openmetadata-ui/src/main/resources/ui/src/components/ServiceInsights/AgentsStatusWidget/AgentsStatusWidget.interface.ts
@ -12,12 +12,14 @@
 */
 import { ReactNode } from 'react';
 import { AgentStatus } from '../../../enums/ServiceInsights.enum';
+import { WorkflowInstance } from '../../../generated/governance/workflows/workflowInstance';
 import { ServiceInsightWidgetCommonProps } from '../ServiceInsightsTab.interface';

 export interface AgentsStatusWidgetProps
  extends ServiceInsightWidgetCommonProps {
  isLoading: boolean;
  agentsInfo: AgentsInfo[];
+  liveAutoPilotStatusData?: WorkflowInstance;
 }

 export interface AgentsInfo {
--- a/openmetadata-ui/src/main/resources/ui/src/components/ServiceInsights/AgentsStatusWidget/AgentsStatusWidget.tsx
+++ b/openmetadata-ui/src/main/resources/ui/src/components/ServiceInsights/AgentsStatusWidget/AgentsStatusWidget.tsx
@ -28,7 +28,7 @@ import './agents-status-widget.less';
 import { AgentsStatusWidgetProps } from './AgentsStatusWidget.interface';

 function AgentsStatusWidget({
-  workflowStatesData,
+  liveAutoPilotStatusData,
  isLoading,
  agentsInfo,
 }: Readonly<AgentsStatusWidgetProps>) {
@ -36,8 +36,12 @@ function AgentsStatusWidget({

  const agentsRunningStatusMessage = useMemo(
    () =>
-      getAgentRunningStatusMessage(isLoading, agentsInfo, workflowStatesData),
-    [workflowStatesData, isLoading, agentsInfo]
+      getAgentRunningStatusMessage(
+        isLoading,
+        agentsInfo,
+        liveAutoPilotStatusData
+      ),
+    [liveAutoPilotStatusData, isLoading, agentsInfo]
  );

  const agentStatusSummary = useMemo(() => {
@ -47,8 +51,11 @@ function AgentsStatusWidget({
  return (
    <Collapse
      className="service-insights-collapse-widget agents-status-widget"
+      data-testid="agent-status-widget"
      expandIcon={() => (
-        <div className="expand-icon-container">
+        <div
+          className="expand-icon-container"
+          data-testid="agent-status-widget-expand-icon">
          {isLoading ? (
            <Skeleton.Input active size="small" />
          ) : (
@ -56,15 +63,20 @@ function AgentsStatusWidget({
              {Object.entries(agentStatusSummary).map(([key, value]) => (
                <div
                  className={classNames('agent-status-summary-item', key)}
+                  data-testid={`agent-status-summary-item-${key}`}
                  key={key}>
                  {getIconFromStatus(key)}
-                  <Typography.Text>{value}</Typography.Text>
+                  <Typography.Text data-testid="pipeline-count">
+                    {value}
+                  </Typography.Text>
                  <Typography.Text>{key}</Typography.Text>
                </div>
              ))}
            </div>
          )}
-          <Typography.Text className="text-primary">
+          <Typography.Text
+            className="text-primary"
+            data-testid="agent-status-widget-view-more">
            {t('label.view-more')}
          </Typography.Text>
          <ArrowSvg className="text-primary" height={14} width={14} />
@ -112,7 +124,8 @@ function AgentsStatusWidget({
                    className={classNames(
                      'agent-status-card',
                      agent.isCollateAgent ? 'collate-agent' : ''
-                    )}>
+                    )}
+                    data-testid={`agent-status-card-${agent.label}`}>
                    <Space align="center" size={8}>
                      {agent.agentIcon}
                      <Typography.Text>{agent.label}</Typography.Text>
--- a/openmetadata-ui/src/main/resources/ui/src/components/ServiceInsights/AgentsStatusWidget/agents-status-widget.less
+++ b/openmetadata-ui/src/main/resources/ui/src/components/ServiceInsights/AgentsStatusWidget/agents-status-widget.less
@ -85,6 +85,19 @@
    }
  }

+  .RUNNING {
+    color: @blue-31;
+  }
+
+  .FINISHED {
+    color: @green-10;
+  }
+
+  .EXCEPTION,
+  .FAILURE {
+    color: @red-10;
+  }
+
  .ant-collapse-arrow.expand-icon-container {
    display: flex;
    align-items: center;
--- a/openmetadata-ui/src/main/resources/ui/src/components/ServiceInsights/PlatformInsightsWidget/PlatformInsightsWidget.test.tsx
+++ b/openmetadata-ui/src/main/resources/ui/src/components/ServiceInsights/PlatformInsightsWidget/PlatformInsightsWidget.test.tsx
@ -0,0 +1,506 @@
+/*
+ *  Copyright 2025 Collate.
+ *  Licensed 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 { render, screen } from '@testing-library/react';
+import { MemoryRouter } from 'react-router-dom';
+import { SystemChartType } from '../../../enums/DataInsight.enum';
+import { ServiceCategory } from '../../../enums/service.enum';
+import { getTitleByChartType } from '../../../utils/ServiceInsightsTabUtils';
+import { getReadableCountString } from '../../../utils/ServiceUtils';
+import { useRequiredParams } from '../../../utils/useRequiredParams';
+import PlatformInsightsWidget from './PlatformInsightsWidget';
+import { PlatformInsightsWidgetProps } from './PlatformInsightsWidget.interface';
+
+// Mock dependencies
+jest.mock('../../../utils/useRequiredParams', () => ({
+  useRequiredParams: jest.fn(),
+}));
+
+jest.mock('../../../utils/ServiceInsightsTabUtils', () => ({
+  getTitleByChartType: jest.fn(),
+}));
+
+jest.mock('../../../utils/ServiceUtils', () => ({
+  getReadableCountString: jest.fn(),
+}));
+
+// Mock SVG components
+jest.mock('../../../assets/svg/ic-arrow-down.svg', () => {
+  return {
+    __esModule: true,
+    ReactComponent: () => <div data-testid="arrow-down-icon" />,
+  };
+});
+
+jest.mock('../../../assets/svg/ic-trend-up.svg', () => {
+  return {
+    __esModule: true,
+    ReactComponent: () => <div data-testid="arrow-up-icon" />,
+  };
+});
+
+const mockUseRequiredParams = useRequiredParams as jest.MockedFunction<any>;
+
+const mockGetTitleByChartType = getTitleByChartType as jest.MockedFunction<
+  typeof getTitleByChartType
+>;
+
+const mockGetReadableCountString =
+  getReadableCountString as jest.MockedFunction<typeof getReadableCountString>;
+
+describe('PlatformInsightsWidget', () => {
+  const mockServiceDetails = {
+    id: 'test-service-id',
+    name: 'test-service',
+    serviceType: 'Mysql' as any,
+    fullyQualifiedName: 'test-service-fqn',
+  } as any;
+
+  const mockChartsData = [
+    {
+      chartType: SystemChartType.DescriptionCoverage,
+      currentPercentage: 85,
+      percentageChange: 5,
+      isIncreased: true,
+      numberOfDays: 7,
+    },
+    {
+      chartType: SystemChartType.PIICoverage,
+      currentPercentage: 60,
+      percentageChange: -2,
+      isIncreased: false,
+      numberOfDays: 7,
+    },
+    {
+      chartType: SystemChartType.TierCoverage,
+      currentPercentage: 75,
+      percentageChange: 0,
+      isIncreased: true,
+      numberOfDays: 1,
+    },
+    {
+      chartType: SystemChartType.OwnersCoverage,
+      currentPercentage: 90,
+      percentageChange: 10,
+      isIncreased: true,
+      numberOfDays: 30,
+    },
+  ];
+
+  const defaultProps: PlatformInsightsWidgetProps = {
+    chartsData: mockChartsData,
+    isLoading: false,
+    serviceDetails: mockServiceDetails,
+  };
+
+  beforeEach(() => {
+    jest.clearAllMocks();
+    mockUseRequiredParams.mockReturnValue({
+      serviceCategory: ServiceCategory.DATABASE_SERVICES,
+    });
+    mockGetTitleByChartType.mockImplementation(
+      (chartType) => `Title for ${chartType}`
+    );
+    mockGetReadableCountString.mockImplementation((value) => value.toString());
+  });
+
+  const renderComponent = (props = {}) => {
+    return render(
+      <MemoryRouter>
+        <PlatformInsightsWidget {...defaultProps} {...props} />
+      </MemoryRouter>
+    );
+  };
+
+  describe('Component Rendering', () => {
+    it('should render the component with correct header', () => {
+      renderComponent();
+
+      expect(
+        screen.getByText('label.entity-insight-plural')
+      ).toBeInTheDocument();
+      expect(
+        screen.getByText('message.platform-insight-description')
+      ).toBeInTheDocument();
+    });
+
+    it('should render expand/collapse functionality', () => {
+      renderComponent();
+
+      const expandIcon = screen.getByText('label.view-more');
+
+      expect(expandIcon).toBeInTheDocument();
+    });
+
+    it('should render the export class for platform insights chart', () => {
+      renderComponent();
+
+      const exportContainer = document.querySelector(
+        '.export-platform-insights-chart'
+      );
+
+      expect(exportContainer).toBeInTheDocument();
+    });
+  });
+
+  describe('Loading State', () => {
+    it('should render skeleton loaders when isLoading is true', () => {
+      renderComponent({ isLoading: true });
+
+      // Should render 5 skeleton cards for database services (includes HealthyDataAssets)
+      const skeletonCards = document.querySelectorAll('.ant-skeleton');
+
+      expect(skeletonCards).toHaveLength(5);
+    });
+
+    it('should render skeleton loaders for non-database services', () => {
+      mockUseRequiredParams.mockReturnValue({
+        serviceCategory: ServiceCategory.MESSAGING_SERVICES,
+      });
+      renderComponent({ isLoading: true });
+
+      // Should render 4 skeleton cards (excludes HealthyDataAssets for non-database services)
+      const skeletonCards = document.querySelectorAll('.ant-skeleton');
+
+      expect(skeletonCards).toHaveLength(4);
+    });
+  });
+
+  describe('Chart Data Rendering', () => {
+    it('should render all chart cards when data is available', () => {
+      renderComponent();
+
+      // Should render 4 chart cards for database services
+      const chartCards = screen.getAllByText(/Title for/);
+
+      expect(chartCards).toHaveLength(4);
+    });
+
+    it('should render chart cards for non-database services (excluding HealthyDataAssets)', () => {
+      mockUseRequiredParams.mockReturnValue({
+        serviceCategory: ServiceCategory.MESSAGING_SERVICES,
+      });
+      renderComponent();
+
+      // Should render 4 chart cards (excludes HealthyDataAssets for non-database services)
+      const chartCards = screen.getAllByText(/Title for/);
+
+      expect(chartCards).toHaveLength(4);
+    });
+
+    it('should render chart titles correctly', () => {
+      renderComponent();
+
+      expect(mockGetTitleByChartType).toHaveBeenCalledWith(
+        SystemChartType.DescriptionCoverage
+      );
+      expect(mockGetTitleByChartType).toHaveBeenCalledWith(
+        SystemChartType.PIICoverage
+      );
+      expect(mockGetTitleByChartType).toHaveBeenCalledWith(
+        SystemChartType.TierCoverage
+      );
+      expect(mockGetTitleByChartType).toHaveBeenCalledWith(
+        SystemChartType.OwnersCoverage
+      );
+    });
+
+    it('should render current percentage values', () => {
+      renderComponent();
+
+      expect(screen.getByText('85%')).toBeInTheDocument();
+      expect(screen.getByText('60%')).toBeInTheDocument();
+      expect(screen.getByText('75%')).toBeInTheDocument();
+      expect(screen.getByText('90%')).toBeInTheDocument();
+    });
+  });
+
+  describe('Percentage Change Display', () => {
+    it('should render positive percentage change with green color and up arrow', () => {
+      renderComponent();
+
+      const positiveChange = screen.getByText('5%');
+
+      expect(positiveChange).toBeInTheDocument();
+      expect(positiveChange).toHaveStyle({ color: 'rgb(6, 118, 71)' }); // GREEN_1
+    });
+
+    it('should render negative percentage change with red color and down arrow', () => {
+      renderComponent();
+
+      const negativeChange = screen.getByText('-2%');
+
+      expect(negativeChange).toBeInTheDocument();
+      expect(negativeChange).toHaveStyle({ color: 'rgb(240, 68, 56)' }); // RED_1
+    });
+
+    it('should render percentage change text but not icon when value is 0', () => {
+      const chartsDataWithZeroChange = [
+        {
+          chartType: SystemChartType.DescriptionCoverage,
+          currentPercentage: 85,
+          percentageChange: 0,
+          isIncreased: true,
+          numberOfDays: 7,
+        },
+      ];
+
+      renderComponent({ chartsData: chartsDataWithZeroChange });
+
+      // Should show the percentage change text
+      expect(screen.getByText('0%')).toBeInTheDocument();
+
+      // Should not show the icon (since showIcon = false when percentageChange === 0)
+      const icons = screen.getAllByTestId('arrow-up-icon');
+
+      expect(icons).toHaveLength(1); // Only the one from the header
+    });
+
+    it('should render percentage change icon correctly', () => {
+      renderComponent();
+
+      const upIcons = screen.getAllByTestId('arrow-up-icon');
+
+      expect(upIcons.length).toBeGreaterThan(0);
+    });
+  });
+
+  describe('Time Period Display', () => {
+    it('should render "in the last day" for single day', () => {
+      renderComponent();
+
+      expect(screen.getByText('label.in-the-last-day')).toBeInTheDocument();
+    });
+
+    it('should render "in last X days" for multiple days', () => {
+      renderComponent();
+
+      expect(screen.getAllByText('label.in-last-number-of-days')).toHaveLength(
+        3
+      ); // Multiple charts have this text
+    });
+  });
+
+  describe('Service Category Filtering', () => {
+    it('should include HealthyDataAssets chart for database services', () => {
+      mockUseRequiredParams.mockReturnValue({
+        serviceCategory: ServiceCategory.DATABASE_SERVICES,
+      });
+
+      const chartsDataWithHealthyAssets = [
+        ...mockChartsData,
+        {
+          chartType: SystemChartType.HealthyDataAssets,
+          currentPercentage: 95,
+          percentageChange: 3,
+          isIncreased: true,
+          numberOfDays: 7,
+        },
+      ];
+
+      renderComponent({ chartsData: chartsDataWithHealthyAssets });
+
+      expect(mockGetTitleByChartType).toHaveBeenCalledWith(
+        SystemChartType.HealthyDataAssets
+      );
+    });
+
+    it('should exclude HealthyDataAssets chart for non-database services', () => {
+      mockUseRequiredParams.mockReturnValue({
+        serviceCategory: ServiceCategory.MESSAGING_SERVICES,
+      });
+
+      const chartsDataWithHealthyAssets = [
+        ...mockChartsData,
+        {
+          chartType: SystemChartType.HealthyDataAssets,
+          currentPercentage: 95,
+          percentageChange: 3,
+          isIncreased: true,
+          numberOfDays: 7,
+        },
+      ];
+
+      renderComponent({ chartsData: chartsDataWithHealthyAssets });
+
+      expect(mockGetTitleByChartType).not.toHaveBeenCalledWith(
+        SystemChartType.HealthyDataAssets
+      );
+    });
+  });
+
+  describe('Empty State', () => {
+    it('should handle empty charts data', () => {
+      renderComponent({ chartsData: [] });
+
+      const chartCards = screen.queryAllByText(/Title for/);
+
+      expect(chartCards).toHaveLength(0);
+    });
+
+    it('should handle undefined percentage change', () => {
+      const chartsDataWithUndefinedChange = [
+        {
+          chartType: SystemChartType.DescriptionCoverage,
+          currentPercentage: 85,
+          percentageChange: undefined,
+          isIncreased: true,
+          numberOfDays: 7,
+        },
+      ];
+
+      renderComponent({ chartsData: chartsDataWithUndefinedChange });
+
+      // Should not render percentage change section
+      expect(
+        screen.queryByText('label.in-last-number-of-days')
+      ).not.toBeInTheDocument();
+    });
+  });
+
+  describe('CSS Classes and Styling', () => {
+    it('should apply correct container class for 5 charts (database services)', () => {
+      renderComponent();
+
+      const colElement = document.querySelector('.other-charts-container');
+
+      expect(colElement).toBeInTheDocument();
+      expect(colElement).not.toHaveClass('four-chart-container');
+    });
+
+    it('should apply correct container class for non-4 charts', () => {
+      const limitedChartsData = mockChartsData.slice(0, 2);
+      renderComponent({ chartsData: limitedChartsData });
+
+      const colElement = document.querySelector('.four-chart-container');
+
+      expect(colElement).not.toBeInTheDocument();
+    });
+
+    it('should apply correct CSS classes to chart cards', () => {
+      renderComponent();
+
+      const chartCards = screen.getAllByText(/Title for/);
+      chartCards.forEach((card) => {
+        const cardElement = card.closest('.widget-info-card');
+
+        expect(cardElement).toHaveClass('other-charts-card');
+      });
+    });
+  });
+
+  describe('Utility Function Calls', () => {
+    it('should call getReadableCountString for current percentage', () => {
+      renderComponent();
+
+      expect(mockGetReadableCountString).toHaveBeenCalledWith(85);
+      expect(mockGetReadableCountString).toHaveBeenCalledWith(60);
+      expect(mockGetReadableCountString).toHaveBeenCalledWith(75);
+      expect(mockGetReadableCountString).toHaveBeenCalledWith(90);
+    });
+
+    it('should call getReadableCountString for percentage change', () => {
+      renderComponent();
+
+      expect(mockGetReadableCountString).toHaveBeenCalledWith(5);
+      expect(mockGetReadableCountString).toHaveBeenCalledWith(-2);
+      expect(mockGetReadableCountString).toHaveBeenCalledWith(0);
+      expect(mockGetReadableCountString).toHaveBeenCalledWith(10);
+    });
+
+    it('should call getTitleByChartType for each chart', () => {
+      renderComponent();
+
+      expect(mockGetTitleByChartType).toHaveBeenCalledTimes(4);
+    });
+  });
+
+  describe('Edge Cases', () => {
+    it('should handle charts data with noRecords flag', () => {
+      const chartsDataWithNoRecords = [
+        {
+          chartType: SystemChartType.DescriptionCoverage,
+          currentPercentage: 85,
+          percentageChange: 5,
+          isIncreased: true,
+          numberOfDays: 7,
+          noRecords: true,
+        },
+      ];
+
+      renderComponent({ chartsData: chartsDataWithNoRecords });
+
+      // Should still render the chart card
+      expect(
+        screen.getByText('Title for assets_with_description')
+      ).toBeInTheDocument();
+    });
+
+    it('should handle very large percentage values', () => {
+      const chartsDataWithLargeValues = [
+        {
+          chartType: SystemChartType.DescriptionCoverage,
+          currentPercentage: 999999,
+          percentageChange: 999999,
+          isIncreased: true,
+          numberOfDays: 7,
+        },
+      ];
+
+      renderComponent({ chartsData: chartsDataWithLargeValues });
+
+      expect(screen.getAllByText('999999%')).toHaveLength(2); // Both current and percentage change
+    });
+
+    it('should handle zero current percentage', () => {
+      const chartsDataWithZeroPercentage = [
+        {
+          chartType: SystemChartType.DescriptionCoverage,
+          currentPercentage: 0,
+          percentageChange: 5,
+          isIncreased: true,
+          numberOfDays: 7,
+        },
+      ];
+
+      renderComponent({ chartsData: chartsDataWithZeroPercentage });
+
+      expect(screen.getByText('0%')).toBeInTheDocument();
+    });
+  });
+
+  describe('Component Integration', () => {
+    it('should integrate with useRequiredParams hook', () => {
+      renderComponent();
+
+      expect(mockUseRequiredParams).toHaveBeenCalled();
+    });
+
+    it('should handle route parameter changes', () => {
+      const { rerender } = renderComponent();
+
+      // Change service category
+      mockUseRequiredParams.mockReturnValue({
+        serviceCategory: ServiceCategory.MESSAGING_SERVICES,
+      });
+
+      rerender(
+        <MemoryRouter>
+          <PlatformInsightsWidget {...defaultProps} />
+        </MemoryRouter>
+      );
+
+      expect(mockUseRequiredParams).toHaveBeenCalled();
+    });
+  });
+});
--- a/openmetadata-ui/src/main/resources/ui/src/components/ServiceInsights/PlatformInsightsWidget/PlatformInsightsWidget.tsx
+++ b/openmetadata-ui/src/main/resources/ui/src/components/ServiceInsights/PlatformInsightsWidget/PlatformInsightsWidget.tsx
@ -11,14 +11,20 @@
 *  limitations under the License.
 */
 import { Card, Col, Collapse, Row, Skeleton, Typography } from 'antd';
+import classNames from 'classnames';
 import { isUndefined } from 'lodash';
+import { ServiceTypes } from 'Models';
+import { useMemo } from 'react';
 import { useTranslation } from 'react-i18next';
 import { ReactComponent as ArrowSvg } from '../../../assets/svg/ic-arrow-down.svg';
 import { ReactComponent as ArrowUp } from '../../../assets/svg/ic-trend-up.svg';
 import { GREEN_1, RED_1 } from '../../../constants/Color.constants';
 import { PLATFORM_INSIGHTS_CHARTS } from '../../../constants/ServiceInsightsTab.constants';
+import { SystemChartType } from '../../../enums/DataInsight.enum';
+import { ServiceCategory } from '../../../enums/service.enum';
 import { getTitleByChartType } from '../../../utils/ServiceInsightsTabUtils';
 import { getReadableCountString } from '../../../utils/ServiceUtils';
+import { useRequiredParams } from '../../../utils/useRequiredParams';
 import './platform-insights-widget.less';
 import { PlatformInsightsWidgetProps } from './PlatformInsightsWidget.interface';

@ -26,8 +32,28 @@ function PlatformInsightsWidget({
  chartsData,
  isLoading,
 }: Readonly<PlatformInsightsWidgetProps>) {
+  const { serviceCategory } =
+    useRequiredParams<{ serviceCategory: ServiceTypes }>();
  const { t } = useTranslation();

+  const { filteredCharts, filteredChartsData, containerClassName } =
+    useMemo(() => {
+      const filteredCharts = PLATFORM_INSIGHTS_CHARTS.filter((chart) =>
+        chart === SystemChartType.HealthyDataAssets
+          ? serviceCategory === ServiceCategory.DATABASE_SERVICES
+          : true
+      );
+
+      return {
+        filteredCharts,
+        filteredChartsData: chartsData.filter((chart) =>
+          filteredCharts.includes(chart.chartType)
+        ),
+        containerClassName:
+          filteredCharts.length === 4 ? 'four-chart-container' : '',
+      };
+    }, [serviceCategory, chartsData]);
+
  return (
    <Collapse
      className="service-insights-collapse-widget platform-insights-card"
@ -57,9 +83,11 @@ function PlatformInsightsWidget({
        key="1">
        {/* Don't remove this class name, it is used for exporting the platform insights chart */}
        <Row className="export-platform-insights-chart" gutter={16}>
-          <Col className="other-charts-container" span={24}>
+          <Col
+            className={classNames('other-charts-container', containerClassName)}
+            span={24}>
            {isLoading
-              ? PLATFORM_INSIGHTS_CHARTS.map((chartType) => (
+              ? filteredCharts.map((chartType) => (
                  <Card
                    className="widget-info-card other-charts-card"
                    key={chartType}>
@ -70,7 +98,7 @@ function PlatformInsightsWidget({
                    />
                  </Card>
                ))
-              : chartsData.map((chart) => {
+              : filteredChartsData.map((chart) => {
                  const icon = chart.isIncreased ? (
                    <ArrowUp color={GREEN_1} height={11} width={11} />
                  ) : (
@ -91,13 +119,13 @@ function PlatformInsightsWidget({
                      <Typography.Text className="font-semibold text-sm">
                        {getTitleByChartType(chart.chartType)}
                      </Typography.Text>
-                      <Row align="bottom" className="m-t-xs" gutter={8}>
+                      <Row align="top" className="m-t-xs" gutter={8}>
                        <Col span={12}>
-                          <Typography.Title level={3}>
+                          <Typography.Text className="current-percentage">
                            {`${getReadableCountString(
                              chart.currentPercentage
                            )}%`}
-                          </Typography.Title>
+                          </Typography.Text>
                        </Col>
                        {!isUndefined(chart.percentageChange) && (
                          <Col
@ -106,7 +134,7 @@ function PlatformInsightsWidget({
                            <div className="percent-change-tag">
                              {showIcon && icon}
                              <Typography.Text
-                                className="font-medium text-sm"
+                                className="font-medium text-xs"
                                style={{
                                  color: chart.isIncreased ? GREEN_1 : RED_1,
                                }}>
--- a/openmetadata-ui/src/main/resources/ui/src/components/ServiceInsights/PlatformInsightsWidget/platform-insights-widget.less
+++ b/openmetadata-ui/src/main/resources/ui/src/components/ServiceInsights/PlatformInsightsWidget/platform-insights-widget.less
@ -19,6 +19,10 @@
    gap: 16px;
  }

+  .other-charts-container.four-chart-container {
+    grid-template-columns: repeat(4, 1fr);
+  }
+
  .percent-change-tag {
    display: flex;
    align-items: center;
@ -43,6 +47,12 @@
      &::after {
        display: none;
      }
+
+      .current-percentage {
+        line-height: 20px;
+        font-size: 20px;
+        font-weight: 600;
+      }
    }
  }
 }
--- a/openmetadata-ui/src/main/resources/ui/src/components/ServiceInsights/ServiceInsightsTab.interface.ts
+++ b/openmetadata-ui/src/main/resources/ui/src/components/ServiceInsights/ServiceInsightsTab.interface.ts
@ -73,6 +73,5 @@ export interface CollateAgentLiveInfo
 export interface TotalAssetsCount {
  name: string;
  value: number;
-  fill: string;
  icon: JSX.Element;
 }
--- a/openmetadata-ui/src/main/resources/ui/src/components/ServiceInsights/ServiceInsightsTab.tsx
+++ b/openmetadata-ui/src/main/resources/ui/src/components/ServiceInsights/ServiceInsightsTab.tsx
@ -14,7 +14,7 @@
 import { Col, Row } from 'antd';
 import { AxiosError } from 'axios';
 import { isEmpty, isUndefined } from 'lodash';
-import { ServiceTypes } from 'Models';
+import { Bucket, ServiceTypes } from 'Models';
 import { useCallback, useEffect, useRef, useState } from 'react';
 import { SOCKET_EVENTS } from '../../constants/constants';
 import {
@ -22,12 +22,14 @@ import {
  PLATFORM_INSIGHTS_CHARTS,
  PLATFORM_INSIGHTS_LIVE_CHARTS,
 } from '../../constants/ServiceInsightsTab.constants';
-import { totalDataAssetsWidgetColors } from '../../constants/TotalDataAssetsWidget.constants';
 import { useWebSocketConnector } from '../../context/WebSocketProvider/WebSocketProvider';
 import { SystemChartType } from '../../enums/DataInsight.enum';
 import { SearchIndex } from '../../enums/search.enum';
 import { AppRunRecord } from '../../generated/entity/applications/appRunRecord';
-import { WorkflowStatus } from '../../generated/governance/workflows/workflowInstance';
+import {
+  WorkflowInstance,
+  WorkflowStatus,
+} from '../../generated/governance/workflows/workflowInstance';
 import { getAgentRuns } from '../../rest/applicationAPI';
 import {
  getMultiChartsPreviewByName,
@ -44,7 +46,7 @@ import {
  getCurrentMillis,
  getDayAgoStartGMTinMillis,
 } from '../../utils/date-time/DateTimeUtils';
-import { getEntityNameLabel } from '../../utils/EntityUtils';
+import { getEntityFeedLink, getEntityNameLabel } from '../../utils/EntityUtils';
 import {
  filterDistributionChartItem,
  getAssetsByServiceType,
@ -53,7 +55,10 @@ import {
  getPlatformInsightsChartDataFormattingMethod,
 } from '../../utils/ServiceInsightsTabUtils';
 import serviceUtilClassBase from '../../utils/ServiceUtilClassBase';
-import { getServiceNameQueryFilter } from '../../utils/ServiceUtils';
+import {
+  getEntityTypeFromServiceCategory,
+  getServiceNameQueryFilter,
+} from '../../utils/ServiceUtils';
 import { getEntityIcon } from '../../utils/TableUtils';
 import { showErrorToast } from '../../utils/ToastUtils';
 import { useRequiredParams } from '../../utils/useRequiredParams';
@ -82,6 +87,9 @@ const ServiceInsightsTab = ({
  const [isLoading, setIsLoading] = useState(false);
  const [totalAssetsCount, setTotalAssetsCount] =
    useState<Array<TotalAssetsCount>>();
+  const [liveAutoPilotStatusData, setLiveAutoPilotStatusData] = useState<
+    WorkflowInstance | undefined
+  >(workflowStatesData?.mainInstanceState);
  const sessionIdRef = useRef<string>();

  const serviceName = serviceDetails.name;
@ -97,14 +105,22 @@ const ServiceInsightsTab = ({

      const assets = getAssetsByServiceType(serviceCategory);

-      const buckets = response.aggregations['entityType'].buckets.filter(
-        (bucket) => assets.includes(bucket.key)
-      );
+      // Arrange the buckets in the order of the assets
+      const buckets = assets.reduce((acc, curr) => {
+        const bucket = response.aggregations['entityType'].buckets.find(
+          (bucket) => bucket.key === curr
+        );

-      const entityCountsArray = buckets.map((bucket, index) => ({
+        if (!isUndefined(bucket)) {
+          return [...acc, bucket];
+        }
+
+        return acc;
+      }, [] as Bucket[]);
+
+      const entityCountsArray = buckets.map((bucket) => ({
        name: getEntityNameLabel(bucket.key),
        value: bucket.doc_count ?? 0,
-        fill: totalDataAssetsWidgetColors[index],
        icon: getEntityIcon(bucket.key, '', { height: 16, width: 16 }) ?? <></>,
      }));

@ -178,11 +194,16 @@ const ServiceInsightsTab = ({

  const triggerSocketConnection = useCallback(async () => {
    if (isUndefined(sessionIdRef.current)) {
+      const entityType = getEntityTypeFromServiceCategory(serviceCategory);
      const { sessionId } = await setChartDataStreamConnection({
        chartNames: LIVE_CHARTS_LIST,
        serviceName,
        startTime: getCurrentDayStartGMTinMillis(),
        endTime: getCurrentDayStartGMTinMillis() + 360000000,
+        entityLink: getEntityFeedLink(
+          entityType,
+          serviceDetails.fullyQualifiedName
+        ),
      });

      sessionIdRef.current = sessionId;
@ -251,10 +272,13 @@ const ServiceInsightsTab = ({

          setTotalAssetsCount(
            getFormattedTotalAssetsDataFromSocketData(
-              data?.data?.total_data_assets_live
+              data?.data?.total_data_assets_live,
+              serviceCategory
            )
          );

+          setLiveAutoPilotStatusData(data.workflowInstances?.[0]);
+
          setChartsResults((prev) => ({
            platformInsightsChart,
            piiDistributionChart: prev?.piiDistributionChart ?? [],
@ -263,7 +287,7 @@ const ServiceInsightsTab = ({
        }
      }
    },
-    [serviceDetails]
+    [serviceDetails, serviceCategory]
  );

  useEffect(() => {
@ -317,7 +341,7 @@ const ServiceInsightsTab = ({
                isCollateAIagentsLoading ||
                isIngestionPipelineLoading
              }
-              workflowStatesData={workflowStatesData}
+              liveAutoPilotStatusData={liveAutoPilotStatusData}
            />
          </Col>
          <Col span={24}>
--- a/openmetadata-ui/src/main/resources/ui/src/components/ServiceInsights/TotalDataAssetsWidget/TotalDataAssetsWidget.tsx
+++ b/openmetadata-ui/src/main/resources/ui/src/components/ServiceInsights/TotalDataAssetsWidget/TotalDataAssetsWidget.tsx
@ -49,7 +49,9 @@ function TotalDataAssetsWidget({
  );

  return (
-    <Card className="widget-info-card total-data-assets-widget">
+    <Card
+      className="widget-info-card total-data-assets-widget"
+      data-testid="total-data-assets-widget">
      <div className="flex items-center gap-2">
        <div className="p-0 icon-container">
          <PieChartIcon height={16} width={16} />
@ -69,18 +71,14 @@ function TotalDataAssetsWidget({
                className="flex items-center justify-between"
                key={entity.name}>
                <div className="flex items-center gap-3">
-                  <div
-                    className="bullet"
-                    style={{
-                      backgroundColor: entity.fill,
-                    }}
-                  />
                  <div className="p-0 icon-container">{entity.icon}</div>

                  <Typography.Text>{entity.name}</Typography.Text>
                </div>

-                <Typography.Text className="font-bold">
+                <Typography.Text
+                  className="font-bold"
+                  data-testid={`${entity.name}-count`}>
                  {getReadableCountString(entity.value)}
                </Typography.Text>
              </div>
--- a/openmetadata-ui/src/main/resources/ui/src/components/ServiceInsights/TotalDataAssetsWidget/total-data-assets-widget.less
+++ b/openmetadata-ui/src/main/resources/ui/src/components/ServiceInsights/TotalDataAssetsWidget/total-data-assets-widget.less
@ -15,11 +15,6 @@
 .total-data-assets-widget.ant-card {
  height: 100%;

-  .bullet {
-    width: 8px;
-    height: 8px;
-    border-radius: 4px;
-  }
  .icon-container {
    display: flex;
    align-items: center;
@ -35,7 +30,7 @@
    display: flex;
    flex-direction: column;
    gap: 8px;
-    background-color: @grey-1;
+    background-color: @grey-50;
    border-radius: @border-rad-xs;
    height: 100%;
  }
--- a/openmetadata-ui/src/main/resources/ui/src/constants/TotalDataAssetsWidget.constants.ts
+++ b/openmetadata-ui/src/main/resources/ui/src/constants/TotalDataAssetsWidget.constants.ts
@ -1,35 +0,0 @@
-/*
- *  Copyright 2025 Collate.
- *  Licensed 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 {
-  BLUE_2,
-  DESERT,
-  ELECTRIC_VIOLET,
-  LEMON_ZEST,
-  MY_SIN,
-  PINK_SALMON,
-  RIPTIDE,
-  SAN_MARINO,
-  SILVER_TREE,
-} from './Color.constants';
-
-export const totalDataAssetsWidgetColors = [
-  BLUE_2,
-  PINK_SALMON,
-  SILVER_TREE,
-  SAN_MARINO,
-  RIPTIDE,
-  MY_SIN,
-  DESERT,
-  ELECTRIC_VIOLET,
-  LEMON_ZEST,
-];
--- a/openmetadata-ui/src/main/resources/ui/src/locale/languages/de-de.json
+++ b/openmetadata-ui/src/main/resources/ui/src/locale/languages/de-de.json
@ -1610,6 +1610,7 @@
    "three-dots-symbol": "•••",
    "thursday": "Donnerstag",
    "tier": "Stufe",
+    "tier-label-type": "Stufen-Etikettentyp",
    "tier-number": "Stufe {{tier}}",
    "tier-plural": "Tiers",
    "tier-plural-lowercase": "stufen",
--- a/openmetadata-ui/src/main/resources/ui/src/locale/languages/en-us.json
+++ b/openmetadata-ui/src/main/resources/ui/src/locale/languages/en-us.json
@ -1610,6 +1610,7 @@
    "three-dots-symbol": "•••",
    "thursday": "Thursday",
    "tier": "Tier",
+    "tier-label-type": "Tier Label Type",
    "tier-number": "Tier{{tier}}",
    "tier-plural": "Tiers",
    "tier-plural-lowercase": "tiers",
--- a/openmetadata-ui/src/main/resources/ui/src/locale/languages/es-es.json
+++ b/openmetadata-ui/src/main/resources/ui/src/locale/languages/es-es.json
@ -1610,6 +1610,7 @@
    "three-dots-symbol": "•••",
    "thursday": "Jueves",
    "tier": "Nivel",
+    "tier-label-type": "Tipo de Etiqueta de Nivel",
    "tier-number": "Nivel {{tier}}",
    "tier-plural": "Niveles",
    "tier-plural-lowercase": "niveles",
--- a/openmetadata-ui/src/main/resources/ui/src/locale/languages/fr-fr.json
+++ b/openmetadata-ui/src/main/resources/ui/src/locale/languages/fr-fr.json
@ -1610,6 +1610,7 @@
    "three-dots-symbol": "•••",
    "thursday": "Jeudi",
    "tier": "Niveau",
+    "tier-label-type": "Type d'Étiquette de Niveau",
    "tier-number": "Niveau {{tier}}",
    "tier-plural": "Niveaux",
    "tier-plural-lowercase": "niveaux",
--- a/openmetadata-ui/src/main/resources/ui/src/locale/languages/gl-es.json
+++ b/openmetadata-ui/src/main/resources/ui/src/locale/languages/gl-es.json
@ -1610,6 +1610,7 @@
    "three-dots-symbol": "•••",
    "thursday": "Xoves",
    "tier": "Nivel",
+    "tier-label-type": "Tipo de Etiqueta de Nivel",
    "tier-number": "Nivel{{tier}}",
    "tier-plural": "Niveles",
    "tier-plural-lowercase": "niveis",
--- a/openmetadata-ui/src/main/resources/ui/src/locale/languages/he-he.json
+++ b/openmetadata-ui/src/main/resources/ui/src/locale/languages/he-he.json
@ -1610,6 +1610,7 @@
    "three-dots-symbol": "•••",
    "thursday": "יום חמישי",
    "tier": "שכבת מידע",
+    "tier-label-type": "סוג תווית שכבת מידע",
    "tier-number": "שכבת מידע {{tier}}",
    "tier-plural": "רמות",
    "tier-plural-lowercase": "שכבות",
--- a/openmetadata-ui/src/main/resources/ui/src/locale/languages/ja-jp.json
+++ b/openmetadata-ui/src/main/resources/ui/src/locale/languages/ja-jp.json
@ -1610,6 +1610,7 @@
    "three-dots-symbol": "•••",
    "thursday": "木曜日",
    "tier": "ティア",
+    "tier-label-type": "ティアラベルタイプ",
    "tier-number": "ティア{{tier}}",
    "tier-plural": "ティア",
    "tier-plural-lowercase": "ティア",
--- a/openmetadata-ui/src/main/resources/ui/src/locale/languages/ko-kr.json
+++ b/openmetadata-ui/src/main/resources/ui/src/locale/languages/ko-kr.json
@ -1610,6 +1610,7 @@
    "three-dots-symbol": "•••",
    "thursday": "목요일",
    "tier": "계층",
+    "tier-label-type": "계층 라벨 유형",
    "tier-number": "계층{{tier}}",
    "tier-plural": "계층들",
    "tier-plural-lowercase": "계층들",
--- a/openmetadata-ui/src/main/resources/ui/src/locale/languages/mr-in.json
+++ b/openmetadata-ui/src/main/resources/ui/src/locale/languages/mr-in.json
@ -1610,6 +1610,7 @@
    "three-dots-symbol": "•••",
    "thursday": "गुरुवार",
    "tier": "स्तर",
+    "tier-label-type": "स्तर लेबल प्रकार",
    "tier-number": "स्तर{{tier}}",
    "tier-plural": "टियर्स",
    "tier-plural-lowercase": "स्तर",
--- a/openmetadata-ui/src/main/resources/ui/src/locale/languages/nl-nl.json
+++ b/openmetadata-ui/src/main/resources/ui/src/locale/languages/nl-nl.json
@ -1610,6 +1610,7 @@
    "three-dots-symbol": "•••",
    "thursday": "donderdag",
    "tier": "Niveau",
+    "tier-label-type": "Niveau-etikettype",
    "tier-number": "Niveau {{tier}}",
    "tier-plural": "Lagen",
    "tier-plural-lowercase": "niveaus",
--- a/openmetadata-ui/src/main/resources/ui/src/locale/languages/pr-pr.json
+++ b/openmetadata-ui/src/main/resources/ui/src/locale/languages/pr-pr.json
@ -1610,6 +1610,7 @@
    "three-dots-symbol": "•••",
    "thursday": "پنج‌شنبه",
    "tier": "سطح",
+    "tier-label-type": "نوع برچسب سطح",
    "tier-number": "سطح {{tier}}",
    "tier-plural": "طبقات",
    "tier-plural-lowercase": "سطوح",
--- a/openmetadata-ui/src/main/resources/ui/src/locale/languages/pt-br.json
+++ b/openmetadata-ui/src/main/resources/ui/src/locale/languages/pt-br.json
@ -1610,6 +1610,7 @@
    "three-dots-symbol": "•••",
    "thursday": "Quinta-feira",
    "tier": "Camada",
+    "tier-label-type": "Tipo de Rótulo de Camada",
    "tier-number": "Camada{{tier}}",
    "tier-plural": "Camadas",
    "tier-plural-lowercase": "camadas",
--- a/openmetadata-ui/src/main/resources/ui/src/locale/languages/pt-pt.json
+++ b/openmetadata-ui/src/main/resources/ui/src/locale/languages/pt-pt.json
@ -1610,6 +1610,7 @@
    "three-dots-symbol": "•••",
    "thursday": "Quinta-feira",
    "tier": "Camada",
+    "tier-label-type": "Tipo de Etiqueta de Camada",
    "tier-number": "Camada{{tier}}",
    "tier-plural": "Camadas",
    "tier-plural-lowercase": "camadas",
--- a/openmetadata-ui/src/main/resources/ui/src/locale/languages/ru-ru.json
+++ b/openmetadata-ui/src/main/resources/ui/src/locale/languages/ru-ru.json
@ -1610,6 +1610,7 @@
    "three-dots-symbol": "•••",
    "thursday": "Четверг",
    "tier": "Уровень",
+    "tier-label-type": "Тип метки уровня",
    "tier-number": "Уровень{{tier}}",
    "tier-plural": "Уровни",
    "tier-plural-lowercase": "уровни",
--- a/openmetadata-ui/src/main/resources/ui/src/locale/languages/th-th.json
+++ b/openmetadata-ui/src/main/resources/ui/src/locale/languages/th-th.json
@ -1610,6 +1610,7 @@
    "three-dots-symbol": "•••",
    "thursday": "วันพฤหัสบดี",
    "tier": "ระดับ",
+    "tier-label-type": "ประเภทป้ายระดับ",
    "tier-number": "ระดับ {{tier}}",
    "tier-plural": "ชั้น",
    "tier-plural-lowercase": "ระดับหลายรายการ",
--- a/openmetadata-ui/src/main/resources/ui/src/locale/languages/tr-tr.json
+++ b/openmetadata-ui/src/main/resources/ui/src/locale/languages/tr-tr.json
@ -1610,6 +1610,7 @@
    "three-dots-symbol": "•••",
    "thursday": "Perşembe",
    "tier": "Katman",
+    "tier-label-type": "Katman Etiketi Türü",
    "tier-number": "Katman{{tier}}",
    "tier-plural": "Katmanlar",
    "tier-plural-lowercase": "katmanlar",
--- a/openmetadata-ui/src/main/resources/ui/src/locale/languages/zh-cn.json
+++ b/openmetadata-ui/src/main/resources/ui/src/locale/languages/zh-cn.json
@ -1610,6 +1610,7 @@
    "three-dots-symbol": "•••",
    "thursday": "星期四",
    "tier": "分级",
+    "tier-label-type": "分级标签类型",
    "tier-number": "{{tier}}级",
    "tier-plural": "层",
    "tier-plural-lowercase": "分级",
--- a/openmetadata-ui/src/main/resources/ui/src/rest/DataInsightAPI.ts
+++ b/openmetadata-ui/src/main/resources/ui/src/rest/DataInsightAPI.ts
@ -76,6 +76,7 @@ export const setChartDataStreamConnection = async (params: {
  serviceName: string;
  startTime: number;
  endTime: number;
+  entityLink: string;
 }) => {
  const response = await APIClient.post<StartChartDataStreamConnectionResponse>(
    `/analytics/dataInsights/system/charts/stream`,
--- a/openmetadata-ui/src/main/resources/ui/src/utils/AgentsStatusWidgetUtils.tsx
+++ b/openmetadata-ui/src/main/resources/ui/src/utils/AgentsStatusWidgetUtils.tsx
@ -29,7 +29,6 @@ import { AgentsInfo } from '../components/ServiceInsights/AgentsStatusWidget/Age
 import {
  AgentsLiveInfo,
  CollateAgentLiveInfo,
-  WorkflowStatesData,
 } from '../components/ServiceInsights/ServiceInsightsTab.interface';
 import {
  AUTOPILOT_AGENTS_ORDERED_LIST,
@ -53,7 +52,10 @@ import {
  PipelineType,
  ProviderType,
 } from '../generated/entity/services/ingestionPipelines/ingestionPipeline';
-import { WorkflowStatus } from '../generated/governance/workflows/workflowInstance';
+import {
+  WorkflowInstance,
+  WorkflowStatus,
+} from '../generated/governance/workflows/workflowInstance';
 import { t } from './i18next/LocalUtil';

 export const getAgentLabelFromType = (agentType: string) => {
@ -300,7 +302,7 @@ export const getIconFromStatus = (status?: string) => {
 export const getAgentRunningStatusMessage = (
  isLoading: boolean,
  agentsInfo: AgentsInfo[],
-  workflowStatesData?: WorkflowStatesData
+  liveAutoPilotStatusData?: WorkflowInstance
 ) => {
  if (isLoading) {
    return (
@ -309,22 +311,28 @@ export const getAgentRunningStatusMessage = (
  }

  let message = '';
+  let Icon: SvgComponent = () => null;
+  const status = liveAutoPilotStatusData?.status ?? '';

-  switch (workflowStatesData?.mainInstanceState?.status) {
+  switch (status) {
    case WorkflowStatus.Running:
      message = t('message.auto-pilot-agents-running-message');
+      Icon = RunningIcon;

      break;
    case WorkflowStatus.Failure:
      message = t('message.auto-pilot-agents-failed-message');
+      Icon = ErrorIcon;

      break;
    case WorkflowStatus.Finished:
      message = t('message.auto-pilot-agents-finished-message');
+      Icon = CheckIcon;

      break;
    case WorkflowStatus.Exception:
      message = t('message.auto-pilot-agents-exception-message');
+      Icon = ErrorIcon;

      break;
  }
@ -334,10 +342,14 @@ export const getAgentRunningStatusMessage = (
  }

  return (
-    <Typography.Text
-      className="text-grey-muted text-sm"
-      data-testid="agents-status-message">
-      {message}
-    </Typography.Text>
+    <div className="flex items-center gap-1">
+      <Icon className={status} height={14} width={14} />
+
+      <Typography.Text
+        className="text-grey-muted text-sm"
+        data-testid="agents-status-message">
+        {message}
+      </Typography.Text>
+    </div>
  );
 };
--- a/openmetadata-ui/src/main/resources/ui/src/utils/ServiceInsightsTabUtils.tsx
+++ b/openmetadata-ui/src/main/resources/ui/src/utils/ServiceInsightsTabUtils.tsx
@ -11,7 +11,15 @@
 *  limitations under the License.
 */
 import { Typography } from 'antd';
-import { first, isEmpty, last, round, sortBy, toLower } from 'lodash';
+import {
+  first,
+  isEmpty,
+  isUndefined,
+  last,
+  round,
+  sortBy,
+  toLower,
+} from 'lodash';
 import { ServiceTypes } from 'Models';
 import { ReactComponent as DescriptionPlaceholderIcon } from '../assets/svg/ic-flat-doc.svg';
 import { ReactComponent as TablePlaceholderIcon } from '../assets/svg/ic-large-table.svg';
@ -22,7 +30,6 @@ import { ReactComponent as PiiPlaceholderIcon } from '../assets/svg/security-saf
 import ErrorPlaceHolder from '../components/common/ErrorWithPlaceholder/ErrorPlaceHolder';
 import { ChartsResults } from '../components/ServiceInsights/ServiceInsightsTab.interface';
 import { SERVICE_AUTOPILOT_AGENT_TYPES } from '../constants/Services.constant';
-import { totalDataAssetsWidgetColors } from '../constants/TotalDataAssetsWidget.constants';
 import { ERROR_PLACEHOLDER_TYPE, SIZE } from '../enums/common.enum';
 import { SystemChartType } from '../enums/DataInsight.enum';
 import { EntityType } from '../enums/entity.enum';
@ -173,12 +180,24 @@ export const getPlatformInsightsChartDataFormattingMethod =
  };

 export const getFormattedTotalAssetsDataFromSocketData = (
-  socketData: DataInsightCustomChartResult
+  socketData: DataInsightCustomChartResult,
+  serviceCategory: ServiceTypes
 ) => {
-  const entityCountsArray = socketData.results.map((result, index) => ({
+  const assets = getAssetsByServiceType(serviceCategory);
+
+  const buckets = assets.reduce((acc, curr) => {
+    const bucket = socketData.results.find((bucket) => bucket.group === curr);
+
+    if (!isUndefined(bucket)) {
+      return [...acc, bucket];
+    }
+
+    return acc;
+  }, [] as DataInsightCustomChartResult['results']);
+
+  const entityCountsArray = buckets.map((result) => ({
    name: getEntityNameLabel(result.group),
    value: result.count ?? 0,
-    fill: totalDataAssetsWidgetColors[index],
    icon: getEntityIcon(result.group, '', { height: 16, width: 16 }) ?? <></>,
  }));