From 55b26ed59067d1f1b959e47387fb493703032e94 Mon Sep 17 00:00:00 2001
From: dotta <dotta@example.com>
Date: Mon, 23 Mar 2026 17:18:17 -0500
Subject: [PATCH] Address Greptile review on agent runtime PR

Co-Authored-By: Paperclip <noreply@paperclip.ing>
---
 ui/src/pages/InstanceSettings.tsx | 64 +++++++++++++++++++++++--------
 1 file changed, 48 insertions(+), 16 deletions(-)

diff --git a/ui/src/pages/InstanceSettings.tsx b/ui/src/pages/InstanceSettings.tsx
index 1a48bf0a..18faff5e 100644
--- a/ui/src/pages/InstanceSettings.tsx
+++ b/ui/src/pages/InstanceSettings.tsx
@@ -80,25 +80,50 @@ export function InstanceSettings() {
   const disableAllMutation = useMutation({
     mutationFn: async (agentRows: InstanceSchedulerHeartbeatAgent[]) => {
       const enabled = agentRows.filter((a) => a.heartbeatEnabled);
-      for (const agentRow of enabled) {
-        const agent = await agentsApi.get(agentRow.id, agentRow.companyId);
-        const runtimeConfig = asRecord(agent.runtimeConfig) ?? {};
-        const heartbeat = asRecord(runtimeConfig.heartbeat) ?? {};
-        await agentsApi.update(
-          agentRow.id,
-          {
-            runtimeConfig: {
-              ...runtimeConfig,
-              heartbeat: { ...heartbeat, enabled: false },
+      if (enabled.length === 0) return enabled;
+
+      const results = await Promise.allSettled(
+        enabled.map(async (agentRow) => {
+          const agent = await agentsApi.get(agentRow.id, agentRow.companyId);
+          const runtimeConfig = asRecord(agent.runtimeConfig) ?? {};
+          const heartbeat = asRecord(runtimeConfig.heartbeat) ?? {};
+          await agentsApi.update(
+            agentRow.id,
+            {
+              runtimeConfig: {
+                ...runtimeConfig,
+                heartbeat: { ...heartbeat, enabled: false },
+              },
             },
-          },
-          agentRow.companyId,
+            agentRow.companyId,
+          );
+        }),
+      );
+
+      const failures = results.filter((result): result is PromiseRejectedResult => result.status === "rejected");
+      if (failures.length > 0) {
+        const firstError = failures[0]?.reason;
+        const detail = firstError instanceof Error ? firstError.message : "Unknown error";
+        throw new Error(
+          failures.length === 1
+            ? `Failed to disable 1 timer heartbeat: ${detail}`
+            : `Failed to disable ${failures.length} of ${enabled.length} timer heartbeats. First error: ${detail}`,
         );
       }
+      return enabled;
     },
-    onSuccess: async () => {
+    onSuccess: async (updatedRows) => {
       setActionError(null);
-      await queryClient.invalidateQueries({ queryKey: queryKeys.instance.schedulerHeartbeats });
+      const companies = new Set(updatedRows.map((row) => row.companyId));
+      await Promise.all([
+        queryClient.invalidateQueries({ queryKey: queryKeys.instance.schedulerHeartbeats }),
+        ...Array.from(companies, (companyId) =>
+          queryClient.invalidateQueries({ queryKey: queryKeys.agents.list(companyId) }),
+        ),
+        ...updatedRows.map((row) =>
+          queryClient.invalidateQueries({ queryKey: queryKeys.agents.detail(row.id) }),
+        ),
+      ]);
     },
     onError: (error) => {
       setActionError(error instanceof Error ? error.message : "Failed to disable all heartbeats.");
@@ -108,7 +133,8 @@ export function InstanceSettings() {
   const agents = heartbeatsQuery.data ?? [];
   const activeCount = agents.filter((agent) => agent.schedulerActive).length;
   const disabledCount = agents.length - activeCount;
-  const anyEnabled = agents.some((a) => a.heartbeatEnabled);
+  const enabledCount = agents.filter((agent) => agent.heartbeatEnabled).length;
+  const anyEnabled = enabledCount > 0;
 
   const grouped = useMemo(() => {
     const map = new Map<string, { companyName: string; agents: InstanceSchedulerHeartbeatAgent[] }>();
@@ -159,7 +185,13 @@ export function InstanceSettings() {
             size="sm"
             className="ml-auto h-7 text-xs"
             disabled={disableAllMutation.isPending}
-            onClick={() => disableAllMutation.mutate(agents)}
+            onClick={() => {
+              const noun = enabledCount === 1 ? "agent" : "agents";
+              if (!window.confirm(`Disable timer heartbeats for all ${enabledCount} enabled ${noun}?`)) {
+                return;
+              }
+              disableAllMutation.mutate(agents);
+            }}
           >
             {disableAllMutation.isPending ? "Disabling..." : "Disable All"}
           </Button>