feat(desktop): auto-recover in-progress desktop tasks on client lifecycle events

- reset stale in_progress tasks owned by a client on startup, offline, and lease expiry - monitor tasks re-queue for another pick; publish tasks move to unknown for manual reconcile - send startup=true flag on first heartbeat so the server can trigger recovery - finalize the linked desktop task when a phase2 monitor callback arrives via desktop_tasks path - short-circuit the desktop monitor attempt write after the callback already finalized it Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
2026-04-23 09:11:23 +08:00
parent f91d2645d3
commit ce028c56bf
6 changed files with 491 additions and 38 deletions
@@ -13,6 +13,7 @@ import (

 	"github.com/google/uuid"
 	"github.com/jackc/pgx/v5"
+	"github.com/jackc/pgx/v5/pgtype"
 	"github.com/jackc/pgx/v5/pgxpool"
 	"go.uber.org/zap"
 	"golang.org/x/net/publicsuffix"
@@ -415,6 +416,16 @@ func (s *MonitoringCallbackService) handleTaskResultForInstallation(
 		QueueEnqueuedAt: &enqueuedAt,
 	})

+	if task.ExecutionOwner == "desktop_tasks" {
+		if finalizeErr := s.completeLinkedDesktopMonitorTask(ctx, task, req); finalizeErr != nil && s.logger != nil {
+			s.logger.Warn("phase2 monitor desktop task finalize failed",
+				zap.Int64("monitor_task_id", task.ID),
+				zap.String("ai_platform_id", task.AIPlatformID),
+				zap.Error(finalizeErr),
+			)
+		}
+	}
+
 	return &MonitoringTaskResultResponse{
 		TaskID:               task.ID,
 		TaskStatus:           "received",
@@ -1757,6 +1768,190 @@ func (s *MonitoringCallbackService) processTaskResult(ctx context.Context, task
 	}, nil
 }

+func buildDesktopMonitorTaskResultValue(task *monitoringCollectTask, req MonitoringTaskResultRequest, runStatus string) map[string]any {
+	result := map[string]any{
+		"monitoring_task_id": task.ID,
+		"status":             runStatus,
+	}
+
+	if task != nil {
+		if strings.TrimSpace(task.AIPlatformID) != "" {
+			result["platform"] = strings.TrimSpace(task.AIPlatformID)
+		}
+		if strings.TrimSpace(task.QuestionText) != "" {
+			result["question_text"] = strings.TrimSpace(task.QuestionText)
+		}
+	}
+	if value := strings.TrimSpace(optionalStringValue(req.ProviderModel)); value != "" {
+		result["provider_model"] = value
+	}
+	if value := strings.TrimSpace(optionalStringValue(req.ProviderRequestID)); value != "" {
+		result["provider_request_id"] = value
+	}
+	if value := strings.TrimSpace(optionalStringValue(req.RequestID)); value != "" {
+		result["request_id"] = value
+	}
+	if value := strings.TrimSpace(optionalStringValue(req.Answer)); value != "" {
+		result["answer"] = value
+	}
+	if len(req.Citations) > 0 {
+		result["citations"] = req.Citations
+	}
+	if len(req.SearchResults) > 0 {
+		result["search_results"] = req.SearchResults
+	}
+	if len(req.RawResponseJSON) > 0 {
+		result["raw_response_json"] = req.RawResponseJSON
+	}
+	return result
+}
+
+func buildDesktopMonitorTaskErrorValue(task *monitoringCollectTask, req MonitoringTaskResultRequest) map[string]any {
+	message := strings.TrimSpace(optionalStringValue(req.ErrorMessage))
+	if message == "" {
+		message = "monitoring task failed"
+	}
+
+	result := map[string]any{
+		"code":               "monitoring_task_failed",
+		"message":            message,
+		"monitoring_task_id": task.ID,
+	}
+	if task != nil && strings.TrimSpace(task.AIPlatformID) != "" {
+		result["platform"] = strings.TrimSpace(task.AIPlatformID)
+	}
+	if len(req.RawResponseJSON) > 0 {
+		result["raw_response_json"] = req.RawResponseJSON
+	}
+	return result
+}
+
+func (s *MonitoringCallbackService) completeLinkedDesktopMonitorTask(
+	ctx context.Context,
+	task *monitoringCollectTask,
+	req MonitoringTaskResultRequest,
+) error {
+	if s == nil || s.businessPool == nil || task == nil || task.ExecutionOwner != "desktop_tasks" {
+		return nil
+	}
+
+	runStatus := normalizeRunStatus(req.Status)
+	if runStatus == "" {
+		return nil
+	}
+
+	tx, err := s.businessPool.Begin(ctx)
+	if err != nil {
+		return response.ErrInternal(50041, "desktop_task_complete_begin_failed", "failed to start phase2 desktop task completion")
+	}
+	defer tx.Rollback(ctx)
+
+	var (
+		desktopID       uuid.UUID
+		workspaceID     int64
+		activeAttemptID pgtype.UUID
+	)
+	if err := tx.QueryRow(ctx, `
+		SELECT desktop_id, workspace_id, active_attempt_id
+		FROM desktop_tasks
+		WHERE kind = 'monitor'
+		  AND monitor_task_id = $1
+		  AND status IN ('in_progress', 'unknown')
+		ORDER BY created_at DESC, id DESC
+		LIMIT 1
+		FOR UPDATE
+	`, task.ID).Scan(&desktopID, &workspaceID, &activeAttemptID); err != nil {
+		if err == pgx.ErrNoRows {
+			return nil
+		}
+		return response.ErrInternal(50041, "desktop_task_lookup_failed", "failed to inspect phase2 desktop task state")
+	}
+
+	resultJSON, err := marshalOptionalJSON(buildDesktopMonitorTaskResultValue(task, req, runStatus))
+	if err != nil {
+		return response.ErrInternal(50041, "desktop_task_payload_encode_failed", "failed to encode phase2 desktop task result payload")
+	}
+
+	var errorJSON []byte
+	if runStatus == "failed" {
+		errorJSON, err = marshalOptionalJSON(buildDesktopMonitorTaskErrorValue(task, req))
+		if err != nil {
+			return response.ErrInternal(50041, "desktop_task_payload_encode_failed", "failed to encode phase2 desktop task error payload")
+		}
+	}
+
+	if _, err := tx.Exec(ctx, `
+		UPDATE desktop_tasks
+		SET status = $2,
+		    result = $3,
+		    error = $4,
+		    active_attempt_id = NULL,
+		    lease_token_hash = NULL,
+		    lease_expires_at = NULL,
+		    updated_at = NOW()
+		WHERE desktop_id = $1
+	`, desktopID, runStatus, resultJSON, errorJSON); err != nil {
+		return response.ErrInternal(50041, "desktop_task_complete_failed", "failed to finalize phase2 desktop task")
+	}
+
+	taskRepo := repository.NewDesktopTaskRepository(tx)
+	if activeAttemptID.Valid {
+		resolvedAttemptID, convErr := uuid.FromBytes(activeAttemptID.Bytes[:])
+		if convErr != nil {
+			return response.ErrInternal(50041, "desktop_task_attempt_decode_failed", "failed to decode phase2 desktop task attempt id")
+		}
+		if finishErr := taskRepo.FinishAttempt(ctx, repository.FinishDesktopTaskAttemptParams{
+			TaskID:      desktopID,
+			AttemptID:   resolvedAttemptID,
+			FinalStatus: &runStatus,
+			Error:       errorJSON,
+		}); finishErr != nil {
+			return response.ErrInternal(50041, "desktop_task_attempt_finish_failed", "failed to finalize phase2 desktop task attempt")
+		}
+	}
+
+	finalizedTask, err := taskRepo.GetByDesktopID(ctx, desktopID, workspaceID)
+	if err != nil {
+		return response.ErrInternal(50041, "desktop_task_reload_failed", "failed to reload finalized phase2 desktop task")
+	}
+
+	if err := tx.Commit(ctx); err != nil {
+		return response.ErrInternal(50041, "desktop_task_complete_commit_failed", "failed to commit phase2 desktop task completion")
+	}
+
+	title, businessDate, schedulerGroupKey, questionText := deriveDesktopTaskEventMetadataFromPayload(
+		finalizedTask.Kind,
+		finalizedTask.Payload,
+	)
+	if publishErr := publishDesktopTaskEvent(ctx, s.rabbitMQ, DesktopTaskEvent{
+		Type:                "task_completed",
+		TaskID:              finalizedTask.DesktopID.String(),
+		JobID:               finalizedTask.JobID.String(),
+		WorkspaceID:         finalizedTask.WorkspaceID,
+		TargetAccountID:     finalizedTask.TargetAccountID.String(),
+		TargetClientID:      finalizedTask.TargetClientID.String(),
+		Platform:            finalizedTask.Platform,
+		Title:               title,
+		BusinessDate:        businessDate,
+		SchedulerGroupKey:   schedulerGroupKey,
+		QuestionText:        questionText,
+		Status:              finalizedTask.Status,
+		Kind:                finalizedTask.Kind,
+		Priority:            finalizedTask.Priority,
+		Lane:                finalizedTask.Lane,
+		InterruptGeneration: finalizedTask.InterruptGeneration,
+		UpdatedAt:           finalizedTask.UpdatedAt,
+	}); publishErr != nil && s.logger != nil {
+		s.logger.Warn("phase2 desktop task completion event publish failed",
+			zap.Int64("monitor_task_id", task.ID),
+			zap.String("desktop_task_id", finalizedTask.DesktopID.String()),
+			zap.Error(publishErr),
+		)
+	}
+
+	return nil
+}
+
 func (s *MonitoringCallbackService) replaceCitationFacts(ctx context.Context, tx pgx.Tx, tenantID, runID int64) error {
 	if _, err := tx.Exec(ctx, `
 		DELETE FROM monitoring_citation_facts