kandev/apps/backend/internal/agent/runtime/lifecycle/execution_store.go

// Package lifecycle manages agent execution lifecycles including tracking,
// state transitions, and cleanup.
package lifecycle

import (
	"errors"
	"fmt"
	"sync"

	v1 "github.com/kandev/kandev/pkg/api/v1"
)

// ErrExecutionNotFound is returned when an execution doesn't exist in the store.
var ErrExecutionNotFound = errors.New("execution not found")

// ErrExecutionAlreadyExistsForSession is returned by Add when the session
// already maps to a different execution. The previous behavior was to silently
// overwrite the bySession index, which orphaned the prior execution: the
// agentctl/subprocess kept running but was unreachable via GetBySessionID, so
// nothing ever cleaned it up. Callers must Remove the prior execution before
// adding a replacement.
var ErrExecutionAlreadyExistsForSession = errors.New("execution already exists for session")

// ErrAgentAlreadyRunning is returned by LaunchAgent when an agent execution is
// already tracked for the requested session. The error fires both when the
// execution is live (a concurrent caller raced us) and when it is stale (a
// prior registration whose process never started or exited without cleanup);
// callers must probe IsAgentRunningForSession to decide. Wrapped at producer
// sites with %w so callers can use errors.Is rather than string-match the
// surrounding context.
var ErrAgentAlreadyRunning = errors.New("session already has an agent running")

// ErrAgentReported is returned when the agent subprocess publishes an error
// event during a prompt turn. Distinguished from infrastructure errors
// (network, timeout) so callers can suppress duplicate error-message creation
// — the agent failure path already records the error in the session.
var ErrAgentReported = errors.New("agent error")

// ExecutionStore provides thread-safe storage and retrieval of agent executions.
// It maintains three indexes for efficient lookup by execution ID, session ID, and container ID.
type ExecutionStore struct {
	executions  map[string]*AgentExecution
	bySession   map[string]string // sessionID -> executionID
	byContainer map[string]string // containerID -> executionID
	mu          sync.RWMutex
}

// NewExecutionStore creates a new ExecutionStore with initialized maps.
func NewExecutionStore() *ExecutionStore {
	return &ExecutionStore{
		executions:  make(map[string]*AgentExecution),
		bySession:   make(map[string]string),
		byContainer: make(map[string]string),
	}
}

// Add adds an agent execution to all tracking maps.
// The execution must have a valid ID. SessionID and ContainerID are optional
// but will be indexed if present.
//
// Returns ErrExecutionAlreadyExistsForSession if the session is already mapped
// to a different execution. Callers must explicitly Remove the prior execution
// before adding a replacement — otherwise the prior subprocess is orphaned.
func (s *ExecutionStore) Add(execution *AgentExecution) error {
	if execution == nil || execution.ID == "" {
		return fmt.Errorf("execution and execution.ID are required")
	}

	s.mu.Lock()
	defer s.mu.Unlock()

	if execution.SessionID != "" {
		if existingID, exists := s.bySession[execution.SessionID]; exists && existingID != execution.ID {
			return fmt.Errorf("%w: session=%s existing=%s new=%s",
				ErrExecutionAlreadyExistsForSession,
				execution.SessionID, existingID, execution.ID)
		}
	}

	s.executions[execution.ID] = execution

	if execution.SessionID != "" {
		s.bySession[execution.SessionID] = execution.ID
	}

	if execution.ContainerID != "" {
		s.byContainer[execution.ContainerID] = execution.ID
	}
	return nil
}

// Remove removes an agent execution from all tracking maps.
func (s *ExecutionStore) Remove(executionID string) {
	s.mu.Lock()
	defer s.mu.Unlock()

	execution, exists := s.executions[executionID]
	if !exists {
		return
	}

	// Remove from secondary indexes
	if execution.SessionID != "" {
		delete(s.bySession, execution.SessionID)
	}
	if execution.ContainerID != "" {
		delete(s.byContainer, execution.ContainerID)
	}

	// Remove from primary map
	delete(s.executions, executionID)
}

// Get returns an agent execution by its ID.
func (s *ExecutionStore) Get(executionID string) (*AgentExecution, bool) {
	s.mu.RLock()
	defer s.mu.RUnlock()

	execution, exists := s.executions[executionID]
	return execution, exists
}

// GetBySessionID returns the agent execution associated with a session ID.
func (s *ExecutionStore) GetBySessionID(sessionID string) (*AgentExecution, bool) {
	s.mu.RLock()
	defer s.mu.RUnlock()

	executionID, exists := s.bySession[sessionID]
	if !exists {
		return nil, false
	}

	execution, exists := s.executions[executionID]
	return execution, exists
}

// OwnsPromptGeneration reports whether executionID and generation still own
// sessionID's active lifecycle prompt.
func (s *ExecutionStore) OwnsPromptGeneration(sessionID, executionID string, generation uint64) bool {
	s.mu.RLock()
	defer s.mu.RUnlock()

	currentExecutionID, exists := s.bySession[sessionID]
	if !exists || currentExecutionID != executionID {
		return false
	}
	execution, exists := s.executions[currentExecutionID]
	return exists && generation != 0 && execution.promptGeneration == generation
}

// GetByTaskEnvironmentID returns any execution associated with a task environment ID.
func (s *ExecutionStore) GetByTaskEnvironmentID(taskEnvironmentID string) (*AgentExecution, bool) {
	s.mu.RLock()
	defer s.mu.RUnlock()

	for _, execution := range s.executions {
		if execution.TaskEnvironmentID == taskEnvironmentID {
			return execution, true
		}
	}
	return nil, false
}

// GetByContainerID returns the agent execution associated with a container ID.
func (s *ExecutionStore) GetByContainerID(containerID string) (*AgentExecution, bool) {
	s.mu.RLock()
	defer s.mu.RUnlock()

	executionID, exists := s.byContainer[containerID]
	if !exists {
		return nil, false
	}

	execution, exists := s.executions[executionID]
	return execution, exists
}

// List returns all tracked agent executions.
func (s *ExecutionStore) List() []*AgentExecution {
	s.mu.RLock()
	defer s.mu.RUnlock()

	result := make([]*AgentExecution, 0, len(s.executions))
	for _, execution := range s.executions {
		result = append(result, execution)
	}
	return result
}

// UpdateStatus updates the status of an agent execution.
func (s *ExecutionStore) UpdateStatus(executionID string, status v1.AgentStatus) {
	s.mu.Lock()
	defer s.mu.Unlock()

	if execution, exists := s.executions[executionID]; exists {
		execution.Status = status
	}
}

// BeginPrompt assigns a new immutable identity to the next prompt dispatch.
func (s *ExecutionStore) BeginPrompt(executionID string) (uint64, error) {
	execution, exists := s.Get(executionID)
	if !exists {
		return 0, ErrExecutionNotFound
	}
	execution.promptLifecycleMu.Lock()
	defer execution.promptLifecycleMu.Unlock()

	s.mu.Lock()
	defer s.mu.Unlock()

	current, exists := s.executions[executionID]
	if !exists || current != execution {
		return 0, ErrExecutionNotFound
	}
	return beginExecutionPrompt(current), nil
}

func beginExecutionPrompt(execution *AgentExecution) uint64 {
	execution.promptGeneration++
	execution.Status = v1.AgentStatusRunning
	return execution.promptGeneration
}

// UpdateError updates the error message of an agent execution and sets its status to failed.
func (s *ExecutionStore) UpdateError(executionID string, errorMsg string) {
	s.mu.Lock()
	defer s.mu.Unlock()

	if execution, exists := s.executions[executionID]; exists {
		execution.ErrorMessage = errorMsg
		execution.Status = v1.AgentStatusFailed
	}
}

// WithLock executes a function with the store lock held, providing access to the execution.
// Returns ErrExecutionNotFound if the execution doesn't exist.
// The function should be fast to avoid blocking other operations.
func (s *ExecutionStore) WithLock(executionID string, fn func(execution *AgentExecution)) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	execution, exists := s.executions[executionID]
	if !exists {
		return ErrExecutionNotFound
	}
	fn(execution)
	return nil
}

// WithRLock executes a function with the store read lock held, providing access to the execution.
// Returns ErrExecutionNotFound if the execution doesn't exist.
func (s *ExecutionStore) WithRLock(executionID string, fn func(execution *AgentExecution)) error {
	s.mu.RLock()
	defer s.mu.RUnlock()

	execution, exists := s.executions[executionID]
	if !exists {
		return ErrExecutionNotFound
	}
	fn(execution)
	return nil
}