fix: harden desktop monitoring recovery

server/internal/tenant/app/desktop_task_service.go

@@ -1901,7 +1901,8 @@ type PublishLeaseRecoveryResult struct {
 }
 
 type MonitorLeaseRecoveryResult struct {
-	Failed int `json:"failed"`
+	Requeued int `json:"requeued"`
+	Failed   int `json:"failed"`
 }
 
 // resolvePublishRecoveryOutcome decides what to do with a publish task whose lease was lost.
@@ -1923,7 +1924,7 @@ func resolvePublishRecoveryOutcome(submitStarted bool, attempts int) (status str
 }
 
 func resolveMonitorRecoveryOutcome(mode desktopTaskRecoveryMode, attempts int) (status string, payloadKind string) {
-	if mode == desktopTaskRecoveryModeLeaseExpiry || attempts >= 2 {
+	if mode != desktopTaskRecoveryModeLeaseExpiry && attempts >= 2 {
 		return "failed", "monitor_final"
 	}
 	return "queued", "monitor"
@@ -2026,12 +2027,21 @@ func (s *DesktopTaskService) recoverClientTasks(
 						continue
 					}
 				}
+			} else if item.MonitorTaskID.Valid {
+				requeuedMonitorTaskIDs, requeueErr := s.requeueMonitorCollectTasksForDesktopRecovery(ctx, []int64{item.MonitorTaskID.Int64}, monitorErrorForTask)
+				if requeueErr != nil {
+					return requeueErr
+				}
+				if len(requeuedMonitorTaskIDs) == 0 {
+					continue
+				}
 			}
 			if _, execErr := tx.Exec(ctx, `
-				UPDATE desktop_tasks
-				SET status = $2,
-				    error = $3,
-				    active_attempt_id = NULL,
+					UPDATE desktop_tasks
+					SET status = $2,
+					    error = $3,
+					    result = CASE WHEN $2 = 'queued' THEN NULL ELSE result END,
+					    active_attempt_id = NULL,
 				    lease_token_hash = NULL,
 				    lease_expires_at = NULL,
 				    interrupted_at = COALESCE(interrupted_at, NOW()),
@@ -2355,7 +2365,7 @@ func (s *DesktopTaskService) RecoverExpiredMonitorTasks(ctx context.Context, lim
 		limit = 1000
 	}
 
-	errorJSON, reason, _, err := desktopTaskRecoveryPayload(desktopTaskRecoveryModeLeaseExpiry, "monitor_final")
+	errorJSON, reason, _, err := desktopTaskRecoveryPayload(desktopTaskRecoveryModeLeaseExpiry, "monitor")
 	if err != nil {
 		return result, response.ErrInternal(50195, "desktop_task_recovery_payload_failed", "failed to encode monitor desktop task recovery payload")
 	}
@@ -2390,7 +2400,7 @@ func (s *DesktopTaskService) RecoverExpiredMonitorTasks(ctx context.Context, lim
 			s.logWarn("expired desktop monitor recovery scan failed", scanErr)
 			return result, response.ErrInternal(50198, "desktop_task_recovery_scan_failed", "failed to parse recovered desktop tasks")
 		}
-		item.Status = "failed"
+		item.Status = "queued"
 		item.ErrorJSON = errorJSON
 		recovered = append(recovered, item)
 		if item.MonitorTaskID.Valid {
@@ -2402,18 +2412,18 @@ func (s *DesktopTaskService) RecoverExpiredMonitorTasks(ctx context.Context, lim
 		return result, response.ErrInternal(50198, "desktop_task_recovery_scan_failed", "failed to iterate recovered desktop tasks")
 	}
 
-	failedMonitorTaskIDs, err := s.failMonitorCollectTasksForDesktopRecovery(ctx, monitorTaskIDs, errorJSON)
+	requeuedMonitorTaskIDs, err := s.requeueMonitorCollectTasksForDesktopRecovery(ctx, monitorTaskIDs, errorJSON)
 	if err != nil {
 		return result, err
 	}
-	failedMonitorTaskSet := int64Set(failedMonitorTaskIDs)
+	requeuedMonitorTaskSet := int64Set(requeuedMonitorTaskIDs)
 	recoverableDesktopIDs := make([]uuid.UUID, 0, len(recovered))
 	recoverable := make([]recoveredDesktopTaskLease, 0, len(recovered))
 	for _, item := range recovered {
 		if !item.MonitorTaskID.Valid {
 			continue
 		}
-		if _, ok := failedMonitorTaskSet[item.MonitorTaskID.Int64]; !ok {
+		if _, ok := requeuedMonitorTaskSet[item.MonitorTaskID.Int64]; !ok {
 			continue
 		}
 		recoverableDesktopIDs = append(recoverableDesktopIDs, item.TaskID)
@@ -2423,7 +2433,7 @@ func (s *DesktopTaskService) RecoverExpiredMonitorTasks(ctx context.Context, lim
 		return result, nil
 	}
 
-	updateRows, err := tx.Query(ctx, failExpiredMonitorDesktopTasksSQL(), recoverableDesktopIDs, errorJSON, reason)
+	updateRows, err := tx.Query(ctx, requeueExpiredMonitorDesktopTasksSQL(), recoverableDesktopIDs, errorJSON, reason)
 	if err != nil {
 		s.logWarn("expired desktop monitor recovery update failed", err)
 		return result, response.ErrInternal(50197, "desktop_task_recovery_query_failed", "failed to select recoverable desktop tasks")
@@ -2462,7 +2472,7 @@ func (s *DesktopTaskService) RecoverExpiredMonitorTasks(ctx context.Context, lim
 		if finishErr := repo.FinishAttempt(ctx, repository.FinishDesktopTaskAttemptParams{
 			TaskID:      item.TaskID,
 			AttemptID:   attemptID,
-			FinalStatus: literalStringPtr("failed"),
+			FinalStatus: literalStringPtr("aborted"),
 			Error:       errorJSON,
 		}); finishErr != nil {
 			s.logWarn("expired desktop monitor recovery attempt finish failed", finishErr, zap.String("task_id", item.TaskID.String()))
@@ -2473,7 +2483,7 @@ func (s *DesktopTaskService) RecoverExpiredMonitorTasks(ctx context.Context, lim
 		return result, response.ErrInternal(50199, "desktop_task_recovery_commit_failed", "failed to commit desktop task recovery")
 	}
 
-	result.Failed = len(updatedDesktopIDs)
+	result.Requeued = len(updatedDesktopIDs)
 
 	for _, item := range recoverable {
 		if _, ok := updatedDesktopIDs[item.TaskID]; !ok {
@@ -2484,12 +2494,12 @@ func (s *DesktopTaskService) RecoverExpiredMonitorTasks(ctx context.Context, lim
 			s.logWarn("expired desktop monitor recovery reload after commit failed", getErr, zap.String("task_id", item.TaskID.String()))
 			continue
 		}
-		s.publishTaskEvent(ctx, task, "task_completed")
+		s.publishTaskEvent(ctx, task, "task_available")
 	}
 
-	if result.Failed > 0 && s.logger != nil {
+	if result.Requeued > 0 && s.logger != nil {
 		s.logger.Info("expired desktop monitor tasks recovered",
-			zap.Int("failed", result.Failed),
+			zap.Int("requeued", result.Requeued),
 		)
 	}
 
@@ -2519,16 +2529,18 @@ func expiredMonitorDesktopTasksCandidateSQL() string {
 	`
 }
 
-func failExpiredMonitorDesktopTasksSQL() string {
+func requeueExpiredMonitorDesktopTasksSQL() string {
 	return `
 		UPDATE desktop_tasks
-		SET status = 'failed',
+		SET status = 'queued',
 		    error = $2,
+		    result = NULL,
 		    active_attempt_id = NULL,
 		    lease_token_hash = NULL,
 		    lease_expires_at = NULL,
 		    interrupted_at = COALESCE(interrupted_at, NOW()),
 		    interrupt_reason = COALESCE(interrupt_reason, $3),
+		    enqueued_at = NOW(),
 		    updated_at = NOW()
 		WHERE desktop_id = ANY($1::uuid[])
 		  AND kind = 'monitor'
@@ -2537,6 +2549,69 @@ func failExpiredMonitorDesktopTasksSQL() string {
 	`
 }
 
+func (s *DesktopTaskService) requeueMonitorCollectTasksForDesktopRecovery(ctx context.Context, monitorTaskIDs []int64, errorJSON []byte) ([]int64, error) {
+	if s == nil || s.monitoringPool == nil || len(monitorTaskIDs) == 0 {
+		return nil, nil
+	}
+	monitorTaskIDs = uniqueInt64s(monitorTaskIDs)
+	message := desktopTaskRecoveryMessage(errorJSON)
+	requestPayloadJSON, err := json.Marshal(map[string]any{
+		"runtime_error": json.RawMessage(errorJSON),
+	})
+	if err != nil {
+		return nil, response.ErrInternal(50195, "desktop_task_recovery_payload_failed", "failed to encode monitor desktop task recovery payload")
+	}
+	rows, err := s.monitoringPool.Query(ctx, `
+		WITH input AS (
+			SELECT unnest($1::bigint[]) AS id
+		),
+		updated AS (
+			UPDATE monitoring_collect_tasks
+			SET status = 'pending',
+			    lease_token_hash = NULL,
+			    leased_to_executor = NULL,
+			    leased_at = NULL,
+			    lease_expires_at = NULL,
+			    callback_received_at = NULL,
+			    completed_at = NULL,
+			    skip_reason = NULL,
+			    dispatch_after = NOW() + interval '30 seconds',
+			    error_message = NULLIF($2, ''),
+			    request_payload_json = (
+					CASE
+						WHEN request_payload_json IS NULL OR jsonb_typeof(request_payload_json) <> 'object'
+							THEN '{}'::jsonb
+						ELSE request_payload_json
+					END
+				) || $3::jsonb,
+			    updated_at = NOW()
+			WHERE id IN (SELECT id FROM input)
+				  AND COALESCE(execution_owner, 'legacy') = 'desktop_tasks'
+				  AND callback_received_at IS NULL
+				  AND status IN ('pending', 'leased', 'expired')
+				RETURNING id
+		)
+		SELECT id FROM updated
+	`, monitorTaskIDs, message, requestPayloadJSON)
+	if err != nil {
+		return nil, response.ErrInternal(50041, "task_update_failed", "failed to requeue expired monitor desktop tasks")
+	}
+	defer rows.Close()
+
+	requeuedIDs := make([]int64, 0, len(monitorTaskIDs))
+	for rows.Next() {
+		var id int64
+		if scanErr := rows.Scan(&id); scanErr != nil {
+			return nil, response.ErrInternal(50041, "task_update_failed", "failed to parse requeued monitor desktop tasks")
+		}
+		requeuedIDs = append(requeuedIDs, id)
+	}
+	if err := rows.Err(); err != nil {
+		return nil, response.ErrInternal(50041, "task_update_failed", "failed to iterate requeued monitor desktop tasks")
+	}
+	return requeuedIDs, nil
+}
+
 func (s *DesktopTaskService) failMonitorCollectTasksForDesktopRecovery(ctx context.Context, monitorTaskIDs []int64, errorJSON []byte) ([]int64, error) {
 	if s == nil || s.monitoringPool == nil || len(monitorTaskIDs) == 0 {
 		return nil, nil
@@ -2586,8 +2661,8 @@ func (s *DesktopTaskService) failMonitorCollectTasksForDesktopRecovery(ctx conte
 			  AND t.status = 'failed'
 			  AND (
 				t.error_message = $2
-				OR t.request_payload_json #>> '{runtime_error,source}' = 'desktop_task_lease_expiry'
-				OR t.request_payload_json #>> '{runtime_error,reason}' = 'lease_expired'
+				OR t.request_payload_json #>> '{runtime_error,source}' = 'desktop_task_recovery'
+				OR t.request_payload_json #>> '{runtime_error,source}' = 'desktop_client_offline'
 			  )
 		)
 		SELECT id FROM updated
@@ -2595,7 +2670,7 @@ func (s *DesktopTaskService) failMonitorCollectTasksForDesktopRecovery(ctx conte
 		SELECT id FROM already_failed
 	`, monitorTaskIDs, message, requestPayloadJSON)
 	if err != nil {
-		return nil, response.ErrInternal(50041, "task_update_failed", "failed to fail expired monitor desktop tasks")
+		return nil, response.ErrInternal(50041, "task_update_failed", "failed to fail recovered monitor desktop tasks")
 	}
 	defer rows.Close()
 
@@ -2700,6 +2775,8 @@ func desktopTaskRecoveryPayload(mode desktopTaskRecoveryMode, kind string) ([]by
 		message = "desktop publish task may have already been submitted to the platform before the lease was lost; kept as unknown for manual reconcile to avoid duplicate publishing"
 	} else if isMonitorFinal {
 		message = "desktop monitor task lost its lease; task has been marked failed to avoid repeated platform retries"
+	} else if kind == "monitor" {
+		message = "desktop monitor task lost its lease; task has been re-queued"
 	}
 
 	switch mode {
@@ -2735,8 +2812,10 @@ func desktopTaskRecoveryPayload(mode desktopTaskRecoveryMode, kind string) ([]by
 			message = fmt.Sprintf("desktop task lease expired before publish completion; task exceeded %d attempts and has been marked failed", desktopPublishMaxAttempts)
 		} else if isPublishUncertain {
 			message = "desktop task lease expired while the publish task was mid-submit; kept as unknown for manual reconcile to avoid duplicate publishing"
-		} else if isMonitorFinal || kind == "monitor" {
+		} else if isMonitorFinal {
 			message = "desktop task lease expired before monitor completion; task has been marked failed"
+		} else if kind == "monitor" {
+			message = "desktop task lease expired before monitor completion; task has been re-queued"
 		} else {
 			message = "desktop task lease expired before publish completion; task has been re-queued for retry"
 		}