Update .gitignore to include log files and database journal files. Modify go.mod to update dependencies for go-sqlite3 and cloud.google.com/go/compute/metadata. Enhance Makefile to include logging options for manager and runner commands. Introduce new job token handling in auth package and implement database migration scripts. Refactor manager and runner components to improve job processing and metadata extraction. Add support for video preview in frontend components and enhance WebSocket management for channel subscriptions.

This commit is contained in:
2026-01-02 13:55:19 -06:00
parent edc8ea160c
commit 94490237fe
44 changed files with 9463 additions and 7875 deletions

View File

@@ -0,0 +1,333 @@
package api
import (
"fmt"
"log"
"strings"
"sync"
"time"
"jiggablend/pkg/types"
"github.com/gorilla/websocket"
)
// JobConnection wraps a WebSocket connection for job communication.
type JobConnection struct {
conn *websocket.Conn
writeMu sync.Mutex
stopPing chan struct{}
stopHeartbeat chan struct{}
isConnected bool
connMu sync.RWMutex
}
// NewJobConnection creates a new job connection wrapper.
func NewJobConnection() *JobConnection {
return &JobConnection{}
}
// Connect establishes a WebSocket connection for a job (no runnerID needed).
func (j *JobConnection) Connect(managerURL, jobPath, jobToken string) error {
wsPath := jobPath + "/ws"
wsURL := strings.Replace(managerURL, "http://", "ws://", 1)
wsURL = strings.Replace(wsURL, "https://", "wss://", 1)
wsURL += wsPath
log.Printf("Connecting to job WebSocket: %s", wsPath)
dialer := websocket.Dialer{
HandshakeTimeout: 10 * time.Second,
}
conn, _, err := dialer.Dial(wsURL, nil)
if err != nil {
return fmt.Errorf("failed to connect job WebSocket: %w", err)
}
j.conn = conn
// Send auth message
authMsg := map[string]interface{}{
"type": "auth",
"job_token": jobToken,
}
if err := conn.WriteJSON(authMsg); err != nil {
conn.Close()
return fmt.Errorf("failed to send auth: %w", err)
}
// Wait for auth_ok
conn.SetReadDeadline(time.Now().Add(30 * time.Second))
var authResp map[string]string
if err := conn.ReadJSON(&authResp); err != nil {
conn.Close()
return fmt.Errorf("failed to read auth response: %w", err)
}
if authResp["type"] == "error" {
conn.Close()
return fmt.Errorf("auth failed: %s", authResp["message"])
}
if authResp["type"] != "auth_ok" {
conn.Close()
return fmt.Errorf("unexpected auth response: %s", authResp["type"])
}
// Clear read deadline after auth
conn.SetReadDeadline(time.Time{})
// Set up ping/pong handler for keepalive
conn.SetPongHandler(func(string) error {
conn.SetReadDeadline(time.Now().Add(90 * time.Second))
return nil
})
// Start ping goroutine
j.stopPing = make(chan struct{})
j.connMu.Lock()
j.isConnected = true
j.connMu.Unlock()
go j.pingLoop()
// Start WebSocket heartbeat goroutine
j.stopHeartbeat = make(chan struct{})
go j.heartbeatLoop()
return nil
}
// pingLoop sends periodic pings to keep the WebSocket connection alive.
func (j *JobConnection) pingLoop() {
defer func() {
if rec := recover(); rec != nil {
log.Printf("Ping loop panicked: %v", rec)
}
}()
ticker := time.NewTicker(30 * time.Second)
defer ticker.Stop()
for {
select {
case <-j.stopPing:
return
case <-ticker.C:
j.writeMu.Lock()
if j.conn != nil {
deadline := time.Now().Add(10 * time.Second)
if err := j.conn.WriteControl(websocket.PingMessage, []byte{}, deadline); err != nil {
log.Printf("Failed to send ping, closing connection: %v", err)
j.connMu.Lock()
j.isConnected = false
if j.conn != nil {
j.conn.Close()
j.conn = nil
}
j.connMu.Unlock()
}
}
j.writeMu.Unlock()
}
}
}
// Heartbeat sends a heartbeat message over WebSocket to keep runner online.
func (j *JobConnection) Heartbeat() {
if j.conn == nil {
return
}
j.writeMu.Lock()
defer j.writeMu.Unlock()
msg := map[string]interface{}{
"type": "runner_heartbeat",
"timestamp": time.Now().Unix(),
}
if err := j.conn.WriteJSON(msg); err != nil {
log.Printf("Failed to send WebSocket heartbeat: %v", err)
// Handle connection failure
j.connMu.Lock()
j.isConnected = false
if j.conn != nil {
j.conn.Close()
j.conn = nil
}
j.connMu.Unlock()
}
}
// heartbeatLoop sends periodic heartbeat messages over WebSocket.
func (j *JobConnection) heartbeatLoop() {
defer func() {
if rec := recover(); rec != nil {
log.Printf("WebSocket heartbeat loop panicked: %v", rec)
}
}()
ticker := time.NewTicker(30 * time.Second)
defer ticker.Stop()
for {
select {
case <-j.stopHeartbeat:
return
case <-ticker.C:
j.Heartbeat()
}
}
}
// Close closes the WebSocket connection.
func (j *JobConnection) Close() {
j.connMu.Lock()
j.isConnected = false
j.connMu.Unlock()
// Stop heartbeat goroutine
if j.stopHeartbeat != nil {
close(j.stopHeartbeat)
j.stopHeartbeat = nil
}
// Stop ping goroutine
if j.stopPing != nil {
close(j.stopPing)
j.stopPing = nil
}
if j.conn != nil {
j.conn.Close()
j.conn = nil
}
}
// IsConnected returns true if the connection is established.
func (j *JobConnection) IsConnected() bool {
j.connMu.RLock()
defer j.connMu.RUnlock()
return j.isConnected && j.conn != nil
}
// Log sends a log entry to the manager.
func (j *JobConnection) Log(taskID int64, level types.LogLevel, message string) {
if j.conn == nil {
return
}
j.writeMu.Lock()
defer j.writeMu.Unlock()
msg := map[string]interface{}{
"type": "log_entry",
"data": map[string]interface{}{
"task_id": taskID,
"log_level": string(level),
"message": message,
},
"timestamp": time.Now().Unix(),
}
if err := j.conn.WriteJSON(msg); err != nil {
log.Printf("Failed to send job log, connection may be broken: %v", err)
// Close the connection on write error
j.connMu.Lock()
j.isConnected = false
if j.conn != nil {
j.conn.Close()
j.conn = nil
}
j.connMu.Unlock()
}
}
// Progress sends a progress update to the manager.
func (j *JobConnection) Progress(taskID int64, progress float64) {
if j.conn == nil {
return
}
j.writeMu.Lock()
defer j.writeMu.Unlock()
msg := map[string]interface{}{
"type": "progress",
"data": map[string]interface{}{
"task_id": taskID,
"progress": progress,
},
"timestamp": time.Now().Unix(),
}
if err := j.conn.WriteJSON(msg); err != nil {
log.Printf("Failed to send job progress, connection may be broken: %v", err)
// Close the connection on write error
j.connMu.Lock()
j.isConnected = false
if j.conn != nil {
j.conn.Close()
j.conn = nil
}
j.connMu.Unlock()
}
}
// OutputUploaded notifies that an output file was uploaded.
func (j *JobConnection) OutputUploaded(taskID int64, fileName string) {
if j.conn == nil {
return
}
j.writeMu.Lock()
defer j.writeMu.Unlock()
msg := map[string]interface{}{
"type": "output_uploaded",
"data": map[string]interface{}{
"task_id": taskID,
"file_name": fileName,
},
"timestamp": time.Now().Unix(),
}
if err := j.conn.WriteJSON(msg); err != nil {
log.Printf("Failed to send output uploaded, connection may be broken: %v", err)
// Close the connection on write error
j.connMu.Lock()
j.isConnected = false
if j.conn != nil {
j.conn.Close()
j.conn = nil
}
j.connMu.Unlock()
}
}
// Complete sends task completion to the manager.
func (j *JobConnection) Complete(taskID int64, success bool, errorMsg error) {
if j.conn == nil {
log.Printf("Cannot send task complete: WebSocket connection is nil")
return
}
j.writeMu.Lock()
defer j.writeMu.Unlock()
msg := map[string]interface{}{
"type": "task_complete",
"data": map[string]interface{}{
"task_id": taskID,
"success": success,
"error": errorMsg,
},
"timestamp": time.Now().Unix(),
}
if err := j.conn.WriteJSON(msg); err != nil {
log.Printf("Failed to send task complete, connection may be broken: %v", err)
// Close the connection on write error
j.connMu.Lock()
j.isConnected = false
if j.conn != nil {
j.conn.Close()
j.conn = nil
}
j.connMu.Unlock()
}
}

View File

@@ -0,0 +1,421 @@
// Package api provides HTTP and WebSocket communication with the manager server.
package api
import (
"bytes"
"encoding/json"
"fmt"
"io"
"mime/multipart"
"net/http"
"net/url"
"os"
"path/filepath"
"strings"
"time"
"jiggablend/pkg/types"
)
// ManagerClient handles all HTTP communication with the manager server.
type ManagerClient struct {
baseURL string
apiKey string
runnerID int64
httpClient *http.Client // Standard timeout for quick requests
longClient *http.Client // No timeout for large file transfers
}
// NewManagerClient creates a new manager client.
func NewManagerClient(baseURL string) *ManagerClient {
return &ManagerClient{
baseURL: strings.TrimSuffix(baseURL, "/"),
httpClient: &http.Client{Timeout: 30 * time.Second},
longClient: &http.Client{Timeout: 0}, // No timeout for large transfers
}
}
// SetCredentials sets the API key and runner ID after registration.
func (m *ManagerClient) SetCredentials(runnerID int64, apiKey string) {
m.runnerID = runnerID
m.apiKey = apiKey
}
// GetRunnerID returns the registered runner ID.
func (m *ManagerClient) GetRunnerID() int64 {
return m.runnerID
}
// GetAPIKey returns the API key.
func (m *ManagerClient) GetAPIKey() string {
return m.apiKey
}
// GetBaseURL returns the base URL.
func (m *ManagerClient) GetBaseURL() string {
return m.baseURL
}
// Request performs an authenticated HTTP request with standard timeout.
func (m *ManagerClient) Request(method, path string, body []byte) (*http.Response, error) {
return m.doRequest(method, path, body, m.httpClient)
}
// RequestLong performs an authenticated HTTP request with no timeout.
// Use for large file uploads/downloads.
func (m *ManagerClient) RequestLong(method, path string, body []byte) (*http.Response, error) {
return m.doRequest(method, path, body, m.longClient)
}
func (m *ManagerClient) doRequest(method, path string, body []byte, client *http.Client) (*http.Response, error) {
if m.apiKey == "" {
return nil, fmt.Errorf("not authenticated")
}
fullURL := m.baseURL + path
req, err := http.NewRequest(method, fullURL, bytes.NewReader(body))
if err != nil {
return nil, err
}
req.Header.Set("Authorization", "Bearer "+m.apiKey)
if len(body) > 0 {
req.Header.Set("Content-Type", "application/json")
}
return client.Do(req)
}
// RequestWithToken performs an authenticated HTTP request using a specific token.
func (m *ManagerClient) RequestWithToken(method, path, token string, body []byte) (*http.Response, error) {
return m.doRequestWithToken(method, path, token, body, m.httpClient)
}
// RequestLongWithToken performs a long-running request with a specific token.
func (m *ManagerClient) RequestLongWithToken(method, path, token string, body []byte) (*http.Response, error) {
return m.doRequestWithToken(method, path, token, body, m.longClient)
}
func (m *ManagerClient) doRequestWithToken(method, path, token string, body []byte, client *http.Client) (*http.Response, error) {
fullURL := m.baseURL + path
req, err := http.NewRequest(method, fullURL, bytes.NewReader(body))
if err != nil {
return nil, err
}
req.Header.Set("Authorization", "Bearer "+token)
if len(body) > 0 {
req.Header.Set("Content-Type", "application/json")
}
return client.Do(req)
}
// RegisterRequest is the request body for runner registration.
type RegisterRequest struct {
Name string `json:"name"`
Hostname string `json:"hostname"`
Capabilities string `json:"capabilities"`
APIKey string `json:"api_key"`
Fingerprint string `json:"fingerprint,omitempty"`
}
// RegisterResponse is the response from runner registration.
type RegisterResponse struct {
ID int64 `json:"id"`
}
// Register registers the runner with the manager.
func (m *ManagerClient) Register(name, hostname string, capabilities map[string]interface{}, registrationToken, fingerprint string) (int64, error) {
capsJSON, err := json.Marshal(capabilities)
if err != nil {
return 0, fmt.Errorf("failed to marshal capabilities: %w", err)
}
reqBody := RegisterRequest{
Name: name,
Hostname: hostname,
Capabilities: string(capsJSON),
APIKey: registrationToken,
}
// Only send fingerprint for non-fixed API keys
if !strings.HasPrefix(registrationToken, "jk_r0_") {
reqBody.Fingerprint = fingerprint
}
body, _ := json.Marshal(reqBody)
resp, err := m.httpClient.Post(
m.baseURL+"/api/runner/register",
"application/json",
bytes.NewReader(body),
)
if err != nil {
return 0, fmt.Errorf("connection error: %w", err)
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusCreated {
bodyBytes, _ := io.ReadAll(resp.Body)
errorBody := string(bodyBytes)
// Check for token-related errors (should not retry)
if resp.StatusCode == http.StatusUnauthorized || resp.StatusCode == http.StatusBadRequest {
errorLower := strings.ToLower(errorBody)
if strings.Contains(errorLower, "invalid") ||
strings.Contains(errorLower, "expired") ||
strings.Contains(errorLower, "already used") ||
strings.Contains(errorLower, "token") {
return 0, fmt.Errorf("token error: %s", errorBody)
}
}
return 0, fmt.Errorf("registration failed (status %d): %s", resp.StatusCode, errorBody)
}
var result RegisterResponse
if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
return 0, fmt.Errorf("failed to decode response: %w", err)
}
m.runnerID = result.ID
m.apiKey = registrationToken
return result.ID, nil
}
// NextJobResponse represents the response from the next-job endpoint.
type NextJobResponse struct {
JobToken string `json:"job_token"`
JobPath string `json:"job_path"`
Task NextJobTaskInfo `json:"task"`
}
// NextJobTaskInfo contains task information from the next-job response.
type NextJobTaskInfo struct {
TaskID int64 `json:"task_id"`
JobID int64 `json:"job_id"`
JobName string `json:"job_name"`
Frame int `json:"frame"`
TaskType string `json:"task_type"`
Metadata *types.BlendMetadata `json:"metadata,omitempty"`
}
// PollNextJob polls the manager for the next available job.
// Returns nil, nil if no job is available.
func (m *ManagerClient) PollNextJob() (*NextJobResponse, error) {
if m.runnerID == 0 || m.apiKey == "" {
return nil, fmt.Errorf("runner not authenticated")
}
path := fmt.Sprintf("/api/runner/workers/%d/next-job", m.runnerID)
resp, err := m.Request("GET", path, nil)
if err != nil {
return nil, fmt.Errorf("failed to poll for job: %w", err)
}
defer resp.Body.Close()
if resp.StatusCode == http.StatusNoContent {
return nil, nil // No job available
}
if resp.StatusCode != http.StatusOK {
body, _ := io.ReadAll(resp.Body)
return nil, fmt.Errorf("unexpected status %d: %s", resp.StatusCode, string(body))
}
var job NextJobResponse
if err := json.NewDecoder(resp.Body).Decode(&job); err != nil {
return nil, fmt.Errorf("failed to decode job response: %w", err)
}
return &job, nil
}
// DownloadContext downloads the job context tar file.
func (m *ManagerClient) DownloadContext(contextPath, jobToken string) (io.ReadCloser, error) {
resp, err := m.RequestLongWithToken("GET", contextPath, jobToken, nil)
if err != nil {
return nil, fmt.Errorf("failed to download context: %w", err)
}
if resp.StatusCode != http.StatusOK {
body, _ := io.ReadAll(resp.Body)
resp.Body.Close()
return nil, fmt.Errorf("context download failed with status %d: %s", resp.StatusCode, string(body))
}
return resp.Body, nil
}
// UploadFile uploads a file to the manager.
func (m *ManagerClient) UploadFile(uploadPath, jobToken, filePath string) error {
file, err := os.Open(filePath)
if err != nil {
return fmt.Errorf("failed to open file: %w", err)
}
defer file.Close()
// Create multipart form
body := &bytes.Buffer{}
writer := multipart.NewWriter(body)
part, err := writer.CreateFormFile("file", filepath.Base(filePath))
if err != nil {
return fmt.Errorf("failed to create form file: %w", err)
}
if _, err := io.Copy(part, file); err != nil {
return fmt.Errorf("failed to copy file to form: %w", err)
}
writer.Close()
fullURL := m.baseURL + uploadPath
req, err := http.NewRequest("POST", fullURL, body)
if err != nil {
return err
}
req.Header.Set("Authorization", "Bearer "+jobToken)
req.Header.Set("Content-Type", writer.FormDataContentType())
resp, err := m.longClient.Do(req)
if err != nil {
return fmt.Errorf("failed to upload file: %w", err)
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusCreated && resp.StatusCode != http.StatusOK {
respBody, _ := io.ReadAll(resp.Body)
return fmt.Errorf("upload failed with status %d: %s", resp.StatusCode, string(respBody))
}
return nil
}
// GetJobMetadata retrieves job metadata from the manager.
func (m *ManagerClient) GetJobMetadata(jobID int64) (*types.BlendMetadata, error) {
path := fmt.Sprintf("/api/runner/jobs/%d/metadata?runner_id=%d", jobID, m.runnerID)
resp, err := m.Request("GET", path, nil)
if err != nil {
return nil, err
}
defer resp.Body.Close()
if resp.StatusCode == http.StatusNotFound {
return nil, nil // No metadata found
}
if resp.StatusCode != http.StatusOK {
body, _ := io.ReadAll(resp.Body)
return nil, fmt.Errorf("failed to get job metadata: %s", string(body))
}
var metadata types.BlendMetadata
if err := json.NewDecoder(resp.Body).Decode(&metadata); err != nil {
return nil, err
}
return &metadata, nil
}
// JobFile represents a file associated with a job.
type JobFile struct {
ID int64 `json:"id"`
JobID int64 `json:"job_id"`
FileType string `json:"file_type"`
FilePath string `json:"file_path"`
FileName string `json:"file_name"`
FileSize int64 `json:"file_size"`
}
// GetJobFiles retrieves the list of files for a job.
func (m *ManagerClient) GetJobFiles(jobID int64) ([]JobFile, error) {
path := fmt.Sprintf("/api/runner/jobs/%d/files?runner_id=%d", jobID, m.runnerID)
resp, err := m.Request("GET", path, nil)
if err != nil {
return nil, err
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
body, _ := io.ReadAll(resp.Body)
return nil, fmt.Errorf("failed to get job files: %s", string(body))
}
var files []JobFile
if err := json.NewDecoder(resp.Body).Decode(&files); err != nil {
return nil, err
}
return files, nil
}
// DownloadFrame downloads a frame file from the manager.
func (m *ManagerClient) DownloadFrame(jobID int64, fileName, destPath string) error {
encodedFileName := url.PathEscape(fileName)
path := fmt.Sprintf("/api/runner/files/%d/%s?runner_id=%d", jobID, encodedFileName, m.runnerID)
resp, err := m.RequestLong("GET", path, nil)
if err != nil {
return err
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
body, _ := io.ReadAll(resp.Body)
return fmt.Errorf("download failed: %s", string(body))
}
file, err := os.Create(destPath)
if err != nil {
return err
}
defer file.Close()
_, err = io.Copy(file, resp.Body)
return err
}
// SubmitMetadata submits extracted metadata to the manager.
func (m *ManagerClient) SubmitMetadata(jobID int64, metadata types.BlendMetadata) error {
metadataJSON, err := json.Marshal(metadata)
if err != nil {
return fmt.Errorf("failed to marshal metadata: %w", err)
}
path := fmt.Sprintf("/api/runner/jobs/%d/metadata?runner_id=%d", jobID, m.runnerID)
fullURL := m.baseURL + path
req, err := http.NewRequest("POST", fullURL, bytes.NewReader(metadataJSON))
if err != nil {
return fmt.Errorf("failed to create request: %w", err)
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Authorization", "Bearer "+m.apiKey)
resp, err := m.httpClient.Do(req)
if err != nil {
return fmt.Errorf("failed to submit metadata: %w", err)
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
body, _ := io.ReadAll(resp.Body)
return fmt.Errorf("metadata submission failed: %s", string(body))
}
return nil
}
// DownloadBlender downloads a Blender version from the manager.
func (m *ManagerClient) DownloadBlender(version string) (io.ReadCloser, error) {
path := fmt.Sprintf("/api/runner/blender/download?version=%s&runner_id=%d", version, m.runnerID)
resp, err := m.RequestLong("GET", path, nil)
if err != nil {
return nil, fmt.Errorf("failed to download blender from manager: %w", err)
}
if resp.StatusCode != http.StatusOK {
body, _ := io.ReadAll(resp.Body)
resp.Body.Close()
return nil, fmt.Errorf("failed to download blender: status %d, body: %s", resp.StatusCode, string(body))
}
return resp.Body, nil
}