agent-team/internal/api/server.go

package api

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"github.com/gorilla/websocket"
	"github.com/labstack/echo/v4"
	"github.com/sdaduanbilei/agent-team/internal/hub"
	"github.com/sdaduanbilei/agent-team/internal/room"
	"github.com/sdaduanbilei/agent-team/internal/skill"
)

type Server struct {
	e          *echo.Echo
	agentsDir  string
	skillsDir  string
	roomsDir   string
	rooms      map[string]*room.Room
	mu         sync.RWMutex
	clients    map[string]map[*websocket.Conn]bool // roomID -> conns
	clientsMu  sync.Mutex
	upgrader   websocket.Upgrader
}

func New(agentsDir, skillsDir, roomsDir string) *Server {
	s := &Server{
		e:         echo.New(),
		agentsDir: agentsDir,
		skillsDir: skillsDir,
		roomsDir:  roomsDir,
		rooms:     make(map[string]*room.Room),
		clients:   make(map[string]map[*websocket.Conn]bool),
		upgrader:  websocket.Upgrader{CheckOrigin: func(r *http.Request) bool { return true }},
	}
	s.loadRooms()
	s.routes()
	return s
}

func (s *Server) Start(addr string) error {
	return s.e.Start(addr)
}

func (s *Server) routes() {
	s.e.Static("/", "web/dist")

	g := s.e.Group("/api")
	g.GET("/rooms", s.listRooms)
	g.POST("/rooms", s.createRoom)
	g.GET("/agents", s.listAgents)
	g.GET("/agents/:name/files/:file", s.readAgentFile)
	g.PUT("/agents/:name/files/:file", s.writeAgentFile)
	g.POST("/agents", s.createAgent)
	g.DELETE("/agents/:name", s.deleteAgent)
	g.GET("/skills", s.listSkills)
	g.GET("/skills/:name", s.getSkill)
	g.POST("/skills", s.createSkill)
	g.POST("/hub/install", s.hubInstall)
	g.GET("/rooms/:id/workspace", s.listWorkspace)
	g.GET("/rooms/:id/workspace/:file", s.getWorkspaceFile)
	g.GET("/rooms/:id/tasks", s.getTasks)
	g.GET("/rooms/:id/history", s.listHistory)
	g.GET("/rooms/:id/messages", s.getMessages)

	s.e.GET("/ws/:roomID", s.wsHandler)
}

func (s *Server) loadRooms() {
	entries, _ := os.ReadDir(s.roomsDir)
	for _, e := range entries {
		if !e.IsDir() {
			continue
		}
		r, err := room.Load(filepath.Join(s.roomsDir, e.Name()), s.agentsDir, s.skillsDir)
		if err != nil {
			continue
		}
		r.Broadcast = func(ev room.Event) { s.broadcast(ev.RoomID, ev) }
		s.rooms[e.Name()] = r
	}
}

func (s *Server) broadcast(roomID string, ev room.Event) {
	s.clientsMu.Lock()
	defer s.clientsMu.Unlock()
	data, _ := json.Marshal(ev)
	for conn := range s.clients[roomID] {
		conn.WriteMessage(websocket.TextMessage, data)
	}
}

func (s *Server) wsHandler(c echo.Context) error {
	roomID := c.Param("roomID")
	conn, err := s.upgrader.Upgrade(c.Response(), c.Request(), nil)
	if err != nil {
		return err
	}
	s.clientsMu.Lock()
	if s.clients[roomID] == nil {
		s.clients[roomID] = make(map[*websocket.Conn]bool)
	}
	s.clients[roomID][conn] = true
	s.clientsMu.Unlock()

	defer func() {
		s.clientsMu.Lock()
		delete(s.clients[roomID], conn)
		s.clientsMu.Unlock()
		conn.Close()
	}()

	for {
		_, msg, err := conn.ReadMessage()
		if err != nil {
			break
		}
		var ev struct {
			Type    string `json:"type"`
			Content string `json:"content"`
		}
		if json.Unmarshal(msg, &ev) != nil || ev.Type != "user_message" {
			continue
		}
		s.mu.RLock()
		r := s.rooms[roomID]
		s.mu.RUnlock()
		if r == nil {
			continue
		}
		go r.Handle(context.Background(), ev.Content)
	}
	return nil
}

// --- REST handlers ---

func (s *Server) listRooms(c echo.Context) error {
	s.mu.RLock()
	defer s.mu.RUnlock()
	type roomInfo struct {
		ID     string      `json:"id"`
		Name   string      `json:"name"`
		Type   string      `json:"type"`
		Status room.Status `json:"status"`
		Master string      `json:"master"`
	}
	var list []roomInfo
	for id, r := range s.rooms {
		list = append(list, roomInfo{ID: id, Name: r.Config.Name, Type: string(r.Config.Type), Status: r.Status, Master: r.Config.Master})
	}
	return c.JSON(200, list)
}

func (s *Server) createRoom(c echo.Context) error {
	var cfg room.Config
	if err := c.Bind(&cfg); err != nil {
		return err
	}
	dir := filepath.Join(s.roomsDir, cfg.Name)
	os.MkdirAll(filepath.Join(dir, "workspace"), 0755)
	os.MkdirAll(filepath.Join(dir, "history"), 0755)

	content := "---\nname: " + cfg.Name + "\ntype: " + string(cfg.Type) + "\nmaster: " + cfg.Master + "\nmembers:\n"
	for _, m := range cfg.Members {
		content += "  - " + m + "\n"
	}
	content += "---\n"
	os.WriteFile(filepath.Join(dir, "room.md"), []byte(content), 0644)

	r, err := room.Load(dir, s.agentsDir, s.skillsDir)
	if err != nil {
		return c.JSON(500, map[string]string{"error": err.Error()})
	}
	r.Broadcast = func(ev room.Event) { s.broadcast(ev.RoomID, ev) }
	s.mu.Lock()
	s.rooms[cfg.Name] = r
	s.mu.Unlock()
	return c.JSON(201, map[string]string{"id": cfg.Name})
}

func (s *Server) listAgents(c echo.Context) error {
	entries, _ := os.ReadDir(s.agentsDir)
	type agentInfo struct {
		Name string `json:"name"`
	}
	var list []agentInfo
	for _, e := range entries {
		if e.IsDir() {
			list = append(list, agentInfo{Name: e.Name()})
		}
	}
	return c.JSON(200, list)
}

func (s *Server) readAgentFile(c echo.Context) error {
	name := c.Param("name")
	file := c.Param("file") // AGENT.md or SOUL.md
	data, err := os.ReadFile(filepath.Join(s.agentsDir, name, file))
	if err != nil {
		return c.JSON(404, map[string]string{"error": "not found"})
	}
	return c.JSON(200, map[string]string{"content": string(data)})
}

func (s *Server) writeAgentFile(c echo.Context) error {
	name := c.Param("name")
	file := c.Param("file")
	var body struct {
		Content string `json:"content"`
	}
	if err := c.Bind(&body); err != nil {
		return err
	}
	dir := filepath.Join(s.agentsDir, name)
	os.MkdirAll(dir, 0755)
	return os.WriteFile(filepath.Join(dir, file), []byte(body.Content), 0644)
}

func (s *Server) createAgent(c echo.Context) error {
	var body struct {
		Name string `json:"name"`
	}
	if err := c.Bind(&body); err != nil {
		return err
	}
	dir := filepath.Join(s.agentsDir, body.Name)
	os.MkdirAll(filepath.Join(dir, "memory"), 0755)
	agentMD := "---\nname: " + body.Name + "\ndescription: \nprovider: deepseek\nmodel: deepseek-chat\napi_key_env: DEEPSEEK_API_KEY\nskills: []\n---\n"
	os.WriteFile(filepath.Join(dir, "AGENT.md"), []byte(agentMD), 0644)
	os.WriteFile(filepath.Join(dir, "SOUL.md"), []byte("You are "+body.Name+"."), 0644)
	return c.JSON(201, map[string]string{"name": body.Name})
}

func (s *Server) deleteAgent(c echo.Context) error {
	name := c.Param("name")
	return os.RemoveAll(filepath.Join(s.agentsDir, name))
}

func (s *Server) listSkills(c echo.Context) error {
	metas, _ := skill.Discover(s.skillsDir)
	return c.JSON(200, metas)
}

func (s *Server) getSkill(c echo.Context) error {
	name := c.Param("name")
	sk, err := skill.Load(filepath.Join(s.skillsDir, name))
	if err != nil {
		return c.JSON(404, map[string]string{"error": "not found"})
	}
	return c.JSON(200, sk)
}

func (s *Server) hubInstall(c echo.Context) error {
	var body struct {
		Repo string `json:"repo"`
	}
	if err := c.Bind(&body); err != nil {
		return err
	}
	if err := hub.Install(body.Repo, s.agentsDir, s.skillsDir); err != nil {
		return c.JSON(500, map[string]string{"error": err.Error()})
	}
	s.loadRooms()
	return c.JSON(200, map[string]string{"status": "installed"})
}

func (s *Server) listWorkspace(c echo.Context) error {
	id := c.Param("id")
	dir := filepath.Join(s.roomsDir, id, "workspace")
	entries, _ := os.ReadDir(dir)
	var files []string
	for _, e := range entries {
		if !e.IsDir() {
			files = append(files, e.Name())
		}
	}
	return c.JSON(200, files)
}

func (s *Server) getTasks(c echo.Context) error {
	id := c.Param("id")
	data, _ := os.ReadFile(filepath.Join(s.roomsDir, id, "tasks.md"))
	return c.JSON(200, map[string]string{"content": string(data)})
}

func (s *Server) listHistory(c echo.Context) error {
	id := c.Param("id")
	entries, _ := os.ReadDir(filepath.Join(s.roomsDir, id, "history"))
	var files []string
	for _, e := range entries {
		files = append(files, e.Name())
	}
	return c.JSON(200, files)
}

func (s *Server) createSkill(c echo.Context) error {
	var body struct {
		Name    string `json:"name"`
		Content string `json:"content"`
	}
	if err := c.Bind(&body); err != nil {
		return err
	}
	dir := filepath.Join(s.skillsDir, body.Name)
	os.MkdirAll(dir, 0755)
	content := body.Content
	if content == "" {
		content = "---\nname: " + body.Name + "\ndescription: \n---\n\n# " + body.Name + "\n\n描述这个 skill 的用途和使用步骤。\n"
	}
	return os.WriteFile(filepath.Join(dir, "SKILL.md"), []byte(content), 0644)
}

func (s *Server) getWorkspaceFile(c echo.Context) error {
	id := c.Param("id")
	file := c.Param("file")
	data, err := os.ReadFile(filepath.Join(s.roomsDir, id, "workspace", file))
	if err != nil {
		return c.JSON(404, map[string]string{"error": "not found"})
	}
	return c.JSON(200, map[string]string{"content": string(data)})
}

func (s *Server) getMessages(c echo.Context) error {
	id := c.Param("id")
	// Read today's history file and parse into messages
	historyFile := filepath.Join(s.roomsDir, id, "history", time.Now().Format("2006-01-02")+".md")
	data, err := os.ReadFile(historyFile)
	if err != nil {
		return c.JSON(200, []interface{}{})
	}
	type msg struct {
		ID      string `json:"id"`
		Agent   string `json:"agent"`
		Role    string `json:"role"`
		Content string `json:"content"`
	}
	var msgs []msg
	lines := strings.Split(string(data), "\n## ")
	for i, block := range lines {
		if block == "" {
			continue
		}
		// Parse "**[HH:MM:SS] agentName** (role)\n\ncontent"
		parts := strings.SplitN(block, "\n\n", 2)
		if len(parts) < 2 {
			continue
		}
		header := parts[0]
		content := strings.TrimSpace(parts[1])
		// Extract role and agent from header like "**[15:04:05] agentName** (role)"
		var agentName, role string
		fmt.Sscanf(header, "**[%*s] %s** (%s)", &agentName, &role)
		agentName = strings.TrimSuffix(agentName, "**")
		role = strings.TrimSuffix(role, ")")
		if agentName == "" {
			agentName = "unknown"
		}
		if role == "" {
			role = "member"
		}
		msgs = append(msgs, msg{ID: fmt.Sprintf("%d", i), Agent: agentName, Role: role, Content: content})
	}
	return c.JSON(200, msgs)
}