Agent-Squad/adapters/sqlite/queue.go

package platform

import (
	"fmt"
	"regexp"
	"strings"
	"time"
)

// Issue is the durable work record. A task is one execution of it.
type Issue struct {
	ID, WorkspaceID, Title, Body, Status, ParentID, Stage, AssigneeKind, AssigneeID string
}

func (s *Store) CreateIssue(issue Issue) (Issue, error) {
	if issue.ID == "" {
		issue.ID = newID()
	}
	if issue.Status == "" {
		issue.Status = "todo"
	}
	_, err := s.db.Exec(`INSERT INTO issue(id,workspace_id,title,body,status,parent_id,stage,assignee_kind,assignee_id,created_at)
		VALUES(?,?,?,?,?,?,?,?,?,?)`,
		issue.ID, issue.WorkspaceID, issue.Title, issue.Body, issue.Status, issue.ParentID, issue.Stage, issue.AssigneeKind, issue.AssigneeID, now())
	return issue, err
}

func (s *Store) GetIssue(id string) (Issue, error) {
	var issue Issue
	err := s.db.QueryRow(`SELECT id,workspace_id,title,body,status,IFNULL(parent_id,''),IFNULL(stage,''),IFNULL(assignee_kind,''),IFNULL(assignee_id,'') FROM issue WHERE id=?`, id).
		Scan(&issue.ID, &issue.WorkspaceID, &issue.Title, &issue.Body, &issue.Status, &issue.ParentID, &issue.Stage, &issue.AssigneeKind, &issue.AssigneeID)
	return issue, err
}

// AssignIssue records the owner. Backlog, or start=false, does not enqueue.
// A squad assignee enqueues only the leader.
func (s *Store) AssignIssue(issueID, kind, id string, start bool) error {
	issue, err := s.GetIssue(issueID)
	if err != nil {
		return err
	}
	if _, err := s.db.Exec(`UPDATE issue SET assignee_kind=?, assignee_id=? WHERE id=?`, kind, id, issueID); err != nil {
		return err
	}
	issue.AssigneeKind, issue.AssigneeID = kind, id
	if !start || issue.Status == "backlog" {
		return nil
	}
	return s.enqueueAssignee(issue)
}

func (s *Store) enqueueAssignee(issue Issue) error {
	switch issue.AssigneeKind {
	case "agent":
		_, err := s.enqueue(issue, issue.AssigneeID, false, "")
		return err
	case "squad":
		var leader, squadID string
		if err := s.db.QueryRow(`SELECT id, leader_agent_id FROM squad WHERE id=?`, issue.AssigneeID).Scan(&squadID, &leader); err != nil {
			return err
		}
		_, err := s.enqueue(issue, leader, true, squadID)
		return err
	default:
		return nil
	}
}

// Comment is one note on an issue.
type Comment struct {
	ID, IssueID, AuthorKind, AuthorID, Body, ParentID string
}

var mentionRE = regexp.MustCompile(`mention://(agent|member)/([0-9a-fA-F]+)`)

// AddComment stores the comment and routes mentions. /note and a bare @all do not enqueue.
func (s *Store) AddComment(c Comment) (Comment, error) {
	if c.ID == "" {
		c.ID = newID()
	}
	if _, err := s.db.Exec(`INSERT INTO comment(id,issue_id,author_kind,author_id,body,parent_id,created_at) VALUES(?,?,?,?,?,?,?)`,
		c.ID, c.IssueID, c.AuthorKind, c.AuthorID, c.Body, c.ParentID, now()); err != nil {
		return Comment{}, err
	}
	if err := s.routeComment(c); err != nil {
		return Comment{}, err
	}
	return c, nil
}

func (s *Store) routeComment(c Comment) error {
	trimmed := strings.TrimSpace(c.Body)
	if strings.HasPrefix(trimmed, "/note") {
		return nil
	}
	if strings.Contains(trimmed, "@all") && !mentionRE.MatchString(c.Body) {
		return nil
	}
	issue, err := s.GetIssue(c.IssueID)
	if err != nil {
		return err
	}
	if issue.Status == "done" || issue.Status == "cancelled" || issue.Status == "backlog" {
		return nil
	}
	matches := mentionRE.FindAllStringSubmatch(c.Body, -1)
	for _, m := range matches {
		if m[1] != "agent" {
			continue
		}
		if _, err := s.enqueue(issue, m[2], false, issueSquad(issue)); err != nil {
			return err
		}
	}
	if c.ParentID != "" && len(matches) == 0 {
		var kind, author string
		err := s.db.QueryRow(`SELECT author_kind, author_id FROM comment WHERE id=?`, c.ParentID).Scan(&kind, &author)
		if err == nil && kind == "agent" {
			_, err = s.enqueue(issue, author, false, issueSquad(issue))
			return err
		}
	}
	return nil
}

func issueSquad(issue Issue) string {
	if issue.AssigneeKind == "squad" {
		return issue.AssigneeID
	}
	return ""
}

// Task is one execution attempt.
type Task struct {
	ID, WorkspaceID, IssueID, AgentID, RuntimeID, Status, ReasonCode, LeaseUntil, FireAt, SquadID, Token string
	Attempts                                                                                             int
	IsLeader                                                                                             bool
}

func (s *Store) enqueue(issue Issue, agentID string, leader bool, squadID string) (Task, error) {
	var existing string
	err := s.db.QueryRow(`SELECT id FROM task WHERE issue_id=? AND agent_id=? AND status IN ('queued','dispatched')`, issue.ID, agentID).Scan(&existing)
	if err == nil && existing != "" {
		return s.getTask(existing)
	}
	leaderFlag := 0
	if leader {
		leaderFlag = 1
	}
	t := Task{
		ID: newID(), WorkspaceID: issue.WorkspaceID, IssueID: issue.ID, AgentID: agentID,
		Status: "queued", SquadID: squadID, Token: newID(), IsLeader: leader,
	}
	_, err = s.db.Exec(`INSERT INTO task(id,workspace_id,issue_id,agent_id,runtime_id,status,reason_code,attempts,lease_until,fire_at,is_leader,squad_id,token,created_at,updated_at)
		VALUES(?,?,?,?,?,?,?,0,?,?,?,?,?,?,?)`,
		t.ID, t.WorkspaceID, t.IssueID, t.AgentID, "", t.Status, "", "", "", leaderFlag, t.SquadID, t.Token, now(), now())
	return t, err
}

// DeferTask parks a task until fireAt.
func (s *Store) DeferTask(issue Issue, agentID string, fireAt time.Time) (Task, error) {
	t := Task{ID: newID(), WorkspaceID: issue.WorkspaceID, IssueID: issue.ID, AgentID: agentID, Status: "deferred", Token: newID(), FireAt: fireAt.UTC().Format(time.RFC3339Nano)}
	_, err := s.db.Exec(`INSERT INTO task(id,workspace_id,issue_id,agent_id,runtime_id,status,reason_code,attempts,lease_until,fire_at,is_leader,squad_id,token,created_at,updated_at)
		VALUES(?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)`,
		t.ID, t.WorkspaceID, t.IssueID, t.AgentID, "", "deferred", "", 0, "", t.FireAt, 0, "", t.Token, now(), now())
	if err != nil {
		return Task{}, err
	}
	_, err = s.db.Exec(`INSERT INTO wakeup(id,task_id,kind,fire_at,created_at) VALUES(?,?,?,?,?)`, newID(), t.ID, "deferred", t.FireAt, now())
	return t, err
}

func (s *Store) getTask(id string) (Task, error) {
	var t Task
	var leader int
	err := s.db.QueryRow(`SELECT id,workspace_id,IFNULL(issue_id,''),agent_id,IFNULL(runtime_id,''),status,IFNULL(reason_code,''),attempts,IFNULL(lease_until,''),IFNULL(fire_at,''),is_leader,IFNULL(squad_id,''),token FROM task WHERE id=?`, id).
		Scan(&t.ID, &t.WorkspaceID, &t.IssueID, &t.AgentID, &t.RuntimeID, &t.Status, &t.ReasonCode, &t.Attempts, &t.LeaseUntil, &t.FireAt, &leader, &t.SquadID, &t.Token)
	t.IsLeader = leader == 1
	return t, err
}

// PromoteDeferred moves due deferred tasks to queued.
func (s *Store) PromoteDeferred(at time.Time) (int, error) {
	res, err := s.db.Exec(`UPDATE task SET status='queued', updated_at=? WHERE status='deferred' AND fire_at<=?`, now(), at.UTC().Format(time.RFC3339Nano))
	if err != nil {
		return 0, err
	}
	n, _ := res.RowsAffected()
	return int(n), nil
}

// claimSQL is the whole claim: one atomic conditional UPDATE whose two
// guarded counts bound the agent by its own max_concurrency and the runtime by
// max_concurrency over every in-flight task of that runtime (an agent bound to
// it, or a task this runtime already leased). A count that cannot be read
// fails the statement, so the claim is refused instead of allowed.
const claimSQL = `UPDATE task SET status='dispatched', runtime_id=?, lease_until=?, updated_at=?
	WHERE id=? AND status='queued'
	AND (SELECT COUNT(*) FROM task WHERE agent_id=? AND status IN ('dispatched','running')) < ?
	AND (SELECT COUNT(*) FROM task t WHERE t.status IN ('dispatched','running')
		AND (t.runtime_id=? OR EXISTS (SELECT 1 FROM agent a WHERE a.id=t.agent_id AND a.runtime_id=?))) < ?`

// claimCapacitySQL counts the agent's own and the runtime's in-flight tasks in
// one read, so the pre-check below matches the guards of claimSQL.
const claimCapacitySQL = `SELECT
	(SELECT COUNT(*) FROM task WHERE agent_id=? AND status IN ('dispatched','running')),
	(SELECT COUNT(*) FROM task t WHERE t.status IN ('dispatched','running')
		AND (t.runtime_id=? OR EXISTS (SELECT 1 FROM agent a WHERE a.id=t.agent_id AND a.runtime_id=?)))`

// Claim takes one queued task whose agent and runtime are under their caps.
// The task becomes dispatched with a five-minute lease.
//
// Both caps live inside the single conditional UPDATE, so concurrent claims
// can overshoot neither the agent cap nor the runtime total: the agent count
// is bounded by agent.max_concurrency and the runtime count (every in-flight
// task whose agent belongs to this runtime, plus every task this runtime
// already leased) by runtime.max_concurrency. A read error is a refusal, never
// a silent allow: the candidate is skipped, and no task is handed out.
func (s *Store) Claim(runtimeID string, at time.Time) (Task, bool, error) {
	rt, err := s.getRuntime(runtimeID)
	if err != nil {
		return Task{}, false, err
	}
	rows, err := s.db.Query(`SELECT id, agent_id FROM task WHERE status='queued' AND workspace_id=? ORDER BY created_at`, rt.WorkspaceID)
	if err != nil {
		return Task{}, false, err
	}
	type candidate struct{ id, agentID string }
	var candidates []candidate
	for rows.Next() {
		var id, agentID string
		if err := rows.Scan(&id, &agentID); err != nil {
			rows.Close()
			return Task{}, false, err
		}
		candidates = append(candidates, candidate{id, agentID})
	}
	if err := rows.Err(); err != nil {
		rows.Close()
		return Task{}, false, err
	}
	rows.Close()
	for _, cand := range candidates {
		agent, err := s.GetAgent(cand.agentID)
		if err != nil || agent.Archived {
			continue
		}
		if agent.RuntimeID != "" && agent.RuntimeID != runtimeID {
			continue
		}
		// Pre-check only: the UPDATE below re-counts both caps under the write
		// lock. A count that cannot be read means unknown capacity, so refuse.
		var agentBusy, runtimeBusy int
		err = s.db.QueryRow(claimCapacitySQL, cand.agentID, runtimeID, runtimeID).Scan(&agentBusy, &runtimeBusy)
		if err != nil {
			return Task{}, false, fmt.Errorf("claim capacity count: %w", err)
		}
		if agentBusy >= agent.MaxConcurrency || runtimeBusy >= rt.MaxConcurrency {
			continue
		}
		lease := at.UTC().Add(5 * time.Minute).Format(time.RFC3339Nano)
		res, err := s.db.Exec(claimSQL,
			runtimeID, lease, now(), cand.id, cand.agentID, agent.MaxConcurrency,
			runtimeID, runtimeID, rt.MaxConcurrency)
		if err != nil {
			return Task{}, false, err
		}
		n, _ := res.RowsAffected()
		if n == 1 {
			t, err := s.getTask(cand.id)
			return t, true, err
		}
	}
	return Task{}, false, nil
}

func (s *Store) getRuntime(id string) (Runtime, error) {
	var rt Runtime
	err := s.db.QueryRow(`SELECT id,workspace_id,name,command,protocol,max_concurrency,IFNULL(last_heartbeat,'') FROM runtime WHERE id=?`, id).
		Scan(&rt.ID, &rt.WorkspaceID, &rt.Name, &rt.Command, &rt.Protocol, &rt.MaxConcurrency, &rt.LastHeartbeat)
	return rt, err
}

// MarkRunning starts the lease clock of an already claimed task.
func (s *Store) MarkRunning(taskID string) error {
	_, err := s.db.Exec(`UPDATE task SET status='running', updated_at=? WHERE id=?`, now(), taskID)
	return err
}

// TransientReason reports whether a failure may be retried.
func TransientReason(code string) bool {
	switch code {
	case "runtime_offline", "daemon_restart", "timeout", "skill_bundle", "tool_network":
		return true
	default:
		return false
	}
}

func retryLimit(code string) int {
	if code == "tool_network" {
		return 3
	}
	if TransientReason(code) {
		return 2
	}
	return 0
}

// FinishTask closes a task. Transient failures requeue until the attempt cap.
// Autopilot run_only tasks are not retried. A failed issue task rolls
// in_progress back to todo when nothing else is still active.
func (s *Store) FinishTask(taskID, status, reason string) error {
	t, err := s.getTask(taskID)
	if err != nil {
		return err
	}
	if status == "failed" && TransientReason(reason) && t.IssueID != "" && t.Attempts+1 < retryLimit(reason) {
		_, err = s.db.Exec(`UPDATE task SET status='queued', reason_code=?, attempts=attempts+1, runtime_id='', lease_until='', updated_at=? WHERE id=?`,
			reason, now(), taskID)
		return err
	}
	if status != "completed" && status != "failed" && status != "cancelled" {
		return fmt.Errorf("unsupported terminal status %s", status)
	}
	if _, err := s.db.Exec(`UPDATE task SET status=?, reason_code=?, updated_at=? WHERE id=?`, status, reason, now(), taskID); err != nil {
		return err
	}
	if status == "failed" && t.IssueID != "" {
		var active int
		_ = s.db.QueryRow(`SELECT COUNT(*) FROM task WHERE issue_id=? AND status IN ('deferred','queued','dispatched','waiting_local_directory','running')`, t.IssueID).Scan(&active)
		if active == 0 {
			_, _ = s.db.Exec(`UPDATE issue SET status='todo' WHERE id=? AND status='in_progress'`, t.IssueID)
		}
	}
	return nil
}

// RecoverOrphans fails dispatched tasks whose lease expired and requeues
// running tasks whose runtime heartbeat is older than three minutes.
func (s *Store) RecoverOrphans(at time.Time) (int, error) {
	n := 0
	rows, err := s.db.Query(`SELECT id FROM task WHERE status='dispatched' AND lease_until < ?`, at.UTC().Format(time.RFC3339Nano))
	if err != nil {
		return 0, err
	}
	ids, err := scanStrings(rows)
	if err != nil {
		return 0, err
	}
	for _, id := range ids {
		if err := s.FinishTask(id, "failed", "timeout"); err != nil {
			return n, err
		}
		n++
	}
	cutoff := at.UTC().Add(-3 * time.Minute).Format(time.RFC3339Nano)
	rows, err = s.db.Query(`SELECT task.id FROM task JOIN runtime ON runtime.id=task.runtime_id WHERE task.status='running' AND runtime.last_heartbeat < ?`, cutoff)
	if err != nil {
		return n, err
	}
	ids, err = scanStrings(rows)
	if err != nil {
		return n, err
	}
	for _, id := range ids {
		if err := s.FinishTask(id, "failed", "runtime_offline"); err != nil {
			return n, err
		}
		n++
	}
	return n, nil
}

// AddTaskMessage records a streamed progress line.
func (s *Store) AddTaskMessage(taskID, kind, body string) error {
	_, err := s.db.Exec(`INSERT INTO task_message(id,task_id,kind,body,created_at) VALUES(?,?,?,?,?)`, newID(), taskID, kind, body, now())
	return err
}

// AddUsage records token counts for one model.
func (s *Store) AddUsage(taskID, model string, input, output, cacheRead, cacheWrite int) error {
	_, err := s.db.Exec(`INSERT INTO usage_(id,task_id,model,input_tokens,output_tokens,cache_read,cache_write,created_at) VALUES(?,?,?,?,?,?,?,?)`,
		newID(), taskID, model, input, output, cacheRead, cacheWrite, now())
	return err
}

// TaskByToken resolves the agent-side credential.
func (s *Store) TaskByToken(token string) (Task, error) {
	var id string
	err := s.db.QueryRow(`SELECT id FROM task WHERE token=?`, token).Scan(&id)
	if err != nil {
		return Task{}, err
	}
	return s.getTask(id)
}