agent-team/internal/room/room.go

package room

import (
	"bytes"
	"context"
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/sdaduanbilei/agent-team/internal/agent"
	"github.com/sdaduanbilei/agent-team/internal/llm"
	"github.com/sdaduanbilei/agent-team/internal/skill"
	"gopkg.in/yaml.v3"
)

type RoomType string

const (
	TypeDept   RoomType = "dept"
	TypeLeader RoomType = "leader"
)

type Status string

const (
	StatusPending  Status = "pending"
	StatusThinking Status = "thinking"
	StatusWorking  Status = "working"
)

type Config struct {
	Name    string   `yaml:"name"`
	Type    RoomType `yaml:"type"`
	Master  string   `yaml:"master"`  // agent name
	Members []string `yaml:"members"` // agent names
}

type Room struct {
	Config    Config
	Dir       string
	master    *agent.Agent
	members   map[string]*agent.Agent
	skillMeta []skill.Meta
	Status    Status
	ActiveAgent string // for working status display
	Broadcast func(Event) // set by api layer
}

type EventType string

const (
	EvtAgentMessage EventType = "agent_message"
	EvtTaskAssign   EventType = "task_assign"
	EvtReview       EventType = "review"
	EvtRoomStatus   EventType = "room_status"
	EvtTasksUpdate  EventType = "tasks_update"
	EvtWorkspaceFile EventType = "workspace_file"
)

type Event struct {
	Type        EventType `json:"type"`
	RoomID      string    `json:"room_id"`
	Agent       string    `json:"agent,omitempty"`
	Role        string    `json:"role,omitempty"` // master | member
	Content     string    `json:"content,omitempty"`
	Streaming   bool      `json:"streaming,omitempty"`
	From        string    `json:"from,omitempty"`
	To          string    `json:"to,omitempty"`
	Task        string    `json:"task,omitempty"`
	Feedback    string    `json:"feedback,omitempty"`
	Status      Status    `json:"status,omitempty"`
	ActiveAgent string    `json:"active_agent,omitempty"`
	Action      string    `json:"action,omitempty"`
	Filename    string    `json:"filename,omitempty"`
}

func Load(roomDir string, agentsDir string, skillsDir string) (*Room, error) {
	data, err := os.ReadFile(filepath.Join(roomDir, "room.md"))
	if err != nil {
		return nil, err
	}
	cfg, err := parseRoomConfig(data)
	if err != nil {
		return nil, err
	}

	r := &Room{Config: cfg, Dir: roomDir, members: make(map[string]*agent.Agent)}

	r.master, err = agent.Load(filepath.Join(agentsDir, cfg.Master))
	if err != nil {
		return nil, fmt.Errorf("load master %s: %w", cfg.Master, err)
	}
	for _, name := range cfg.Members {
		a, err := agent.Load(filepath.Join(agentsDir, name))
		if err != nil {
			return nil, fmt.Errorf("load member %s: %w", name, err)
		}
		r.members[name] = a
	}

	r.skillMeta, _ = skill.Discover(skillsDir)
	return r, nil
}

func (r *Room) emit(e Event) {
	e.RoomID = r.Config.Name
	if r.Broadcast != nil {
		r.Broadcast(e)
	}
}

func (r *Room) setStatus(s Status, activeAgent, action string) {
	r.Status = s
	r.ActiveAgent = activeAgent
	r.emit(Event{Type: EvtRoomStatus, Status: s, ActiveAgent: activeAgent, Action: action})
}

// AppendHistory persists a message to today's history file.
func (r *Room) AppendHistory(role, agentName, content string) {
	dir := filepath.Join(r.Dir, "history")
	os.MkdirAll(dir, 0755)
	filename := filepath.Join(dir, time.Now().Format("2006-01-02")+".md")
	line := fmt.Sprintf("\n**[%s] %s** (%s)\n\n%s\n", time.Now().Format("15:04:05"), agentName, role, content)
	f, err := os.OpenFile(filename, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
	if err != nil {
		return
	}
	defer f.Close()
	f.WriteString(line)
}

// Handle processes a user message through master orchestration.
func (r *Room) Handle(ctx context.Context, userMsg string) error {
	r.AppendHistory("user", "user", userMsg)
	r.setStatus(StatusThinking, "", "")

	// Build master context
	teamXML := r.buildTeamXML()
	skillXML := skill.ToXML(r.skillMeta)
	systemPrompt := r.master.BuildSystemPrompt(teamXML + "\n\n" + skillXML)

	masterMsgs := []llm.Message{
		llm.NewMsg("system", systemPrompt+"\n\nYou are the master of this team. When you need a team member to do something, output a line like: ASSIGN:<member_name>:<task description>. When you are done reviewing and satisfied, output DONE:<summary>."),
		llm.NewMsg("user", userMsg),
	}

	// Master planning loop
	for iteration := 0; iteration < 5; iteration++ {
		var masterReply strings.Builder
		_, err := r.master.Chat(ctx, masterMsgs, func(token string) {
			masterReply.WriteString(token)
			r.emit(Event{Type: EvtAgentMessage, Agent: r.master.Config.Name, Role: "master", Content: token, Streaming: true})
		})
		if err != nil {
			return err
		}
		reply := masterReply.String()
		masterMsgs = append(masterMsgs, llm.NewMsg("assistant", reply))
		r.AppendHistory("master", r.master.Config.Name, reply)

		// Parse assignments
		assignments := parseAssignments(reply)
		if len(assignments) == 0 {
			// No assignments, master is done
			break
		}

		// Execute assignments
		var results strings.Builder
		for memberName, task := range assignments {
			member, ok := r.members[memberName]
			if !ok {
				continue
			}
			r.setStatus(StatusWorking, member.Config.Name, task)
			r.emit(Event{Type: EvtTaskAssign, From: r.master.Config.Name, To: memberName, Task: task})

			memberSystem := member.BuildSystemPrompt(skillXML)
			memberMsgs := []llm.Message{
				llm.NewMsg("system", memberSystem),
				llm.NewMsg("user", task),
			}
			var memberReply strings.Builder
			_, err := member.Chat(ctx, memberMsgs, func(token string) {
				memberReply.WriteString(token)
				r.emit(Event{Type: EvtAgentMessage, Agent: memberName, Role: "member", Content: token, Streaming: true})
			})
			if err != nil {
				results.WriteString(fmt.Sprintf("[%s] error: %v\n", memberName, err))
				continue
			}
			result := memberReply.String()
			results.WriteString(fmt.Sprintf("[%s] %s\n", memberName, result))
			r.AppendHistory("member", memberName, result)

			// Save workspace file if member produced a document
			if strings.Contains(result, "# ") {
				filename := fmt.Sprintf("%s-%s.md", memberName, time.Now().Format("20060102-150405"))
				r.saveWorkspace(filename, result)
				r.emit(Event{Type: EvtWorkspaceFile, Filename: filename, Content: result})
			}
		}

		// Feed results back to master for review
		r.setStatus(StatusThinking, "", "")
		masterMsgs = append(masterMsgs, llm.NewMsg("user", "Team results:\n"+results.String()+"\nPlease review. If satisfied output DONE:<summary>, otherwise output ASSIGN instructions for revisions."))

		// Update tasks
		r.updateTasks(masterMsgs)

		if strings.Contains(reply, "DONE:") {
			break
		}
	}

	r.setStatus(StatusPending, "", "")

	// Auto-update master memory after task completion
	go r.updateMasterMemory(context.Background(), userMsg, masterMsgs)

	return nil
}

func (r *Room) updateMasterMemory(ctx context.Context, task string, msgs []llm.Message) {
	summaryPrompt := fmt.Sprintf("Based on this task: %q\nSummarize key learnings and patterns in 3-5 bullet points for future reference. Be concise.", task)
	memMsgs := append(msgs, llm.NewMsg("user", summaryPrompt))
	summary, err := r.master.Chat(ctx, memMsgs, nil)
	if err != nil || summary == "" {
		return
	}
	filename := time.Now().Format("2006-01") + ".md"
	existing, _ := os.ReadFile(filepath.Join(r.master.Dir, "memory", filename))
	content := string(existing) + fmt.Sprintf("\n## %s — %s\n\n%s\n", time.Now().Format("2006-01-02"), task[:min(50, len(task))], summary)
	r.master.SaveMemory(filename, content)
}

func min(a, b int) int {
	if a < b {
		return a
	}
	return b
}

func parseAssignments(text string) map[string]string {
	result := make(map[string]string)
	for _, line := range strings.Split(text, "\n") {
		if strings.HasPrefix(line, "ASSIGN:") {
			parts := strings.SplitN(strings.TrimPrefix(line, "ASSIGN:"), ":", 2)
			if len(parts) == 2 {
				result[strings.TrimSpace(parts[0])] = strings.TrimSpace(parts[1])
			}
		}
	}
	return result
}

func (r *Room) buildTeamXML() string {
	var sb strings.Builder
	sb.WriteString("<team_members>\n")
	for name, a := range r.members {
		fmt.Fprintf(&sb, "  <member>\n    <name>%s</name>\n    <description>%s</description>\n  </member>\n", name, a.Config.Description)
	}
	sb.WriteString("</team_members>")
	return sb.String()
}

func (r *Room) saveWorkspace(filename, content string) {
	dir := filepath.Join(r.Dir, "workspace")
	os.MkdirAll(dir, 0755)
	os.WriteFile(filepath.Join(dir, filename), []byte(content), 0644)
}

func (r *Room) updateTasks(msgs []llm.Message) {
	// Extract task list from conversation and save
	var tasks strings.Builder
	tasks.WriteString("# Tasks\n\n")
	for _, m := range msgs {
		if m.Role == "assistant" && strings.Contains(m.Content, "ASSIGN:") {
			for _, line := range strings.Split(m.Content, "\n") {
				if strings.HasPrefix(line, "ASSIGN:") {
					parts := strings.SplitN(strings.TrimPrefix(line, "ASSIGN:"), ":", 2)
					if len(parts) == 2 {
						tasks.WriteString(fmt.Sprintf("- [ ] [%s] %s\n", strings.TrimSpace(parts[0]), strings.TrimSpace(parts[1])))
					}
				}
			}
		}
	}
	content := tasks.String()
	os.WriteFile(filepath.Join(r.Dir, "tasks.md"), []byte(content), 0644)
	r.emit(Event{Type: EvtTasksUpdate, Content: content})
}

func parseRoomConfig(data []byte) (Config, error) {
	var cfg Config
	if !bytes.HasPrefix(data, []byte("---")) {
		return cfg, fmt.Errorf("missing frontmatter")
	}
	parts := bytes.SplitN(data, []byte("---"), 3)
	if len(parts) < 3 {
		return cfg, fmt.Errorf("invalid frontmatter")
	}
	return cfg, yaml.Unmarshal(parts[1], &cfg)
}