Agent-Squad/platform/skills.go

package platform

import (
	"archive/zip"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"
)

// Skill is one imported SKILL.md plus attached files.
type Skill struct {
	ID, WorkspaceID, Name, Description, Source string
}

// ImportSkill loads a skill from a directory or a zip file.
// conflict is fail, overwrite, rename, or skip.
func (s *Store) ImportSkill(workspaceID, src, conflict string) (Skill, error) {
	name, desc, files, err := readSkillSource(src)
	if err != nil {
		return Skill{}, err
	}
	return s.saveSkill(workspaceID, name, desc, src, files, conflict)
}

func (s *Store) saveSkill(workspaceID, name, desc, source string, files map[string]string, conflict string) (Skill, error) {
	var existing string
	_ = s.db.QueryRow(`SELECT id FROM skill WHERE workspace_id=? AND name=?`, workspaceID, name).Scan(&existing)
	if existing != "" {
		switch conflict {
		case "skip":
			return s.getSkill(existing)
		case "overwrite":
			if _, err := s.db.Exec(`DELETE FROM skill_file WHERE skill_id=?`, existing); err != nil {
				return Skill{}, err
			}
			if _, err := s.db.Exec(`UPDATE skill SET description=?, source=? WHERE id=?`, desc, source, existing); err != nil {
				return Skill{}, err
			}
			return s.writeSkillFiles(Skill{ID: existing, WorkspaceID: workspaceID, Name: name, Description: desc, Source: source}, files)
		case "rename":
			name = name + "-" + newID()[:4]
		case "fail", "":
			return Skill{}, fmt.Errorf("skill %q already exists", name)
		default:
			return Skill{}, fmt.Errorf("unknown conflict policy %q", conflict)
		}
	}
	sk := Skill{ID: newID(), WorkspaceID: workspaceID, Name: name, Description: desc, Source: source}
	if _, err := s.db.Exec(`INSERT INTO skill(id,workspace_id,name,description,source,created_at) VALUES(?,?,?,?,?,?)`,
		sk.ID, sk.WorkspaceID, sk.Name, sk.Description, sk.Source, now()); err != nil {
		return Skill{}, err
	}
	return s.writeSkillFiles(sk, files)
}

func (s *Store) writeSkillFiles(sk Skill, files map[string]string) (Skill, error) {
	for path, body := range files {
		if _, err := s.db.Exec(`INSERT INTO skill_file(id,skill_id,path,body) VALUES(?,?,?,?)`, newID(), sk.ID, path, body); err != nil {
			return Skill{}, err
		}
	}
	return sk, nil
}

func (s *Store) getSkill(id string) (Skill, error) {
	var sk Skill
	err := s.db.QueryRow(`SELECT id,workspace_id,name,description,source FROM skill WHERE id=?`, id).
		Scan(&sk.ID, &sk.WorkspaceID, &sk.Name, &sk.Description, &sk.Source)
	return sk, err
}

// RefreshSkill re-reads the original source path.
func (s *Store) RefreshSkill(id string) (Skill, error) {
	sk, err := s.getSkill(id)
	if err != nil {
		return Skill{}, err
	}
	return s.ImportSkill(sk.WorkspaceID, sk.Source, "overwrite")
}

// BindSkill attaches a skill to an agent. enabled can be toggled later.
func (s *Store) BindSkill(agentID, skillID string, enabled bool) error {
	flag := 0
	if enabled {
		flag = 1
	}
	_, err := s.db.Exec(`INSERT INTO agent_skill(agent_id,skill_id,enabled) VALUES(?,?,?)
		ON CONFLICT(agent_id, skill_id) DO UPDATE SET enabled=excluded.enabled`, agentID, skillID, flag)
	if err != nil {
		_, err = s.db.Exec(`DELETE FROM agent_skill WHERE agent_id=? AND skill_id=?`, agentID, skillID)
		if err != nil {
			return err
		}
		_, err = s.db.Exec(`INSERT INTO agent_skill(agent_id,skill_id,enabled) VALUES(?,?,?)`, agentID, skillID, flag)
	}
	return err
}

func readSkillSource(src string) (string, string, map[string]string, error) {
	info, err := os.Stat(src)
	if err != nil {
		return "", "", nil, err
	}
	files := map[string]string{}
	if info.IsDir() {
		err = filepath.Walk(src, func(path string, fi os.FileInfo, err error) error {
			if err != nil || fi.IsDir() {
				return err
			}
			rel, _ := filepath.Rel(src, path)
			raw, err := os.ReadFile(path)
			if err != nil {
				return err
			}
			files[rel] = string(raw)
			return nil
		})
	} else if strings.HasSuffix(src, ".zip") {
		zr, err := zip.OpenReader(src)
		if err != nil {
			return "", "", nil, err
		}
		defer zr.Close()
		for _, f := range zr.File {
			if f.FileInfo().IsDir() {
				continue
			}
			rc, err := f.Open()
			if err != nil {
				return "", "", nil, err
			}
			raw, err := io.ReadAll(rc)
			_ = rc.Close()
			if err != nil {
				return "", "", nil, err
			}
			files[f.Name] = string(raw)
		}
	} else {
		return "", "", nil, fmt.Errorf("skill source must be a directory or zip")
	}
	if err != nil {
		return "", "", nil, err
	}
	body := ""
	for path, content := range files {
		if filepath.Base(path) == "SKILL.md" {
			body = content
			break
		}
	}
	if body == "" {
		return "", "", nil, fmt.Errorf("SKILL.md not found in %s", src)
	}
	name := skillName(src, body)
	return name, skillDescription(body), files, nil
}

func skillName(src, body string) string {
	if info, err := os.Stat(src); err == nil && info.IsDir() {
		return filepath.Base(src)
	}
	return strings.TrimSuffix(filepath.Base(src), ".zip")
}

func skillDescription(body string) string {
	for _, line := range strings.Split(body, "\n") {
		line = strings.TrimSpace(line)
		if strings.HasPrefix(line, "description:") {
			return strings.TrimSpace(strings.TrimPrefix(line, "description:"))
		}
	}
	return ""
}