feat: enhance agent SSE handling and worker restart logic

- Add sseRestartListener to handle vGPU restart events separately from config updates.
- Update sseConfigListener to improve logging and connection handling.
- Modify restartWorker to include actual worker info for better process management.
- Introduce a retry mechanism for failed worker restarts to improve resilience.
- Update settings.local.json to include additional Bash commands for enhanced functionality.

Author: andystime

.claude/settings.local.json

@@ -4,7 +4,12 @@
       "Bash(go build:*)",
       "Bash(go env:*)",
       "Bash(curl:*)",
-      "Bash(echo exit: $?:*)"
+      "Bash(echo exit: $?:*)",
+      "Bash(npm install:*)",
+      "Bash(where node:*)",
+      "Bash(winget install:*)",
+      "Bash(cmd:*)",
+      "Bash(powershell:*)"
     ]
   }
 }

internal/agent/agent.go

@@ -255,9 +255,10 @@ func (a *Agent) Start() error {
 	}
 
 	// Start background tasks
-	a.wg.Add(2)
+	a.wg.Add(3)
 	go a.statusReportLoop()
 	go a.sseConfigListener()
+	go a.sseRestartListener()
 
 	klog.Infof("Agent started: agent_id=%s pid=%d", a.agentID, os.Getpid())
 

internal/agent/sse.go

@@ -19,8 +19,9 @@ const (
 	sseVGPURestartTopicSuffix = "_vgpu_restart"
 )
 
-// sseConfigListener connects to the SSE endpoint and triggers config re-fetch on new events.
-// It reconnects automatically with exponential backoff.
+// sseConfigListener connects to the SSE endpoint for config-update events and
+// triggers config re-fetch on new events. It reconnects automatically with
+// exponential backoff.
 func (a *Agent) sseConfigListener() {
 	defer a.wg.Done()
 
@@ -35,22 +36,52 @@ func (a *Agent) sseConfigListener() {
 
 		err := a.listenSSE()
 		if err != nil {
-			klog.Warningf("SSE connection error: %v", err)
+			klog.Warningf("SSE config connection error: %v", err)
 		}
 
-		// Check if context is done before reconnecting
 		select {
 		case <-a.ctx.Done():
 			return
 		case <-time.After(backoff):
 		}
 
-		// Exponential backoff
 		backoff = min(backoff*2, sseReconnectMax)
 	}
 }
 
-// listenSSE opens a single SSE connection and processes events until disconnected.
+// sseRestartListener connects to the SSE endpoint for vGPU restart events and
+// routes them to the reconciler. It runs on a dedicated connection so that
+// messages can be attributed to the restart topic without relying on the
+// broker setting the SSE `event:` field (which sse.tensor-fusion.ai does not).
+func (a *Agent) sseRestartListener() {
+	defer a.wg.Done()
+
+	backoff := sseReconnectMin
+
+	for {
+		select {
+		case <-a.ctx.Done():
+			return
+		default:
+		}
+
+		err := a.listenSSERestart()
+		if err != nil {
+			klog.Warningf("SSE restart connection error: %v", err)
+		}
+
+		select {
+		case <-a.ctx.Done():
+			return
+		case <-time.After(backoff):
+		}
+
+		backoff = min(backoff*2, sseReconnectMax)
+	}
+}
+
+// listenSSE opens a single SSE connection for config-update events (topic =
+// agentID) and triggers a debounced pullConfig on every received frame.
 func (a *Agent) listenSSE() error {
 	ctx, cancel := context.WithCancel(a.ctx)
 	defer cancel()
@@ -59,7 +90,11 @@ func (a *Agent) listenSSE() error {
 	if err != nil {
 		return err
 	}
-	req.Header.Set(sseTopicHeader, strings.Join([]string{a.agentID, a.vgpuRestartTopic()}, ","))
+	// Subscribe only to the config-update topic.
+	// Restart events are handled by a separate connection (listenSSERestart)
+	// because the SSE broker does not include the topic name in the frame,
+	// making it impossible to route events from a multi-topic subscription.
+	req.Header.Set(sseTopicHeader, a.agentID)
 	req.Header.Set("Accept", "text/event-stream")
 	req.Header.Set("Cache-Control", "no-cache")
 
@@ -71,25 +106,31 @@ func (a *Agent) listenSSE() error {
 	defer func() { _ = resp.Body.Close() }()
 
 	if resp.StatusCode != http.StatusOK {
-		klog.Warningf("SSE endpoint returned status %d", resp.StatusCode)
+		klog.Warningf("SSE config endpoint returned status %d", resp.StatusCode)
 		return nil
 	}
 
-	klog.Infof("SSE connection established: topics=%s,%s", a.agentID, a.vgpuRestartTopic())
+	klog.Infof("SSE config connection established: topic=%s", a.agentID)
 
 	scanner := bufio.NewScanner(resp.Body)
 	var debounceTimer *time.Timer
-	var eventType string
 	var eventDataLines []string
 
 	flushEvent := func() {
 		if len(eventDataLines) == 0 {
-			eventType = ""
 			return
 		}
-		a.handleSSEEvent(eventType, eventDataLines, &debounceTimer)
-		eventType = ""
 		eventDataLines = nil
+		// Debounce config pull so event bursts result in one pullConfig call.
+		if debounceTimer != nil {
+			debounceTimer.Stop()
+		}
+		debounceTimer = time.AfterFunc(sseDebounceDelay, func() {
+			klog.Infof("SSE config event received, triggering config re-fetch")
+			if err := a.pullConfig(); err != nil {
+				klog.Errorf("Failed to pull config after SSE event: %v", err)
+			}
+		})
 	}
 
 	for scanner.Scan() {
@@ -103,12 +144,9 @@ func (a *Agent) listenSSE() error {
 		}
 
 		line := scanner.Text()
-
 		switch {
 		case line == "":
 			flushEvent()
-		case strings.HasPrefix(line, "event:"):
-			eventType = strings.TrimSpace(strings.TrimPrefix(line, "event:"))
 		case strings.HasPrefix(line, "data:"):
 			eventDataLines = append(eventDataLines, strings.TrimSpace(strings.TrimPrefix(line, "data:")))
 		}
@@ -123,30 +161,77 @@ func (a *Agent) listenSSE() error {
 		return err
 	}
 
-	klog.Infof("SSE connection closed by server, will reconnect")
+	klog.Infof("SSE config connection closed by server, will reconnect")
 	return nil
 }
 
-func (a *Agent) vgpuRestartTopic() string {
-	return a.agentID + sseVGPURestartTopicSuffix
-}
+// listenSSERestart opens a single SSE connection for vGPU restart events
+// (topic = agentID + "_vgpu_restart"). Every received frame is forwarded
+// directly to handleVGPURestartEvent.
+func (a *Agent) listenSSERestart() error {
+	ctx, cancel := context.WithCancel(a.ctx)
+	defer cancel()
 
-func (a *Agent) handleSSEEvent(eventType string, dataLines []string, debounceTimer **time.Timer) {
-	if strings.TrimSpace(eventType) == a.vgpuRestartTopic() {
-		a.handleVGPURestartEvent(dataLines)
-		return
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, sseEndpoint, nil)
+	if err != nil {
+		return err
 	}
+	req.Header.Set(sseTopicHeader, a.vgpuRestartTopic())
+	req.Header.Set("Accept", "text/event-stream")
+	req.Header.Set("Cache-Control", "no-cache")
 
-	// Debounce config pull so event bursts result in one pullConfig call.
-	if *debounceTimer != nil {
-		(*debounceTimer).Stop()
+	client := &http.Client{}
+	resp, err := client.Do(req)
+	if err != nil {
+		return err
 	}
-	*debounceTimer = time.AfterFunc(sseDebounceDelay, func() {
-		klog.Infof("SSE event received, triggering config re-fetch")
-		if err := a.pullConfig(); err != nil {
-			klog.Errorf("Failed to pull config after SSE event: %v", err)
+	defer func() { _ = resp.Body.Close() }()
+
+	if resp.StatusCode != http.StatusOK {
+		klog.Warningf("SSE restart endpoint returned status %d", resp.StatusCode)
+		return nil
+	}
+
+	klog.Infof("SSE restart connection established: topic=%s", a.vgpuRestartTopic())
+
+	scanner := bufio.NewScanner(resp.Body)
+	var eventDataLines []string
+
+	flushEvent := func() {
+		if len(eventDataLines) == 0 {
+			return
+		}
+		a.handleVGPURestartEvent(eventDataLines)
+		eventDataLines = nil
+	}
+
+	for scanner.Scan() {
+		select {
+		case <-ctx.Done():
+			return nil
+		default:
 		}
-	})
+
+		line := scanner.Text()
+		switch {
+		case line == "":
+			flushEvent()
+		case strings.HasPrefix(line, "data:"):
+			eventDataLines = append(eventDataLines, strings.TrimSpace(strings.TrimPrefix(line, "data:")))
+		}
+	}
+	flushEvent()
+
+	if err := scanner.Err(); err != nil {
+		return err
+	}
+
+	klog.Infof("SSE restart connection closed by server, will reconnect")
+	return nil
+}
+
+func (a *Agent) vgpuRestartTopic() string {
+	return a.agentID + sseVGPURestartTopicSuffix
 }
 
 func (a *Agent) handleVGPURestartEvent(dataLines []string) int {

internal/hypervisor/reconciler.go

@@ -177,7 +177,7 @@ func (r *Reconciler) reconcile() {
 				klog.Infof("Force restart requested for worker: worker_id=%s", workerID)
 			}
 			// Structural change (GPU allocation, executable, args) requires restart
-			if err := r.restartWorker(desiredInfo); err != nil {
+			if err := r.restartWorker(desiredInfo, actualWorker); err != nil {
 				klog.Errorf("Failed to restart worker: worker_id=%s error=%v", workerID, err)
 				if forceRestart {
 					retryRestarts[workerID] = struct{}{}
@@ -213,6 +213,12 @@ func (r *Reconciler) reconcile() {
 			r.forceRestarts[workerID] = struct{}{}
 		}
 		r.mu.Unlock()
+		// Schedule a reconcile soon so failed restarts are retried quickly
+		// rather than waiting for the 30-second ticker.
+		go func() {
+			time.Sleep(5 * time.Second)
+			r.TriggerReconcile()
+		}()
 	}
 
 	if added > 0 || removed > 0 || updated > 0 {
@@ -248,17 +254,38 @@ func (r *Reconciler) stopWorker(workerID string) error {
 	return nil
 }
 
-func (r *Reconciler) restartWorker(info *api.WorkerInfo) error {
+func (r *Reconciler) restartWorker(desired *api.WorkerInfo, actual *api.WorkerInfo) error {
 	// Stop first
-	if err := r.stopWorker(info.WorkerUID); err != nil {
-		klog.Warningf("Failed to stop worker during restart: worker_id=%s error=%v", info.WorkerUID, err)
+	if err := r.stopWorker(desired.WorkerUID); err != nil {
+		klog.Warningf("Failed to stop worker during restart: worker_id=%s error=%v", desired.WorkerUID, err)
 	}
 
-	// Small delay to ensure cleanup
-	time.Sleep(100 * time.Millisecond)
+	// StopWorker sends a signal but does not wait for the process to exit.
+	// Poll until the old process releases the port, then start the new one.
+	// Without this wait, StartWorker fails with "port already in use".
+	if actual != nil && actual.WorkerRunningInfo != nil && actual.WorkerRunningInfo.PID > 0 {
+		pid := int(actual.WorkerRunningInfo.PID)
+		if isProcessRunning(pid) {
+			deadline := time.Now().Add(10 * time.Second)
+			for time.Now().Before(deadline) {
+				if !isProcessRunning(pid) {
+					break
+				}
+				time.Sleep(100 * time.Millisecond)
+			}
+			if isProcessRunning(pid) {
+				klog.Warningf("Worker process did not exit within 10s, force killing: worker_id=%s pid=%d", desired.WorkerUID, pid)
+				forceKillWorkerProcess(pid)
+				time.Sleep(200 * time.Millisecond) // Allow OS to reclaim port after force kill
+			}
+		}
+	} else {
+		// PID unknown: short delay to let any in-flight socket closure finish
+		time.Sleep(500 * time.Millisecond)
+	}
 
 	// Start with new config
-	return r.startWorker(info)
+	return r.startWorker(desired)
 }
 
 // needsRestart checks if structural config changed (GPU allocation, executable, args)