diff --git a/server/cmd/scheduler/main.go b/server/cmd/scheduler/main.go index 06b80b0..e4aa002 100644 --- a/server/cmd/scheduler/main.go +++ b/server/cmd/scheduler/main.go @@ -76,7 +76,7 @@ func main() { app.Config.MonitoringDispatch, app.Config.BrandLibrary, app.Logger, - ) + ).WithRedis(app.Redis) ctx, cancel := context.WithCancel(context.Background()) defer cancel() diff --git a/server/cmd/tenant-api/main.go b/server/cmd/tenant-api/main.go index 454ee0c..76902d9 100644 --- a/server/cmd/tenant-api/main.go +++ b/server/cmd/tenant-api/main.go @@ -32,7 +32,7 @@ func main() { app.DesktopTaskStreams.Run(workerCtx) app.DesktopDispatch.Run(workerCtx) - monitoringService := tenantapp.NewMonitoringService(app.DB, app.MonitoringDB, app.RabbitMQ, app.Config.MonitoringDispatch, app.Config.BrandLibrary, app.Logger) + monitoringService := tenantapp.NewMonitoringService(app.DB, app.MonitoringDB, app.RabbitMQ, app.Config.MonitoringDispatch, app.Config.BrandLibrary, app.Logger).WithRedis(app.Redis) monitoringCallbackService := tenantapp.NewMonitoringCallbackService(app.DB, app.MonitoringDB, app.RabbitMQ, app.LLM, app.Logger) tenantapp.NewMonitoringCollectOutboxWorker(monitoringService, app.Logger).Start(workerCtx) tenantapp.NewMonitoringResultIngestWorker(monitoringCallbackService, app.RabbitMQ, app.Logger, app.Config.MonitoringWorkers).Start(workerCtx) diff --git a/server/internal/tenant/app/desktop_task_service.go b/server/internal/tenant/app/desktop_task_service.go index c5e50ea..18681f4 100644 --- a/server/internal/tenant/app/desktop_task_service.go +++ b/server/internal/tenant/app/desktop_task_service.go @@ -11,6 +11,7 @@ import ( "github.com/jackc/pgx/v5" "github.com/jackc/pgx/v5/pgtype" "github.com/jackc/pgx/v5/pgxpool" + goredis "github.com/redis/go-redis/v9" "go.uber.org/zap" "github.com/geo-platform/tenant-api/internal/shared/auth" @@ -26,6 +27,7 @@ type DesktopTaskService struct { monitoringPool *pgxpool.Pool repo repository.DesktopTaskRepository cache sharedcache.Cache + redis *goredis.Client messaging *rabbitmq.Client logger *zap.Logger } @@ -45,6 +47,13 @@ func (s *DesktopTaskService) WithCache(c sharedcache.Cache) *DesktopTaskService return s } +func (s *DesktopTaskService) WithRedis(redis *goredis.Client) *DesktopTaskService { + if s != nil { + s.redis = redis + } + return s +} + type DesktopTaskView struct { ID string `json:"id"` JobID string `json:"job_id"` @@ -163,7 +172,9 @@ func (s *DesktopTaskService) Lease(ctx context.Context, client *repository.Deskt var task *repository.DesktopTask switch { case taskID != nil: - task, err = s.repo.LeaseQueuedByDesktopID(ctx, *taskID, client.ID, leaseParams) + task, err = s.leaseQueuedDesktopTaskByID(ctx, client, *taskID, leaseParams) + case req.Kind != nil && strings.TrimSpace(*req.Kind) == "monitor": + task, err = s.leaseNextQueuedMonitorTask(ctx, client, leaseParams) default: task, err = s.repo.LeaseNextQueued(ctx, client.ID, req.Kind, leaseParams) } @@ -219,6 +230,288 @@ func (s *DesktopTaskService) Lease(ctx context.Context, client *repository.Deskt }, nil } +const desktopTaskRepositoryReturningColumns = ` + t.desktop_id, + t.job_id, + t.tenant_id, + t.workspace_id, + t.target_account_id, + t.target_client_id, + t.platform_id, + t.kind, + t.payload, + t.status, + t.priority, + t.lane, + t.lane_weight, + t.source, + t.scheduler_group_key, + t.monitor_task_id, + t.supersedes_task_id, + t.control_flags, + t.interrupt_generation, + t.dedup_key, + t.active_attempt_id, + t.lease_expires_at, + t.attempts, + t.result, + t.error, + t.started_at, + t.interrupted_at, + t.interrupt_reason, + t.enqueued_at, + t.created_at, + t.updated_at` + +type desktopTaskRepositoryScanner interface { + Scan(dest ...any) error +} + +func (s *DesktopTaskService) leaseNextQueuedMonitorTask( + ctx context.Context, + 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 + + row := s.pool.QueryRow(ctx, ` + 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 ( + dt.target_client_id = $3 + OR ($4::boolean AND dt.target_account_id = ANY($5::uuid[])) + ) + 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 + LIMIT 1 + FOR UPDATE SKIP LOCKED + ) + UPDATE desktop_tasks AS t + SET target_client_id = $3, + active_attempt_id = $6, + lease_token_hash = $7, + lease_expires_at = now() + interval '10 minutes', + status = 'in_progress', + attempts = t.attempts + 1, + updated_at = now() + FROM candidate + WHERE t.desktop_id = candidate.desktop_id + RETURNING `+desktopTaskRepositoryReturningColumns, + client.TenantID, + client.WorkspaceID, + client.ID, + hasAccountIDs, + accountIDs, + params.AttemptID, + params.LeaseTokenHash, + ) + return scanRepositoryDesktopTask(row) +} + +func (s *DesktopTaskService) leaseQueuedDesktopTaskByID( + ctx context.Context, + client *repository.DesktopClient, + 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 + + row := s.pool.QueryRow(ctx, ` + WITH candidate AS ( + SELECT dt.desktop_id + FROM desktop_tasks AS dt + WHERE dt.desktop_id = $1 + AND dt.workspace_id = $2 + AND dt.status = 'queued' + 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 + ) + ) + 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, + lease_expires_at = now() + interval '10 minutes', + status = 'in_progress', + attempts = t.attempts + 1, + updated_at = now() + FROM candidate + WHERE t.desktop_id = candidate.desktop_id + RETURNING `+desktopTaskRepositoryReturningColumns, + desktopID, + client.WorkspaceID, + client.TenantID, + client.ID, + hasAccountIDs, + accountIDs, + params.AttemptID, + params.LeaseTokenHash, + ) + return scanRepositoryDesktopTask(row) +} + +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 + schedulerGroupKey pgtype.Text + monitorTaskID pgtype.Int8 + supersedesTaskID pgtype.UUID + dedupKey pgtype.Text + activeAttemptID pgtype.UUID + leaseExpiresAt pgtype.Timestamptz + startedAt pgtype.Timestamptz + interruptedAt pgtype.Timestamptz + interruptReason pgtype.Text + enqueuedAt pgtype.Timestamptz + ) + if err := row.Scan( + &task.DesktopID, + &task.JobID, + &task.TenantID, + &task.WorkspaceID, + &task.TargetAccountID, + &task.TargetClientID, + &task.Platform, + &task.Kind, + &task.Payload, + &task.Status, + &task.Priority, + &task.Lane, + &task.LaneWeight, + &task.Source, + &schedulerGroupKey, + &monitorTaskID, + &supersedesTaskID, + &task.ControlFlags, + &task.InterruptGeneration, + &dedupKey, + &activeAttemptID, + &leaseExpiresAt, + &task.Attempts, + &task.Result, + &task.Error, + &startedAt, + &interruptedAt, + &interruptReason, + &enqueuedAt, + &task.CreatedAt, + &task.UpdatedAt, + ); err != nil { + return nil, err + } + + task.SchedulerGroupKey = desktopTaskNullableText(schedulerGroupKey) + task.MonitorTaskID = desktopTaskNullableInt64(monitorTaskID) + task.SupersedesTaskID = desktopTaskNullableUUID(supersedesTaskID) + task.DedupKey = desktopTaskNullableText(dedupKey) + task.ActiveAttemptID = desktopTaskNullableUUID(activeAttemptID) + task.LeaseExpiresAt = desktopTaskNullableTime(leaseExpiresAt) + task.StartedAt = desktopTaskNullableTime(startedAt) + task.InterruptedAt = desktopTaskNullableTime(interruptedAt) + task.InterruptReason = desktopTaskNullableText(interruptReason) + if enqueuedAt.Valid { + task.EnqueuedAt = enqueuedAt.Time + } + return &task, nil +} + +func desktopTaskNullableText(value pgtype.Text) *string { + if !value.Valid { + return nil + } + text := value.String + return &text +} + +func desktopTaskNullableInt64(value pgtype.Int8) *int64 { + if !value.Valid { + return nil + } + number := value.Int64 + return &number +} + +func desktopTaskNullableUUID(value pgtype.UUID) *uuid.UUID { + if !value.Valid { + return nil + } + resolved, err := uuid.FromBytes(value.Bytes[:]) + if err != nil { + return nil + } + return &resolved +} + +func desktopTaskNullableTime(value pgtype.Timestamptz) *time.Time { + if !value.Valid { + return nil + } + timestamp := value.Time + return ×tamp +} + func (s *DesktopTaskService) Extend(ctx context.Context, client *repository.DesktopClient, desktopID uuid.UUID, req ExtendDesktopTaskRequest) (*DesktopTaskView, error) { if client == nil { return nil, response.ErrUnauthorized(40108, "desktop_client_missing", "desktop client context is required") @@ -731,7 +1024,7 @@ func (s *DesktopTaskService) classifyLeaseError(ctx context.Context, client *rep return response.ErrInternal(50095, "desktop_task_lookup_failed", "failed to inspect desktop task state") } - if task.TargetClientID != client.ID { + if task.TargetClientID != client.ID && !s.clientCanLeaseMonitorTask(ctx, client, task) { return response.ErrForbidden(40381, "desktop_task_not_target_client", "desktop task is assigned to another desktop client") } @@ -742,6 +1035,22 @@ func (s *DesktopTaskService) classifyLeaseError(ctx context.Context, client *rep return response.ErrConflict(40991, "desktop_task_not_queued", "desktop task is not queued") } +func (s *DesktopTaskService) clientCanLeaseMonitorTask(ctx context.Context, client *repository.DesktopClient, task *repository.DesktopTask) bool { + 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 { + return false + } + for _, accountID := range accountIDs { + if accountID == task.TargetAccountID { + return true + } + } + return false +} + func buildDesktopTaskView(task *repository.DesktopTask) DesktopTaskView { var activeAttemptID *string if task.ActiveAttemptID != nil { diff --git a/server/internal/tenant/app/monitoring_daily_task_worker.go b/server/internal/tenant/app/monitoring_daily_task_worker.go index 9515b9a..bcf638d 100644 --- a/server/internal/tenant/app/monitoring_daily_task_worker.go +++ b/server/internal/tenant/app/monitoring_daily_task_worker.go @@ -374,7 +374,7 @@ func (s *MonitoringService) generateDailyMonitoringTasksForBrand( if materializeLimit <= 0 && dispatchLimit <= 0 { return 0, 0, 0, 0, nil } - if executionOwner != "desktop_tasks" || plan.PrimaryClientID == nil { + if executionOwner != "desktop_tasks" { return 0, 0, 0, 0, nil } @@ -435,7 +435,6 @@ func (s *MonitoringService) generateDailyMonitoringTasksForBrand( businessDay, now.Add(defaultMonitoringDailyLookahead), now.Add(-defaultMonitoringDailyDesktopRetryCooldown), - *plan.PrimaryClientID, dispatchLimit, ) if err != nil { @@ -577,7 +576,9 @@ func (s *MonitoringService) loadDailyMonitoringPrimaryClientPlans( ORDER BY c.tenant_id ASC, c.workspace_id ASC `, freeBrandLimit, paidBrandLimit, questionLimit, supportedPlatforms) if err != nil { - return nil, response.ErrInternal(50041, "query_failed", "failed to load daily monitoring plans") + appErr := response.ErrInternal(50041, "query_failed", "failed to load daily monitoring plans") + appErr.Cause = err + return nil, appErr } defer rows.Close() @@ -661,7 +662,9 @@ func (s *MonitoringService) loadDailyMonitoringPlanPlatformJSON( ORDER BY p.tenant_id ASC, p.workspace_id ASC, p.client_id ASC `, tenantIDs, workspaceIDs, clientIDs, supportedPlatforms) if err != nil { - return nil, response.ErrInternal(50041, "query_failed", "failed to load daily monitoring plan platforms") + appErr := response.ErrInternal(50041, "query_failed", "failed to load daily monitoring plan platforms") + appErr.Cause = err + return nil, appErr } defer rows.Close() @@ -865,15 +868,15 @@ func (s *MonitoringService) ensureDailyMonitoringCollectTasks( 'pending', 100, 'normal', - $5, + NULL, input.dispatch_after, - $6, - $7 + $5, + $6 FROM unnest( - $8::bigint[], - $9::bytea[], - $10::text[], - $11::timestamptz[] + $7::bigint[], + $8::bytea[], + $9::text[], + $10::timestamptz[] ) AS input(question_id, question_hash, ai_platform_id, dispatch_after) ON CONFLICT ( tenant_id, brand_id, question_id, ai_platform_id, collector_type, run_mode, business_date @@ -884,7 +887,6 @@ func (s *MonitoringService) ensureDailyMonitoringCollectTasks( brand.BrandID, monitoringCollectorType, businessDate, - nullableUUID(plan.PrimaryClientID), shardKey, executionOwner, questionIDs, @@ -906,7 +908,6 @@ func (s *MonitoringService) loadDueDailyMonitorDesktopTaskSpecs( businessDate time.Time, cutoff time.Time, retryBefore time.Time, - targetClientID uuid.UUID, limit int, ) ([]monitorDesktopTaskSpec, error) { if q == nil || limit <= 0 { @@ -925,11 +926,10 @@ func (s *MonitoringService) loadDueDailyMonitorDesktopTaskSpecs( AND t.trigger_source = 'automatic' AND t.run_mode = 'desktop_standard' AND COALESCE(t.execution_owner, 'legacy') = 'desktop_tasks' - AND t.target_client_id = $5 - AND (t.dispatch_after IS NULL OR t.dispatch_after <= $6) - AND (t.last_dispatched_at IS NULL OR t.last_dispatched_at <= $7) + AND (t.dispatch_after IS NULL OR t.dispatch_after <= $5) + AND (t.last_dispatched_at IS NULL OR t.last_dispatched_at <= $6) ORDER BY t.dispatch_after ASC NULLS FIRST, t.dispatch_priority DESC, t.id ASC - LIMIT $8 + LIMIT $7 FOR UPDATE SKIP LOCKED ), claimed_tasks AS ( @@ -964,9 +964,9 @@ func (s *MonitoringService) loadDueDailyMonitorDesktopTaskSpecs( ORDER BY s.projected_at DESC, s.id DESC LIMIT 1 ), '') AS question_text_snapshot - FROM claimed_tasks t - ORDER BY t.dispatch_after ASC NULLS FIRST, t.dispatch_priority DESC, t.id ASC - `, tenantID, brandID, monitoringCollectorType, monitoringBusinessDateText(businessDate), targetClientID, cutoff.UTC(), retryBefore.UTC(), limit) + FROM claimed_tasks t + ORDER BY t.dispatch_after ASC NULLS FIRST, t.dispatch_priority DESC, t.id ASC + `, tenantID, brandID, monitoringCollectorType, monitoringBusinessDateText(businessDate), cutoff.UTC(), retryBefore.UTC(), limit) if err != nil { return nil, response.ErrInternal(50041, "query_failed", "failed to load due daily monitor desktop task specs") } @@ -996,7 +996,6 @@ func (s *MonitoringService) loadDueDailyMonitorDesktopTaskSpecs( return nil, response.ErrInternal(50041, "scan_failed", "failed to parse due daily monitor desktop task specs") } spec.WorkspaceID = workspaceID - spec.TargetClientID = targetClientID spec.BusinessDate = businessDay.Format("2006-01-02") spec.QuestionHash = encodeQuestionHash(questionHash) spec.SchedulerGroupKey = monitoringSchedulerGroupKey(spec.QuestionHash, spec.QuestionID, spec.QuestionText) diff --git a/server/internal/tenant/app/monitoring_phase2_desktop_tasks.go b/server/internal/tenant/app/monitoring_phase2_desktop_tasks.go index 9070470..4ea9759 100644 --- a/server/internal/tenant/app/monitoring_phase2_desktop_tasks.go +++ b/server/internal/tenant/app/monitoring_phase2_desktop_tasks.go @@ -26,6 +26,7 @@ type monitorDesktopTaskSpec struct { MonitorTaskID int64 TenantID int64 WorkspaceID int64 + TargetAccountID uuid.UUID TargetClientID uuid.UUID PlatformID string BusinessDate string @@ -40,6 +41,19 @@ type monitorDesktopTaskSpec struct { InterruptGeneration int } +type monitorDesktopAccountCandidate struct { + PlatformID string + AccountID uuid.UUID + DBClientID *uuid.UUID + HealthRank int + RecordOrder int +} + +type monitorDesktopTaskTarget struct { + AccountID uuid.UUID + ClientID uuid.UUID +} + type monitorDesktopInterruptTarget struct { TaskID string MonitorTaskID int64 @@ -181,7 +195,7 @@ func (s *MonitoringService) dispatchMonitorDesktopTasks( return nil, nil, nil } - accountMap, err := s.loadMonitorDesktopAccountMap(ctx, specs[0].TenantID, specs[0].WorkspaceID, specs[0].TargetClientID) + targets, err := s.loadMonitorDesktopTaskTargets(ctx, specs[0].TenantID, specs[0].WorkspaceID, specs) if err != nil { return nil, nil, err } @@ -196,7 +210,15 @@ func (s *MonitoringService) dispatchMonitorDesktopTasks( publishableTasks := make([]*repository.DesktopTask, 0, len(specs)) fallbackLegacyTaskIDs := make([]int64, 0) for _, spec := range specs { - task, shouldPublish, fallbackLegacy, upsertErr := s.upsertMonitorDesktopTask(ctx, tx, repo, spec, accountMap) + target, ok := targets[normalizeMonitoringPlatformID(spec.PlatformID)] + if !ok || target.AccountID == uuid.Nil || target.ClientID == uuid.Nil { + fallbackLegacyTaskIDs = append(fallbackLegacyTaskIDs, spec.MonitorTaskID) + continue + } + spec.TargetAccountID = target.AccountID + spec.TargetClientID = target.ClientID + + task, shouldPublish, fallbackLegacy, upsertErr := s.upsertMonitorDesktopTask(ctx, tx, repo, spec) if upsertErr != nil { return nil, nil, upsertErr } @@ -215,20 +237,32 @@ func (s *MonitoringService) dispatchMonitorDesktopTasks( return publishableTasks, fallbackLegacyTaskIDs, nil } -func (s *MonitoringService) loadMonitorDesktopAccountMap( +func (s *MonitoringService) loadMonitorDesktopTaskTargets( ctx context.Context, tenantID, workspaceID int64, - targetClientID uuid.UUID, -) (map[string]uuid.UUID, error) { + specs []monitorDesktopTaskSpec, +) (map[string]monitorDesktopTaskTarget, error) { + platformIDs := uniqueMonitorDesktopTaskPlatformIDs(specs) + if len(platformIDs) == 0 { + return nil, nil + } + rows, err := s.businessPool.Query(ctx, ` - SELECT DISTINCT ON (platform_id) + SELECT platform_id, - desktop_id + desktop_id, + COALESCE(client_id::text, '') AS client_id, + CASE health + WHEN 'live' THEN 0 + WHEN 'risk' THEN 1 + WHEN 'captcha' THEN 2 + ELSE 3 + END AS health_rank FROM platform_accounts WHERE tenant_id = $1 AND workspace_id = $2 - AND client_id = $3 AND deleted_at IS NULL + AND platform_id = ANY($3::text[]) ORDER BY platform_id ASC, CASE health WHEN 'live' THEN 0 @@ -239,35 +273,158 @@ func (s *MonitoringService) loadMonitorDesktopAccountMap( verified_at DESC NULLS LAST, updated_at DESC, created_at DESC - `, tenantID, workspaceID, targetClientID) + `, tenantID, workspaceID, platformIDs) if err != nil { return nil, response.ErrInternal(50117, "desktop_account_lookup_failed", "failed to resolve monitor desktop accounts") } defer rows.Close() - result := make(map[string]uuid.UUID) + candidates := make([]monitorDesktopAccountCandidate, 0) for rows.Next() { var ( - platformID string - desktopID uuid.UUID + item monitorDesktopAccountCandidate + clientIDText string ) - if scanErr := rows.Scan(&platformID, &desktopID); scanErr != nil { + if scanErr := rows.Scan(&item.PlatformID, &item.AccountID, &clientIDText, &item.HealthRank); scanErr != nil { return nil, response.ErrInternal(50117, "desktop_account_lookup_failed", "failed to parse monitor desktop accounts") } - result[platformID] = desktopID + if parsedClientID, parseErr := uuid.Parse(strings.TrimSpace(clientIDText)); parseErr == nil { + item.DBClientID = &parsedClientID + } + item.PlatformID = normalizeMonitoringPlatformID(item.PlatformID) + item.RecordOrder = len(candidates) + candidates = append(candidates, item) } if err := rows.Err(); err != nil { return nil, response.ErrInternal(50117, "desktop_account_lookup_failed", "failed to iterate monitor desktop accounts") } + + accountIDs := monitorDesktopAccountCandidateAccountIDs(candidates) + accountPresence := loadDesktopAccountPresence(ctx, s.redis, accountIDs) + clientIDs := monitorDesktopAccountCandidateClientIDs(candidates, accountPresence) + onlineClients, err := s.loadOnlineMonitorDesktopClients(ctx, workspaceID, clientIDs) + if err != nil { + return nil, err + } + return selectMonitorDesktopTaskTargets(candidates, accountPresence, onlineClients), nil +} + +func uniqueMonitorDesktopTaskPlatformIDs(specs []monitorDesktopTaskSpec) []string { + if len(specs) == 0 { + return nil + } + seen := make(map[string]struct{}, len(specs)) + result := make([]string, 0, len(specs)) + for _, spec := range specs { + platformID := normalizeMonitoringPlatformID(spec.PlatformID) + if platformID == "" { + continue + } + if _, ok := seen[platformID]; ok { + continue + } + seen[platformID] = struct{}{} + result = append(result, platformID) + } + return result +} + +func monitorDesktopAccountCandidateAccountIDs(candidates []monitorDesktopAccountCandidate) []uuid.UUID { + result := make([]uuid.UUID, 0, len(candidates)) + for _, candidate := range candidates { + if candidate.AccountID != uuid.Nil { + result = append(result, candidate.AccountID) + } + } + return uniqueUUIDs(result) +} + +func monitorDesktopAccountCandidateClientIDs(candidates []monitorDesktopAccountCandidate, accountPresence map[uuid.UUID]uuid.UUID) []uuid.UUID { + result := make([]uuid.UUID, 0, len(candidates)*2) + for _, candidate := range candidates { + if candidate.DBClientID != nil && *candidate.DBClientID != uuid.Nil { + result = append(result, *candidate.DBClientID) + } + if accountPresence != nil { + if clientID, ok := accountPresence[candidate.AccountID]; ok && clientID != uuid.Nil { + result = append(result, clientID) + } + } + } + return uniqueUUIDs(result) +} + +func (s *MonitoringService) loadOnlineMonitorDesktopClients(ctx context.Context, workspaceID int64, clientIDs []uuid.UUID) (map[uuid.UUID]bool, error) { + result := make(map[uuid.UUID]bool) + clientIDs = uniqueUUIDs(clientIDs) + if len(clientIDs) == 0 { + return result, nil + } + + for clientID, online := range loadDesktopClientPresence(ctx, s.redis, clientIDs) { + if online { + result[clientID] = true + } + } + + rows, err := s.businessPool.Query(ctx, ` + SELECT id + FROM desktop_clients + WHERE workspace_id = $1 + AND id = ANY($2::uuid[]) + AND revoked_at IS NULL + AND last_seen_at IS NOT NULL + AND last_seen_at >= NOW() - $3::interval + `, workspaceID, clientIDs, formatPgInterval(desktopClientPresenceTTL)) + if err != nil { + return nil, response.ErrInternal(50117, "desktop_client_lookup_failed", "failed to resolve online monitor desktop clients") + } + defer rows.Close() + + for rows.Next() { + var clientID uuid.UUID + if scanErr := rows.Scan(&clientID); scanErr != nil { + return nil, response.ErrInternal(50117, "desktop_client_lookup_failed", "failed to parse online monitor desktop clients") + } + result[clientID] = true + } + if err := rows.Err(); err != nil { + return nil, response.ErrInternal(50117, "desktop_client_lookup_failed", "failed to iterate online monitor desktop clients") + } return result, nil } +func selectMonitorDesktopTaskTargets( + candidates []monitorDesktopAccountCandidate, + accountPresence map[uuid.UUID]uuid.UUID, + onlineClients map[uuid.UUID]bool, +) map[string]monitorDesktopTaskTarget { + result := make(map[string]monitorDesktopTaskTarget) + for _, candidate := range candidates { + platformID := normalizeMonitoringPlatformID(candidate.PlatformID) + if platformID == "" || candidate.AccountID == uuid.Nil { + continue + } + if _, exists := result[platformID]; exists { + continue + } + + if clientID, ok := accountPresence[candidate.AccountID]; ok && clientID != uuid.Nil && onlineClients[clientID] { + result[platformID] = monitorDesktopTaskTarget{AccountID: candidate.AccountID, ClientID: clientID} + continue + } + if candidate.DBClientID != nil && *candidate.DBClientID != uuid.Nil && onlineClients[*candidate.DBClientID] { + result[platformID] = monitorDesktopTaskTarget{AccountID: candidate.AccountID, ClientID: *candidate.DBClientID} + } + } + return result +} + func (s *MonitoringService) upsertMonitorDesktopTask( ctx context.Context, tx pgx.Tx, repo repository.DesktopTaskRepository, spec monitorDesktopTaskSpec, - accountMap map[string]uuid.UUID, ) (*repository.DesktopTask, bool, bool, error) { var ( existingDesktopID uuid.UUID @@ -297,9 +454,9 @@ func (s *MonitoringService) upsertMonitorDesktopTask( if existingStatus != "queued" && !expiredInProgress { return task, false, false, nil } - targetAccountID := task.TargetAccountID - if override, ok := accountMap[spec.PlatformID]; ok { - targetAccountID = override + targetAccountID := spec.TargetAccountID + if targetAccountID == uuid.Nil { + targetAccountID = task.TargetAccountID } payloadJSON, payloadErr := marshalMonitorDesktopTaskPayload(spec) if payloadErr != nil { @@ -372,8 +529,8 @@ func (s *MonitoringService) upsertMonitorDesktopTask( return nil, false, false, response.ErrInternal(50118, "desktop_task_lookup_failed", "failed to inspect active phase2 monitor desktop task") } - targetAccountID, ok := accountMap[spec.PlatformID] - if !ok { + targetAccountID := spec.TargetAccountID + if targetAccountID == uuid.Nil || spec.TargetClientID == uuid.Nil { return nil, false, true, nil } payloadJSON, payloadErr := marshalMonitorDesktopTaskPayload(spec) diff --git a/server/internal/tenant/app/monitoring_phase2_desktop_tasks_test.go b/server/internal/tenant/app/monitoring_phase2_desktop_tasks_test.go new file mode 100644 index 0000000..8f518c7 --- /dev/null +++ b/server/internal/tenant/app/monitoring_phase2_desktop_tasks_test.go @@ -0,0 +1,78 @@ +package app + +import ( + "testing" + + "github.com/google/uuid" + "github.com/stretchr/testify/assert" +) + +func TestSelectMonitorDesktopTaskTargetsPrefersPresentAccountClient(t *testing.T) { + accountID := mustParseDailyMonitoringUUID(t, "00000000-0000-0000-0000-000000000101") + staleClientID := mustParseDailyMonitoringUUID(t, "00000000-0000-0000-0000-000000000201") + liveClientID := mustParseDailyMonitoringUUID(t, "00000000-0000-0000-0000-000000000202") + + targets := selectMonitorDesktopTaskTargets( + []monitorDesktopAccountCandidate{ + { + PlatformID: "qwen", + AccountID: accountID, + DBClientID: &staleClientID, + }, + }, + map[uuid.UUID]uuid.UUID{ + accountID: liveClientID, + }, + map[uuid.UUID]bool{ + liveClientID: true, + }, + ) + + assert.Equal(t, monitorDesktopTaskTarget{ + AccountID: accountID, + ClientID: liveClientID, + }, targets["qwen"]) +} + +func TestSelectMonitorDesktopTaskTargetsFallsBackToOnlineDBClient(t *testing.T) { + accountID := mustParseDailyMonitoringUUID(t, "00000000-0000-0000-0000-000000000102") + clientID := mustParseDailyMonitoringUUID(t, "00000000-0000-0000-0000-000000000203") + + targets := selectMonitorDesktopTaskTargets( + []monitorDesktopAccountCandidate{ + { + PlatformID: "deepseek", + AccountID: accountID, + DBClientID: &clientID, + }, + }, + nil, + map[uuid.UUID]bool{ + clientID: true, + }, + ) + + assert.Equal(t, monitorDesktopTaskTarget{ + AccountID: accountID, + ClientID: clientID, + }, targets["deepseek"]) +} + +func TestSelectMonitorDesktopTaskTargetsSkipsOfflineClients(t *testing.T) { + accountID := mustParseDailyMonitoringUUID(t, "00000000-0000-0000-0000-000000000103") + clientID := mustParseDailyMonitoringUUID(t, "00000000-0000-0000-0000-000000000204") + + targets := selectMonitorDesktopTaskTargets( + []monitorDesktopAccountCandidate{ + { + PlatformID: "kimi", + AccountID: accountID, + DBClientID: &clientID, + }, + }, + nil, + map[uuid.UUID]bool{}, + ) + + assert.Empty(t, targets) +} diff --git a/server/internal/tenant/app/monitoring_service.go b/server/internal/tenant/app/monitoring_service.go index f8e1621..fca725c 100644 --- a/server/internal/tenant/app/monitoring_service.go +++ b/server/internal/tenant/app/monitoring_service.go @@ -14,6 +14,7 @@ import ( "github.com/google/uuid" "github.com/jackc/pgx/v5" "github.com/jackc/pgx/v5/pgxpool" + goredis "github.com/redis/go-redis/v9" "go.uber.org/zap" "github.com/geo-platform/tenant-api/internal/shared/auth" @@ -40,6 +41,7 @@ type MonitoringService struct { businessPool *pgxpool.Pool monitoringPool *pgxpool.Pool rabbitMQ *rabbitmq.Client + redis *goredis.Client logger *zap.Logger desktopTasksRolloutPercentage int brandLibraryConfig config.BrandLibraryConfig @@ -62,6 +64,13 @@ func NewMonitoringService( } } +func (s *MonitoringService) WithRedis(redis *goredis.Client) *MonitoringService { + if s != nil { + s.redis = redis + } + return s +} + type MonitoringDashboardCompositeResponse struct { Overview MonitoringOverview `json:"overview"` Runtime MonitoringDashboardRuntime `json:"runtime"` diff --git a/server/internal/tenant/transport/desktop_client_handler.go b/server/internal/tenant/transport/desktop_client_handler.go index 458382e..476705c 100644 --- a/server/internal/tenant/transport/desktop_client_handler.go +++ b/server/internal/tenant/transport/desktop_client_handler.go @@ -20,7 +20,7 @@ func NewDesktopClientHandler(a *bootstrap.App) *DesktopClientHandler { a.MonitoringDB, a.RabbitMQ, a.Logger, - ).WithCache(a.Cache) + ).WithCache(a.Cache).WithRedis(a.Redis) return &DesktopClientHandler{ svc: app.NewDesktopClientService( repository.NewDesktopClientRepository(a.DB), diff --git a/server/internal/tenant/transport/desktop_task_handler.go b/server/internal/tenant/transport/desktop_task_handler.go index 6bc63a3..74e6596 100644 --- a/server/internal/tenant/transport/desktop_task_handler.go +++ b/server/internal/tenant/transport/desktop_task_handler.go @@ -27,7 +27,7 @@ func NewDesktopTaskHandler(a *bootstrap.App) *DesktopTaskHandler { a.MonitoringDB, a.RabbitMQ, a.Logger, - ).WithCache(a.Cache), + ).WithCache(a.Cache).WithRedis(a.Redis), publishSvc: app.NewPublishJobService( a.DB, a.RabbitMQ, diff --git a/server/internal/tenant/transport/monitoring_handler.go b/server/internal/tenant/transport/monitoring_handler.go index 3319569..1d43381 100644 --- a/server/internal/tenant/transport/monitoring_handler.go +++ b/server/internal/tenant/transport/monitoring_handler.go @@ -26,7 +26,7 @@ type monitoringCollectNowRequest struct { func NewMonitoringHandler(a *bootstrap.App) *MonitoringHandler { return &MonitoringHandler{ - svc: app.NewMonitoringService(a.DB, a.MonitoringDB, a.RabbitMQ, a.Config.MonitoringDispatch, a.Config.BrandLibrary, a.Logger), + svc: app.NewMonitoringService(a.DB, a.MonitoringDB, a.RabbitMQ, a.Config.MonitoringDispatch, a.Config.BrandLibrary, a.Logger).WithRedis(a.Redis), } }