AIRouter/internal/task/task.go

package task

import (
	"errors"
	"fmt"
	"sync"
	"time"

	"github.com/edgeai/gateway/internal/observability"
)

// TaskState represents the lifecycle state of a task.
type TaskState string

const (
	StateQueued    TaskState = "QUEUED"
	StateRunning   TaskState = "RUNNING"
	StateStreaming TaskState = "STREAMING"
	StateCompleted TaskState = "COMPLETED"
	StateFailed    TaskState = "FAILED"
	StateCancelled TaskState = "CANCELLED"
)

// TaskPriority levels (P0 highest, P4 lowest).
type TaskPriority int

const (
	PriorityRealtime   TaskPriority = 0 // P0
	PriorityHigh       TaskPriority = 1 // P1
	PriorityNormal     TaskPriority = 2 // P2 (default)
	PriorityLow        TaskPriority = 3 // P3
	PriorityBackground TaskPriority = 4 // P4
)

// Task represents an inference task in the system.
type Task struct {
	ID           string
	RequestID    string
	SessionID    string
	AppID        string
	TenantID     string
	LogicalModel string
	ActualModel  string
	Priority     TaskPriority
	State        TaskState
	Stream       bool
	CreatedAt    time.Time
	StartedAt    *time.Time
	CompletedAt  *time.Time
	CancelReason string
	ErrorMessage string
	InputTokens  int
	OutputTokens int
	NodeID       string
	Degraded     bool
	cancelCh     chan struct{}
	cancelOnce   sync.Once
	mu           sync.RWMutex
}

// NewTask creates a new task in QUEUED state.
func NewTask(id, requestID, appID, tenantID, logicalModel string, priority TaskPriority, stream bool) *Task {
	return &Task{
		ID:           id,
		RequestID:    requestID,
		AppID:        appID,
		TenantID:     tenantID,
		LogicalModel: logicalModel,
		Priority:     priority,
		State:        StateQueued,
		Stream:       stream,
		CreatedAt:    time.Now(),
		cancelCh:     make(chan struct{}),
	}
}

// AllowedTransitions defines valid state transitions.
var allowedTransitions = map[TaskState][]TaskState{
	StateQueued:    {StateRunning, StateFailed, StateCancelled},
	StateRunning:   {StateStreaming, StateCompleted, StateFailed, StateCancelled},
	StateStreaming: {StateCompleted, StateFailed, StateCancelled},
	StateCompleted: {},
	StateFailed:    {},
	StateCancelled: {},
}

// Transition changes the task state if the transition is valid.
func (t *Task) Transition(to TaskState) error {
	t.mu.Lock()
	defer t.mu.Unlock()

	allowed, ok := allowedTransitions[t.State]
	if !ok {
		return fmt.Errorf("unknown current state: %s", t.State)
	}

	valid := false
	for _, s := range allowed {
		if s == to {
			valid = true
			break
		}
	}
	if !valid {
		return fmt.Errorf("invalid transition: %s -> %s", t.State, to)
	}

	from := t.State
	t.State = to
	now := time.Now()

	switch to {
	case StateRunning:
		t.StartedAt = &now
	case StateCompleted, StateFailed, StateCancelled:
		t.CompletedAt = &now
	}

	_ = from
	return nil
}

// Cancel signals task cancellation and transitions to CANCELLED if possible.
func (t *Task) Cancel(reason string) error {
	t.cancelOnce.Do(func() {
		close(t.cancelCh)
	})

	t.mu.Lock()
	defer t.mu.Unlock()

	if t.State == StateCompleted || t.State == StateFailed || t.State == StateCancelled {
		return errors.New("task already in terminal state")
	}

	t.CancelReason = reason
	t.State = StateCancelled
	now := time.Now()
	t.CompletedAt = &now
	return nil
}

// Cancelled returns a channel that's closed when the task is cancelled.
func (t *Task) Cancelled() <-chan struct{} {
	return t.cancelCh
}

// IsCancelled returns true if the task has been cancelled.
func (t *Task) IsCancelled() bool {
	select {
	case <-t.cancelCh:
		return true
	default:
		return false
	}
}

// GetState returns the current state (thread-safe).
func (t *Task) GetState() TaskState {
	t.mu.RLock()
	defer t.mu.RUnlock()
	return t.State
}

// IsTerminal returns true if the task is in a terminal state.
func (t *Task) IsTerminal() bool {
	s := t.GetState()
	return s == StateCompleted || s == StateFailed || s == StateCancelled
}

// StateMachineLogger logs state transitions.
type StateMachineLogger struct {
	logger *observability.Logger
}

func NewStateMachineLogger(logger *observability.Logger) *StateMachineLogger {
	return &StateMachineLogger{logger: logger}
}

// LogTransition logs a state transition.
func (sml *StateMachineLogger) LogTransition(task *Task, from, to TaskState, reason string) {
	sml.logger.Info("task state transition",
		observability.F().
			Event("state_transition").
			TaskID(task.ID).
			Set("from_state", string(from)).
			Set("to_state", string(to)).
			Reason(reason))
	_ = from // used in log field above
}