MuyueWorkspace/internal/api/server.go

package api

import (
	"encoding/json"
	"fmt"
	"net/http"
	"os/exec"
	"strings"
	"sync/atomic"

	"github.com/muyue/muyue/internal/agent"
	"github.com/muyue/muyue/internal/config"
	"github.com/muyue/muyue/internal/installer"
	"github.com/muyue/muyue/internal/scanner"
	"github.com/muyue/muyue/internal/workflow"
)

type Server struct {
	config              *config.MuyueConfig
	scanResult          *scanner.ScanResult
	mux                 *http.ServeMux
	convStore           *ConversationStore
	shellConvStore      *ShellConvStore
	consumption         *consumptionStore
	agentRegistry       *agent.Registry
	agentToolsJSON      json.RawMessage
	shellAgentRegistry  *agent.Registry
	shellAgentToolsJSON json.RawMessage
	workflowEngine      *workflow.Engine
	browserTestStore    *BrowserTestStore
	activeCrushAgents   atomic.Int32
	activeClaudeAgents  atomic.Int32
}

func NewServer(cfg *config.MuyueConfig) *Server {
	s := &Server{
		mux: http.NewServeMux(),
	}
	// Auto-initialize config if nil or if no config file exists on disk
	if cfg == nil || !config.Exists() {
		defaultCfg := config.Default()
		if cfg != nil {
			// Preserve any user-provided settings from cfg
			defaultCfg.Profile = cfg.Profile
			defaultCfg.AI = cfg.AI
			defaultCfg.Tools = cfg.Tools
			defaultCfg.BMAD = cfg.BMAD
			defaultCfg.Terminal = cfg.Terminal
		}
		// Save initial config to establish the file for first-time usage
		if err := config.Save(defaultCfg); err != nil {
			_ = err
		}
		cfg = defaultCfg
	}
	s.config = cfg
	s.scanResult = scanner.ScanSystem()
	s.convStore = NewConversationStore()
	s.shellConvStore = NewShellConvStore()
	s.consumption = newConsumptionStore()
	s.agentRegistry = agent.DefaultRegistry()
	s.browserTestStore = NewBrowserTestStore()
	if err := RegisterBrowserTestTool(s.agentRegistry, s.browserTestStore); err != nil {
		// Tool registration only fails for duplicate names — non-fatal
		_ = err
	}
	tools := s.agentRegistry.OpenAITools()
	toolsJSON, _ := json.Marshal(tools)
	s.agentToolsJSON = json.RawMessage(toolsJSON)

	s.shellAgentRegistry = agent.NewRegistry()
	terminalTool, _ := agent.NewTerminalTool()
	s.shellAgentRegistry.Register(terminalTool)
	shellTools := s.shellAgentRegistry.OpenAITools()
	shellToolsJSON, _ := json.Marshal(shellTools)
	s.shellAgentToolsJSON = json.RawMessage(shellToolsJSON)

	s.workflowEngine, _ = workflow.NewEngine(s.agentRegistry)
	s.initStarship()
	s.routes()
	return s
}

func (s *Server) routes() {
	s.mux.HandleFunc("/api/info", s.handleInfo)
	s.mux.HandleFunc("/api/system", s.handleSystem)
	s.mux.HandleFunc("/api/tools", s.handleTools)
	s.mux.HandleFunc("/api/config", s.handleConfig)
	s.mux.HandleFunc("/api/providers", s.handleProviders)
	s.mux.HandleFunc("/api/skills", s.handleSkills)
	s.mux.HandleFunc("/api/lsp", s.handleLSP)
	s.mux.HandleFunc("/api/mcp", s.handleMCP)
	s.mux.HandleFunc("/api/updates", s.handleUpdates)
	s.mux.HandleFunc("/api/install", s.handleInstall)
	s.mux.HandleFunc("/api/scan", s.handleScan)
	s.mux.HandleFunc("/api/editors", s.handleEditors)
	s.mux.HandleFunc("/api/preferences", s.handleUpdatePreferences)
	s.mux.HandleFunc("/api/terminal", s.handleTerminal)
	s.mux.HandleFunc("/api/ws/terminal", s.handleTerminalWS)
	s.mux.HandleFunc("/api/terminal/sessions", s.handleTerminalSessions)
	s.mux.HandleFunc("/api/terminal/sessions/", s.handleTerminalSessionsDelete)
	s.mux.HandleFunc("/api/terminal/themes", s.handleGetTerminalThemes)
	s.mux.HandleFunc("/api/terminal/settings", s.handleSaveTerminalSettings)
	s.mux.HandleFunc("/api/mcp/configure", s.handleMCPConfigure)
	s.mux.HandleFunc("/api/config/profile", s.handleSaveProfile)
	s.mux.HandleFunc("/api/config/provider", s.handleSaveProvider)
	s.mux.HandleFunc("/api/config/reset", s.handleResetConfig)
	s.mux.HandleFunc("/api/starship/apply-theme", s.handleApplyStarshipTheme)
	s.mux.HandleFunc("/api/providers/validate", s.handleValidateProvider)
	s.mux.HandleFunc("/api/update/run", s.handleRunUpdate)
	s.mux.HandleFunc("/api/images/", s.handleServeImage)
	s.mux.HandleFunc("/api/chat", s.handleChat)
	s.mux.HandleFunc("/api/chat/history", s.handleChatHistory)
	s.mux.HandleFunc("/api/chat/clear", s.handleChatClear)
	s.mux.HandleFunc("/api/chat/summarize", s.handleChatSummarize)
	s.mux.HandleFunc("/api/tool/call", s.handleToolCall)
	s.mux.HandleFunc("/api/tools/list", s.handleToolList)
	s.mux.HandleFunc("/api/shell/chat", s.handleShellChat)
	s.mux.HandleFunc("/api/shell/chat/history", s.handleShellChatHistory)
	s.mux.HandleFunc("/api/shell/chat/clear", s.handleShellChatClear)
	s.mux.HandleFunc("/api/shell/analyze", s.handleShellAnalyze)
	s.mux.HandleFunc("/api/shell/analysis", s.handleShellAnalysisGet)
	s.mux.HandleFunc("/api/workflow", s.handleWorkflowCreate)
	s.mux.HandleFunc("/api/workflow/list", s.handleWorkflowList)
	s.mux.HandleFunc("/api/workflow/", s.handleWorkflowGet)
	s.mux.HandleFunc("/api/workflow/plan", s.handleWorkflowPlan)
	s.mux.HandleFunc("/api/workflow/execute/", s.handleWorkflowExecute)
	s.mux.HandleFunc("/api/workflow/approve/", s.handleWorkflowApprove)
	s.mux.HandleFunc("/api/conversations", s.handleListConversations)
	s.mux.HandleFunc("/api/conversations/search", s.handleSearchConversations)
	s.mux.HandleFunc("/api/conversations/export", s.handleExportConversation)
	s.mux.HandleFunc("/api/conversations/", s.handleDeleteConversation)
	s.mux.HandleFunc("/api/lsp/install", s.handleLSPInstall)
	s.mux.HandleFunc("/api/skills/deploy", s.handleSkillsDeploy)
	s.mux.HandleFunc("/api/skills/undeploy", s.handleSkillsUndeploy)
	s.mux.HandleFunc("/api/ssh/connections", s.handleSSHConnections)
	s.mux.HandleFunc("/api/ssh/test", s.handleSSHTest)

	s.mux.HandleFunc("/api/mcp/status", s.handleMCPStatus)
	s.mux.HandleFunc("/api/mcp/registry", s.handleMCPRegistry)
	s.mux.HandleFunc("/api/lsp/health", s.handleLSPHealth)
	s.mux.HandleFunc("/api/lsp/auto-install", s.handleLSPAutoInstall)
	s.mux.HandleFunc("/api/lsp/editor-config", s.handleLSPEditorConfig)
	s.mux.HandleFunc("/api/skills/validate", s.handleSkillValidate)
	s.mux.HandleFunc("/api/skills/test", s.handleSkillTest)
	s.mux.HandleFunc("/api/skills/export", s.handleSkillExport)
	s.mux.HandleFunc("/api/skills/import", s.handleSkillImport)
	s.mux.HandleFunc("/api/dashboard/status", s.handleDashboardStatus)
	s.mux.HandleFunc("/api/ai/task", s.handleAITask)
	s.mux.HandleFunc("/api/providers/quota", s.handleProvidersQuota)
	s.mux.HandleFunc("/api/providers/consumption", s.handleProvidersConsumption)
	s.mux.HandleFunc("/api/recent-commands", s.handleRecentCommands)
	s.mux.HandleFunc("/api/running-processes", s.handleRunningProcesses)
	s.mux.HandleFunc("/api/system/metrics", s.handleSystemMetrics)

	s.mux.HandleFunc("/api/test/snippet", s.handleBrowserTestSnippet)
	s.mux.HandleFunc("/api/test/sessions", s.handleBrowserTestSessions)
	s.mux.HandleFunc("/api/test/console/", s.handleBrowserTestConsole)
	s.mux.HandleFunc("/api/ws/browser-test", s.handleBrowserTestWS)
}

func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if strings.HasPrefix(r.URL.Path, "/api/ws/") || strings.HasPrefix(r.URL.Path, "/api/images/") {
		s.mux.ServeHTTP(w, r)
		return
	}
	w.Header().Set("Content-Type", "application/json")
	if origin := r.Header.Get("Origin"); isAllowedOrigin(origin) {
		w.Header().Set("Access-Control-Allow-Origin", origin)
		w.Header().Set("Vary", "Origin")
	}
	w.Header().Set("Access-Control-Allow-Methods", "GET, POST, PUT, DELETE, OPTIONS")
	w.Header().Set("Access-Control-Allow-Headers", "Content-Type")
	if r.Method == "OPTIONS" {
		w.WriteHeader(http.StatusOK)
		return
	}
	s.mux.ServeHTTP(w, r)
}

func isAllowedOrigin(origin string) bool {
	if origin == "" {
		return false
	}
	switch {
	case strings.HasPrefix(origin, "http://127.0.0.1"),
		strings.HasPrefix(origin, "http://localhost"),
		strings.HasPrefix(origin, "http://[::1]"),
		strings.HasPrefix(origin, "https://127.0.0.1"),
		strings.HasPrefix(origin, "https://localhost"),
		strings.HasPrefix(origin, "https://[::1]"):
		return true
	}
	return false
}

const maxCrushAgents = 2
const maxClaudeAgents = 2

func (s *Server) AcquireAgentSlot(toolName string) (release func(), err error) {
	var counter *atomic.Int32
	var max int32
	switch toolName {
	case "crush_run":
		counter = &s.activeCrushAgents
		max = maxCrushAgents
	case "claude_run":
		counter = &s.activeClaudeAgents
		max = maxClaudeAgents
	default:
		return func() {}, nil
	}
	current := counter.Add(1)
	if current > max {
		counter.Add(-1)
		return nil, fmt.Errorf("Limite de %d agents %s atteinte", max, toolName)
	}
	return func() { counter.Add(-1) }, nil
}

func (s *Server) initStarship() {
	if _, err := exec.LookPath("starship"); err != nil {
		inst := installer.New(s.config)
		if result := inst.InstallTool("starship"); !result.Success {
			return
		}
	}
	ApplyStarshipTheme(s.config.Terminal.PromptTheme)
}