Agent-Squad/adapters/sqlite/platform_test.go

package platform

import (
	"context"
	"crypto/hmac"
	"crypto/sha256"
	"encoding/hex"
	"os"
	"path/filepath"
	"strings"
	"testing"
	"time"

	agentsquad "agent-squad/squad"
)

func TestOrchestrationFlow(t *testing.T) {
	s := openTest(t)
	ws, err := s.CreateWorkspace("lab")
	if err != nil {
		t.Fatal(err)
	}
	leader, err := s.CreateAgent(Agent{WorkspaceID: ws.ID, Name: "leader", MaxConcurrency: 2})
	if err != nil {
		t.Fatal(err)
	}
	member, err := s.CreateAgent(Agent{WorkspaceID: ws.ID, Name: "writer", MaxConcurrency: 1})
	if err != nil {
		t.Fatal(err)
	}
	human, err := s.CreateMember(ws.ID, "ada")
	if err != nil {
		t.Fatal(err)
	}
	copied, err := s.CopyAgent(leader.ID)
	if err != nil || copied.CustomEnv != "" || copied.RuntimeID != "" {
		t.Fatalf("copy leaked config: %+v %v", copied, err)
	}
	squad, err := s.CreateSquad(Squad{
		WorkspaceID: ws.ID, Name: "docs", LeaderAgentID: leader.ID, Instructions: "ship the note",
	}, []SquadMember{{Kind: "agent", ID: member.ID, RoleDescription: "writes"}, {Kind: "human", ID: human.ID, RoleDescription: "reviews"}})
	if err != nil {
		t.Fatal(err)
	}
	brief, err := s.LeaderBrief(squad.ID)
	if err != nil || !strings.Contains(brief, "Squad Operating Protocol") || !strings.Contains(brief, "mention://agent/") || !strings.Contains(brief, "ship the note") {
		t.Fatalf("brief: %s %v", brief, err)
	}
	issue, err := s.CreateIssue(Issue{WorkspaceID: ws.ID, Title: "write the note", Body: "three facts", Status: "todo"})
	if err != nil {
		t.Fatal(err)
	}
	if err := s.AssignIssue(issue.ID, "squad", squad.ID, false); err != nil {
		t.Fatal(err)
	}
	if _, ok, err := claimAny(t, s, ws.ID); err != nil || ok {
		t.Fatal("start=false must not enqueue")
	}
	if err := s.AssignIssue(issue.ID, "squad", squad.ID, true); err != nil {
		t.Fatal(err)
	}
	rts, err := s.RegisterProbed(ws.ID)
	if err != nil || rts[0].Protocol != "agent-native" {
		t.Fatalf("runtimes %+v %v", rts, err)
	}
	leader.RuntimeID = rts[0].ID
	if err := s.UpdateAgent(leader); err != nil {
		t.Fatal(err)
	}
	member.RuntimeID = rts[0].ID
	if err := s.UpdateAgent(member); err != nil {
		t.Fatal(err)
	}
	task, ok, err := s.Claim(rts[0].ID, time.Now())
	if err != nil || !ok || !task.IsLeader || task.AgentID != leader.ID {
		t.Fatalf("claim %+v %v %v", task, ok, err)
	}
	if err := s.TouchStatus(issue.ID, leader.ID, false); err != nil {
		t.Fatal(err)
	}
	if err := s.RecordActivity(squad.ID, issue.ID, leader.ID, "action", "delegate"); err != nil {
		t.Fatal(err)
	}
	stages, err := s.PlanStages(issue.ID, []string{"draft"})
	if err != nil || len(stages) != 1 {
		t.Fatal(err)
	}
	if _, err := s.DelegateComment(issue.ID, leader.ID, member.ID, "please draft"); err != nil {
		t.Fatal(err)
	}
	memberTask, ok, err := s.Claim(rts[0].ID, time.Now())
	if err != nil || !ok || memberTask.AgentID != member.ID {
		t.Fatalf("member claim %+v %v %v", memberTask, ok, err)
	}
	if err := s.AgentCLI(memberTask.Token, []string{"comment", "draft ready"}); err != nil {
		t.Fatal(err)
	}
	if err := s.FinishTask(task.ID, "completed", ""); err != nil {
		t.Fatal(err)
	}
	if err := s.CloseIssue(stages[0].ID); err != nil {
		t.Fatal(err)
	}
	woke, ok, err := s.Claim(rts[0].ID, time.Now())
	if err != nil || !ok || woke.AgentID != leader.ID {
		t.Fatalf("child_done wake %+v %v %v", woke, ok, err)
	}
	if err := s.TouchStatus(issue.ID, leader.ID, true); err != nil {
		t.Fatal(err)
	}
	got, err := s.GetIssue(issue.ID)
	if err != nil || got.Status != "in_review" {
		t.Fatalf("status %+v %v", got, err)
	}
	work, err := s.PrepareWorkdir(t.TempDir(), task)
	if err != nil {
		t.Fatal(err)
	}
	text, err := s.RunInProcess(context.Background(), woke, work, stubExecutor{result: &agentsquad.ExecResult{
		Success: true, Logs: "done",
		Patches: map[string][]byte{"output.md": []byte("done")},
	}})
	if err != nil || text != "done" {
		t.Fatalf("in-process %q %v", text, err)
	}

	script := filepath.Join(t.TempDir(), "fake-cli")
	if err := os.WriteFile(script, []byte("#!/bin/sh\necho ok\n"), 0o755); err != nil {
		t.Fatal(err)
	}
	fakeTask, err := s.enqueue(issue, leader.ID, false, squad.ID)
	if err != nil {
		t.Fatal(err)
	}
	if err := s.RunCommand(context.Background(), fakeTask, script, work); err != nil {
		t.Fatal(err)
	}
	finished, err := s.getTask(fakeTask.ID)
	if err != nil || finished.Status != "completed" {
		t.Fatalf("fake cli %+v %v", finished, err)
	}
}

func TestAutopilotSkillAndQueueEdges(t *testing.T) {
	s := openTest(t)
	ws, _ := s.CreateWorkspace("lab")
	agent, _ := s.CreateAgent(Agent{WorkspaceID: ws.ID, Name: "bot", MaxConcurrency: 1})
	dir := t.TempDir()
	if err := os.MkdirAll(filepath.Join(dir, "review"), 0o755); err != nil {
		t.Fatal(err)
	}
	if err := os.WriteFile(filepath.Join(dir, "review", "SKILL.md"), []byte("description: careful\nbody"), 0o644); err != nil {
		t.Fatal(err)
	}
	sk, err := s.ImportSkill(ws.ID, filepath.Join(dir, "review"), "fail")
	if err != nil {
		t.Fatal(err)
	}
	if _, err := s.ImportSkill(ws.ID, filepath.Join(dir, "review"), "fail"); err == nil {
		t.Fatal("conflict fail")
	}
	if _, err := s.ImportSkill(ws.ID, filepath.Join(dir, "review"), "skip"); err != nil {
		t.Fatal(err)
	}
	if err := s.BindSkill(agent.ID, sk.ID, true); err != nil {
		t.Fatal(err)
	}
	ap, err := s.CreateAutopilot(Autopilot{
		WorkspaceID: ws.ID, Title: "daily", Runbook: "check", AssigneeKind: "agent", AssigneeID: agent.ID, ExecutionMode: "create_issue",
	})
	if err != nil {
		t.Fatal(err)
	}
	if err := s.AddSchedule(ap.ID, "* * * * *", "UTC"); err != nil {
		t.Fatal(err)
	}
	n, err := s.FireSchedules(time.Now())
	if err != nil || n != 1 {
		t.Fatalf("cron %d %v", n, err)
	}
	token, err := s.AddWebhook(ap.ID, "secret")
	if err != nil {
		t.Fatal(err)
	}
	body := []byte(`{"event":"push"}`)
	mac := sign(body, "secret")
	if _, err := s.AcceptWebhook(token, body, map[string]string{"X-Hub-Signature-256": mac, "Idempotency-Key": "k1"}); err != nil {
		t.Fatal(err)
	}
	got, err := s.AcceptWebhook(token, body, map[string]string{"X-Hub-Signature-256": mac, "Idempotency-Key": "k1"})
	if err != nil || got != "duplicate" {
		t.Fatalf("dedupe %q %v", got, err)
	}
	issue, _ := s.CreateIssue(Issue{WorkspaceID: ws.ID, Title: "later", Status: "todo"})
	if _, err := s.DeferTask(issue, agent.ID, time.Now().Add(-time.Minute)); err != nil {
		t.Fatal(err)
	}
	moved, err := s.PromoteDeferred(time.Now())
	if err != nil || moved < 1 {
		t.Fatalf("promote %d %v", moved, err)
	}
	if _, err := s.AddComment(Comment{IssueID: issue.ID, AuthorKind: "human", AuthorID: "h", Body: "/note secret"}); err != nil {
		t.Fatal(err)
	}
	chat, err := s.CreateChat(ChatSession{WorkspaceID: ws.ID, OwnerID: "h"})
	if err != nil {
		t.Fatal(err)
	}
	if err := s.AddChatMessage(chat.ID, "user", "hi"); err != nil {
		t.Fatal(err)
	}
	proj, err := s.CreateProject(Project{WorkspaceID: ws.ID, Name: "repo"})
	if err != nil {
		t.Fatal(err)
	}
	if err := s.AddResource(proj.ID, "local_directory", "/tmp/src", "worktree"); err != nil {
		t.Fatal(err)
	}
	if err := s.ArchiveAgent(agent.ID); err != nil {
		t.Fatal(err)
	}
}

func openTest(t *testing.T) *Store {
	t.Helper()
	s, err := Open(filepath.Join(t.TempDir(), "orch.db"))
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = s.Close() })
	return s
}

func TestR2JournalModeWAL(t *testing.T) {
	s := openTest(t)
	var mode string
	if err := s.db.QueryRow(`PRAGMA journal_mode`).Scan(&mode); err != nil {
		t.Fatal(err)
	}
	if mode != "wal" {
		t.Fatalf("journal_mode %q", mode)
	}
}

func claimAny(t *testing.T, s *Store, workspaceID string) (Task, bool, error) {
	t.Helper()
	rt, err := s.RegisterRuntime(Runtime{WorkspaceID: workspaceID, Name: "tmp", Command: "in-process", Protocol: "agent-native"})
	if err != nil {
		return Task{}, false, err
	}
	return s.Claim(rt.ID, time.Now())
}

func sign(body []byte, secret string) string {
	mac := hmac.New(sha256.New, []byte(secret))
	_, _ = mac.Write(body)
	return "sha256=" + hex.EncodeToString(mac.Sum(nil))
}