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>
This commit is contained in:
@@ -132,6 +132,9 @@ func (s *DesktopTaskService) Lease(ctx context.Context, client *repository.Deskt
|
||||
if err := s.dropStaleMonitorTasksForClient(ctx, client); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := s.recoverExpiredClientTasks(ctx, client); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
rawToken, tokenHash, err := newDesktopClientToken()
|
||||
if err != nil {
|
||||
@@ -783,6 +786,245 @@ func normalizeDesktopTaskTerminalStatus(status string) string {
|
||||
}
|
||||
}
|
||||
|
||||
type desktopTaskRecoveryMode string
|
||||
|
||||
const (
|
||||
desktopTaskRecoveryModeStartup desktopTaskRecoveryMode = "startup"
|
||||
desktopTaskRecoveryModeDisconnect desktopTaskRecoveryMode = "disconnect"
|
||||
desktopTaskRecoveryModeLeaseExpiry desktopTaskRecoveryMode = "lease_expired"
|
||||
)
|
||||
|
||||
type recoveredDesktopTaskLease struct {
|
||||
TaskID uuid.UUID
|
||||
WorkspaceID int64
|
||||
Kind string
|
||||
Status string
|
||||
ActiveAttempt pgtype.UUID
|
||||
}
|
||||
|
||||
func (s *DesktopTaskService) recoverClientTasksOnStartup(ctx context.Context, client *repository.DesktopClient) error {
|
||||
return s.recoverClientTasks(ctx, client, desktopTaskRecoveryModeStartup)
|
||||
}
|
||||
|
||||
func (s *DesktopTaskService) recoverClientTasksOnDisconnect(ctx context.Context, client *repository.DesktopClient) error {
|
||||
return s.recoverClientTasks(ctx, client, desktopTaskRecoveryModeDisconnect)
|
||||
}
|
||||
|
||||
func (s *DesktopTaskService) recoverExpiredClientTasks(ctx context.Context, client *repository.DesktopClient) error {
|
||||
return s.recoverClientTasks(ctx, client, desktopTaskRecoveryModeLeaseExpiry)
|
||||
}
|
||||
|
||||
func (s *DesktopTaskService) recoverClientTasks(
|
||||
ctx context.Context,
|
||||
client *repository.DesktopClient,
|
||||
mode desktopTaskRecoveryMode,
|
||||
) error {
|
||||
if s == nil || client == nil || s.pool == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
publishErrorJSON, publishReason, publishMessage, err := desktopTaskRecoveryPayload(mode, "publish")
|
||||
if err != nil {
|
||||
return response.ErrInternal(50195, "desktop_task_recovery_payload_failed", "failed to encode publish desktop task recovery payload")
|
||||
}
|
||||
monitorErrorJSON, monitorReason, _, err := desktopTaskRecoveryPayload(mode, "monitor")
|
||||
if err != nil {
|
||||
return response.ErrInternal(50195, "desktop_task_recovery_payload_failed", "failed to encode monitor desktop task recovery payload")
|
||||
}
|
||||
|
||||
tx, err := s.pool.Begin(ctx)
|
||||
if err != nil {
|
||||
return response.ErrInternal(50196, "desktop_task_recovery_begin_failed", "failed to start desktop task recovery")
|
||||
}
|
||||
defer tx.Rollback(ctx)
|
||||
|
||||
repo := repository.NewDesktopTaskRepository(tx)
|
||||
query := `
|
||||
SELECT
|
||||
desktop_id,
|
||||
workspace_id,
|
||||
kind,
|
||||
active_attempt_id
|
||||
FROM desktop_tasks
|
||||
WHERE target_client_id = $1
|
||||
AND workspace_id = $2
|
||||
AND status = 'in_progress'
|
||||
`
|
||||
queryArgs := []any{client.ID, client.WorkspaceID}
|
||||
if mode == desktopTaskRecoveryModeLeaseExpiry {
|
||||
query += `
|
||||
AND lease_expires_at IS NOT NULL
|
||||
AND lease_expires_at < NOW()
|
||||
`
|
||||
}
|
||||
query += `
|
||||
FOR UPDATE
|
||||
`
|
||||
|
||||
rows, err := tx.Query(ctx, query, queryArgs...)
|
||||
if err != nil {
|
||||
s.logWarn("desktop task recovery query failed", err, zap.String("client_id", client.ID.String()), zap.String("mode", string(mode)))
|
||||
return response.ErrInternal(50197, "desktop_task_recovery_query_failed", "failed to select recoverable desktop tasks")
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
recovered := make([]recoveredDesktopTaskLease, 0)
|
||||
for rows.Next() {
|
||||
var item recoveredDesktopTaskLease
|
||||
if scanErr := rows.Scan(
|
||||
&item.TaskID,
|
||||
&item.WorkspaceID,
|
||||
&item.Kind,
|
||||
&item.Status,
|
||||
&item.ActiveAttempt,
|
||||
); scanErr != nil {
|
||||
return response.ErrInternal(50198, "desktop_task_recovery_scan_failed", "failed to parse recovered desktop tasks")
|
||||
}
|
||||
recovered = append(recovered, item)
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
return response.ErrInternal(50198, "desktop_task_recovery_scan_failed", "failed to iterate recovered desktop tasks")
|
||||
}
|
||||
|
||||
for index := range recovered {
|
||||
item := &recovered[index]
|
||||
if item.Kind == "monitor" {
|
||||
if _, execErr := tx.Exec(ctx, `
|
||||
UPDATE desktop_tasks
|
||||
SET status = 'queued',
|
||||
error = $2,
|
||||
active_attempt_id = NULL,
|
||||
lease_token_hash = NULL,
|
||||
lease_expires_at = NULL,
|
||||
interrupted_at = COALESCE(interrupted_at, NOW()),
|
||||
interrupt_reason = $3,
|
||||
enqueued_at = NOW(),
|
||||
updated_at = NOW()
|
||||
WHERE desktop_id = $1
|
||||
`, item.TaskID, monitorErrorJSON, monitorReason); execErr != nil {
|
||||
s.logWarn("desktop monitor task recovery update failed", execErr, zap.String("task_id", item.TaskID.String()), zap.String("mode", string(mode)))
|
||||
return response.ErrInternal(50197, "desktop_task_recovery_query_failed", "failed to select recoverable desktop tasks")
|
||||
}
|
||||
item.Status = "queued"
|
||||
continue
|
||||
}
|
||||
|
||||
if _, execErr := tx.Exec(ctx, `
|
||||
UPDATE desktop_tasks
|
||||
SET status = 'unknown',
|
||||
error = $2,
|
||||
active_attempt_id = NULL,
|
||||
lease_token_hash = NULL,
|
||||
lease_expires_at = NULL,
|
||||
interrupted_at = COALESCE(interrupted_at, NOW()),
|
||||
interrupt_reason = COALESCE(interrupt_reason, $3),
|
||||
updated_at = NOW()
|
||||
WHERE desktop_id = $1
|
||||
`, item.TaskID, publishErrorJSON, publishReason); execErr != nil {
|
||||
s.logWarn("desktop publish task recovery update failed", execErr, zap.String("task_id", item.TaskID.String()), zap.String("mode", string(mode)))
|
||||
return response.ErrInternal(50197, "desktop_task_recovery_query_failed", "failed to select recoverable desktop tasks")
|
||||
}
|
||||
item.Status = "unknown"
|
||||
}
|
||||
|
||||
for _, item := range recovered {
|
||||
if !item.ActiveAttempt.Valid {
|
||||
continue
|
||||
}
|
||||
|
||||
attemptID, convErr := uuid.FromBytes(item.ActiveAttempt.Bytes[:])
|
||||
if convErr != nil {
|
||||
s.logWarn("desktop recovery attempt id decode failed", convErr, zap.String("task_id", item.TaskID.String()))
|
||||
continue
|
||||
}
|
||||
|
||||
finalStatus := literalStringPtr(item.Status)
|
||||
errorJSON := publishErrorJSON
|
||||
if item.Kind == "monitor" {
|
||||
finalStatus = literalStringPtr("aborted")
|
||||
errorJSON = monitorErrorJSON
|
||||
}
|
||||
|
||||
if finishErr := repo.FinishAttempt(ctx, repository.FinishDesktopTaskAttemptParams{
|
||||
TaskID: item.TaskID,
|
||||
AttemptID: attemptID,
|
||||
FinalStatus: finalStatus,
|
||||
Error: errorJSON,
|
||||
}); finishErr != nil {
|
||||
s.logWarn("desktop recovery attempt finish failed", finishErr, zap.String("task_id", item.TaskID.String()))
|
||||
}
|
||||
}
|
||||
|
||||
if err := tx.Commit(ctx); err != nil {
|
||||
return response.ErrInternal(50199, "desktop_task_recovery_commit_failed", "failed to commit desktop task recovery")
|
||||
}
|
||||
|
||||
for _, item := range recovered {
|
||||
task, getErr := s.repo.GetByDesktopID(ctx, item.TaskID, item.WorkspaceID)
|
||||
if getErr != nil {
|
||||
s.logWarn("desktop recovery reload failed", getErr, zap.String("task_id", item.TaskID.String()))
|
||||
continue
|
||||
}
|
||||
|
||||
eventType := "task_completed"
|
||||
if item.Kind == "monitor" {
|
||||
eventType = "task_available"
|
||||
}
|
||||
s.publishTaskEvent(ctx, task, eventType)
|
||||
}
|
||||
|
||||
if len(recovered) > 0 && s.logger != nil {
|
||||
s.logger.Info("desktop tasks recovered",
|
||||
zap.String("client_id", client.ID.String()),
|
||||
zap.String("mode", string(mode)),
|
||||
zap.Int("count", len(recovered)),
|
||||
zap.String("publish_message", publishMessage),
|
||||
)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func desktopTaskRecoveryPayload(mode desktopTaskRecoveryMode, kind string) ([]byte, string, string, error) {
|
||||
reason := strings.TrimSpace(string(mode))
|
||||
message := "desktop task was recovered after the client lost the active lease"
|
||||
source := "desktop_task_recovery"
|
||||
|
||||
switch mode {
|
||||
case desktopTaskRecoveryModeStartup:
|
||||
source = "desktop_client_startup"
|
||||
if kind == "monitor" {
|
||||
message = "desktop client restarted while the monitor task was in progress; task has been re-queued"
|
||||
} else {
|
||||
message = "desktop client restarted while the publish task was in progress; task has been moved to unknown for manual reconcile"
|
||||
}
|
||||
case desktopTaskRecoveryModeDisconnect:
|
||||
source = "desktop_client_offline"
|
||||
if kind == "monitor" {
|
||||
message = "desktop client went offline while the monitor task was in progress; task has been re-queued"
|
||||
} else {
|
||||
message = "desktop client went offline while the publish task was in progress; task has been moved to unknown for manual reconcile"
|
||||
}
|
||||
case desktopTaskRecoveryModeLeaseExpiry:
|
||||
source = "desktop_task_lease_expiry"
|
||||
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 moved to unknown for manual reconcile"
|
||||
}
|
||||
}
|
||||
|
||||
payload, err := marshalOptionalJSON(map[string]any{
|
||||
"reason": reason,
|
||||
"message": message,
|
||||
"source": source,
|
||||
})
|
||||
if err != nil {
|
||||
return nil, "", "", err
|
||||
}
|
||||
return payload, reason, message, nil
|
||||
}
|
||||
|
||||
func (s *DesktopTaskService) dropStaleMonitorTasksForClient(ctx context.Context, client *repository.DesktopClient) error {
|
||||
if s == nil || client == nil || s.pool == nil || s.monitoringPool == nil {
|
||||
return nil
|
||||
|
||||
Reference in New Issue
Block a user