server/agent/manager.go

package agent

import (
	"fmt"
	"log/slog"
	"time"

	"github.com/iomodo/staff/config"
	"github.com/iomodo/staff/llm"
	"github.com/iomodo/staff/task"
	"github.com/iomodo/staff/tm"
)

// Manager manages multiple AI agents with Git operations and task processing
type Manager struct {
	config       *config.Config
	agents       map[string]*Agent
	taskManager  tm.TaskManager
	autoAssigner *task.AutoAssigner
	isRunning    map[string]bool
	roles        []string
	logger       *slog.Logger
}

// NewManager creates a new agent manager
func NewManager(cfg *config.Config, taskManager tm.TaskManager, logger *slog.Logger) (*Manager, error) {
	autoAssigner := task.NewAutoAssigner(cfg.Agents)

	agentRoles := make([]string, 0, len(cfg.Agents))
	for _, agentConfig := range cfg.Agents {
		agentRoles = append(agentRoles, agentConfig.Role)
	}

	manager := &Manager{
		config:       cfg,
		agents:       make(map[string]*Agent),
		taskManager:  taskManager,
		autoAssigner: autoAssigner,
		isRunning:    make(map[string]bool),
		roles:        agentRoles,
		logger:       logger,
	}

	// Initialize agents
	if err := manager.initializeAgents(); err != nil {
		return nil, fmt.Errorf("failed to initialize agents: %w", err)
	}

	return manager, nil
}

// initializeAgents creates agent instances from configuration
func (m *Manager) initializeAgents() error {
	llmConfig := llm.Config{
		Provider: llm.ProviderFake, // Use fake provider for testing
		APIKey:   m.config.OpenAI.APIKey,
		BaseURL:  m.config.OpenAI.BaseURL,
		Timeout:  m.config.OpenAI.Timeout,
	}
	for _, agentConfig := range m.config.Agents {
		agent, err := NewAgent(agentConfig, llmConfig, m.taskManager, m.roles, m.logger)
		if err != nil {
			return fmt.Errorf("failed to create agent %s: %w", agentConfig.Name, err)
		}
		m.agents[agentConfig.Name] = agent
	}
	return nil
}

func (m *Manager) StartAllAgents() {
	// Start all configured agents with a default loop interval
	defaultInterval := 1 * time.Second

	for _, a := range m.agents {
		m.logger.Info("Starting agent",
			slog.String("name", a.Name),
			slog.String("role", a.Role),
			slog.String("model", a.Model))
		if err := a.Start(defaultInterval); err != nil {
			m.logger.Error("Failed to start agent",
				slog.String("agent", a.Name),
				slog.String("error", err.Error()))
			continue
		}
		m.isRunning[a.Name] = true
	}
}

func (m *Manager) StartAgent(agentName string, loopInterval time.Duration) error {
	agent, exists := m.agents[agentName]
	if !exists {
		return fmt.Errorf("agent %s not found", agentName)
	}

	if m.isRunning[agentName] {
		return fmt.Errorf("agent %s is already running", agentName)
	}

	agent.Start(loopInterval)
	m.isRunning[agentName] = true
	return nil
}

// StopAgent stops a running agent
func (m *Manager) StopAgent(agentName string) error {
	agent, exists := m.agents[agentName]
	if !exists {
		return fmt.Errorf("agent %s not found", agentName)
	}
	if !m.isRunning[agentName] {
		return fmt.Errorf("agent %s is not running", agentName)
	}

	agent.Stop()
	m.isRunning[agentName] = false

	m.logger.Info("Stopped agent", slog.String("name", agentName))
	return nil
}

// AutoAssignTask automatically assigns a task to the best matching agent
func (m *Manager) AutoAssignTask(taskID string) error {
	task, err := m.taskManager.GetTask(taskID)
	if err != nil {
		return fmt.Errorf("failed to get task: %w", err)
	}

	agentName, err := m.autoAssigner.AssignTask(task)
	if err != nil {
		return fmt.Errorf("failed to auto-assign task: %w", err)
	}

	task.Assignee = agentName
	if err := m.taskManager.UpdateTask(task); err != nil {
		return fmt.Errorf("failed to update task assignment: %w", err)
	}

	explanation := m.autoAssigner.GetRecommendationExplanation(task, agentName)
	m.logger.Info("Auto-assigned task to agent",
		slog.String("task_id", taskID),
		slog.String("agent", agentName),
		slog.String("explanation", explanation))

	return nil
}

// IsAgentRunning checks if an agent is currently running
func (m *Manager) IsAgentRunning(agentName string) bool {
	return m.isRunning[agentName]
}

// CompleteTaskForAgent finds the agent performing a task and resets its state
func (m *Manager) CompleteTaskForAgent(taskID string) error {
	// Get the task to find the assignee
	task, err := m.taskManager.GetTask(taskID)
	if err != nil {
		return fmt.Errorf("failed to get task %s: %w", taskID, err)
	}

	// Find the agent assigned to this task
	agent, exists := m.agents[task.Assignee]
	if !exists {
		return fmt.Errorf("agent %s not found for task %s", task.Assignee, taskID)
	}

	// Reset the agent's current task and running state
	agent.CurrentTask = nil
	agent.IsRunning = false
	m.isRunning[agent.Name] = false

	m.logger.Info("Completed task for agent",
		slog.String("task_id", taskID),
		slog.String("agent", agent.Name))

	return nil
}

// Close shuts down the agent manager
func (m *Manager) Close() error {
	// Stop all running agents
	for agentName := range m.isRunning {
		if m.isRunning[agentName] {
			m.StopAgent(agentName)
		}
	}

	// Close all LLM providers
	for _, agent := range m.agents {
		if err := agent.Provider.Close(); err != nil {
			m.logger.Error("Error closing provider for agent",
				slog.String("agent", agent.Name),
				slog.String("error", err.Error()))
		}
	}

	return m.taskManager.Close()
}