package transport import ( "context" "encoding/json" "net/http" "time" "github.com/gin-gonic/gin" "github.com/gorilla/websocket" "go.uber.org/zap" "github.com/geo-platform/tenant-api/internal/bootstrap" "github.com/geo-platform/tenant-api/internal/shared/stream" tenantapp "github.com/geo-platform/tenant-api/internal/tenant/app" "github.com/geo-platform/tenant-api/internal/tenant/repository" ) const ( dispatchWriteWait = 10 * time.Second dispatchPongWait = 60 * time.Second dispatchPingPeriod = 25 * time.Second dispatchReadLimitBytes = 2048 dispatchCloseGracePause = 500 * time.Millisecond dispatchReplayLimit = 20 ) type DesktopDispatchHandler struct { app *bootstrap.App upgrader websocket.Upgrader } func NewDesktopDispatchHandler(a *bootstrap.App) *DesktopDispatchHandler { return &DesktopDispatchHandler{ app: a, upgrader: websocket.Upgrader{ ReadBufferSize: 4096, WriteBufferSize: 4096, CheckOrigin: func(_ *http.Request) bool { // Origin enforcement happens at the gateway / LB layer; this // endpoint is already gated by the bearer token middleware. return true }, }, } } // Dispatch upgrades the HTTP request into a WebSocket and streams dispatch // events for the current desktop client. The contract intentionally mirrors // the existing SSE events endpoint: connect with a bearer token, receive JSON // frames whose shape matches DesktopDispatchEvent, and reconnect on drop. func (h *DesktopDispatchHandler) Dispatch(c *gin.Context) { client := MustDesktopClient(c.Request.Context()) hub := h.app.DesktopDispatch if hub == nil { c.AbortWithStatus(http.StatusServiceUnavailable) return } conn, err := h.upgrader.Upgrade(c.Writer, c.Request, nil) if err != nil { if h.app.Logger != nil { h.app.Logger.Warn("desktop dispatch upgrade failed", zap.String("client_id", client.ID.String()), zap.Error(err), ) } return } subscriber, cancel := hub.Subscribe(client.ID.String()) defer cancel() tenantapp.MarkDesktopTaskConsumerPresent(c.Request.Context(), h.app.Redis, client.ID) conn.SetReadLimit(dispatchReadLimitBytes) _ = conn.SetReadDeadline(time.Now().Add(dispatchPongWait)) conn.SetPongHandler(func(string) error { _ = conn.SetReadDeadline(time.Now().Add(dispatchPongWait)) return nil }) // Reader goroutine: drains control frames and signals disconnect. We do // not expect application data from the client over this socket. readerDone := make(chan struct{}) go func() { defer close(readerDone) for { if _, _, err := conn.ReadMessage(); err != nil { return } } }() serverTime := time.Now().UTC() connected := stream.DesktopDispatchEvent{ Type: "connected", TargetClientID: client.ID.String(), WorkspaceID: client.WorkspaceID, UpdatedAt: serverTime, ServerTime: &serverTime, } if err := writeDispatchFrame(conn, connected); err != nil { _ = conn.Close() return } h.replayQueuedPublishTasks(c.Request.Context(), client, conn) pingTicker := time.NewTicker(dispatchPingPeriod) defer pingTicker.Stop() defer func() { _ = conn.SetWriteDeadline(time.Now().Add(dispatchWriteWait)) _ = conn.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, "")) time.Sleep(dispatchCloseGracePause) _ = conn.Close() }() outbound := subscriber.Outbound() for { select { case <-c.Request.Context().Done(): return case <-readerDone: return case event, ok := <-outbound: if !ok { return } if err := writeDispatchFrame(conn, event); err != nil { return } case <-pingTicker.C: _ = conn.SetWriteDeadline(time.Now().Add(dispatchWriteWait)) if err := conn.WriteMessage(websocket.PingMessage, nil); err != nil { return } } } } func writeDispatchFrame(conn *websocket.Conn, event stream.DesktopDispatchEvent) error { body, err := json.Marshal(event) if err != nil { return err } _ = conn.SetWriteDeadline(time.Now().Add(dispatchWriteWait)) return conn.WriteMessage(websocket.TextMessage, body) } func (h *DesktopDispatchHandler) replayQueuedPublishTasks(ctx context.Context, client *repository.DesktopClient, conn *websocket.Conn) { if h == nil || h.app == nil || h.app.DB == nil || client == nil || conn == nil { return } rows, err := h.app.DB.Query(ctx, ` SELECT dt.desktop_id, dt.job_id, dt.tenant_id, dt.workspace_id, dt.target_account_id, dt.target_client_id, dt.platform_id, dt.kind, dt.payload, dt.status, dt.priority, dt.lane, dt.lane_weight, dt.source, dt.scheduler_group_key, dt.monitor_task_id, dt.supersedes_task_id, dt.control_flags, dt.interrupt_generation, dt.dedup_key, dt.active_attempt_id, dt.lease_expires_at, dt.attempts, dt.result, dt.error, dt.started_at, dt.interrupted_at, dt.interrupt_reason, dt.enqueued_at, dt.created_at, dt.updated_at FROM desktop_tasks AS dt WHERE dt.tenant_id = $1 AND dt.workspace_id = $2 AND dt.target_client_id = $3 AND dt.kind = 'publish' AND dt.status = 'queued' AND dt.attempts < 3 AND EXISTS ( SELECT 1 FROM platform_accounts AS pa WHERE pa.desktop_id = dt.target_account_id AND pa.workspace_id = dt.workspace_id AND pa.client_id = dt.target_client_id AND pa.deleted_at IS NULL AND pa.delete_requested_at IS NULL ) AND NOT EXISTS ( SELECT 1 FROM desktop_publish_jobs AS j WHERE j.desktop_id = dt.job_id AND j.status <> 'queued' ) ORDER BY dt.priority DESC, COALESCE(dt.enqueued_at, dt.created_at) ASC, dt.created_at ASC, dt.desktop_id ASC LIMIT $4 `, client.TenantID, client.WorkspaceID, client.ID, dispatchReplayLimit) if err != nil { if h.app.Logger != nil { h.app.Logger.Warn("desktop dispatch replay query failed", zap.String("client_id", client.ID.String()), zap.Error(err), ) } return } defer rows.Close() replayed := 0 for rows.Next() { task, scanErr := scanDispatchReplayDesktopTask(rows) if scanErr != nil { if h.app.Logger != nil { h.app.Logger.Warn("desktop dispatch replay scan failed", zap.String("client_id", client.ID.String()), zap.Error(scanErr), ) } return } if err := writeDispatchFrame(conn, tenantapp.DesktopDispatchEventFromTask(task, "task_available")); err != nil { if h.app.Logger != nil { h.app.Logger.Warn("desktop dispatch replay write failed", zap.String("client_id", client.ID.String()), zap.String("task_id", task.DesktopID.String()), zap.Error(err), ) } return } replayed++ } if err := rows.Err(); err != nil && h.app.Logger != nil { h.app.Logger.Warn("desktop dispatch replay rows failed", zap.String("client_id", client.ID.String()), zap.Error(err), ) } if replayed > 0 && h.app.Logger != nil { h.app.Logger.Info("desktop dispatch replayed queued publish tasks", zap.String("client_id", client.ID.String()), zap.Int("count", replayed), ) } } type dispatchReplayDesktopTaskScanner interface { Scan(dest ...any) error } func scanDispatchReplayDesktopTask(row dispatchReplayDesktopTaskScanner) (*repository.DesktopTask, error) { var task repository.DesktopTask 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, &task.SchedulerGroupKey, &task.MonitorTaskID, &task.SupersedesTaskID, &task.ControlFlags, &task.InterruptGeneration, &task.DedupKey, &task.ActiveAttemptID, &task.LeaseExpiresAt, &task.Attempts, &task.Result, &task.Error, &task.StartedAt, &task.InterruptedAt, &task.InterruptReason, &task.EnqueuedAt, &task.CreatedAt, &task.UpdatedAt, ) if err != nil { return nil, err } return &task, nil }