fix: lease monitor tasks by account identity

server/internal/tenant/app/desktop_task_service.go

@@ -28,6 +28,7 @@ import (
 const (
 	desktopPublishMaxAttempts                     = 3
 	maxConcurrentMonitorPlatformsPerDesktopClient = 2
+	monitorLeaseCandidateScanLimitSQL             = "16"
 	monitorSucceededTaskCooldownSQL               = "2 seconds"
 	monitorUnhealthyTaskCooldownSQL               = "30 seconds"
 )
@@ -519,12 +520,6 @@ func (s *DesktopTaskService) leaseNextQueuedMonitorTask(
 	client *repository.DesktopClient,
 	params repository.DesktopTaskLeaseParams,
 ) (*repository.DesktopTask, error) {
-	accountIDs, err := s.loadMonitorLeaseAccountIDs(ctx, client)
-	if err != nil {
-		return nil, err
-	}
-	hasAccountIDs := len(accountIDs) > 0
-
 	tx, err := beginDesktopMonitorLeaseTx(ctx, s.pool, client.ID)
 	if err != nil {
 		return nil, err
@@ -535,8 +530,6 @@ func (s *DesktopTaskService) leaseNextQueuedMonitorTask(
 		client.TenantID,
 		client.WorkspaceID,
 		client.ID,
-		hasAccountIDs,
-		accountIDs,
 		params.AttemptID,
 		params.LeaseTokenHash,
 		maxConcurrentMonitorPlatformsPerDesktopClient,
@@ -553,37 +546,71 @@ func (s *DesktopTaskService) leaseNextQueuedMonitorTask(
 
 func leaseNextQueuedMonitorTaskSQL() string {
 	return `
-		WITH candidate AS (
-			SELECT dt.desktop_id
-			FROM desktop_tasks AS dt
-			WHERE dt.tenant_id = $1
-			  AND dt.workspace_id = $2
-			  AND dt.kind = 'monitor'
-			  AND dt.status = 'queued'
-			  AND NOT EXISTS (
+		WITH current_accounts AS (
+			SELECT desktop_id, user_id, platform_id, platform_uid, account_fingerprint
+			FROM platform_accounts
+			WHERE tenant_id = $1
+			  AND workspace_id = $2
+			  AND client_id = $3
+			  AND deleted_at IS NULL
+			  AND delete_requested_at IS NULL
+		),
+		candidate_pool AS (
+			SELECT dt.desktop_id,
+				ca.desktop_id AS lease_account_id,
+				dt.lane_weight,
+				dt.priority,
+				dt.enqueued_at,
+				dt.created_at,
+				` + monitorAccountLeaseLockSQL("target_account") + ` AS account_slot_lock_key
+			FROM current_accounts AS ca
+			JOIN platform_accounts AS target_account
+			  ON target_account.tenant_id = $1
+			 AND target_account.workspace_id = $2
+			 AND target_account.user_id = ca.user_id
+			 AND target_account.platform_id = ca.platform_id
+			 AND ` + monitorAccountIdentityMatchSQL("ca", "target_account") + `
+			JOIN desktop_tasks AS dt
+			  ON dt.target_account_id = target_account.desktop_id
+			 AND dt.tenant_id = $1
+			 AND dt.workspace_id = $2
+			 AND dt.platform_id = ca.platform_id
+			 AND dt.kind = 'monitor'
+			 AND dt.status = 'queued'
+			WHERE NOT EXISTS (
 				SELECT 1
 				FROM desktop_tasks AS active
+				JOIN platform_accounts AS active_account
+				  ON active_account.desktop_id = active.target_account_id
+				 AND active_account.tenant_id = active.tenant_id
+				 AND active_account.workspace_id = active.workspace_id
 				WHERE active.tenant_id = dt.tenant_id
 				  AND active.workspace_id = dt.workspace_id
 				  AND active.kind = 'monitor'
-				  AND active.target_client_id = $3
 				  AND active.platform_id = dt.platform_id
 				  AND active.status = 'in_progress'
 				  AND active.desktop_id <> dt.desktop_id
+				  AND active_account.user_id = target_account.user_id
+				  AND ` + monitorAccountIdentityMatchSQL("active_account", "target_account") + `
 			  )
 			  AND NOT EXISTS (
 				SELECT 1
 				FROM desktop_tasks AS recent
+				JOIN platform_accounts AS recent_account
+				  ON recent_account.desktop_id = recent.target_account_id
+				 AND recent_account.tenant_id = recent.tenant_id
+				 AND recent_account.workspace_id = recent.workspace_id
 				WHERE recent.tenant_id = dt.tenant_id
 				  AND recent.workspace_id = dt.workspace_id
 				  AND recent.kind = 'monitor'
-				  AND recent.target_client_id = $3
 				  AND recent.platform_id = dt.platform_id
 				  AND (
 					(recent.status = 'succeeded' AND recent.updated_at >= now() - interval '` + monitorSucceededTaskCooldownSQL + `')
 					OR (recent.status IN ('failed', 'unknown', 'aborted') AND recent.updated_at >= now() - interval '` + monitorUnhealthyTaskCooldownSQL + `')
 				  )
 				  AND recent.desktop_id <> dt.desktop_id
+				  AND recent_account.user_id = target_account.user_id
+				  AND ` + monitorAccountIdentityMatchSQL("recent_account", "target_account") + `
 			  )
 			  AND (
 				SELECT COUNT(DISTINCT active_platform.platform_id)
@@ -594,23 +621,44 @@ func leaseNextQueuedMonitorTaskSQL() string {
 				  AND active_platform.target_client_id = $3
 				  AND active_platform.status = 'in_progress'
 				  AND active_platform.desktop_id <> dt.desktop_id
-			  ) < $8::integer
-			  AND (
-				dt.target_client_id = $3
-				OR ($4::boolean AND dt.target_account_id = ANY($5::uuid[]))
+			  ) < $6::integer
+			  AND NOT EXISTS (
+				SELECT 1
+				FROM desktop_tasks AS active_client_platform
+				WHERE active_client_platform.tenant_id = dt.tenant_id
+				  AND active_client_platform.workspace_id = dt.workspace_id
+				  AND active_client_platform.kind = 'monitor'
+				  AND active_client_platform.target_client_id = $3
+				  AND active_client_platform.platform_id = dt.platform_id
+				  AND active_client_platform.status = 'in_progress'
+				  AND active_client_platform.desktop_id <> dt.desktop_id
 			  )
 			ORDER BY dt.lane_weight DESC,
 				dt.priority DESC,
 				COALESCE(dt.enqueued_at, dt.created_at) ASC,
 				dt.created_at ASC,
-				dt.desktop_id ASC
+				dt.desktop_id ASC,
+				ca.desktop_id ASC
+			LIMIT ` + monitorLeaseCandidateScanLimitSQL + `
+			FOR UPDATE OF dt SKIP LOCKED
+		),
+		candidate AS (
+			SELECT desktop_id, lease_account_id
+			FROM candidate_pool
+			WHERE pg_try_advisory_xact_lock(account_slot_lock_key)
+			ORDER BY lane_weight DESC,
+				priority DESC,
+				COALESCE(enqueued_at, created_at) ASC,
+				created_at ASC,
+				desktop_id ASC,
+				lease_account_id ASC
 			LIMIT 1
-			FOR UPDATE SKIP LOCKED
 		)
 		UPDATE desktop_tasks AS t
 		SET target_client_id = $3,
-		    active_attempt_id = $6,
-		    lease_token_hash = $7,
+		    target_account_id = candidate.lease_account_id,
+		    active_attempt_id = $4,
+		    lease_token_hash = $5,
 		    lease_expires_at = now() + interval '10 minutes',
 		    status = 'in_progress',
 		    attempts = t.attempts + 1,
@@ -620,6 +668,51 @@ func leaseNextQueuedMonitorTaskSQL() string {
 		RETURNING ` + desktopTaskRepositoryReturningColumns
 }
 
+func monitorAccountLeaseLockSQL(alias string) string {
+	return fmt.Sprintf(`hashtextextended(
+				concat_ws(
+					':',
+					'monitor_desktop_account_slot',
+					%s.tenant_id::text,
+					%s.workspace_id::text,
+					%s.user_id::text,
+					%s.platform_id,
+					CASE
+						WHEN btrim(COALESCE(%s.account_fingerprint, '')) <> '' THEN 'fp:' || %s.account_fingerprint
+						WHEN btrim(COALESCE(%s.platform_uid, '')) <> '' THEN 'uid:' || %s.platform_uid
+						ELSE 'id:' || %s.desktop_id::text
+					END
+				),
+				0
+			)`,
+		alias,
+		alias,
+		alias,
+		alias,
+		alias, alias,
+		alias, alias,
+		alias,
+	)
+}
+
+func monitorAccountIdentityMatchSQL(leftAlias, rightAlias string) string {
+	return fmt.Sprintf(`(
+				%s.desktop_id = %s.desktop_id
+				OR (
+					btrim(COALESCE(%s.account_fingerprint, '')) <> ''
+					AND %s.account_fingerprint = %s.account_fingerprint
+				)
+				OR (
+					btrim(COALESCE(%s.platform_uid, '')) <> ''
+					AND %s.platform_uid = %s.platform_uid
+				)
+			)`,
+		leftAlias, rightAlias,
+		leftAlias, leftAlias, rightAlias,
+		leftAlias, leftAlias, rightAlias,
+	)
+}
+
 func (s *DesktopTaskService) leaseNextQueuedPublishTask(
 	ctx context.Context,
 	client *repository.DesktopClient,
@@ -683,12 +776,6 @@ func (s *DesktopTaskService) leaseQueuedDesktopTaskByID(
 	desktopID uuid.UUID,
 	params repository.DesktopTaskLeaseParams,
 ) (*repository.DesktopTask, error) {
-	accountIDs, err := s.loadMonitorLeaseAccountIDs(ctx, client)
-	if err != nil {
-		return nil, err
-	}
-	hasAccountIDs := len(accountIDs) > 0
-
 	tx, err := beginDesktopMonitorLeaseTx(ctx, s.pool, client.ID)
 	if err != nil {
 		return nil, err
@@ -700,8 +787,6 @@ func (s *DesktopTaskService) leaseQueuedDesktopTaskByID(
 		client.WorkspaceID,
 		client.TenantID,
 		client.ID,
-		hasAccountIDs,
-		accountIDs,
 		params.AttemptID,
 		params.LeaseTokenHash,
 		desktopPublishMaxAttempts,
@@ -719,13 +804,35 @@ func (s *DesktopTaskService) leaseQueuedDesktopTaskByID(
 
 func leaseQueuedDesktopTaskByIDSQL() string {
 	return `
-		WITH candidate AS (
-			SELECT dt.desktop_id
+		WITH current_accounts AS (
+			SELECT desktop_id, user_id, platform_id, platform_uid, account_fingerprint
+			FROM platform_accounts
+			WHERE tenant_id = $3
+			  AND workspace_id = $2
+			  AND client_id = $4
+			  AND deleted_at IS NULL
+			  AND delete_requested_at IS NULL
+		),
+		candidate AS (
+			SELECT dt.desktop_id,
+				dt.kind,
+				ca.desktop_id AS lease_account_id,
+				` + monitorAccountLeaseLockSQL("target_account") + ` AS account_slot_lock_key
 			FROM desktop_tasks AS dt
+			LEFT JOIN platform_accounts AS target_account
+			  ON dt.kind = 'monitor'
+			 AND target_account.desktop_id = dt.target_account_id
+			 AND target_account.tenant_id = dt.tenant_id
+			 AND target_account.workspace_id = dt.workspace_id
+			LEFT JOIN current_accounts AS ca
+			  ON dt.kind = 'monitor'
+			 AND ca.user_id = target_account.user_id
+			 AND ca.platform_id = dt.platform_id
+			 AND ` + monitorAccountIdentityMatchSQL("ca", "target_account") + `
 			WHERE dt.desktop_id = $1
 			  AND dt.workspace_id = $2
 			  AND dt.status = 'queued'
-			  AND (dt.kind <> 'publish' OR dt.attempts < $9)
+			  AND (dt.kind <> 'publish' OR dt.attempts < $7)
 			  AND (
 				dt.kind <> 'publish'
 				OR EXISTS (
@@ -738,18 +845,27 @@ func leaseQueuedDesktopTaskByIDSQL() string {
 					  AND pa.delete_requested_at IS NULL
 				)
 			  )
+			  AND (
+				dt.kind <> 'monitor'
+				OR ca.desktop_id IS NOT NULL
+			  )
 			  AND (
 				dt.kind <> 'monitor'
 				OR NOT EXISTS (
 					SELECT 1
 					FROM desktop_tasks AS active
+					JOIN platform_accounts AS active_account
+					  ON active_account.desktop_id = active.target_account_id
+					 AND active_account.tenant_id = active.tenant_id
+					 AND active_account.workspace_id = active.workspace_id
 					WHERE active.tenant_id = dt.tenant_id
 					  AND active.workspace_id = dt.workspace_id
 					  AND active.kind = 'monitor'
-					  AND active.target_client_id = $4
 					  AND active.platform_id = dt.platform_id
 					  AND active.status = 'in_progress'
 					  AND active.desktop_id <> dt.desktop_id
+					  AND active_account.user_id = target_account.user_id
+					  AND ` + monitorAccountIdentityMatchSQL("active_account", "target_account") + `
 				)
 			  )
 			  AND (
@@ -757,16 +873,21 @@ func leaseQueuedDesktopTaskByIDSQL() string {
 				OR NOT EXISTS (
 					SELECT 1
 					FROM desktop_tasks AS recent
+					JOIN platform_accounts AS recent_account
+					  ON recent_account.desktop_id = recent.target_account_id
+					 AND recent_account.tenant_id = recent.tenant_id
+					 AND recent_account.workspace_id = recent.workspace_id
 					WHERE recent.tenant_id = dt.tenant_id
 					  AND recent.workspace_id = dt.workspace_id
 					  AND recent.kind = 'monitor'
-					  AND recent.target_client_id = $4
 					  AND recent.platform_id = dt.platform_id
 					  AND (
 						(recent.status = 'succeeded' AND recent.updated_at >= now() - interval '` + monitorSucceededTaskCooldownSQL + `')
 						OR (recent.status IN ('failed', 'unknown', 'aborted') AND recent.updated_at >= now() - interval '` + monitorUnhealthyTaskCooldownSQL + `')
 					  )
 					  AND recent.desktop_id <> dt.desktop_id
+					  AND recent_account.user_id = target_account.user_id
+					  AND ` + monitorAccountIdentityMatchSQL("recent_account", "target_account") + `
 				)
 			  )
 			  AND (
@@ -780,36 +901,55 @@ func leaseQueuedDesktopTaskByIDSQL() string {
 					  AND active_platform.target_client_id = $4
 					  AND active_platform.status = 'in_progress'
 					  AND active_platform.desktop_id <> dt.desktop_id
-				) < $10::integer
+				) < $8::integer
+			  )
+			  AND (
+				dt.kind <> 'monitor'
+				OR NOT EXISTS (
+					SELECT 1
+					FROM desktop_tasks AS active_client_platform
+					WHERE active_client_platform.tenant_id = dt.tenant_id
+					  AND active_client_platform.workspace_id = dt.workspace_id
+					  AND active_client_platform.kind = 'monitor'
+					  AND active_client_platform.target_client_id = $4
+					  AND active_client_platform.platform_id = dt.platform_id
+					  AND active_client_platform.status = 'in_progress'
+					  AND active_client_platform.desktop_id <> dt.desktop_id
+				)
 			  )
 			  AND (
 				(
 					dt.kind = 'monitor'
 					AND dt.tenant_id = $3
-					AND (
-						dt.target_client_id = $4
-						OR ($5::boolean AND dt.target_account_id = ANY($6::uuid[]))
-					)
 				)
 				OR (
 					dt.kind <> 'monitor'
 					AND dt.target_client_id = $4
 				)
 			  )
+			ORDER BY ca.desktop_id ASC NULLS LAST
+			LIMIT 1
+			FOR UPDATE OF dt SKIP LOCKED
+		),
+		locked_candidate AS (
+			SELECT desktop_id, lease_account_id
+			FROM candidate
+			WHERE kind <> 'monitor'
+			   OR pg_try_advisory_xact_lock(account_slot_lock_key)
 			LIMIT 1
-			FOR UPDATE SKIP LOCKED
 		)
 		UPDATE desktop_tasks AS t
 		SET target_client_id = CASE WHEN t.kind = 'monitor' THEN $4 ELSE t.target_client_id END,
-		    active_attempt_id = $7,
-		    lease_token_hash = $8,
+		    target_account_id = CASE WHEN t.kind = 'monitor' THEN locked_candidate.lease_account_id ELSE t.target_account_id END,
+		    active_attempt_id = $5,
+		    lease_token_hash = $6,
 		    lease_expires_at = now() + CASE WHEN t.kind = 'publish' THEN interval '3 minutes' ELSE interval '10 minutes' END,
 		    status = 'in_progress',
 		    attempts = t.attempts + 1,
 		    started_at = COALESCE(t.started_at, now()),
 		    updated_at = now()
-		FROM candidate
-		WHERE t.desktop_id = candidate.desktop_id
+		FROM locked_candidate
+		WHERE t.desktop_id = locked_candidate.desktop_id
 		RETURNING ` + desktopTaskRepositoryReturningColumns
 }
 
@@ -828,39 +968,6 @@ func beginDesktopMonitorLeaseTx(ctx context.Context, pool *pgxpool.Pool, clientI
 	return tx, nil
 }
 
-func (s *DesktopTaskService) loadMonitorLeaseAccountIDs(ctx context.Context, client *repository.DesktopClient) ([]uuid.UUID, error) {
-	if s == nil || client == nil || s.pool == nil {
-		return nil, nil
-	}
-
-	accountIDs := loadTrackedDesktopClientAccountIDs(ctx, s.redis, client.ID)
-	rows, err := s.pool.Query(ctx, `
-		SELECT desktop_id
-		FROM platform_accounts
-		WHERE tenant_id = $1
-		  AND workspace_id = $2
-		  AND client_id = $3
-		  AND deleted_at IS NULL
-		  AND platform_id = ANY($4::text[])
-	`, client.TenantID, client.WorkspaceID, client.ID, monitoringPlatformIDs(defaultMonitoringPlatforms))
-	if err != nil {
-		return nil, response.ErrInternal(50117, "desktop_account_lookup_failed", "failed to resolve monitor lease accounts")
-	}
-	defer rows.Close()
-
-	for rows.Next() {
-		var accountID uuid.UUID
-		if scanErr := rows.Scan(&accountID); scanErr != nil {
-			return nil, response.ErrInternal(50117, "desktop_account_lookup_failed", "failed to parse monitor lease accounts")
-		}
-		accountIDs = append(accountIDs, accountID)
-	}
-	if err := rows.Err(); err != nil {
-		return nil, response.ErrInternal(50117, "desktop_account_lookup_failed", "failed to iterate monitor lease accounts")
-	}
-	return uniqueUUIDs(accountIDs), nil
-}
-
 func scanRepositoryDesktopTask(row desktopTaskRepositoryScanner) (*repository.DesktopTask, error) {
 	var (
 		task              repository.DesktopTask
@@ -1662,16 +1769,28 @@ func (s *DesktopTaskService) clientCanLeaseMonitorTask(ctx context.Context, clie
 	if s == nil || client == nil || task == nil || task.Kind != "monitor" || task.TargetAccountID == uuid.Nil {
 		return false
 	}
-	accountIDs, err := s.loadMonitorLeaseAccountIDs(ctx, client)
-	if err != nil {
+	var canLease bool
+	if err := s.pool.QueryRow(ctx, `
+		SELECT EXISTS (
+			SELECT 1
+			FROM platform_accounts AS target_account
+			JOIN platform_accounts AS client_account
+			  ON client_account.tenant_id = target_account.tenant_id
+			 AND client_account.workspace_id = target_account.workspace_id
+			 AND client_account.user_id = target_account.user_id
+			 AND client_account.platform_id = target_account.platform_id
+			 AND client_account.client_id = $4
+			 AND client_account.deleted_at IS NULL
+			 AND client_account.delete_requested_at IS NULL
+			 AND `+monitorAccountIdentityMatchSQL("client_account", "target_account")+`
+			WHERE target_account.desktop_id = $1
+			  AND target_account.tenant_id = $2
+			  AND target_account.workspace_id = $3
+		)
+	`, task.TargetAccountID, task.TenantID, task.WorkspaceID, client.ID).Scan(&canLease); err != nil {
 		return false
 	}
-	for _, accountID := range accountIDs {
-		if accountID == task.TargetAccountID {
-			return true
-		}
-	}
-	return false
+	return canLease
 }
 
 func buildDesktopTaskView(task *repository.DesktopTask) DesktopTaskView {

server/internal/tenant/app/desktop_task_service_test.go

@@ -239,19 +239,28 @@ func TestBuildCountPublishTasksByStatusesQueryHidesDeletedPublishRecords(t *test
 	}
 }
 
-func TestLeaseNextQueuedMonitorTaskSQLSerializesPerClientPlatform(t *testing.T) {
+func TestLeaseNextQueuedMonitorTaskSQLUsesAccountIdentityAndClientSlots(t *testing.T) {
 	t.Parallel()
 
 	normalized := normalizeSQLForDesktopTaskTest(leaseNextQueuedMonitorTaskSQL())
 	for _, fragment := range []string{
+		"WITH current_accounts AS",
+		"client_id = $3",
+		"dt.target_account_id = target_account.desktop_id",
+		"target_account.account_fingerprint",
+		"target_account.platform_uid",
+		"hashtextextended",
+		"pg_try_advisory_xact_lock(account_slot_lock_key)",
+		"SET target_client_id = $3, target_account_id = candidate.lease_account_id",
 		"NOT EXISTS ( SELECT 1 FROM desktop_tasks AS active",
-		"active.target_client_id = $3",
 		"active.platform_id = dt.platform_id",
 		"active.status = 'in_progress'",
 		"active.desktop_id <> dt.desktop_id",
 		"COUNT(DISTINCT active_platform.platform_id)",
 		"active_platform.target_client_id = $3",
-		"< $8::integer",
+		"< $6::integer",
+		"active_client_platform.target_client_id = $3",
+		"active_client_platform.platform_id = dt.platform_id",
 		"recent.status = 'succeeded'",
 		"recent.updated_at >= now() - interval '2 seconds'",
 		"recent.status IN ('failed', 'unknown', 'aborted')",
@@ -264,19 +273,30 @@ func TestLeaseNextQueuedMonitorTaskSQLSerializesPerClientPlatform(t *testing.T)
 	}
 }
 
-func TestLeaseQueuedDesktopTaskByIDSQLSerializesMonitorPerClientPlatform(t *testing.T) {
+func TestLeaseQueuedDesktopTaskByIDSQLUsesAccountIdentityAndClientSlots(t *testing.T) {
 	t.Parallel()
 
 	normalized := normalizeSQLForDesktopTaskTest(leaseQueuedDesktopTaskByIDSQL())
 	for _, fragment := range []string{
+		"WITH current_accounts AS",
+		"client_id = $4",
+		"dt.kind <> 'monitor' OR ca.desktop_id IS NOT NULL",
+		"target_account.account_fingerprint",
+		"target_account.platform_uid",
+		"SELECT dt.desktop_id, dt.kind, ca.desktop_id AS lease_account_id",
+		"locked_candidate AS",
+		"pg_try_advisory_xact_lock(account_slot_lock_key)",
+		"FROM locked_candidate WHERE t.desktop_id = locked_candidate.desktop_id",
+		"target_account_id = CASE WHEN t.kind = 'monitor' THEN locked_candidate.lease_account_id ELSE t.target_account_id END",
 		"dt.kind <> 'monitor' OR NOT EXISTS",
-		"active.target_client_id = $4",
 		"active.platform_id = dt.platform_id",
 		"active.status = 'in_progress'",
 		"active.desktop_id <> dt.desktop_id",
 		"COUNT(DISTINCT active_platform.platform_id)",
 		"active_platform.target_client_id = $4",
-		"< $10::integer",
+		"< $8::integer",
+		"active_client_platform.target_client_id = $4",
+		"active_client_platform.platform_id = dt.platform_id",
 		"recent.status = 'succeeded'",
 		"recent.updated_at >= now() - interval '2 seconds'",
 		"recent.status IN ('failed', 'unknown', 'aborted')",