803677be57
- Add API server with WebSocket support and REST endpoints - Implement authentication system with API keys and permissions - Add task queue system with Redis backend and error handling - Include storage layer with database migrations and schemas - Add comprehensive logging, metrics, and telemetry - Implement security middleware and network utilities - Add experiment management and container orchestration - Include configuration management with smart defaults
302 lines
8.1 KiB
Go
302 lines
8.1 KiB
Go
package controller
|
|
|
|
import (
|
|
"fmt"
|
|
"time"
|
|
|
|
"github.com/charmbracelet/bubbles/key"
|
|
"github.com/charmbracelet/bubbles/list"
|
|
tea "github.com/charmbracelet/bubbletea"
|
|
"github.com/jfraeys/fetch_ml/cmd/tui/internal/config"
|
|
"github.com/jfraeys/fetch_ml/cmd/tui/internal/model"
|
|
"github.com/jfraeys/fetch_ml/cmd/tui/internal/services"
|
|
"github.com/jfraeys/fetch_ml/internal/logging"
|
|
)
|
|
|
|
// Controller handles all business logic and state updates
|
|
type Controller struct {
|
|
config *config.Config
|
|
server *services.MLServer
|
|
taskQueue *services.TaskQueue
|
|
logger *logging.Logger
|
|
}
|
|
|
|
// New creates a new Controller instance
|
|
func New(cfg *config.Config, srv *services.MLServer, tq *services.TaskQueue, logger *logging.Logger) *Controller {
|
|
return &Controller{
|
|
config: cfg,
|
|
server: srv,
|
|
taskQueue: tq,
|
|
logger: logger,
|
|
}
|
|
}
|
|
|
|
// Init initializes the TUI and returns initial commands
|
|
func (c *Controller) Init() tea.Cmd {
|
|
return tea.Batch(
|
|
tea.SetWindowTitle("FetchML"),
|
|
c.loadAllData(),
|
|
tickCmd(),
|
|
)
|
|
}
|
|
|
|
// Update handles all messages and updates the state
|
|
func (c *Controller) Update(msg tea.Msg, m model.State) (model.State, tea.Cmd) {
|
|
var cmds []tea.Cmd
|
|
|
|
switch msg := msg.(type) {
|
|
case tea.KeyMsg:
|
|
// Handle input mode (for queuing jobs with args)
|
|
if m.InputMode {
|
|
switch msg.String() {
|
|
case "enter":
|
|
args := m.Input.Value()
|
|
m.Input.SetValue("")
|
|
m.InputMode = false
|
|
if job := getSelectedJob(m); job != nil {
|
|
cmds = append(cmds, c.queueJob(job.Name, args))
|
|
}
|
|
return m, tea.Batch(cmds...)
|
|
case "esc":
|
|
m.InputMode = false
|
|
m.Input.SetValue("")
|
|
return m, nil
|
|
}
|
|
var cmd tea.Cmd
|
|
m.Input, cmd = m.Input.Update(msg)
|
|
return m, cmd
|
|
}
|
|
|
|
// Handle settings-specific keys
|
|
if m.ActiveView == model.ViewModeSettings {
|
|
switch msg.String() {
|
|
case "up", "k":
|
|
if m.SettingsIndex > 1 { // Skip index 0 (Status)
|
|
m.SettingsIndex--
|
|
cmds = append(cmds, c.updateSettingsContent(m))
|
|
if m.SettingsIndex == 1 {
|
|
m.ApiKeyInput.Focus()
|
|
} else {
|
|
m.ApiKeyInput.Blur()
|
|
}
|
|
}
|
|
case "down", "j":
|
|
if m.SettingsIndex < 2 {
|
|
m.SettingsIndex++
|
|
cmds = append(cmds, c.updateSettingsContent(m))
|
|
if m.SettingsIndex == 1 {
|
|
m.ApiKeyInput.Focus()
|
|
} else {
|
|
m.ApiKeyInput.Blur()
|
|
}
|
|
}
|
|
case "enter":
|
|
if cmd := c.handleSettingsAction(&m); cmd != nil {
|
|
cmds = append(cmds, cmd)
|
|
}
|
|
case "esc":
|
|
m.ActiveView = model.ViewModeJobs
|
|
m.ApiKeyInput.Blur()
|
|
}
|
|
if m.SettingsIndex == 1 { // API Key input field
|
|
var cmd tea.Cmd
|
|
m.ApiKeyInput, cmd = m.ApiKeyInput.Update(msg)
|
|
cmds = append(cmds, cmd)
|
|
// Force update settings view to show typed characters immediately
|
|
cmds = append(cmds, c.updateSettingsContent(m))
|
|
}
|
|
return m, tea.Batch(cmds...)
|
|
}
|
|
|
|
// Handle global keys
|
|
switch {
|
|
case key.Matches(msg, m.Keys.Quit):
|
|
return m, tea.Quit
|
|
case key.Matches(msg, m.Keys.Refresh):
|
|
m.IsLoading = true
|
|
m.Status = "Refreshing all data..."
|
|
m.LastRefresh = time.Now()
|
|
cmds = append(cmds, c.loadAllData())
|
|
case key.Matches(msg, m.Keys.RefreshGPU):
|
|
m.Status = "Refreshing GPU status..."
|
|
cmds = append(cmds, c.loadGPU())
|
|
case key.Matches(msg, m.Keys.Trigger):
|
|
if job := getSelectedJob(m); job != nil {
|
|
cmds = append(cmds, c.queueJob(job.Name, ""))
|
|
}
|
|
case key.Matches(msg, m.Keys.TriggerArgs):
|
|
if job := getSelectedJob(m); job != nil {
|
|
m.InputMode = true
|
|
m.Input.Focus()
|
|
}
|
|
case key.Matches(msg, m.Keys.ViewQueue):
|
|
m.ActiveView = model.ViewModeQueue
|
|
cmds = append(cmds, c.showQueue(m))
|
|
case key.Matches(msg, m.Keys.ViewContainer):
|
|
m.ActiveView = model.ViewModeContainer
|
|
cmds = append(cmds, c.loadContainer())
|
|
case key.Matches(msg, m.Keys.ViewGPU):
|
|
m.ActiveView = model.ViewModeGPU
|
|
cmds = append(cmds, c.loadGPU())
|
|
case key.Matches(msg, m.Keys.ViewJobs):
|
|
m.ActiveView = model.ViewModeJobs
|
|
case key.Matches(msg, m.Keys.ViewSettings):
|
|
m.ActiveView = model.ViewModeSettings
|
|
m.SettingsIndex = 1 // Start at Input field, skip Status
|
|
m.ApiKeyInput.Focus()
|
|
cmds = append(cmds, c.updateSettingsContent(m))
|
|
case key.Matches(msg, m.Keys.ViewExperiments):
|
|
m.ActiveView = model.ViewModeExperiments
|
|
cmds = append(cmds, c.loadExperiments())
|
|
case key.Matches(msg, m.Keys.Cancel):
|
|
if job := getSelectedJob(m); job != nil && job.TaskID != "" {
|
|
cmds = append(cmds, c.cancelTask(job.TaskID))
|
|
}
|
|
case key.Matches(msg, m.Keys.Delete):
|
|
if job := getSelectedJob(m); job != nil && job.Status == model.StatusPending {
|
|
cmds = append(cmds, c.deleteJob(job.Name))
|
|
}
|
|
case key.Matches(msg, m.Keys.MarkFailed):
|
|
if job := getSelectedJob(m); job != nil && job.Status == model.StatusRunning {
|
|
cmds = append(cmds, c.markFailed(job.Name))
|
|
}
|
|
case key.Matches(msg, m.Keys.Help):
|
|
m.ShowHelp = !m.ShowHelp
|
|
}
|
|
|
|
case tea.WindowSizeMsg:
|
|
m.Width = msg.Width
|
|
m.Height = msg.Height
|
|
|
|
// Update component sizes
|
|
h, v := 4, 2 // docStyle.GetFrameSize() approx
|
|
listHeight := msg.Height - v - 8
|
|
m.JobList.SetSize(msg.Width/3-h, listHeight)
|
|
|
|
panelWidth := msg.Width*2/3 - h - 2
|
|
panelHeight := (listHeight - 6) / 3
|
|
|
|
m.GpuView.Width = panelWidth
|
|
m.GpuView.Height = panelHeight
|
|
m.ContainerView.Width = panelWidth
|
|
m.ContainerView.Height = panelHeight
|
|
m.QueueView.Width = panelWidth
|
|
m.QueueView.Height = listHeight - 4
|
|
m.SettingsView.Width = panelWidth
|
|
m.SettingsView.Height = listHeight - 4
|
|
m.ExperimentsView.Width = panelWidth
|
|
m.ExperimentsView.Height = listHeight - 4
|
|
|
|
case JobsLoadedMsg:
|
|
m.Jobs = []model.Job(msg)
|
|
calculateJobStats(&m)
|
|
items := make([]list.Item, len(m.Jobs))
|
|
for i, job := range m.Jobs {
|
|
items[i] = job
|
|
}
|
|
cmds = append(cmds, m.JobList.SetItems(items))
|
|
m.Status = formatStatus(m)
|
|
m.IsLoading = false
|
|
|
|
case TasksLoadedMsg:
|
|
m.QueuedTasks = []*model.Task(msg)
|
|
m.Status = formatStatus(m)
|
|
|
|
case GpuLoadedMsg:
|
|
m.GpuView.SetContent(string(msg))
|
|
m.GpuView.GotoTop()
|
|
|
|
case ContainerLoadedMsg:
|
|
m.ContainerView.SetContent(string(msg))
|
|
m.ContainerView.GotoTop()
|
|
|
|
case QueueLoadedMsg:
|
|
m.QueueView.SetContent(string(msg))
|
|
m.QueueView.GotoTop()
|
|
|
|
case SettingsContentMsg:
|
|
m.SettingsView.SetContent(string(msg))
|
|
|
|
case ExperimentsLoadedMsg:
|
|
m.ExperimentsView.SetContent(string(msg))
|
|
m.ExperimentsView.GotoTop()
|
|
|
|
case SettingsUpdateMsg:
|
|
// Settings content was updated, just trigger a re-render
|
|
|
|
case StatusMsg:
|
|
if msg.Level == "error" {
|
|
m.ErrorMsg = msg.Text
|
|
m.Status = "Error occurred - check status"
|
|
} else {
|
|
m.ErrorMsg = ""
|
|
m.Status = msg.Text
|
|
}
|
|
|
|
case TickMsg:
|
|
var spinCmd tea.Cmd
|
|
m.Spinner, spinCmd = m.Spinner.Update(msg)
|
|
cmds = append(cmds, spinCmd)
|
|
|
|
// Auto-refresh every 10 seconds
|
|
if time.Since(m.LastRefresh) > 10*time.Second && !m.IsLoading {
|
|
m.LastRefresh = time.Now()
|
|
cmds = append(cmds, c.loadAllData())
|
|
}
|
|
cmds = append(cmds, tickCmd())
|
|
|
|
default:
|
|
var spinCmd tea.Cmd
|
|
m.Spinner, spinCmd = m.Spinner.Update(msg)
|
|
cmds = append(cmds, spinCmd)
|
|
}
|
|
|
|
// Update all bubble components
|
|
var cmd tea.Cmd
|
|
m.JobList, cmd = m.JobList.Update(msg)
|
|
cmds = append(cmds, cmd)
|
|
|
|
m.GpuView, cmd = m.GpuView.Update(msg)
|
|
cmds = append(cmds, cmd)
|
|
|
|
m.ContainerView, cmd = m.ContainerView.Update(msg)
|
|
cmds = append(cmds, cmd)
|
|
|
|
m.QueueView, cmd = m.QueueView.Update(msg)
|
|
cmds = append(cmds, cmd)
|
|
|
|
m.ExperimentsView, cmd = m.ExperimentsView.Update(msg)
|
|
cmds = append(cmds, cmd)
|
|
|
|
return m, tea.Batch(cmds...)
|
|
}
|
|
|
|
// ExperimentsLoadedMsg is sent when experiments are loaded
|
|
type ExperimentsLoadedMsg string
|
|
|
|
func (c *Controller) loadExperiments() tea.Cmd {
|
|
return func() tea.Msg {
|
|
commitIDs, err := c.taskQueue.ListExperiments()
|
|
if err != nil {
|
|
return StatusMsg{Level: "error", Text: fmt.Sprintf("Failed to list experiments: %v", err)}
|
|
}
|
|
|
|
if len(commitIDs) == 0 {
|
|
return ExperimentsLoadedMsg("Experiments:\n\nNo experiments found.")
|
|
}
|
|
|
|
var output string
|
|
output += "Experiments:\n\n"
|
|
|
|
for _, commitID := range commitIDs {
|
|
details, err := c.taskQueue.GetExperimentDetails(commitID)
|
|
if err != nil {
|
|
output += fmt.Sprintf("Error loading %s: %v\n\n", commitID, err)
|
|
continue
|
|
}
|
|
output += details + "\n----------------------------------------\n\n"
|
|
}
|
|
|
|
return ExperimentsLoadedMsg(output)
|
|
}
|
|
}
|