// Package helpers provides shared utilities for WebSocket handlers. package helpers import ( "encoding/json" "fmt" "strings" "github.com/jfraeys/fetch_ml/internal/queue" ) // ErrorCode represents WebSocket error codes type ErrorCode byte // TaskErrorMapper maps task errors to error codes type TaskErrorMapper struct{} // NewTaskErrorMapper creates a new task error mapper func NewTaskErrorMapper() *TaskErrorMapper { return &TaskErrorMapper{} } // MapError maps a task error to an error code based on status and error message func (m *TaskErrorMapper) MapError(t *queue.Task, defaultCode ErrorCode) ErrorCode { if t == nil { return defaultCode } status := strings.ToLower(strings.TrimSpace(t.Status)) errStr := strings.ToLower(strings.TrimSpace(t.Error)) if status == "cancelled" { return 0x24 // ErrorCodeJobCancelled } if strings.Contains(errStr, "out of memory") || strings.Contains(errStr, "oom") { return 0x30 // ErrorCodeOutOfMemory } if strings.Contains(errStr, "no space left") || strings.Contains(errStr, "disk full") { return 0x31 // ErrorCodeDiskFull } if strings.Contains(errStr, "rate limit") || strings.Contains(errStr, "too many requests") || strings.Contains(errStr, "throttle") { return 0x33 // ErrorCodeServiceUnavailable } if strings.Contains(errStr, "timed out") || strings.Contains(errStr, "timeout") || strings.Contains(errStr, "deadline") { return 0x14 // ErrorCodeTimeout } if strings.Contains(errStr, "connection refused") || strings.Contains(errStr, "connection reset") || strings.Contains(errStr, "network unreachable") { return 0x12 // ErrorCodeNetworkError } if strings.Contains(errStr, "queue") && strings.Contains(errStr, "not configured") { return 0x32 // ErrorCodeInvalidConfiguration } // Default for worker-side execution failures if status == "failed" { return 0x23 // ErrorCodeJobExecutionFailed } return defaultCode } // MapJupyterError maps Jupyter task errors to error codes func (m *TaskErrorMapper) MapJupyterError(t *queue.Task) ErrorCode { if t == nil { return 0x00 // ErrorCodeUnknownError } status := strings.ToLower(strings.TrimSpace(t.Status)) errStr := strings.ToLower(strings.TrimSpace(t.Error)) if status == "cancelled" { return 0x24 // ErrorCodeJobCancelled } if strings.Contains(errStr, "out of memory") || strings.Contains(errStr, "oom") { return 0x30 // ErrorCodeOutOfMemory } if strings.Contains(errStr, "no space left") || strings.Contains(errStr, "disk full") { return 0x31 // ErrorCodeDiskFull } if strings.Contains(errStr, "rate limit") || strings.Contains(errStr, "too many requests") || strings.Contains(errStr, "throttle") { return 0x33 // ErrorCodeServiceUnavailable } if strings.Contains(errStr, "timed out") || strings.Contains(errStr, "timeout") || strings.Contains(errStr, "deadline") { return 0x14 // ErrorCodeTimeout } if strings.Contains(errStr, "connection refused") || strings.Contains(errStr, "connection reset") || strings.Contains(errStr, "network unreachable") { return 0x12 // ErrorCodeNetworkError } if strings.Contains(errStr, "queue") && strings.Contains(errStr, "not configured") { return 0x32 // ErrorCodeInvalidConfiguration } // Default for worker-side execution failures if status == "failed" { return 0x23 // ErrorCodeJobExecutionFailed } return 0x00 // ErrorCodeUnknownError } // ResourceRequest represents resource requirements type ResourceRequest struct { CPU int MemoryGB int GPU int GPUMemory string } // ParseResourceRequest parses an optional resource request from bytes. // Format: [cpu:1][memory_gb:1][gpu:1][gpu_mem_len:1][gpu_mem:var] // If payload is empty, returns nil. func ParseResourceRequest(payload []byte) (*ResourceRequest, error) { if len(payload) == 0 { return nil, nil } if len(payload) < 4 { return nil, fmt.Errorf("resource payload too short") } cpu := int(payload[0]) mem := int(payload[1]) gpu := int(payload[2]) gpuMemLen := int(payload[3]) if gpuMemLen < 0 || len(payload) < 4+gpuMemLen { return nil, fmt.Errorf("invalid gpu memory length") } gpuMem := "" if gpuMemLen > 0 { gpuMem = string(payload[4 : 4+gpuMemLen]) } return &ResourceRequest{CPU: cpu, MemoryGB: mem, GPU: gpu, GPUMemory: gpuMem}, nil } // JSONResponseBuilder helps build JSON data responses type JSONResponseBuilder struct { data interface{} } // NewJSONResponseBuilder creates a new JSON response builder func NewJSONResponseBuilder(data interface{}) *JSONResponseBuilder { return &JSONResponseBuilder{data: data} } // Build marshals the data to JSON func (b *JSONResponseBuilder) Build() ([]byte, error) { return json.Marshal(b.data) } // BuildOrEmpty marshals the data to JSON or returns empty array on error func (b *JSONResponseBuilder) BuildOrEmpty() []byte { data, err := json.Marshal(b.data) if err != nil { return []byte("[]") } return data } // StringPtr returns a pointer to a string func StringPtr(s string) *string { return &s } // IntPtr returns a pointer to an int func IntPtr(i int) *int { return &i } // MarshalJSONOrEmpty marshals data to JSON or returns empty array on error func MarshalJSONOrEmpty(data interface{}) []byte { b, err := json.Marshal(data) if err != nil { return []byte("[]") } return b } // MarshalJSONBytes marshals data to JSON bytes with error handling func MarshalJSONBytes(data interface{}) ([]byte, error) { return json.Marshal(data) } // IsEmptyJSON checks if JSON data is empty or "null" func IsEmptyJSON(data []byte) bool { if len(data) == 0 { return true } // Check for "null", "[]", "{}" or empty after trimming s := strings.TrimSpace(string(data)) return s == "" || s == "null" || s == "[]" || s == "{}" }