422 lines
12 KiB
Go
422 lines
12 KiB
Go
// 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
|
|
}
|