multica/server/internal/handler/daemon_ws.go

package handler

import (
	"log/slog"
	"net/http"
	"strings"

	"github.com/multica-ai/multica/server/internal/auth"
	"github.com/multica-ai/multica/server/internal/daemonws"
	"github.com/multica-ai/multica/server/internal/middleware"
)

func (h *Handler) DaemonWebSocket(w http.ResponseWriter, r *http.Request) {
	if h.DaemonHub == nil {
		writeErrorCode(w, http.StatusInternalServerError, "daemon_websocket_misconfigured", "daemon websocket unavailable")
		return
	}

	runtimeIDs := parseRuntimeIDs(r)
	userID := requestUserID(r)
	if len(runtimeIDs) == 0 && userID == "" {
		writeError(w, http.StatusBadRequest, "runtime_ids or user identity required")
		return
	}

	identity, ok := h.buildDaemonWebSocketIdentity(w, r, runtimeIDs, userID)
	if !ok {
		return
	}
	h.DaemonHub.HandleWebSocket(w, r, identity)
}

// buildDaemonWebSocketIdentity authenticates the connection's entire runtime
// set with one narrow query and seeds the connection-scoped heartbeat leases.
// Runtime ownership is immutable, so the heartbeat hot path can safely use
// this fixed scope without re-reading agent_runtime every 15 seconds.
func (h *Handler) buildDaemonWebSocketIdentity(w http.ResponseWriter, r *http.Request, runtimeIDs []string, userID string) (daemonws.ClientIdentity, bool) {
	identity := daemonws.ClientIdentity{
		DaemonID:           middleware.DaemonIDFromContext(r.Context()),
		DaemonTokenHash:    daemonTokenHashForRequest(r),
		DaemonTokenOwnerID: middleware.DaemonTokenOwnerIDFromContext(r.Context()),
		UserID:             userID,
		RuntimeIDs:         runtimeIDs,
		RuntimeLeases:      make(map[string]*daemonws.RuntimeLease, len(runtimeIDs)),
		ClientVersion:      r.Header.Get("X-Client-Version"),
		Capabilities:       r.Header.Get("X-Client-Capabilities"),
	}
	if len(runtimeIDs) == 0 {
		return identity, true
	}

	runtimeUUIDs, ok := parseUUIDSliceOrBadRequest(w, runtimeIDs, "runtime_ids")
	if !ok {
		return daemonws.ClientIdentity{}, false
	}
	rows, err := h.Queries.GetAgentRuntimeHeartbeatLeases(r.Context(), runtimeUUIDs)
	if err != nil {
		slog.Warn("load daemon websocket runtime leases failed", "runtimes", len(runtimeIDs), "error", err)
		writeError(w, http.StatusInternalServerError, "failed to load runtimes")
		return daemonws.ClientIdentity{}, false
	}
	byID := make(map[string]int, len(rows))
	for index, row := range rows {
		byID[uuidToString(row.ID)] = index
	}

	workspaceIDs := make([]string, 0, len(runtimeIDs))
	seenWorkspaceIDs := make(map[string]struct{}, len(runtimeIDs))
	for runtimeIndex, runtimeID := range runtimeIDs {
		index, found := byID[uuidToString(runtimeUUIDs[runtimeIndex])]
		if !found {
			writeError(w, http.StatusNotFound, "runtime not found")
			return daemonws.ClientIdentity{}, false
		}
		rt := rows[index]
		// Same runtime-level rule as the HTTP daemon routes (MUL-8154): a
		// daemon token must match the runtime's daemon_id, a user credential
		// must own the runtime. The leases seeded below authorize every later
		// WS heartbeat for this runtime, so this is the only check they get.
		workspaceID := uuidToString(rt.WorkspaceID)
		if !h.requireDaemonRuntimeAuthorized(w, r, workspaceID, rt.DaemonID, rt.OwnerID, "runtime not found") {
			return daemonws.ClientIdentity{}, false
		}
		if workspaceID != "" {
			if _, ok := seenWorkspaceIDs[workspaceID]; !ok {
				seenWorkspaceIDs[workspaceID] = struct{}{}
				workspaceIDs = append(workspaceIDs, workspaceID)
			}
		}
		identity.RuntimeLeases[runtimeID] = daemonws.NewRuntimeLease(
			workspaceID,
			rt.Status,
			rt.LastSeenAt.Time,
			rt.LastSeenAt.Valid,
		)
	}

	primaryWorkspaceID := ""
	if len(workspaceIDs) > 0 {
		primaryWorkspaceID = workspaceIDs[0]
	}
	identity.WorkspaceID = primaryWorkspaceID
	identity.WorkspaceIDs = workspaceIDs
	return identity, true
}

// daemonTokenHashForRequest returns the hash of the mdt_ token behind a
// daemon-token request: the one DaemonAuth recorded, or — for a request that
// carries a daemon context without passing DaemonAuth — the Authorization
// header's. Empty for user credentials.
func daemonTokenHashForRequest(r *http.Request) string {
	if middleware.DaemonWorkspaceIDFromContext(r.Context()) == "" {
		return ""
	}
	if hash := middleware.DaemonTokenHashFromContext(r.Context()); hash != "" {
		return hash
	}
	raw := strings.TrimPrefix(r.Header.Get("Authorization"), "Bearer ")
	if !strings.HasPrefix(raw, "mdt_") {
		return ""
	}
	return auth.HashToken(raw)
}

func parseRuntimeIDs(r *http.Request) []string {
	seen := map[string]struct{}{}
	var out []string
	add := func(raw string) {
		for _, part := range strings.Split(raw, ",") {
			id := strings.TrimSpace(part)
			if id == "" {
				continue
			}
			if _, ok := seen[id]; ok {
				continue
			}
			seen[id] = struct{}{}
			out = append(out, id)
		}
	}
	for _, raw := range r.URL.Query()["runtime_id"] {
		add(raw)
	}
	for _, raw := range r.URL.Query()["runtime_ids"] {
		add(raw)
	}
	return out
}