loop/statemachine.go

package loop

import (
	"context"
	"fmt"
	"log/slog"
	"sync"
	"time"
)

// State represents the possible states of the Agent state machine
type State int

const (
	// StateUnknown is the default state
	StateUnknown State = iota
	// StateReady is the initial state when the agent is initialized and ready to operate
	StateReady
	// StateWaitingForUserInput occurs when the agent is waiting for a user message to start a turn
	StateWaitingForUserInput
	// StateSendingToLLM occurs when the agent is sending message(s) to the LLM
	StateSendingToLLM
	// StateProcessingLLMResponse occurs when the agent is processing a response from the LLM
	StateProcessingLLMResponse
	// StateEndOfTurn occurs when processing is completed without tool use, and the turn ends
	StateEndOfTurn
	// StateToolUseRequested occurs when the LLM has requested to use a tool
	StateToolUseRequested
	// StateCheckingForCancellation occurs when the agent checks if user requested cancellation
	StateCheckingForCancellation
	// StateRunningTool occurs when the agent is executing the requested tool
	StateRunningTool
	// StateCheckingGitCommits occurs when the agent checks for new git commits after tool execution
	StateCheckingGitCommits
	// StateRunningAutoformatters occurs when the agent runs code formatters on new commits
	StateRunningAutoformatters
	// StateCheckingBudget occurs when the agent verifies if budget limits are exceeded
	StateCheckingBudget
	// StateGatheringAdditionalMessages occurs when the agent collects user messages that arrived during tool execution
	StateGatheringAdditionalMessages
	// StateSendingToolResults occurs when the agent sends tool results back to the LLM
	StateSendingToolResults
	// StateCancelled occurs when an operation was cancelled by the user
	StateCancelled
	// StateBudgetExceeded occurs when the budget limit was reached
	StateBudgetExceeded
	// StateError occurs when an error occurred during processing
	StateError
)

// String returns a string representation of the State for logging and debugging
func (s State) String() string {
	switch s {
	case StateUnknown:
		return "Unknown"
	case StateReady:
		return "Ready"
	case StateWaitingForUserInput:
		return "WaitingForUserInput"
	case StateSendingToLLM:
		return "SendingToLLM"
	case StateProcessingLLMResponse:
		return "ProcessingLLMResponse"
	case StateEndOfTurn:
		return "EndOfTurn"
	case StateToolUseRequested:
		return "ToolUseRequested"
	case StateCheckingForCancellation:
		return "CheckingForCancellation"
	case StateRunningTool:
		return "RunningTool"
	case StateCheckingGitCommits:
		return "CheckingGitCommits"
	case StateRunningAutoformatters:
		return "RunningAutoformatters"
	case StateCheckingBudget:
		return "CheckingBudget"
	case StateGatheringAdditionalMessages:
		return "GatheringAdditionalMessages"
	case StateSendingToolResults:
		return "SendingToolResults"
	case StateCancelled:
		return "Cancelled"
	case StateBudgetExceeded:
		return "BudgetExceeded"
	case StateError:
		return "Error"
	default:
		return fmt.Sprintf("Unknown(%d)", int(s))
	}
}

// TransitionEvent represents an event that causes a state transition
type TransitionEvent struct {
	// Description provides a human-readable description of the event
	Description string
	// Data can hold any additional information about the event
	Data interface{}
	// Timestamp is when the event occurred
	Timestamp time.Time
}

// StateTransition represents a transition from one state to another
type StateTransition struct {
	From  State
	To    State
	Event TransitionEvent
}

// StateMachine manages the Agent's states and transitions
type StateMachine struct {
	// mu protects all fields of the StateMachine from concurrent access
	mu sync.RWMutex
	// currentState is the current state of the state machine
	currentState State
	// previousState is the previous state of the state machine
	previousState State
	// stateEnteredAt is when the current state was entered
	stateEnteredAt time.Time
	// transitions maps from states to the states they can transition to
	transitions map[State]map[State]bool
	// history records the history of state transitions
	history []StateTransition
	// maxHistorySize limits the number of transitions to keep in history
	maxHistorySize int
	// eventListeners are notified when state transitions occur
	eventListeners []chan<- StateTransition
	// onTransition is a callback function that's called when a transition occurs
	onTransition func(ctx context.Context, from, to State, event TransitionEvent)
}

// NewStateMachine creates a new state machine initialized to StateReady
func NewStateMachine() *StateMachine {
	sm := &StateMachine{
		currentState:   StateReady,
		previousState:  StateUnknown,
		stateEnteredAt: time.Now(),
		transitions:    make(map[State]map[State]bool),
		maxHistorySize: 100,
		eventListeners: make([]chan<- StateTransition, 0),
	}

	// Initialize valid transitions
	sm.initTransitions()

	return sm
}

// SetMaxHistorySize sets the maximum number of transitions to keep in history
func (sm *StateMachine) SetMaxHistorySize(size int) {
	if size < 1 {
		size = 1 // Ensure we keep at least one entry
	}

	sm.mu.Lock()
	defer sm.mu.Unlock()

	sm.maxHistorySize = size

	// Trim history if needed
	if len(sm.history) > sm.maxHistorySize {
		sm.history = sm.history[len(sm.history)-sm.maxHistorySize:]
	}
}

// AddTransitionListener adds a listener channel that will be notified of state transitions
// Returns a function that can be called to remove the listener
func (sm *StateMachine) AddTransitionListener(listener chan<- StateTransition) func() {
	sm.mu.Lock()
	defer sm.mu.Unlock()

	sm.eventListeners = append(sm.eventListeners, listener)

	// Return a function to remove this listener
	return func() {
		sm.RemoveTransitionListener(listener)
	}
}

// RemoveTransitionListener removes a previously added listener
func (sm *StateMachine) RemoveTransitionListener(listener chan<- StateTransition) {
	sm.mu.Lock()
	defer sm.mu.Unlock()

	for i, l := range sm.eventListeners {
		if l == listener {
			// Remove by swapping with the last element and then truncating
			lastIdx := len(sm.eventListeners) - 1
			sm.eventListeners[i] = sm.eventListeners[lastIdx]
			sm.eventListeners = sm.eventListeners[:lastIdx]
			break
		}
	}
}

// SetTransitionCallback sets a function to be called on every state transition
func (sm *StateMachine) SetTransitionCallback(callback func(ctx context.Context, from, to State, event TransitionEvent)) {
	sm.mu.Lock()
	defer sm.mu.Unlock()

	sm.onTransition = callback
}

// ClearTransitionCallback removes any previously set transition callback
func (sm *StateMachine) ClearTransitionCallback() {
	sm.mu.Lock()
	defer sm.mu.Unlock()

	sm.onTransition = nil
}

// initTransitions initializes the map of valid state transitions
func (sm *StateMachine) initTransitions() {
	// Helper function to add transitions
	addTransition := func(from State, to ...State) {
		// Initialize the map for this state if it doesn't exist
		if _, exists := sm.transitions[from]; !exists {
			sm.transitions[from] = make(map[State]bool)
		}

		// Add all of the 'to' states
		for _, toState := range to {
			sm.transitions[from][toState] = true
		}
	}

	// Define valid transitions based on the state machine diagram

	// Initial state
	addTransition(StateReady, StateWaitingForUserInput)

	// Main flow
	addTransition(StateWaitingForUserInput, StateSendingToLLM, StateError)
	addTransition(StateSendingToLLM, StateProcessingLLMResponse, StateError)
	addTransition(StateProcessingLLMResponse, StateEndOfTurn, StateToolUseRequested, StateError)
	addTransition(StateEndOfTurn, StateWaitingForUserInput)

	// Tool use flow
	addTransition(StateToolUseRequested, StateCheckingForCancellation)
	addTransition(StateCheckingForCancellation, StateRunningTool, StateCancelled)
	addTransition(StateRunningTool, StateCheckingGitCommits, StateError)
	addTransition(StateCheckingGitCommits, StateRunningAutoformatters, StateCheckingBudget)
	addTransition(StateRunningAutoformatters, StateCheckingBudget)
	addTransition(StateCheckingBudget, StateGatheringAdditionalMessages, StateBudgetExceeded)
	addTransition(StateGatheringAdditionalMessages, StateSendingToolResults, StateError)
	addTransition(StateSendingToolResults, StateProcessingLLMResponse, StateError)

	// Terminal states to new turn
	addTransition(StateCancelled, StateWaitingForUserInput)
	addTransition(StateBudgetExceeded, StateWaitingForUserInput)
	addTransition(StateError, StateWaitingForUserInput)
}

// Transition attempts to transition from the current state to the given state
func (sm *StateMachine) Transition(ctx context.Context, newState State, event string) error {
	if sm == nil {
		return fmt.Errorf("nil StateMachine pointer")
	}
	transitionEvent := TransitionEvent{
		Description: event,
		Timestamp:   time.Now(),
	}
	return sm.TransitionWithEvent(ctx, newState, transitionEvent)
}

// TransitionWithEvent attempts to transition from the current state to the given state
// with the provided event information
func (sm *StateMachine) TransitionWithEvent(ctx context.Context, newState State, event TransitionEvent) error {
	// First check if the transition is valid without holding the write lock
	sm.mu.RLock()
	currentState := sm.currentState
	canTransition := false
	if validToStates, exists := sm.transitions[currentState]; exists {
		canTransition = validToStates[newState]
	}
	sm.mu.RUnlock()

	if !canTransition {
		return fmt.Errorf("invalid transition from %s to %s", currentState, newState)
	}

	// Acquire write lock for the actual transition
	sm.mu.Lock()
	defer sm.mu.Unlock()

	// Double-check that the state hasn't changed since we checked
	if sm.currentState != currentState {
		// State changed between our check and lock acquisition
		// Re-check if the transition is still valid
		if validToStates, exists := sm.transitions[sm.currentState]; !exists || !validToStates[newState] {
			return fmt.Errorf("concurrent state change detected: invalid transition from current %s to %s",
				sm.currentState, newState)
		}
	}

	// Calculate duration in current state
	duration := time.Since(sm.stateEnteredAt)

	// Record the transition
	transition := StateTransition{
		From:  sm.currentState,
		To:    newState,
		Event: event,
	}

	// Update state
	sm.previousState = sm.currentState
	sm.currentState = newState
	sm.stateEnteredAt = time.Now()

	// Add to history
	sm.history = append(sm.history, transition)

	// Trim history if it exceeds maximum size
	if len(sm.history) > sm.maxHistorySize {
		sm.history = sm.history[len(sm.history)-sm.maxHistorySize:]
	}

	// Make a local copy of any callback functions to invoke outside the lock
	var onTransition func(ctx context.Context, from, to State, event TransitionEvent)
	var eventListenersCopy []chan<- StateTransition
	if sm.onTransition != nil {
		onTransition = sm.onTransition
	}
	if len(sm.eventListeners) > 0 {
		eventListenersCopy = make([]chan<- StateTransition, len(sm.eventListeners))
		copy(eventListenersCopy, sm.eventListeners)
	}

	// Log the transition
	slog.InfoContext(ctx, "State transition",
		"from", sm.previousState,
		"to", sm.currentState,
		"event", event.Description,
		"duration", duration)

	// Release the lock before notifying listeners to avoid deadlocks
	sm.mu.Unlock()

	// Notify listeners if any
	if onTransition != nil {
		onTransition(ctx, sm.previousState, sm.currentState, event)
	}

	for _, ch := range eventListenersCopy {
		select {
		case ch <- transition:
			// Successfully sent
		default:
			// Channel buffer full or no receiver, log and continue
			slog.WarnContext(ctx, "Failed to notify state transition listener",
				"from", sm.previousState, "to", sm.currentState)
		}
	}

	// Re-acquire the lock that we explicitly released above
	sm.mu.Lock()
	return nil
}

// CurrentState returns the current state
func (sm *StateMachine) CurrentState() State {
	sm.mu.RLock()
	defer sm.mu.RUnlock()
	return sm.currentState
}

// PreviousState returns the previous state
func (sm *StateMachine) PreviousState() State {
	sm.mu.RLock()
	defer sm.mu.RUnlock()
	return sm.previousState
}

// TimeInState returns how long the machine has been in the current state
func (sm *StateMachine) TimeInState() time.Duration {
	sm.mu.RLock()
	enteredAt := sm.stateEnteredAt
	sm.mu.RUnlock()
	return time.Since(enteredAt)
}

// CanTransition returns true if a transition from the from state to the to state is valid
func (sm *StateMachine) CanTransition(from, to State) bool {
	sm.mu.RLock()
	defer sm.mu.RUnlock()
	if validToStates, exists := sm.transitions[from]; exists {
		return validToStates[to]
	}
	return false
}

// History returns the transition history of the state machine
func (sm *StateMachine) History() []StateTransition {
	sm.mu.RLock()
	defer sm.mu.RUnlock()

	// Return a copy to prevent modification
	historyCopy := make([]StateTransition, len(sm.history))
	copy(historyCopy, sm.history)
	return historyCopy
}

// Reset resets the state machine to the initial ready state
func (sm *StateMachine) Reset() {
	sm.mu.Lock()
	defer sm.mu.Unlock()

	sm.currentState = StateReady
	sm.previousState = StateUnknown
	sm.stateEnteredAt = time.Now()
}

// IsInTerminalState returns whether the current state is a terminal state
func (sm *StateMachine) IsInTerminalState() bool {
	sm.mu.RLock()
	defer sm.mu.RUnlock()

	switch sm.currentState {
	case StateEndOfTurn, StateCancelled, StateBudgetExceeded, StateError:
		return true
	default:
		return false
	}
}

// IsInErrorState returns whether the current state is an error state
func (sm *StateMachine) IsInErrorState() bool {
	sm.mu.RLock()
	defer sm.mu.RUnlock()

	switch sm.currentState {
	case StateError, StateCancelled, StateBudgetExceeded:
		return true
	default:
		return false
	}
}

// ForceTransition forces a transition regardless of whether it's valid according to the state machine rules
// This should be used only in critical situations like cancellation or error recovery
func (sm *StateMachine) ForceTransition(ctx context.Context, newState State, reason string) {
	event := TransitionEvent{
		Description: fmt.Sprintf("Forced transition: %s", reason),
		Timestamp:   time.Now(),
	}

	sm.mu.Lock()

	// Calculate duration in current state
	duration := time.Since(sm.stateEnteredAt)

	// Record the transition
	transition := StateTransition{
		From:  sm.currentState,
		To:    newState,
		Event: event,
	}

	// Update state
	sm.previousState = sm.currentState
	sm.currentState = newState
	sm.stateEnteredAt = time.Now()

	// Add to history
	sm.history = append(sm.history, transition)

	// Trim history if it exceeds maximum size
	if len(sm.history) > sm.maxHistorySize {
		sm.history = sm.history[len(sm.history)-sm.maxHistorySize:]
	}

	// Make a local copy of any callback functions to invoke outside the lock
	var onTransition func(ctx context.Context, from, to State, event TransitionEvent)
	var eventListenersCopy []chan<- StateTransition
	if sm.onTransition != nil {
		onTransition = sm.onTransition
	}
	if len(sm.eventListeners) > 0 {
		eventListenersCopy = make([]chan<- StateTransition, len(sm.eventListeners))
		copy(eventListenersCopy, sm.eventListeners)
	}

	// Log the transition
	slog.WarnContext(ctx, "Forced state transition",
		"from", sm.previousState,
		"to", sm.currentState,
		"reason", reason,
		"duration", duration)

	// Release the lock before notifying listeners to avoid deadlocks
	sm.mu.Unlock()

	// Notify listeners if any
	if onTransition != nil {
		onTransition(ctx, sm.previousState, sm.currentState, event)
	}

	for _, ch := range eventListenersCopy {
		select {
		case ch <- transition:
			// Successfully sent
		default:
			// Channel buffer full or no receiver, log and continue
			slog.WarnContext(ctx, "Failed to notify state transition listener for forced transition",
				"from", sm.previousState, "to", sm.currentState)
		}
	}

	// Re-acquire the lock
	sm.mu.Lock()
	defer sm.mu.Unlock()
}