Files
jiggablend/internal/api/runners.go
2025-11-21 17:31:18 -06:00

583 lines
16 KiB
Go

package api
import (
"context"
"database/sql"
"encoding/json"
"fmt"
"io"
"log"
"net/http"
"strings"
"time"
"github.com/go-chi/chi/v5"
"fuego/internal/auth"
"fuego/pkg/types"
)
// handleListRunners lists all runners
func (s *Server) handleListRunners(w http.ResponseWriter, r *http.Request) {
_, err := getUserID(r)
if err != nil {
s.respondError(w, http.StatusUnauthorized, err.Error())
return
}
rows, err := s.db.Query(
`SELECT id, name, hostname, ip_address, status, last_heartbeat, capabilities, created_at
FROM runners ORDER BY created_at DESC`,
)
if err != nil {
s.respondError(w, http.StatusInternalServerError, fmt.Sprintf("Failed to query runners: %v", err))
return
}
defer rows.Close()
runners := []types.Runner{}
for rows.Next() {
var runner types.Runner
err := rows.Scan(
&runner.ID, &runner.Name, &runner.Hostname, &runner.IPAddress,
&runner.Status, &runner.LastHeartbeat, &runner.Capabilities, &runner.CreatedAt,
)
if err != nil {
s.respondError(w, http.StatusInternalServerError, fmt.Sprintf("Failed to scan runner: %v", err))
return
}
runners = append(runners, runner)
}
s.respondJSON(w, http.StatusOK, runners)
}
// runnerAuthMiddleware verifies runner requests using HMAC signatures
func (s *Server) runnerAuthMiddleware(next http.HandlerFunc) http.HandlerFunc {
return func(w http.ResponseWriter, r *http.Request) {
// Get runner ID from query string
runnerIDStr := r.URL.Query().Get("runner_id")
if runnerIDStr == "" {
s.respondError(w, http.StatusBadRequest, "runner_id required in query string")
return
}
var runnerID int64
_, err := fmt.Sscanf(runnerIDStr, "%d", &runnerID)
if err != nil {
s.respondError(w, http.StatusBadRequest, "invalid runner_id")
return
}
// Get runner secret
runnerSecret, err := s.secrets.GetRunnerSecret(runnerID)
if err != nil {
s.respondError(w, http.StatusUnauthorized, "runner not found or not verified")
return
}
// Verify request signature
valid, err := auth.VerifyRequest(r, runnerSecret, 5*time.Minute)
if err != nil || !valid {
s.respondError(w, http.StatusUnauthorized, "invalid signature")
return
}
// Add runner ID to context
ctx := r.Context()
ctx = context.WithValue(ctx, "runner_id", runnerID)
next(w, r.WithContext(ctx))
}
}
// handleRegisterRunner registers a new runner
func (s *Server) handleRegisterRunner(w http.ResponseWriter, r *http.Request) {
var req struct {
types.RegisterRunnerRequest
RegistrationToken string `json:"registration_token"`
}
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
s.respondError(w, http.StatusBadRequest, "Invalid request body")
return
}
if req.Name == "" {
s.respondError(w, http.StatusBadRequest, "Runner name is required")
return
}
if req.RegistrationToken == "" {
s.respondError(w, http.StatusBadRequest, "Registration token is required")
return
}
// Validate registration token
valid, err := s.secrets.ValidateRegistrationToken(req.RegistrationToken)
if err != nil {
s.respondError(w, http.StatusInternalServerError, fmt.Sprintf("Failed to validate token: %v", err))
return
}
if !valid {
s.respondError(w, http.StatusUnauthorized, "Invalid or expired registration token")
return
}
// Get manager secret
managerSecret, err := s.secrets.GetManagerSecret()
if err != nil {
s.respondError(w, http.StatusInternalServerError, "Failed to get manager secret")
return
}
// Generate runner secret
runnerSecret, err := s.secrets.GenerateRunnerSecret()
if err != nil {
s.respondError(w, http.StatusInternalServerError, "Failed to generate runner secret")
return
}
// Register runner
result, err := s.db.Exec(
`INSERT INTO runners (name, hostname, ip_address, status, last_heartbeat, capabilities,
registration_token, runner_secret, manager_secret, verified)
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`,
req.Name, req.Hostname, req.IPAddress, types.RunnerStatusOnline, time.Now(), req.Capabilities,
req.RegistrationToken, runnerSecret, managerSecret, true,
)
if err != nil {
s.respondError(w, http.StatusInternalServerError, fmt.Sprintf("Failed to register runner: %v", err))
return
}
runnerID, _ := result.LastInsertId()
// Return runner info with secrets
s.respondJSON(w, http.StatusCreated, map[string]interface{}{
"id": runnerID,
"name": req.Name,
"hostname": req.Hostname,
"ip_address": req.IPAddress,
"status": types.RunnerStatusOnline,
"runner_secret": runnerSecret,
"manager_secret": managerSecret,
"verified": true,
})
}
// handleRunnerHeartbeat updates runner heartbeat
func (s *Server) handleRunnerHeartbeat(w http.ResponseWriter, r *http.Request) {
runnerID, ok := r.Context().Value("runner_id").(int64)
if !ok {
s.respondError(w, http.StatusBadRequest, "runner_id not found in context")
return
}
_, err := s.db.Exec(
`UPDATE runners SET last_heartbeat = ?, status = ? WHERE id = ?`,
time.Now(), types.RunnerStatusOnline, runnerID,
)
if err != nil {
s.respondError(w, http.StatusInternalServerError, fmt.Sprintf("Failed to update heartbeat: %v", err))
return
}
s.respondJSON(w, http.StatusOK, map[string]string{"message": "Heartbeat updated"})
}
// handleGetRunnerTasks gets pending tasks for a runner
func (s *Server) handleGetRunnerTasks(w http.ResponseWriter, r *http.Request) {
runnerID, ok := r.Context().Value("runner_id").(int64)
if !ok {
s.respondError(w, http.StatusBadRequest, "runner_id not found in context")
return
}
// Get pending tasks
rows, err := s.db.Query(
`SELECT t.id, t.job_id, t.runner_id, t.frame_start, t.frame_end, t.status, t.output_path,
t.created_at, t.started_at, t.completed_at, t.error_message,
j.name as job_name, j.output_format
FROM tasks t
JOIN jobs j ON t.job_id = j.id
WHERE t.status = ? AND j.status != ?
ORDER BY t.created_at ASC
LIMIT 10`,
types.TaskStatusPending, types.JobStatusCancelled,
)
if err != nil {
s.respondError(w, http.StatusInternalServerError, fmt.Sprintf("Failed to query tasks: %v", err))
return
}
defer rows.Close()
tasks := []map[string]interface{}{}
for rows.Next() {
var task types.Task
var runnerID sql.NullInt64
var startedAt, completedAt sql.NullTime
var jobName, outputFormat string
err := rows.Scan(
&task.ID, &task.JobID, &runnerID, &task.FrameStart, &task.FrameEnd,
&task.Status, &task.OutputPath, &task.CreatedAt,
&startedAt, &completedAt, &task.ErrorMessage,
&jobName, &outputFormat,
)
if err != nil {
s.respondError(w, http.StatusInternalServerError, fmt.Sprintf("Failed to scan task: %v", err))
return
}
if runnerID.Valid {
task.RunnerID = &runnerID.Int64
}
if startedAt.Valid {
task.StartedAt = &startedAt.Time
}
if completedAt.Valid {
task.CompletedAt = &completedAt.Time
}
// Get input files for the job
var inputFiles []string
fileRows, err := s.db.Query(
`SELECT file_path FROM job_files WHERE job_id = ? AND file_type = ?`,
task.JobID, types.JobFileTypeInput,
)
if err == nil {
for fileRows.Next() {
var filePath string
if err := fileRows.Scan(&filePath); err == nil {
inputFiles = append(inputFiles, filePath)
}
}
fileRows.Close()
}
tasks = append(tasks, map[string]interface{}{
"task": task,
"job_name": jobName,
"output_format": outputFormat,
"input_files": inputFiles,
})
// Assign task to runner
_, err = s.db.Exec(
`UPDATE tasks SET runner_id = ?, status = ? WHERE id = ?`,
runnerID, types.TaskStatusRunning, task.ID,
)
if err != nil {
s.respondError(w, http.StatusInternalServerError, fmt.Sprintf("Failed to assign task: %v", err))
return
}
}
s.respondJSON(w, http.StatusOK, tasks)
}
// handleCompleteTask marks a task as completed
func (s *Server) handleCompleteTask(w http.ResponseWriter, r *http.Request) {
taskID, err := parseID(r, "id")
if err != nil {
s.respondError(w, http.StatusBadRequest, err.Error())
return
}
var req struct {
OutputPath string `json:"output_path"`
Success bool `json:"success"`
Error string `json:"error,omitempty"`
}
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
s.respondError(w, http.StatusBadRequest, "Invalid request body")
return
}
status := types.TaskStatusCompleted
if !req.Success {
status = types.TaskStatusFailed
}
now := time.Now()
_, err = s.db.Exec(
`UPDATE tasks SET status = ?, output_path = ?, completed_at = ?, error_message = ? WHERE id = ?`,
status, req.OutputPath, now, req.Error, taskID,
)
if err != nil {
s.respondError(w, http.StatusInternalServerError, fmt.Sprintf("Failed to update task: %v", err))
return
}
// Update job progress
var jobID int64
var frameStart, frameEnd int
err = s.db.QueryRow(
`SELECT job_id, frame_start, frame_end FROM tasks WHERE id = ?`,
taskID,
).Scan(&jobID, &frameStart, &frameEnd)
if err == nil {
// Count completed tasks
var totalTasks, completedTasks int
s.db.QueryRow(
`SELECT COUNT(*) FROM tasks WHERE job_id = ?`,
jobID,
).Scan(&totalTasks)
s.db.QueryRow(
`SELECT COUNT(*) FROM tasks WHERE job_id = ? AND status = ?`,
jobID, types.TaskStatusCompleted,
).Scan(&completedTasks)
progress := float64(completedTasks) / float64(totalTasks) * 100.0
// Update job status
var jobStatus string
var outputFormat string
s.db.QueryRow(`SELECT output_format FROM jobs WHERE id = ?`, jobID).Scan(&outputFormat)
if completedTasks == totalTasks {
jobStatus = string(types.JobStatusCompleted)
now := time.Now()
s.db.Exec(
`UPDATE jobs SET status = ?, progress = ?, completed_at = ? WHERE id = ?`,
jobStatus, progress, now, jobID,
)
// For MP4 jobs, create a video generation task
if outputFormat == "MP4" {
go s.generateMP4Video(jobID)
}
} else {
jobStatus = string(types.JobStatusRunning)
var startedAt sql.NullTime
s.db.QueryRow(`SELECT started_at FROM jobs WHERE id = ?`, jobID).Scan(&startedAt)
if !startedAt.Valid {
now := time.Now()
s.db.Exec(`UPDATE jobs SET started_at = ? WHERE id = ?`, now, jobID)
}
s.db.Exec(
`UPDATE jobs SET status = ?, progress = ? WHERE id = ?`,
jobStatus, progress, jobID,
)
}
}
s.respondJSON(w, http.StatusOK, map[string]string{"message": "Task completed"})
}
// handleUpdateTaskProgress updates task progress
func (s *Server) handleUpdateTaskProgress(w http.ResponseWriter, r *http.Request) {
_, err := parseID(r, "id")
if err != nil {
s.respondError(w, http.StatusBadRequest, err.Error())
return
}
var req struct {
Progress float64 `json:"progress"`
}
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
s.respondError(w, http.StatusBadRequest, "Invalid request body")
return
}
// This is mainly for logging/debugging, actual progress is calculated from completed tasks
s.respondJSON(w, http.StatusOK, map[string]string{"message": "Progress updated"})
}
// handleDownloadFileForRunner allows runners to download job files
func (s *Server) handleDownloadFileForRunner(w http.ResponseWriter, r *http.Request) {
jobID, err := parseID(r, "jobId")
if err != nil {
s.respondError(w, http.StatusBadRequest, err.Error())
return
}
fileName := chi.URLParam(r, "fileName")
// Find the file in the database
var filePath string
err = s.db.QueryRow(
`SELECT file_path FROM job_files WHERE job_id = ? AND file_name = ?`,
jobID, fileName,
).Scan(&filePath)
if err == sql.ErrNoRows {
s.respondError(w, http.StatusNotFound, "File not found")
return
}
if err != nil {
s.respondError(w, http.StatusInternalServerError, fmt.Sprintf("Failed to query file: %v", err))
return
}
// Open and serve file
file, err := s.storage.GetFile(filePath)
if err != nil {
s.respondError(w, http.StatusNotFound, "File not found on disk")
return
}
defer file.Close()
w.Header().Set("Content-Type", "application/octet-stream")
w.Header().Set("Content-Disposition", fmt.Sprintf("attachment; filename=%s", fileName))
io.Copy(w, file)
}
// handleUploadFileFromRunner allows runners to upload output files
func (s *Server) handleUploadFileFromRunner(w http.ResponseWriter, r *http.Request) {
jobID, err := parseID(r, "jobId")
if err != nil {
s.respondError(w, http.StatusBadRequest, err.Error())
return
}
err = r.ParseMultipartForm(100 << 20) // 100 MB
if err != nil {
s.respondError(w, http.StatusBadRequest, "Failed to parse form")
return
}
file, header, err := r.FormFile("file")
if err != nil {
s.respondError(w, http.StatusBadRequest, "No file provided")
return
}
defer file.Close()
// Save file
filePath, err := s.storage.SaveOutput(jobID, header.Filename, file)
if err != nil {
s.respondError(w, http.StatusInternalServerError, fmt.Sprintf("Failed to save file: %v", err))
return
}
// Record in database
_, err = s.db.Exec(
`INSERT INTO job_files (job_id, file_type, file_path, file_name, file_size)
VALUES (?, ?, ?, ?, ?)`,
jobID, types.JobFileTypeOutput, filePath, header.Filename, header.Size,
)
if err != nil {
s.respondError(w, http.StatusInternalServerError, fmt.Sprintf("Failed to record file: %v", err))
return
}
s.respondJSON(w, http.StatusCreated, map[string]interface{}{
"file_path": filePath,
"file_name": header.Filename,
})
}
// generateMP4Video generates MP4 video from PNG frames for a completed job
func (s *Server) generateMP4Video(jobID int64) {
// This would be called by a runner or external process
// For now, we'll create a special task that runners can pick up
// In a production system, you might want to use a job queue or have a dedicated video processor
// Get all PNG output files for this job
rows, err := s.db.Query(
`SELECT file_path, file_name FROM job_files
WHERE job_id = ? AND file_type = ? AND file_name LIKE '%.png'
ORDER BY file_name`,
jobID, types.JobFileTypeOutput,
)
if err != nil {
log.Printf("Failed to query PNG files for job %d: %v", jobID, err)
return
}
defer rows.Close()
var pngFiles []string
for rows.Next() {
var filePath, fileName string
if err := rows.Scan(&filePath, &fileName); err == nil {
pngFiles = append(pngFiles, filePath)
}
}
if len(pngFiles) == 0 {
log.Printf("No PNG files found for job %d", jobID)
return
}
// Note: Video generation will be handled by runners when they complete tasks
// Runners can check job status and generate MP4 when all frames are complete
log.Printf("Job %d completed with %d PNG frames - ready for MP4 generation", jobID, len(pngFiles))
}
// handleGetJobStatusForRunner allows runners to check job status
func (s *Server) handleGetJobStatusForRunner(w http.ResponseWriter, r *http.Request) {
jobID, err := parseID(r, "jobId")
if err != nil {
s.respondError(w, http.StatusBadRequest, err.Error())
return
}
var job types.Job
var startedAt, completedAt sql.NullTime
err = s.db.QueryRow(
`SELECT id, user_id, name, status, progress, frame_start, frame_end, output_format,
created_at, started_at, completed_at, error_message
FROM jobs WHERE id = ?`,
jobID,
).Scan(
&job.ID, &job.UserID, &job.Name, &job.Status, &job.Progress,
&job.FrameStart, &job.FrameEnd, &job.OutputFormat,
&job.CreatedAt, &startedAt, &completedAt, &job.ErrorMessage,
)
if err == sql.ErrNoRows {
s.respondError(w, http.StatusNotFound, "Job not found")
return
}
if err != nil {
s.respondError(w, http.StatusInternalServerError, fmt.Sprintf("Failed to query job: %v", err))
return
}
if startedAt.Valid {
job.StartedAt = &startedAt.Time
}
if completedAt.Valid {
job.CompletedAt = &completedAt.Time
}
s.respondJSON(w, http.StatusOK, job)
}
// handleGetJobFilesForRunner allows runners to get job files
func (s *Server) handleGetJobFilesForRunner(w http.ResponseWriter, r *http.Request) {
jobID, err := parseID(r, "jobId")
if err != nil {
s.respondError(w, http.StatusBadRequest, err.Error())
return
}
rows, err := s.db.Query(
`SELECT id, job_id, file_type, file_path, file_name, file_size, created_at
FROM job_files WHERE job_id = ? ORDER BY file_name`,
jobID,
)
if err != nil {
s.respondError(w, http.StatusInternalServerError, fmt.Sprintf("Failed to query files: %v", err))
return
}
defer rows.Close()
files := []types.JobFile{}
for rows.Next() {
var file types.JobFile
err := rows.Scan(
&file.ID, &file.JobID, &file.FileType, &file.FilePath,
&file.FileName, &file.FileSize, &file.CreatedAt,
)
if err != nil {
s.respondError(w, http.StatusInternalServerError, fmt.Sprintf("Failed to scan file: %v", err))
return
}
files = append(files, file)
}
s.respondJSON(w, http.StatusOK, files)
}