Merge pull request #183 from CaitlynByrne/feat/pause-drain

feat: add graceful pause (drain mode) for running agents
2026-03-17 02:43:09 +00:00 · 2026-02-12 07:22:01 +02:00
parent 76dd4b8d80 9721368188
commit c114248b09
12 changed files with 311 additions and 24 deletions
--- a/ui/package-lock.json
+++ b/ui/package-lock.json
@@ -96,6 +96,7 @@
      "integrity": "sha512-e7jT4DxYvIDLk1ZHmU/m/mB19rex9sv0c2ftBtjSBv+kVM/902eh0fINUzD7UwLLNR+jU585GxUJ8/EBfAM5fw==",
      "dev": true,
      "license": "MIT",
+      "peer": true,
      "dependencies": {
        "@babel/code-frame": "^7.27.1",
        "@babel/generator": "^7.28.5",
@@ -2825,6 +2826,7 @@
      "integrity": "sha512-MciR4AKGHWl7xwxkBa6xUGxQJ4VBOmPTF7sL+iGzuahOFaO0jHCsuEfS80pan1ef4gWId1oWOweIhrDEYLuaOw==",
      "dev": true,
      "license": "MIT",
+      "peer": true,
      "dependencies": {
        "undici-types": "~6.21.0"
      }
@@ -2834,6 +2836,7 @@
      "resolved": "https://registry.npmjs.org/@types/react/-/react-19.2.9.tgz",
      "integrity": "sha512-Lpo8kgb/igvMIPeNV2rsYKTgaORYdO1XGVZ4Qz3akwOj0ySGYMPlQWa8BaLn0G63D1aSaAQ5ldR06wCpChQCjA==",
      "license": "MIT",
+      "peer": true,
      "dependencies": {
        "csstype": "^3.2.2"
      }
@@ -2844,6 +2847,7 @@
      "integrity": "sha512-jp2L/eY6fn+KgVVQAOqYItbF0VY/YApe5Mz2F0aykSO8gx31bYCZyvSeYxCHKvzHG5eZjc+zyaS5BrBWya2+kQ==",
      "devOptional": true,
      "license": "MIT",
+      "peer": true,
      "peerDependencies": {
        "@types/react": "^19.2.0"
      }
@@ -2899,6 +2903,7 @@
      "integrity": "sha512-3xP4XzzDNQOIqBMWogftkwxhg5oMKApqY0BAflmLZiFYHqyhSOxv/cd/zPQLTcCXr4AkaKb25joocY0BD1WC6A==",
      "dev": true,
      "license": "MIT",
+      "peer": true,
      "dependencies": {
        "@typescript-eslint/scope-manager": "8.51.0",
        "@typescript-eslint/types": "8.51.0",
@@ -3209,6 +3214,7 @@
      "integrity": "sha512-NZyJarBfL7nWwIq+FDL6Zp/yHEhePMNnnJ0y3qfieCrmNvYct8uvtiV41UvlSe6apAfk0fY1FbWx+NwfmpvtTg==",
      "dev": true,
      "license": "MIT",
+      "peer": true,
      "bin": {
        "acorn": "bin/acorn"
      },
@@ -3340,6 +3346,7 @@
        }
      ],
      "license": "MIT",
+      "peer": true,
      "dependencies": {
        "baseline-browser-mapping": "^2.9.0",
        "caniuse-lite": "^1.0.30001759",
@@ -3611,6 +3618,7 @@
      "resolved": "https://registry.npmjs.org/d3-selection/-/d3-selection-3.0.0.tgz",
      "integrity": "sha512-fmTRWbNMmsmWq6xJV8D19U/gw/bwrHfNXxrIN+HfZgnzqTHp9jOmKMhsTUjXOJnZOdZY9Q28y4yebKzqDKlxlQ==",
      "license": "ISC",
+      "peer": true,
      "engines": {
        "node": ">=12"
      }
@@ -3836,6 +3844,7 @@
      "integrity": "sha512-LEyamqS7W5HB3ujJyvi0HQK/dtVINZvd5mAAp9eT5S/ujByGjiZLCzPcHVzuXbpJDJF/cxwHlfceVUDZ2lnSTw==",
      "dev": true,
      "license": "MIT",
+      "peer": true,
      "dependencies": {
        "@eslint-community/eslint-utils": "^4.8.0",
        "@eslint-community/regexpp": "^4.12.1",
@@ -5836,6 +5845,7 @@
      "integrity": "sha512-5gTmgEY/sqK6gFXLIsQNH19lWb4ebPDLA4SdLP7dsWkIXHWlG66oPuVvXSGFPppYZz8ZDZq0dYYrbHfBCVUb1Q==",
      "dev": true,
      "license": "MIT",
+      "peer": true,
      "engines": {
        "node": ">=12"
      },
@@ -5951,6 +5961,7 @@
      "resolved": "https://registry.npmjs.org/react/-/react-19.2.3.tgz",
      "integrity": "sha512-Ku/hhYbVjOQnXDZFv2+RibmLFGwFdeeKHFcOTlrt7xplBnya5OGn/hIRDsqDiSUcfORsDC7MPxwork8jBwsIWA==",
      "license": "MIT",
+      "peer": true,
      "engines": {
        "node": ">=0.10.0"
      }
@@ -5960,6 +5971,7 @@
      "resolved": "https://registry.npmjs.org/react-dom/-/react-dom-19.2.3.tgz",
      "integrity": "sha512-yELu4WmLPw5Mr/lmeEpox5rw3RETacE++JgHqQzd2dg+YbJuat3jH4ingc+WPZhxaoFzdv9y33G+F7Nl5O0GBg==",
      "license": "MIT",
+      "peer": true,
      "dependencies": {
        "scheduler": "^0.27.0"
      },
@@ -6424,6 +6436,7 @@
      "integrity": "sha512-84MVSjMEHP+FQRPy3pX9sTVV/INIex71s9TL2Gm5FG/WG1SqXeKyZ0k7/blY/4FdOzI12CBy1vGc4og/eus0fw==",
      "dev": true,
      "license": "Apache-2.0",
+      "peer": true,
      "bin": {
        "tsc": "bin/tsc",
        "tsserver": "bin/tsserver"
@@ -6677,6 +6690,7 @@
      "integrity": "sha512-w+N7Hifpc3gRjZ63vYBXA56dvvRlNWRczTdmCBBa+CotUzAPf5b7YMdMR/8CQoeYE5LX3W4wj6RYTgonm1b9DA==",
      "dev": true,
      "license": "MIT",
+      "peer": true,
      "dependencies": {
        "esbuild": "^0.27.0",
        "fdir": "^6.5.0",
--- a/ui/src/components/AgentControl.tsx
+++ b/ui/src/components/AgentControl.tsx
@@ -1,8 +1,10 @@
 import { useState, useEffect, useRef, useCallback } from 'react'
-import { Play, Square, Loader2, GitBranch, Clock } from 'lucide-react'
+import { Play, Square, Loader2, GitBranch, Clock, Pause, PlayCircle } from 'lucide-react'
 import {
  useStartAgent,
  useStopAgent,
+  useGracefulPauseAgent,
+  useGracefulResumeAgent,
  useSettings,
  useUpdateProjectSettings,
 } from '../hooks/useProjects'
@@ -60,12 +62,14 @@ export function AgentControl({ projectName, status, defaultConcurrency = 3 }: Ag

  const startAgent = useStartAgent(projectName)
  const stopAgent = useStopAgent(projectName)
+  const gracefulPause = useGracefulPauseAgent(projectName)
+  const gracefulResume = useGracefulResumeAgent(projectName)
  const { data: nextRun } = useNextScheduledRun(projectName)

  const [showScheduleModal, setShowScheduleModal] = useState(false)

-  const isLoading = startAgent.isPending || stopAgent.isPending
-  const isRunning = status === 'running' || status === 'paused'
+  const isLoading = startAgent.isPending || stopAgent.isPending || gracefulPause.isPending || gracefulResume.isPending
+  const isRunning = status === 'running' || status === 'paused' || status === 'pausing' || status === 'paused_graceful'
  const isLoadingStatus = status === 'loading'
  const isParallel = concurrency > 1

@@ -126,7 +130,7 @@ export function AgentControl({ projectName, status, defaultConcurrency = 3 }: Ag
          </Badge>
        )}

-        {/* Start/Stop button */}
+        {/* Start/Stop/Pause/Resume buttons */}
        {isLoadingStatus ? (
          <Button disabled variant="outline" size="sm">
            <Loader2 size={18} className="animate-spin" />
@@ -146,19 +150,69 @@ export function AgentControl({ projectName, status, defaultConcurrency = 3 }: Ag
            )}
          </Button>
        ) : (
-          <Button
-            onClick={handleStop}
-            disabled={isLoading}
-            variant="destructive"
-            size="sm"
-            title={yoloMode ? 'Stop Agent (YOLO Mode)' : 'Stop Agent'}
-          >
-            {isLoading ? (
-              <Loader2 size={18} className="animate-spin" />
-            ) : (
-              <Square size={18} />
+          <div className="flex items-center gap-1.5">
+            {/* Pausing indicator */}
+            {status === 'pausing' && (
+              <Badge variant="secondary" className="gap-1 animate-pulse">
+                <Loader2 size={12} className="animate-spin" />
+                Pausing...
+              </Badge>
            )}
-          </Button>
+
+            {/* Paused indicator + Resume button */}
+            {status === 'paused_graceful' && (
+              <>
+                <Badge variant="outline" className="gap-1">
+                  Paused
+                </Badge>
+                <Button
+                  onClick={() => gracefulResume.mutate()}
+                  disabled={isLoading}
+                  variant="default"
+                  size="sm"
+                  title="Resume agent"
+                >
+                  {gracefulResume.isPending ? (
+                    <Loader2 size={18} className="animate-spin" />
+                  ) : (
+                    <PlayCircle size={18} />
+                  )}
+                </Button>
+              </>
+            )}
+
+            {/* Graceful pause button (only when running normally) */}
+            {status === 'running' && (
+              <Button
+                onClick={() => gracefulPause.mutate()}
+                disabled={isLoading}
+                variant="outline"
+                size="sm"
+                title="Pause agent (finish current work first)"
+              >
+                {gracefulPause.isPending ? (
+                  <Loader2 size={18} className="animate-spin" />
+                ) : (
+                  <Pause size={18} />
+                )}
+              </Button>
+            )}
+
+            {/* Stop button (always available) */}
+            <Button
+              onClick={handleStop}
+              disabled={isLoading}
+              variant="destructive"
+              size="sm"
+              title="Stop Agent (immediate)"
+            >
+              {stopAgent.isPending ? (
+                <Loader2 size={18} className="animate-spin" />
+              ) : (
+                <Square size={18} />
+              )}
+            </Button>
+          </div>
        )}

        {/* Clock button to open schedule modal */}
--- a/ui/src/components/OrchestratorStatusCard.tsx
+++ b/ui/src/components/OrchestratorStatusCard.tsx
@@ -25,6 +25,10 @@ function getStateText(state: OrchestratorState): string {
      return 'Watching progress...'
    case 'complete':
      return 'Mission accomplished!'
+    case 'draining':
+      return 'Draining agents...'
+    case 'paused':
+      return 'Paused'
    default:
      return 'Orchestrating...'
  }
@@ -42,6 +46,10 @@ function getStateColor(state: OrchestratorState): string {
      return 'text-primary'
    case 'initializing':
      return 'text-yellow-600 dark:text-yellow-400'
+    case 'draining':
+      return 'text-amber-600 dark:text-amber-400'
+    case 'paused':
+      return 'text-muted-foreground'
    default:
      return 'text-muted-foreground'
  }
--- a/ui/src/hooks/useProjects.ts
+++ b/ui/src/hooks/useProjects.ts
@@ -197,6 +197,28 @@ export function useResumeAgent(projectName: string) {
  })
 }

+export function useGracefulPauseAgent(projectName: string) {
+  const queryClient = useQueryClient()
+
+  return useMutation({
+    mutationFn: () => api.gracefulPauseAgent(projectName),
+    onSuccess: () => {
+      queryClient.invalidateQueries({ queryKey: ['agent-status', projectName] })
+    },
+  })
+}
+
+export function useGracefulResumeAgent(projectName: string) {
+  const queryClient = useQueryClient()
+
+  return useMutation({
+    mutationFn: () => api.gracefulResumeAgent(projectName),
+    onSuccess: () => {
+      queryClient.invalidateQueries({ queryKey: ['agent-status', projectName] })
+    },
+  })
+}
+
 // ============================================================================
 // Setup
 // ============================================================================
--- a/ui/src/lib/api.ts
+++ b/ui/src/lib/api.ts
@@ -271,6 +271,18 @@ export async function resumeAgent(projectName: string): Promise<AgentActionRespo
  })
 }

+export async function gracefulPauseAgent(projectName: string): Promise<AgentActionResponse> {
+  return fetchJSON(`/projects/${encodeURIComponent(projectName)}/agent/graceful-pause`, {
+    method: 'POST',
+  })
+}
+
+export async function gracefulResumeAgent(projectName: string): Promise<AgentActionResponse> {
+  return fetchJSON(`/projects/${encodeURIComponent(projectName)}/agent/graceful-resume`, {
+    method: 'POST',
+  })
+}
+
 // ============================================================================
 // Spec Creation API
 // ============================================================================
--- a/ui/src/lib/types.ts
+++ b/ui/src/lib/types.ts
@@ -120,7 +120,7 @@ export interface FeatureUpdate {
 }

 // Agent types
-export type AgentStatus = 'stopped' | 'running' | 'paused' | 'crashed' | 'loading'
+export type AgentStatus = 'stopped' | 'running' | 'paused' | 'crashed' | 'loading' | 'pausing' | 'paused_graceful'

 export interface AgentStatusResponse {
  status: AgentStatus
@@ -216,6 +216,8 @@ export type OrchestratorState =
  | 'spawning'
  | 'monitoring'
  | 'complete'
+  | 'draining'
+  | 'paused'

 // Orchestrator event for recent activity
 export interface OrchestratorEvent {