Agent-Squad/platform/runtime.go

package platform

import (
	"context"
	"encoding/json"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"strings"

	"agent-native/runtime"
	"agent-native/runtime/llm"
)

// CLICommands are the external tools a daemon probes, in order.
var CLICommands = []string{
	"claude", "codex", "opencode", "codearts", "deveco", "openclaw", "hermes", "pi",
	"omp", "cursor-agent", "copilot", "kimi", "reasonix", "dsh", "kiro-cli", "codebuddy",
	"agy", "qodercli", "qodercln", "traecli", "grok", "qwen", "qwenpaw", "dim", "mcode",
}

// ProtocolFor maps a command onto a protocol family.
func ProtocolFor(command string) string {
	switch command {
	case "claude":
		return "claude"
	case "codex":
		return "codex"
	case "pi":
		return "pi"
	case "opencode":
		return "opencode"
	default:
		return "custom"
	}
}

// ProbeCLIs returns commands present on PATH. An empty list is allowed:
// the in-process runtime still runs.
func ProbeCLIs() []string {
	var found []string
	for _, name := range CLICommands {
		if env := os.Getenv("AGENT_SQUAD_" + strings.ToUpper(strings.ReplaceAll(name, "-", "_")) + "_PATH"); env != "" {
			if _, err := os.Stat(env); err == nil {
				found = append(found, env)
				continue
			}
		}
		if path, err := exec.LookPath(name); err == nil {
			found = append(found, path)
		}
	}
	return found
}

// RegisterProbed records every discovered CLI plus the in-process runtime.
func (s *Store) RegisterProbed(workspaceID string) ([]Runtime, error) {
	inproc, err := s.RegisterRuntime(Runtime{
		WorkspaceID: workspaceID, Name: "agent-native", Command: "in-process",
		Protocol: "agent-native", MaxConcurrency: 2,
	})
	if err != nil {
		return nil, err
	}
	out := []Runtime{inproc}
	for _, path := range ProbeCLIs() {
		base := filepath.Base(path)
		rt, err := s.RegisterRuntime(Runtime{
			WorkspaceID: workspaceID, Name: base, Command: path,
			Protocol: ProtocolFor(base), MaxConcurrency: 1,
		})
		if err != nil {
			return nil, err
		}
		out = append(out, rt)
	}
	return out, nil
}

// PrepareWorkdir creates workdir, output, and logs, and writes the brief and skills.
func (s *Store) PrepareWorkdir(root string, task Task) (string, error) {
	base := filepath.Join(root, task.WorkspaceID, task.IssueID)
	for _, sub := range []string{"workdir", "output", "logs"} {
		if err := os.MkdirAll(filepath.Join(base, sub), 0o755); err != nil {
			return "", err
		}
	}
	brief := task.IssueID
	if task.SquadID != "" {
		text, err := s.LeaderBrief(task.SquadID)
		if err != nil {
			return "", err
		}
		brief = text
	}
	work := filepath.Join(base, "workdir")
	for _, name := range []string{"AGENTS.md", "CLAUDE.md"} {
		if err := os.WriteFile(filepath.Join(work, name), []byte(brief), 0o644); err != nil {
			return "", err
		}
	}
	if err := s.materializeSkills(work, task.AgentID); err != nil {
		return "", err
	}
	return work, nil
}

func (s *Store) materializeSkills(work, agentID string) error {
	rows, err := s.db.Query(`SELECT skill.id, skill.name FROM skill JOIN agent_skill ON agent_skill.skill_id=skill.id WHERE agent_skill.agent_id=? AND agent_skill.enabled=1`, agentID)
	if err != nil {
		return err
	}
	type skillRef struct{ id, name string }
	var skills []skillRef
	for rows.Next() {
		var item skillRef
		if err := rows.Scan(&item.id, &item.name); err != nil {
			rows.Close()
			return err
		}
		skills = append(skills, item)
	}
	err = rows.Err()
	rows.Close()
	if err != nil {
		return err
	}
	for _, item := range skills {
		dir := filepath.Join(work, ".agents", "skills", item.name)
		if err := os.MkdirAll(dir, 0o755); err != nil {
			return err
		}
		files, err := s.db.Query(`SELECT path, body FROM skill_file WHERE skill_id=?`, item.id)
		if err != nil {
			return err
		}
		for files.Next() {
			var path, body string
			if err := files.Scan(&path, &body); err != nil {
				files.Close()
				return err
			}
			target := filepath.Join(dir, path)
			if err := os.MkdirAll(filepath.Dir(target), 0o755); err != nil {
				files.Close()
				return err
			}
			if err := os.WriteFile(target, []byte(body), 0o644); err != nil {
				files.Close()
				return err
			}
		}
		if err := files.Err(); err != nil {
			files.Close()
			return err
		}
		files.Close()
	}
	return nil
}

// ExecutionEnv is the environment for one task. custom_env cannot override
// the task token, PATH, or HOME.
func ExecutionEnv(task Task, custom string) []string {
	env := os.Environ()
	reserved := map[string]bool{"AGENT_SQUAD_TOKEN": true, "PATH": true, "HOME": true}
	for _, line := range strings.Split(custom, "\n") {
		if line == "" || !strings.Contains(line, "=") {
			continue
		}
		key := strings.SplitN(line, "=", 2)[0]
		if reserved[key] {
			continue
		}
		env = append(env, line)
	}
	env = append(env, "AGENT_SQUAD_TOKEN="+task.Token, "AGENT_SQUAD_TASK="+task.ID, "AGENT_SQUAD_ISSUE="+task.IssueID)
	return env
}

// RunCommand executes an external CLI in the workdir and records the result.
func (s *Store) RunCommand(ctx context.Context, task Task, command, workdir string) error {
	if err := s.MarkRunning(task.ID); err != nil {
		return err
	}
	cmd := exec.CommandContext(ctx, command)
	cmd.Dir = workdir
	agent, _ := s.GetAgent(task.AgentID)
	cmd.Env = ExecutionEnv(task, agent.CustomEnv)
	out, err := cmd.CombinedOutput()
	_ = s.AddTaskMessage(task.ID, "log", string(out))
	if err != nil {
		return s.FinishTask(task.ID, "failed", "timeout")
	}
	return s.FinishTask(task.ID, "completed", "")
}

// RunInProcess executes the task with AgentSession. stream is required; the
// product supplies a provider, tests supply a script.
func (s *Store) RunInProcess(ctx context.Context, task Task, workdir string, stream llm.StreamFunc) (string, error) {
	if err := s.MarkRunning(task.ID); err != nil {
		return "", err
	}
	issue, _ := s.GetIssue(task.IssueID)
	agent := runtime.NewAgentSession(runtime.Config{
		Model:    "agent-native",
		Stream:   stream,
		MaxTurns: 4,
		OnEvent: func(ev runtime.Event) {
			_ = s.AddTaskMessage(task.ID, ev.Type, ev.Text)
		},
	})
	res, err := agent.Prompt(ctx, llm.Message{Role: llm.RoleUser, Content: issue.Title + "\n" + issue.Body})
	if err != nil {
		_ = s.FinishTask(task.ID, "failed", "config")
		return "", err
	}
	if res.Err != nil {
		_ = s.FinishTask(task.ID, "failed", "model")
		return res.Text, res.Err
	}
	_ = s.AddUsage(task.ID, "agent-native", res.Usage.PromptTokens, res.Usage.CompletionTokens, res.Usage.CacheReadTokens, res.Usage.CacheWriteTokens)
	if err := s.FinishTask(task.ID, "completed", ""); err != nil {
		return res.Text, err
	}
	_ = os.WriteFile(filepath.Join(workdir, "output.md"), []byte(res.Text), 0o644)
	return res.Text, nil
}

// AgentCLI applies a task-token command: comment, status, sub-issue, or activity.
func (s *Store) AgentCLI(token string, args []string) error {
	task, err := s.TaskByToken(token)
	if err != nil {
		return err
	}
	if len(args) == 0 {
		return fmt.Errorf("agent cli: missing command")
	}
	switch args[0] {
	case "comment":
		if len(args) < 2 {
			return fmt.Errorf("comment requires a body")
		}
		_, err = s.AddComment(Comment{IssueID: task.IssueID, AuthorKind: "agent", AuthorID: task.AgentID, Body: args[1]})
		return err
	case "status":
		if len(args) < 2 {
			return fmt.Errorf("status requires a value")
		}
		goal := args[1] == "in_review"
		return s.TouchStatus(task.IssueID, task.AgentID, goal)
	case "sub-issue":
		if len(args) < 2 {
			return fmt.Errorf("sub-issue requires a title")
		}
		_, err = s.CreateIssue(Issue{WorkspaceID: task.WorkspaceID, Title: args[1], Body: args[1], ParentID: task.IssueID, Status: "todo"})
		return err
	case "activity":
		if len(args) < 2 {
			return fmt.Errorf("activity requires an action")
		}
		detail := ""
		if len(args) > 2 {
			detail = args[2]
		}
		return s.RecordActivity(task.SquadID, task.IssueID, task.AgentID, args[1], detail)
	default:
		return fmt.Errorf("unknown agent command %s", args[0])
	}
}

// CheckoutWorktree documents the branch name used for parallel local_directory work.
func WorktreeBranch(agentName, issueID string) string {
	safe := strings.Map(func(r rune) rune {
		if r == '/' || r == ' ' {
			return '-'
		}
		return r
	}, agentName)
	return "agent/" + safe + "/" + issueID
}

// Ensure the JSON import stays referenced for providers that pass raw tool args.
var _ = json.RawMessage(nil)