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")

// ErrPromptActivityNotOwned means a cancellation snapshot no longer belongs
// to the execution and prompt that were captured.
var ErrPromptActivityNotOwned = errors.New("prompt activity no longer owned")

// 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)
	}
	execution.clearRuntimeEnvironment()

	// 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
}

// OwnsPromptActivity reports whether the execution still owns the prompt and
// its activity has not changed since the watchdog captured its snapshot.
func (s *ExecutionStore) OwnsPromptActivity(
	sessionID, executionID string,
	generation, activityEpoch 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]
	if !exists || generation == 0 || activityEpoch == 0 || execution.promptGeneration != generation {
		return false
	}
	return execution.promptActivityEpochSnapshot() == activityEpoch
}

// ClaimPromptActivity validates a captured prompt identity while the store
// lock is held and returns that exact execution for a lifecycle operation.
// Callers must use the returned execution instead of looking it up by session
// again, because a session can acquire a replacement execution while the
// operation waits for the agent to settle.
func (s *ExecutionStore) ClaimPromptActivity(
	sessionID, executionID string,
	generation, activityEpoch uint64,
) (*AgentExecution, error) {
	s.mu.Lock()
	defer s.mu.Unlock()

	currentExecutionID, exists := s.bySession[sessionID]
	if !exists {
		return nil, ErrExecutionNotFound
	}
	execution, exists := s.executions[currentExecutionID]
	if !exists {
		return nil, ErrExecutionNotFound
	}
	if currentExecutionID != executionID || generation == 0 || activityEpoch == 0 || execution.promptGeneration != generation || execution.promptActivityEpochSnapshot() != activityEpoch {
		return nil, ErrPromptActivityNotOwned
	}
	return execution, nil
}

// ActivePromptGeneration returns the generation of the prompt that is dispatched
// and still in flight for executionID, or 0 if there is none. A steer reuses this
// exact generation so the folded completion is attributed to the predecessor's
// still-open waiter rather than discarded as a mismatch. It deliberately returns
// 0 for a generation that is only admitted (not yet dispatched — its buffers are
// mid-reset) or already completed, so a steer never overtakes an about-to-start
// turn or attaches to a finished one; those cases fall back to an ordinary prompt.
func (s *ExecutionStore) ActivePromptGeneration(executionID string) uint64 {
	s.mu.RLock()
	defer s.mu.RUnlock()

	execution, exists := s.executions[executionID]
	if !exists {
		return 0
	}
	gen := execution.promptGeneration
	if gen == 0 || execution.dispatchedPromptGeneration != gen || execution.promptCompletionGeneration == gen {
		return 0
	}
	return gen
}

// ExecutionReference identifies one registered execution without requiring a
// later lookup through the session index. Both IDs are captured from the same
// execution-store snapshot so a caller can keep targeting the original
// execution if a session later acquires a replacement.
type ExecutionReference struct {
	SessionID   string
	ExecutionID string
}

// ListExecutionsForTask returns the session and execution IDs of executions
// registered in-memory under taskID. It is a read-only snapshot, independent
// of any session's persisted database state. Task-scoped stop uses it to
// recover an execution whose session row is already terminal (for example,
// FAILED) but whose prior teardown attempt failed.
func (s *ExecutionStore) ListExecutionsForTask(taskID string) []ExecutionReference {
	if taskID == "" {
		return nil
	}
	s.mu.RLock()
	defer s.mu.RUnlock()

	var references []ExecutionReference
	for _, execution := range s.executions {
		if execution.TaskID == taskID && execution.SessionID != "" && execution.ID != "" {
			references = append(references, ExecutionReference{
				SessionID:   execution.SessionID,
				ExecutionID: execution.ID,
			})
		}
	}
	return references
}

// 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 {
	// A prompt is about to be dispatched through this object, so a
	// recovered-but-not-yet-adopted generation must not later clobber it with
	// a stale pre-restart completion (see recoveredPromptGenerationPending).
	execution.recoveredPromptGenerationPending.Store(false)
	execution.promptGeneration++
	execution.promptCompletionGeneration = 0
	// The new generation is not dispatched until its triggerPrompt succeeds; a
	// steer must not reuse it until then (its buffers are still being reset).
	execution.dispatchedPromptGeneration = 0
	execution.ProviderError = nil
	execution.Status = v1.AgentStatusRunning
	execution.resetActiveTool()
	return execution.promptGeneration
}

// MarkPromptDispatched records that generation reached agentctl and is in
// flight. Ignored if a newer generation has since begun or this one already
// completed, so a steer can only ever reuse a live dispatched turn.
func (s *ExecutionStore) MarkPromptDispatched(executionID string, generation uint64) {
	if generation == 0 {
		return
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	execution, exists := s.executions[executionID]
	if !exists {
		return
	}
	if execution.promptGeneration != generation || execution.promptCompletionGeneration == generation {
		return
	}
	execution.dispatchedPromptGeneration = generation
	execution.armPromptActivity()
}

// 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
}