orca.ai/pkg/actor/system.go

package actor

import (
	"fmt"
	"sync"
	"sync/atomic"

	"github.com/orca/orca/pkg/tool"
)

// System manages the lifecycle of all agents in the Orca actor framework.
//
// It provides centralized agent creation, monitoring, and shutdown
// capabilities. Agents are identified by unique IDs and organized by role.
type System struct {
	mu     sync.RWMutex
	agents map[string]Agent
	nextID int64
}

// NewSystem creates a new empty actor System.
func NewSystem() *System {
	return &System{
		agents: make(map[string]Agent),
	}
}

// AgentInfo holds summary information about a managed agent.
type AgentInfo struct {
	ID     string      `json:"id"`
	Role   string      `json:"role"`
	Status ActorStatus `json:"status"`
}

// CreateOrchestrator creates a new Orchestrator agent and registers it.
func (s *System) CreateOrchestrator(bus interface{}) (*Orchestrator, error) {
	id := s.nextAgentID("orch")
	return s.addOrchestrator(id, bus)
}

// CreateWorker creates a new Worker agent and registers it.
func (s *System) CreateWorker() (*Worker, error) {
	id := s.nextAgentID("worker")
	return s.addWorker(id)
}

// CreateToolWorker creates a new ToolWorker agent with the given tool manager and registers it.
func (s *System) CreateToolWorker(manager *tool.Manager) (*ToolWorker, error) {
	id := s.nextAgentID("tool")
	return s.addToolWorker(id, manager)
}

// nextAgentID generates a unique agent ID with the given prefix.
func (s *System) nextAgentID(prefix string) string {
	n := atomic.AddInt64(&s.nextID, 1)
	return fmt.Sprintf("%s-%d", prefix, n)
}

// addOrchestrator creates and registers an orchestrator.
func (s *System) addOrchestrator(id string, busInterface interface{}) (*Orchestrator, error) {
	mb, ok := busInterface.(interface{ Bus() })
	var orch *Orchestrator
	if ok {
		// If busInterface has a Bus() method, we could extract it here
		_ = mb
	}
	orch = NewOrchestrator(id, nil)

	s.mu.Lock()
	s.agents[id] = orch
	s.mu.Unlock()

	return orch, nil
}

// addWorker creates and registers a worker.
func (s *System) addWorker(id string) (*Worker, error) {
	w := NewWorker(id)

	s.mu.Lock()
	s.agents[id] = w
	s.mu.Unlock()

	return w, nil
}

// addToolWorker creates and registers a tool worker with the given tool manager.
func (s *System) addToolWorker(id string, manager *tool.Manager) (*ToolWorker, error) {
	w := NewToolWorker(id, manager)

	s.mu.Lock()
	s.agents[id] = w
	s.mu.Unlock()

	return w, nil
}

// StopAgent stops and removes a single agent by ID.
func (s *System) StopAgent(id string) error {
	s.mu.Lock()
	agent, ok := s.agents[id]
	if !ok {
		s.mu.Unlock()
		return fmt.Errorf("agent %s not found", id)
	}
	delete(s.agents, id)
	s.mu.Unlock()

	return agent.Stop()
}

// GetAgent retrieves a registered agent by ID.
func (s *System) GetAgent(id string) (Agent, bool) {
	s.mu.RLock()
	defer s.mu.RUnlock()
	agent, ok := s.agents[id]
	return agent, ok
}

// ListAgents returns all registered agents.
func (s *System) ListAgents() []Agent {
	s.mu.RLock()
	defer s.mu.RUnlock()

	agents := make([]Agent, 0, len(s.agents))
	for _, a := range s.agents {
		agents = append(agents, a)
	}
	return agents
}

// AgentInfos returns summary information for all registered agents.
func (s *System) AgentInfos() []AgentInfo {
	s.mu.RLock()
	defer s.mu.RUnlock()

	infos := make([]AgentInfo, 0, len(s.agents))
	for _, a := range s.agents {
		// Try to get status from BaseAgent
		status := StatusIdle
		if ba, ok := a.(*BaseAgent); ok {
			status = ba.Status()
		} else if orch, ok := a.(*Orchestrator); ok {
			status = orch.Status()
		} else if w, ok := a.(*Worker); ok {
			status = w.Status()
		} else if tw, ok := a.(*ToolWorker); ok {
			status = tw.Status()
		}

		infos = append(infos, AgentInfo{
			ID:     a.ID(),
			Role:   a.Role(),
			Status: status,
		})
	}
	return infos
}

// StopAll gracefully stops all registered agents.
func (s *System) StopAll() error {
	s.mu.Lock()
	defer s.mu.Unlock()

	var lastErr error
	for id, agent := range s.agents {
		if err := agent.Stop(); err != nil {
			lastErr = err
		}
		delete(s.agents, id)
	}
	return lastErr
}

// AgentCount returns the number of registered agents.
func (s *System) AgentCount() int {
	s.mu.RLock()
	defer s.mu.RUnlock()
	return len(s.agents)
}