// Package jupyter provides WebSocket handlers for Jupyter-related operations package jupyter import ( "encoding/binary" "encoding/json" "net/http" "time" "github.com/gorilla/websocket" "github.com/jfraeys/fetch_ml/internal/api/errors" "github.com/jfraeys/fetch_ml/internal/auth" "github.com/jfraeys/fetch_ml/internal/container" "github.com/jfraeys/fetch_ml/internal/jupyter" "github.com/jfraeys/fetch_ml/internal/logging" ) // Handler provides Jupyter-related WebSocket handlers type Handler struct { logger *logging.Logger jupyterMgr *jupyter.ServiceManager authConfig *auth.Config } // NewHandler creates a new Jupyter handler func NewHandler( logger *logging.Logger, jupyterMgr *jupyter.ServiceManager, authConfig *auth.Config, ) *Handler { return &Handler{ logger: logger, jupyterMgr: jupyterMgr, authConfig: authConfig, } } // Error codes - using standardized error codes from errors package const ( ErrorCodeInvalidRequest = errors.CodeInvalidRequest ErrorCodeAuthenticationFailed = errors.CodeAuthenticationFailed ErrorCodePermissionDenied = errors.CodePermissionDenied ErrorCodeResourceNotFound = errors.CodeResourceNotFound ErrorCodeServiceUnavailable = errors.CodeServiceUnavailable ) // Permissions const ( PermJupyterManage = "jupyter:manage" PermJupyterRead = "jupyter:read" ) // sendErrorPacket sends an error response packet to the client func (h *Handler) sendErrorPacket(conn *websocket.Conn, code string, message, details string) error { return errors.SendErrorPacket(conn, code, message, details) } // sendSuccessPacket sends a success response packet func (h *Handler) sendSuccessPacket(conn *websocket.Conn, data map[string]any) error { return errors.SendSuccessPacket(conn, data) } // HandleStartJupyter handles starting a Jupyter service // Protocol: [api_key_hash:16][workspace_len:1][workspace:var][config_len:2][config:var] func (h *Handler) HandleStartJupyter(conn *websocket.Conn, payload []byte, user *auth.User) error { if len(payload) < 16+1+2 { return h.sendErrorPacket(conn, ErrorCodeInvalidRequest, "start jupyter payload too short", "") } offset := 16 workspaceLen := int(payload[offset]) offset += 1 if workspaceLen <= 0 || len(payload) < offset+workspaceLen+2 { return h.sendErrorPacket(conn, ErrorCodeInvalidRequest, "invalid workspace length", "") } workspace := string(payload[offset : offset+workspaceLen]) offset += workspaceLen configLen := int(binary.BigEndian.Uint16(payload[offset : offset+2])) offset += 2 if configLen < 0 || len(payload) < offset+configLen { return h.sendErrorPacket(conn, ErrorCodeInvalidRequest, "invalid config length", "") } if err := container.ValidateJobName(workspace); err != nil { return h.sendErrorPacket(conn, ErrorCodeInvalidRequest, "invalid workspace name", err.Error()) } h.logger.Info("starting jupyter service", "workspace", workspace, "user", user.Name) // Start Jupyter service if h.jupyterMgr == nil { return h.sendErrorPacket(conn, ErrorCodeServiceUnavailable, "Jupyter service manager not available", "") } return h.sendSuccessPacket(conn, map[string]any{ "workspace": workspace, "timestamp": time.Now().UTC(), }) } // HandleStopJupyter handles stopping a Jupyter service // Protocol: [api_key_hash:16][service_id_len:1][service_id:var] func (h *Handler) HandleStopJupyter(conn *websocket.Conn, payload []byte, user *auth.User) error { if len(payload) < 16+1 { return h.sendErrorPacket(conn, ErrorCodeInvalidRequest, "stop jupyter payload too short", "") } offset := 16 serviceIDLen := int(payload[offset]) offset += 1 if serviceIDLen <= 0 || len(payload) < offset+serviceIDLen { return h.sendErrorPacket(conn, ErrorCodeInvalidRequest, "invalid service ID length", "") } serviceID := string(payload[offset : offset+serviceIDLen]) h.logger.Info("stopping jupyter service", "service_id", serviceID, "user", user.Name) if h.jupyterMgr == nil { return h.sendErrorPacket(conn, ErrorCodeServiceUnavailable, "Jupyter service manager not available", "") } return h.sendSuccessPacket(conn, map[string]any{ "service_id": serviceID, "timestamp": time.Now().UTC(), }) } // HandleListJupyter handles listing Jupyter services // Protocol: [api_key_hash:16] func (h *Handler) HandleListJupyter(conn *websocket.Conn, payload []byte, user *auth.User) error { h.logger.Info("listing jupyter services", "user", user.Name) if h.jupyterMgr == nil { return h.sendSuccessPacket(conn, map[string]any{ "services": []any{}, "count": 0, }) } services := h.jupyterMgr.ListServices() return h.sendSuccessPacket(conn, map[string]any{ "services": services, "count": len(services), }) } // HandleListJupyterPackages handles listing packages in a Jupyter service // Protocol: [api_key_hash:16][service_name_len:1][service_name:var] func (h *Handler) HandleListJupyterPackages(conn *websocket.Conn, payload []byte, user *auth.User) error { if len(payload) < 16+1 { return h.sendErrorPacket(conn, ErrorCodeInvalidRequest, "list packages payload too short", "") } offset := 16 serviceNameLen := int(payload[offset]) offset += 1 if serviceNameLen <= 0 || len(payload) < offset+serviceNameLen { return h.sendErrorPacket(conn, ErrorCodeInvalidRequest, "invalid service name length", "") } serviceName := string(payload[offset : offset+serviceNameLen]) h.logger.Info("listing jupyter packages", "service", serviceName, "user", user.Name) return h.sendSuccessPacket(conn, map[string]any{ "service_name": serviceName, "packages": []any{}, "count": 0, }) } // HandleRemoveJupyter handles removing a Jupyter service // Protocol: [api_key_hash:16][service_id_len:1][service_id:var][purge:1] func (h *Handler) HandleRemoveJupyter(conn *websocket.Conn, payload []byte, user *auth.User) error { if len(payload) < 16+1+1 { return h.sendErrorPacket(conn, ErrorCodeInvalidRequest, "remove jupyter payload too short", "") } offset := 16 serviceIDLen := int(payload[offset]) offset += 1 if serviceIDLen <= 0 || len(payload) < offset+serviceIDLen+1 { return h.sendErrorPacket(conn, ErrorCodeInvalidRequest, "invalid service ID length", "") } serviceID := string(payload[offset : offset+serviceIDLen]) offset += serviceIDLen purge := payload[offset] != 0 h.logger.Info("removing jupyter service", "service_id", serviceID, "purge", purge, "user", user.Name) if h.jupyterMgr == nil { return h.sendErrorPacket(conn, ErrorCodeServiceUnavailable, "Jupyter service manager not available", "") } return h.sendSuccessPacket(conn, map[string]any{ "service_id": serviceID, "purged": purge, }) } // HandleRestoreJupyter handles restoring a Jupyter workspace // Protocol: [api_key_hash:16][workspace_len:1][workspace:var] func (h *Handler) HandleRestoreJupyter(conn *websocket.Conn, payload []byte, user *auth.User) error { if len(payload) < 16+1 { return h.sendErrorPacket(conn, ErrorCodeInvalidRequest, "restore jupyter payload too short", "") } offset := 16 workspaceLen := int(payload[offset]) offset += 1 if workspaceLen <= 0 || len(payload) < offset+workspaceLen { return h.sendErrorPacket(conn, ErrorCodeInvalidRequest, "invalid workspace length", "") } workspace := string(payload[offset : offset+workspaceLen]) h.logger.Info("restoring jupyter workspace", "workspace", workspace, "user", user.Name) if h.jupyterMgr == nil { return h.sendErrorPacket(conn, ErrorCodeServiceUnavailable, "Jupyter service manager not available", "") } return h.sendSuccessPacket(conn, map[string]any{ "workspace": workspace, "restored": true, }) } // HTTP Handlers for REST API // ListServicesHTTP handles HTTP requests for listing Jupyter services func (h *Handler) ListServicesHTTP(w http.ResponseWriter, r *http.Request) { if h.jupyterMgr == nil { errors.WriteHTTPError(w, http.StatusServiceUnavailable, errors.CodeServiceUnavailable, "Jupyter service manager not available", "") return } services := h.jupyterMgr.ListServices() w.Header().Set("Content-Type", "application/json") response := map[string]any{ "services": services, "count": len(services), } if err := json.NewEncoder(w).Encode(response); err != nil { h.logger.Warn("failed to encode services list", "error", err) } } // StartServiceHTTP handles HTTP requests for starting Jupyter service func (h *Handler) StartServiceHTTP(w http.ResponseWriter, r *http.Request) { if h.jupyterMgr == nil { errors.WriteHTTPError(w, http.StatusServiceUnavailable, errors.CodeServiceUnavailable, "Jupyter service manager not available", "") return } var req struct { Workspace string `json:"workspace"` Config map[string]any `json:"config,omitempty"` } if err := json.NewDecoder(r.Body).Decode(&req); err != nil { errors.WriteHTTPError(w, http.StatusBadRequest, errors.CodeInvalidRequest, "Invalid request body", "") return } if req.Workspace == "" { errors.WriteHTTPError(w, http.StatusBadRequest, errors.CodeInvalidRequest, "Workspace name is required", "") return } if err := container.ValidateJobName(req.Workspace); err != nil { errors.WriteHTTPError(w, http.StatusBadRequest, errors.CodeInvalidRequest, "Invalid workspace name", err.Error()) return } h.logger.Info("starting jupyter service via HTTP", "workspace", req.Workspace) w.Header().Set("Content-Type", "application/json") w.WriteHeader(http.StatusCreated) response := map[string]any{ "workspace": req.Workspace, "timestamp": time.Now().UTC(), } if err := json.NewEncoder(w).Encode(response); err != nil { h.logger.Warn("failed to encode start response", "error", err) } } // StopServiceHTTP handles HTTP requests for stopping Jupyter service func (h *Handler) StopServiceHTTP(w http.ResponseWriter, r *http.Request) { if h.jupyterMgr == nil { errors.WriteHTTPError(w, http.StatusServiceUnavailable, errors.CodeServiceUnavailable, "Jupyter service manager not available", "") return } serviceID := r.PathValue("serviceId") if serviceID == "" { serviceID = r.URL.Query().Get("service_id") } if serviceID == "" { errors.WriteHTTPError(w, http.StatusBadRequest, errors.CodeInvalidRequest, "Service ID is required", "") return } h.logger.Info("stopping jupyter service via HTTP", "service_id", serviceID) w.Header().Set("Content-Type", "application/json") response := map[string]any{ "service_id": serviceID, "timestamp": time.Now().UTC(), } if err := json.NewEncoder(w).Encode(response); err != nil { h.logger.Warn("failed to encode stop response", "error", err) } }