Agent-Squad/adapters/sqlite/db.go

package platform

import (
	"crypto/rand"
	"database/sql"
	"encoding/hex"
	"fmt"
	"strings"
	"time"

	_ "modernc.org/sqlite"
)

// Store is the orchestration database. The goal engine keeps its file workspace.
type Store struct {
	db *sql.DB
}

func Open(path string) (*Store, error) {
	db, err := sql.Open("sqlite", path+"?_pragma=busy_timeout(5000)&_pragma=journal_mode(WAL)")
	if err != nil {
		return nil, err
	}
	if _, err := db.Exec(`PRAGMA foreign_keys = ON`); err != nil {
		_ = db.Close()
		return nil, err
	}
	s := &Store{db: db}
	if err := s.migrate(); err != nil {
		_ = db.Close()
		return nil, err
	}
	return s, nil
}

func (s *Store) Close() error { return s.db.Close() }

func (s *Store) migrate() error {
	_, err := s.db.Exec(schema)
	return err
}

func newID() string {
	var b [16]byte
	_, _ = rand.Read(b[:])
	return hex.EncodeToString(b[:])
}

func now() string { return time.Now().UTC().Format(time.RFC3339Nano) }

const schema = `
CREATE TABLE IF NOT EXISTS workspace (
  id TEXT PRIMARY KEY, name TEXT NOT NULL, created_at TEXT NOT NULL
);
CREATE TABLE IF NOT EXISTS member (
  id TEXT PRIMARY KEY, workspace_id TEXT NOT NULL, name TEXT NOT NULL,
  kind TEXT NOT NULL, created_at TEXT NOT NULL
);
CREATE TABLE IF NOT EXISTS runtime (
  id TEXT PRIMARY KEY, workspace_id TEXT NOT NULL, name TEXT NOT NULL,
  command TEXT NOT NULL, protocol TEXT NOT NULL, max_concurrency INTEGER NOT NULL,
  last_heartbeat TEXT, created_at TEXT NOT NULL
);
CREATE TABLE IF NOT EXISTS agent (
  id TEXT PRIMARY KEY, workspace_id TEXT NOT NULL, name TEXT NOT NULL,
  runtime_id TEXT, custom_env TEXT, custom_args TEXT, mcp_config TEXT,
  max_concurrency INTEGER NOT NULL, archived INTEGER NOT NULL, created_at TEXT NOT NULL
);
CREATE TABLE IF NOT EXISTS squad (
  id TEXT PRIMARY KEY, workspace_id TEXT NOT NULL, name TEXT NOT NULL,
  leader_agent_id TEXT NOT NULL, instructions TEXT NOT NULL, created_at TEXT NOT NULL
);
CREATE TABLE IF NOT EXISTS squad_member (
  squad_id TEXT NOT NULL, member_kind TEXT NOT NULL, member_id TEXT NOT NULL,
  role_description TEXT NOT NULL
);
CREATE TABLE IF NOT EXISTS issue (
  id TEXT PRIMARY KEY, workspace_id TEXT NOT NULL, title TEXT NOT NULL,
  body TEXT NOT NULL, status TEXT NOT NULL, parent_id TEXT, stage TEXT,
  assignee_kind TEXT, assignee_id TEXT, created_at TEXT NOT NULL
);
CREATE TABLE IF NOT EXISTS comment (
  id TEXT PRIMARY KEY, issue_id TEXT NOT NULL, author_kind TEXT NOT NULL,
  author_id TEXT NOT NULL, body TEXT NOT NULL, parent_id TEXT, created_at TEXT NOT NULL
);
CREATE TABLE IF NOT EXISTS task (
  id TEXT PRIMARY KEY, workspace_id TEXT NOT NULL, issue_id TEXT, agent_id TEXT NOT NULL,
  runtime_id TEXT, status TEXT NOT NULL, reason_code TEXT, attempts INTEGER NOT NULL,
  lease_until TEXT, fire_at TEXT, is_leader INTEGER NOT NULL, squad_id TEXT,
  token TEXT, created_at TEXT NOT NULL, updated_at TEXT NOT NULL
);
CREATE TABLE IF NOT EXISTS task_message (
  id TEXT PRIMARY KEY, task_id TEXT NOT NULL, kind TEXT NOT NULL, body TEXT NOT NULL, created_at TEXT NOT NULL
);
CREATE TABLE IF NOT EXISTS skill (
  id TEXT PRIMARY KEY, workspace_id TEXT NOT NULL, name TEXT NOT NULL,
  description TEXT NOT NULL, source TEXT NOT NULL, created_at TEXT NOT NULL
);
CREATE TABLE IF NOT EXISTS skill_file (
  id TEXT PRIMARY KEY, skill_id TEXT NOT NULL, path TEXT NOT NULL, body TEXT NOT NULL
);
CREATE TABLE IF NOT EXISTS agent_skill (
  agent_id TEXT NOT NULL, skill_id TEXT NOT NULL, enabled INTEGER NOT NULL
);
CREATE TABLE IF NOT EXISTS autopilot (
  id TEXT PRIMARY KEY, workspace_id TEXT NOT NULL, title TEXT NOT NULL,
  runbook TEXT NOT NULL, assignee_kind TEXT NOT NULL, assignee_id TEXT NOT NULL,
  status TEXT NOT NULL, execution_mode TEXT NOT NULL, created_at TEXT NOT NULL
);
CREATE TABLE IF NOT EXISTS trigger_ (
  id TEXT PRIMARY KEY, autopilot_id TEXT NOT NULL, kind TEXT NOT NULL,
  cron TEXT, timezone TEXT, token TEXT, secret TEXT
);
CREATE TABLE IF NOT EXISTS autopilot_run (
  id TEXT PRIMARY KEY, autopilot_id TEXT NOT NULL, status TEXT NOT NULL,
  detail TEXT, created_at TEXT NOT NULL
);
CREATE TABLE IF NOT EXISTS webhook_delivery (
  id TEXT PRIMARY KEY, autopilot_id TEXT NOT NULL, dedupe_key TEXT NOT NULL, created_at TEXT NOT NULL,
  UNIQUE(autopilot_id, dedupe_key)
);
CREATE TABLE IF NOT EXISTS chat_session (
  id TEXT PRIMARY KEY, workspace_id TEXT NOT NULL, owner_id TEXT NOT NULL,
  project_id TEXT, created_at TEXT NOT NULL
);
CREATE TABLE IF NOT EXISTS chat_message (
  id TEXT PRIMARY KEY, session_id TEXT NOT NULL, role TEXT NOT NULL, body TEXT NOT NULL, created_at TEXT NOT NULL
);
CREATE TABLE IF NOT EXISTS project (
  id TEXT PRIMARY KEY, workspace_id TEXT NOT NULL, name TEXT NOT NULL, created_at TEXT NOT NULL
);
CREATE TABLE IF NOT EXISTS project_resource (
  id TEXT PRIMARY KEY, project_id TEXT NOT NULL, kind TEXT NOT NULL,
  locator TEXT NOT NULL, mode TEXT NOT NULL
);
CREATE TABLE IF NOT EXISTS wakeup (
  id TEXT PRIMARY KEY, task_id TEXT NOT NULL, kind TEXT NOT NULL, fire_at TEXT, created_at TEXT NOT NULL
);
CREATE TABLE IF NOT EXISTS usage_ (
  id TEXT PRIMARY KEY, task_id TEXT NOT NULL, model TEXT NOT NULL,
  input_tokens INTEGER NOT NULL, output_tokens INTEGER NOT NULL,
  cache_read INTEGER NOT NULL, cache_write INTEGER NOT NULL, created_at TEXT NOT NULL
);
CREATE TABLE IF NOT EXISTS activity (
  id TEXT PRIMARY KEY, squad_id TEXT NOT NULL, issue_id TEXT NOT NULL,
  actor_id TEXT NOT NULL, action TEXT NOT NULL, detail TEXT NOT NULL, created_at TEXT NOT NULL
);
CREATE UNIQUE INDEX IF NOT EXISTS agent_skill_pk ON agent_skill(agent_id, skill_id);
`

// Workspace is a tenant boundary.
type Workspace struct {
	ID   string
	Name string
}

func (s *Store) CreateWorkspace(name string) (Workspace, error) {
	w := Workspace{ID: newID(), Name: name}
	_, err := s.db.Exec(`INSERT INTO workspace(id,name,created_at) VALUES(?,?,?)`, w.ID, w.Name, now())
	return w, err
}

// Member is a human identity. Agents live in their own table.
type Member struct {
	ID, WorkspaceID, Name, Kind string
}

func (s *Store) CreateMember(workspaceID, name string) (Member, error) {
	m := Member{ID: newID(), WorkspaceID: workspaceID, Name: name, Kind: "human"}
	_, err := s.db.Exec(`INSERT INTO member(id,workspace_id,name,kind,created_at) VALUES(?,?,?,?,?)`,
		m.ID, m.WorkspaceID, m.Name, m.Kind, now())
	return m, err
}

// Agent is a durable executor identity.
type Agent struct {
	ID, WorkspaceID, Name, RuntimeID, CustomEnv, CustomArgs, MCPConfig string
	MaxConcurrency                                                     int
	Archived                                                           bool
}

func (s *Store) CreateAgent(a Agent) (Agent, error) {
	if a.ID == "" {
		a.ID = newID()
	}
	if a.MaxConcurrency <= 0 {
		a.MaxConcurrency = 1
	}
	_, err := s.db.Exec(`INSERT INTO agent(id,workspace_id,name,runtime_id,custom_env,custom_args,mcp_config,max_concurrency,archived,created_at)
		VALUES(?,?,?,?,?,?,?,?,0,?)`,
		a.ID, a.WorkspaceID, a.Name, a.RuntimeID, a.CustomEnv, a.CustomArgs, a.MCPConfig, a.MaxConcurrency, now())
	return a, err
}

func (s *Store) GetAgent(id string) (Agent, error) {
	var a Agent
	var archived int
	err := s.db.QueryRow(`SELECT id,workspace_id,name,IFNULL(runtime_id,''),IFNULL(custom_env,''),IFNULL(custom_args,''),IFNULL(mcp_config,''),max_concurrency,archived FROM agent WHERE id=?`, id).
		Scan(&a.ID, &a.WorkspaceID, &a.Name, &a.RuntimeID, &a.CustomEnv, &a.CustomArgs, &a.MCPConfig, &a.MaxConcurrency, &archived)
	a.Archived = archived == 1
	return a, err
}

// ArchiveAgent cancels unfinished tasks for that agent.
func (s *Store) ArchiveAgent(id string) error {
	tx, err := s.db.Begin()
	if err != nil {
		return err
	}
	defer tx.Rollback()
	if _, err := tx.Exec(`UPDATE agent SET archived=1 WHERE id=?`, id); err != nil {
		return err
	}
	if _, err := tx.Exec(`UPDATE task SET status='cancelled', reason_code='archived', updated_at=? WHERE agent_id=? AND status IN ('deferred','queued','dispatched','waiting_local_directory','running')`, now(), id); err != nil {
		return err
	}
	return tx.Commit()
}

// CopyAgent copies identity only. Env, MCP, and runtime config are not copied.
func (s *Store) CopyAgent(id string) (Agent, error) {
	src, err := s.GetAgent(id)
	if err != nil {
		return Agent{}, err
	}
	return s.CreateAgent(Agent{
		WorkspaceID: src.WorkspaceID, Name: src.Name + " copy", MaxConcurrency: src.MaxConcurrency,
	})
}

func (s *Store) UpdateAgent(a Agent) error {
	_, err := s.db.Exec(`UPDATE agent SET name=?, runtime_id=?, custom_env=?, custom_args=?, mcp_config=?, max_concurrency=? WHERE id=?`,
		a.Name, a.RuntimeID, a.CustomEnv, a.CustomArgs, a.MCPConfig, a.MaxConcurrency, a.ID)
	return err
}

// Runtime is one registered executor, including the in-process agent-native runtime.
type Runtime struct {
	ID, WorkspaceID, Name, Command, Protocol string
	MaxConcurrency                           int
	LastHeartbeat                            string
}

func (s *Store) RegisterRuntime(rt Runtime) (Runtime, error) {
	if rt.ID == "" {
		rt.ID = newID()
	}
	if rt.MaxConcurrency <= 0 {
		rt.MaxConcurrency = 1
	}
	if rt.Protocol == "" {
		rt.Protocol = "custom"
	}
	_, err := s.db.Exec(`INSERT INTO runtime(id,workspace_id,name,command,protocol,max_concurrency,last_heartbeat,created_at) VALUES(?,?,?,?,?,?,?,?)`,
		rt.ID, rt.WorkspaceID, rt.Name, rt.Command, rt.Protocol, rt.MaxConcurrency, now(), now())
	return rt, err
}

func (s *Store) Heartbeat(runtimeID string) error {
	_, err := s.db.Exec(`UPDATE runtime SET last_heartbeat=? WHERE id=?`, now(), runtimeID)
	return err
}

func scanStrings(rows *sql.Rows) ([]string, error) {
	defer rows.Close()
	var out []string
	for rows.Next() {
		var v string
		if err := rows.Scan(&v); err != nil {
			return nil, err
		}
		out = append(out, v)
	}
	return out, rows.Err()
}

func placeholder(n int) string {
	parts := make([]string, n)
	for i := range parts {
		parts[i] = "?"
	}
	return strings.Join(parts, ",")
}

func fmtErr(msg string, err error) error {
	if err == nil {
		return nil
	}
	return fmt.Errorf("%s: %w", msg, err)
}