agent-dashboard/internal/server/store.go

package server

import (
	"database/sql"
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"time"

	_ "modernc.org/sqlite"

	"github.com/longbin/agent-notify/internal/remote"
)

const offlineThreshold = 5 * time.Minute

type Store struct {
	db *sql.DB
}

type SessionRow struct {
	SessionKey     string    `json:"session_key"`
	Hostname       string    `json:"hostname"`
	IPs            []string  `json:"ips"`
	Agent          string    `json:"agent"`
	CWD            string    `json:"cwd"`
	Status         string    `json:"status"`
	Event          string    `json:"event,omitempty"`
	ConversationID string    `json:"conversation_id,omitempty"`
	LastUser       string    `json:"last_user,omitempty"`
	LastAgent      string    `json:"last_agent,omitempty"`
	UpdatedAt      time.Time `json:"updated_at"`
	ReceivedAt     time.Time `json:"received_at"`
}

type ListFilters struct {
	Host   string
	CWD    string
	Agent  string
	Status string
}

type MetaResult struct {
	Hosts  []string `json:"hosts"`
	CWDs   []string `json:"cwds"`
	Agents []string `json:"agents"`
}

func Open(path string) (*Store, error) {
	if err := os.MkdirAll(filepath.Dir(path), 0755); err != nil {
		return nil, err
	}
	db, err := sql.Open("sqlite", path)
	if err != nil {
		return nil, err
	}
	s := &Store{db: db}
	if err := s.initSchema(); err != nil {
		db.Close()
		return nil, err
	}
	return s, nil
}

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

func (s *Store) initSchema() error {
	stmts := []string{
		`CREATE TABLE IF NOT EXISTS sessions (
			session_key TEXT PRIMARY KEY,
			hostname TEXT NOT NULL,
			ips TEXT NOT NULL,
			agent TEXT NOT NULL,
			cwd TEXT NOT NULL,
			status TEXT NOT NULL,
			event TEXT NOT NULL DEFAULT '',
			conversation_id TEXT NOT NULL DEFAULT '',
			last_user TEXT NOT NULL DEFAULT '',
			last_agent TEXT NOT NULL DEFAULT '',
			updated_at TEXT NOT NULL,
			received_at TEXT NOT NULL
		)`,
		`CREATE INDEX IF NOT EXISTS idx_sessions_hostname ON sessions(hostname)`,
		`CREATE INDEX IF NOT EXISTS idx_sessions_cwd ON sessions(cwd)`,
		`CREATE INDEX IF NOT EXISTS idx_sessions_agent ON sessions(agent)`,
		`CREATE INDEX IF NOT EXISTS idx_sessions_updated_at ON sessions(updated_at)`,
	}
	for _, stmt := range stmts {
		if _, err := s.db.Exec(stmt); err != nil {
			return err
		}
	}
	return nil
}

func (s *Store) Upsert(report remote.StatusReport) error {
	key := remote.SessionKey(report.Hostname, remote.PrimaryIP(report.IPs), report.CWD, report.Agent)
	ipsJSON, err := json.Marshal(report.IPs)
	if err != nil {
		return err
	}
	receivedAt := time.Now().UTC()
	updatedAt := report.UpdatedAt.UTC()
	if updatedAt.IsZero() {
		updatedAt = receivedAt
	}

	_, err = s.db.Exec(`
		INSERT INTO sessions (
			session_key, hostname, ips, agent, cwd, status, event,
			conversation_id, last_user, last_agent, updated_at, received_at
		) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
		ON CONFLICT(session_key) DO UPDATE SET
			hostname = excluded.hostname,
			ips = excluded.ips,
			agent = excluded.agent,
			cwd = excluded.cwd,
			status = excluded.status,
			event = excluded.event,
			conversation_id = excluded.conversation_id,
			last_user = CASE WHEN excluded.last_user = '' THEN sessions.last_user ELSE excluded.last_user END,
			last_agent = CASE WHEN excluded.last_agent = '' THEN sessions.last_agent ELSE excluded.last_agent END,
			updated_at = excluded.updated_at,
			received_at = excluded.received_at
	`, key, report.Hostname, string(ipsJSON), report.Agent, report.CWD, report.Status,
		report.Event, report.ConversationID, report.LastUser, report.LastAgent,
		updatedAt.Format(time.RFC3339Nano), receivedAt.Format(time.RFC3339Nano))
	return err
}

func (s *Store) List(filters ListFilters) ([]SessionRow, error) {
	query := `SELECT session_key, hostname, ips, agent, cwd, status, event,
		conversation_id, last_user, last_agent, updated_at, received_at
		FROM sessions WHERE 1=1`
	var args []any

	if filters.Host != "" {
		query += ` AND hostname = ?`
		args = append(args, filters.Host)
	}
	if filters.CWD != "" {
		query += ` AND cwd = ?`
		args = append(args, filters.CWD)
	}
	if filters.Agent != "" {
		query += ` AND agent = ?`
		args = append(args, filters.Agent)
	}
	query += ` ORDER BY updated_at DESC`

	rows, err := s.db.Query(query, args...)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var result []SessionRow
	for rows.Next() {
		row, err := scanSessionRow(rows)
		if err != nil {
			return nil, err
		}
		row.Status = displayStatus(row.Status, row.UpdatedAt)
		if filters.Status != "" && row.Status != filters.Status {
			continue
		}
		result = append(result, row)
	}
	return result, rows.Err()
}

func displayStatus(stored string, updatedAt time.Time) string {
	if time.Since(updatedAt) > offlineThreshold {
		return "offline"
	}
	return stored
}

func (s *Store) Meta() (MetaResult, error) {
	var meta MetaResult
	for _, q := range []struct {
		sql  string
		dest *[]string
	}{
		{`SELECT DISTINCT hostname FROM sessions ORDER BY hostname`, &meta.Hosts},
		{`SELECT DISTINCT cwd FROM sessions ORDER BY cwd`, &meta.CWDs},
		{`SELECT DISTINCT agent FROM sessions ORDER BY agent`, &meta.Agents},
	} {
		rows, err := s.db.Query(q.sql)
		if err != nil {
			return MetaResult{}, err
		}
		for rows.Next() {
			var v string
			if err := rows.Scan(&v); err != nil {
				rows.Close()
				return MetaResult{}, err
			}
			*q.dest = append(*q.dest, v)
		}
		if err := rows.Close(); err != nil {
			return MetaResult{}, err
		}
		if err := rows.Err(); err != nil {
			return MetaResult{}, err
		}
	}
	return meta, nil
}

type rowScanner interface {
	Scan(dest ...any) error
}

func scanSessionRow(rows rowScanner) (SessionRow, error) {
	var row SessionRow
	var ipsJSON, updatedAtStr, receivedAtStr string
	if err := rows.Scan(
		&row.SessionKey, &row.Hostname, &ipsJSON, &row.Agent, &row.CWD, &row.Status,
		&row.Event, &row.ConversationID, &row.LastUser, &row.LastAgent,
		&updatedAtStr, &receivedAtStr,
	); err != nil {
		return SessionRow{}, err
	}
	if err := json.Unmarshal([]byte(ipsJSON), &row.IPs); err != nil {
		return SessionRow{}, fmt.Errorf("decode ips: %w", err)
	}
	if row.IPs == nil {
		row.IPs = []string{}
	}
	var err error
	row.UpdatedAt, err = parseTime(updatedAtStr)
	if err != nil {
		return SessionRow{}, fmt.Errorf("updated_at: %w", err)
	}
	row.ReceivedAt, err = parseTime(receivedAtStr)
	if err != nil {
		return SessionRow{}, fmt.Errorf("received_at: %w", err)
	}
	return row, nil
}

func parseTime(s string) (time.Time, error) {
	for _, layout := range []string{time.RFC3339Nano, time.RFC3339} {
		if t, err := time.Parse(layout, s); err == nil {
			return t.UTC(), nil
		}
	}
	return time.Time{}, fmt.Errorf("invalid time %q", s)
}