Agent-Squad/platform/autopilot.go

package platform

import (
	"crypto/hmac"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"strings"
	"time"
)

// Autopilot is a runbook fired by cron, webhook, or API.
type Autopilot struct {
	ID, WorkspaceID, Title, Runbook, AssigneeKind, AssigneeID, Status, ExecutionMode string
}

func (s *Store) CreateAutopilot(a Autopilot) (Autopilot, error) {
	if a.ID == "" {
		a.ID = newID()
	}
	if a.Status == "" {
		a.Status = "active"
	}
	if a.ExecutionMode != "create_issue" && a.ExecutionMode != "run_only" {
		return Autopilot{}, fmt.Errorf("execution mode must be create_issue or run_only")
	}
	_, err := s.db.Exec(`INSERT INTO autopilot(id,workspace_id,title,runbook,assignee_kind,assignee_id,status,execution_mode,created_at) VALUES(?,?,?,?,?,?,?,?,?)`,
		a.ID, a.WorkspaceID, a.Title, a.Runbook, a.AssigneeKind, a.AssigneeID, a.Status, a.ExecutionMode, now())
	return a, err
}

// AddSchedule stores a five-field cron expression and an IANA timezone.
func (s *Store) AddSchedule(autopilotID, cronExpr, timezone string) error {
	if _, err := time.LoadLocation(timezone); err != nil {
		return fmt.Errorf("timezone: %w", err)
	}
	if len(strings.Fields(cronExpr)) != 5 {
		return fmt.Errorf("cron must have 5 fields")
	}
	_, err := s.db.Exec(`INSERT INTO trigger_(id,autopilot_id,kind,cron,timezone) VALUES(?,?,?,?,?)`,
		newID(), autopilotID, "schedule", cronExpr, timezone)
	return err
}

// AddWebhook returns the token clients post to.
func (s *Store) AddWebhook(autopilotID, secret string) (string, error) {
	token := newID()
	_, err := s.db.Exec(`INSERT INTO trigger_(id,autopilot_id,kind,token,secret) VALUES(?,?,?,?,?)`,
		newID(), autopilotID, "webhook", token, secret)
	return token, err
}

// FireSchedules runs every active schedule that matches at.
func (s *Store) FireSchedules(at time.Time) (int, error) {
	rows, err := s.db.Query(`SELECT autopilot.id, trigger_.cron, trigger_.timezone FROM autopilot JOIN trigger_ ON trigger_.autopilot_id=autopilot.id WHERE autopilot.status='active' AND trigger_.kind='schedule'`)
	if err != nil {
		return 0, err
	}
	type due struct{ id, cron, tz string }
	var list []due
	for rows.Next() {
		var item due
		if err := rows.Scan(&item.id, &item.cron, &item.tz); err != nil {
			rows.Close()
			return 0, err
		}
		list = append(list, item)
	}
	err = rows.Err()
	rows.Close()
	if err != nil {
		return 0, err
	}
	n := 0
	for _, item := range list {
		if !cronDue(item.cron, item.tz, at) {
			continue
		}
		if err := s.fire(item.id, "schedule"); err != nil {
			return n, err
		}
		n++
	}
	return n, nil
}

// AcceptWebhook validates, dedupes, and fires. Body must be a JSON object or array up to 256KiB.
func (s *Store) AcceptWebhook(token string, body []byte, headers map[string]string) (string, error) {
	if len(body) > 256*1024 {
		return "", fmt.Errorf("webhook body exceeds 256KiB")
	}
	var probe any
	if err := json.Unmarshal(body, &probe); err != nil {
		return "", fmt.Errorf("webhook body must be JSON")
	}
	switch probe.(type) {
	case map[string]any, []any:
	default:
		return "", fmt.Errorf("webhook body must be a JSON object or array")
	}
	var autopilotID, secret string
	err := s.db.QueryRow(`SELECT autopilot_id, IFNULL(secret,'') FROM trigger_ WHERE kind='webhook' AND token=?`, token).Scan(&autopilotID, &secret)
	if err != nil {
		return "", err
	}
	if secret != "" {
		sig := headers["X-Hub-Signature-256"]
		mac := hmac.New(sha256.New, []byte(secret))
		_, _ = mac.Write(body)
		want := "sha256=" + hex.EncodeToString(mac.Sum(nil))
		if !hmac.Equal([]byte(sig), []byte(want)) {
			return "", fmt.Errorf("webhook signature mismatch")
		}
	}
	key := headers["Idempotency-Key"]
	if key == "" {
		key = headers["X-GitHub-Delivery"]
	}
	if key != "" {
		if _, err := s.db.Exec(`INSERT INTO webhook_delivery(id,autopilot_id,dedupe_key,created_at) VALUES(?,?,?,?)`, newID(), autopilotID, key, now()); err != nil {
			if strings.Contains(err.Error(), "UNIQUE") {
				return "duplicate", nil
			}
			return "", err
		}
	}
	if err := s.fire(autopilotID, eventName(headers, probe)); err != nil {
		return "", err
	}
	return eventName(headers, probe), nil
}

func eventName(headers map[string]string, probe any) string {
	if v := headers["X-GitHub-Event"]; v != "" {
		return v
	}
	if obj, ok := probe.(map[string]any); ok {
		if v, ok := obj["action"].(string); ok && v != "" {
			return v
		}
		if v, ok := obj["event"].(string); ok && v != "" {
			return v
		}
	}
	return "webhook"
}

func (s *Store) fire(autopilotID, detail string) error {
	var a Autopilot
	err := s.db.QueryRow(`SELECT id,workspace_id,title,runbook,assignee_kind,assignee_id,status,execution_mode FROM autopilot WHERE id=?`, autopilotID).
		Scan(&a.ID, &a.WorkspaceID, &a.Title, &a.Runbook, &a.AssigneeKind, &a.AssigneeID, &a.Status, &a.ExecutionMode)
	if err != nil {
		return err
	}
	if a.Status != "active" {
		return nil
	}
	status := "completed"
	if a.ExecutionMode == "create_issue" {
		issue, err := s.CreateIssue(Issue{WorkspaceID: a.WorkspaceID, Title: a.Title, Body: a.Runbook, Status: "todo", AssigneeKind: a.AssigneeKind, AssigneeID: a.AssigneeID})
		if err != nil {
			status = "failed"
			detail = err.Error()
		} else if err := s.AssignIssue(issue.ID, a.AssigneeKind, a.AssigneeID, true); err != nil {
			status = "failed"
			detail = err.Error()
		}
	} else if a.AssigneeKind == "agent" {
		issue := Issue{WorkspaceID: a.WorkspaceID}
		t, err := s.enqueue(issue, a.AssigneeID, false, "")
		if err != nil {
			status = "failed"
		} else {
			_, _ = s.db.Exec(`UPDATE task SET issue_id='' WHERE id=?`, t.ID)
		}
	}
	if _, err := s.db.Exec(`INSERT INTO autopilot_run(id,autopilot_id,status,detail,created_at) VALUES(?,?,?,?,?)`, newID(), autopilotID, status, detail, now()); err != nil {
		return err
	}
	return s.maybePause(autopilotID)
}

func (s *Store) maybePause(autopilotID string) error {
	since := time.Now().UTC().Add(-7 * 24 * time.Hour).Format(time.RFC3339Nano)
	var total, failed int
	_ = s.db.QueryRow(`SELECT COUNT(*) FROM autopilot_run WHERE autopilot_id=? AND created_at>=?`, autopilotID, since).Scan(&total)
	_ = s.db.QueryRow(`SELECT COUNT(*) FROM autopilot_run WHERE autopilot_id=? AND created_at>=? AND status='failed'`, autopilotID, since).Scan(&failed)
	if total >= 50 && failed*10 >= total*9 {
		_, err := s.db.Exec(`UPDATE autopilot SET status='paused' WHERE id=?`, autopilotID)
		return err
	}
	return nil
}

func cronDue(expr, timezone string, at time.Time) bool {
	loc, err := time.LoadLocation(timezone)
	if err != nil {
		return false
	}
	at = at.In(loc)
	fields := strings.Fields(expr)
	if len(fields) != 5 {
		return false
	}
	values := []int{at.Minute(), at.Hour(), at.Day(), int(at.Month()), int(at.Weekday())}
	for i, field := range fields {
		if !cronField(field, values[i]) {
			return false
		}
	}
	return true
}

func cronField(field string, value int) bool {
	if field == "*" {
		return true
	}
	if strings.HasPrefix(field, "*/") {
		var step int
		if _, err := fmt.Sscan(field[2:], &step); err != nil || step <= 0 {
			return false
		}
		return value%step == 0
	}
	var n int
	if _, err := fmt.Sscan(field, &n); err != nil {
		return false
	}
	return n == value
}

// ChatSession is private to its owner and may point at a project.
type ChatSession struct {
	ID, WorkspaceID, OwnerID, ProjectID string
}

func (s *Store) CreateChat(c ChatSession) (ChatSession, error) {
	if c.ID == "" {
		c.ID = newID()
	}
	_, err := s.db.Exec(`INSERT INTO chat_session(id,workspace_id,owner_id,project_id,created_at) VALUES(?,?,?,?,?)`,
		c.ID, c.WorkspaceID, c.OwnerID, c.ProjectID, now())
	return c, err
}

func (s *Store) AddChatMessage(sessionID, role, body string) error {
	_, err := s.db.Exec(`INSERT INTO chat_message(id,session_id,role,body,created_at) VALUES(?,?,?,?,?)`, newID(), sessionID, role, body, now())
	return err
}

// Project groups resources.
type Project struct {
	ID, WorkspaceID, Name string
}

func (s *Store) CreateProject(p Project) (Project, error) {
	if p.ID == "" {
		p.ID = newID()
	}
	_, err := s.db.Exec(`INSERT INTO project(id,workspace_id,name,created_at) VALUES(?,?,?,?)`, p.ID, p.WorkspaceID, p.Name, now())
	return p, err
}

// AddResource records a github_repo or local_directory. local_directory mode is in_place or worktree.
func (s *Store) AddResource(projectID, kind, locator, mode string) error {
	switch kind {
	case "github_repo":
		mode = ""
	case "local_directory":
		if mode != "in_place" && mode != "worktree" {
			return fmt.Errorf("local_directory mode must be in_place or worktree")
		}
	default:
		return fmt.Errorf("unknown resource kind %s", kind)
	}
	_, err := s.db.Exec(`INSERT INTO project_resource(id,project_id,kind,locator,mode) VALUES(?,?,?,?,?)`, newID(), projectID, kind, locator, mode)
	return err
}