teabot/internal/state/state.go

// Package state persists which Gitea events teabot has already handled so a
// restart does not re-trigger work. State is a single JSON file under the
// configured state directory (default ~/.local/state/teabot/state.json).
package state

import (
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"sync"
	"time"
)

// StateFileName is the JSON file holding processed-event bookkeeping.
const StateFileName = "state.json"

// RepoState tracks what has been handled for one repository.
type RepoState struct {
	// ProcessedIssues holds issue indexes already dispatched.
	ProcessedIssues map[int64]bool `json:"processed_issues"`
	// ProcessedPulls holds pull-request indexes already dispatched.
	ProcessedPulls map[int64]bool `json:"processed_pulls"`
	// ProcessedComments holds comment IDs already dispatched.
	ProcessedComments map[int64]bool `json:"processed_comments"`
	// ActedIssues/ActedPulls record which issues/PRs teabot ran a session
	// for, so comment follow-ups only fire on threads the bot engaged with.
	ActedIssues map[int64]bool `json:"acted_issues"`
	ActedPulls  map[int64]bool `json:"acted_pulls"`
	// Seeded is set the first time a repo is polled: existing open issues/PRs
	// and recent comments are recorded as processed WITHOUT dispatching, so a
	// fresh install does not stampede every open item.
	Seeded bool `json:"seeded"`
	// LastPoll is the time of the last completed poll, used to bound `since`
	// queries on subsequent cycles.
	LastPoll time.Time `json:"last_poll"`
}

func newRepoState() *RepoState {
	return &RepoState{
		ProcessedIssues:   map[int64]bool{},
		ProcessedPulls:    map[int64]bool{},
		ProcessedComments: map[int64]bool{},
		ActedIssues:       map[int64]bool{},
		ActedPulls:        map[int64]bool{},
	}
}

// data is the on-disk document.
type data struct {
	Repos map[string]*RepoState `json:"repos"`
}

// Store is a thread-safe, file-backed processed-event tracker.
type Store struct {
	path string
	mu   sync.Mutex
	d    *data
}

// New loads the store from dir, creating an empty one if the file is absent.
func New(dir string) (*Store, error) {
	s := &Store{
		path: filepath.Join(dir, StateFileName),
		d:    &data{Repos: map[string]*RepoState{}},
	}
	raw, err := os.ReadFile(s.path)
	if err != nil {
		if os.IsNotExist(err) {
			return s, nil
		}
		return nil, fmt.Errorf("reading state %s: %w", s.path, err)
	}
	if len(raw) == 0 {
		return s, nil
	}
	if err := json.Unmarshal(raw, s.d); err != nil {
		return nil, fmt.Errorf("parsing state %s: %w", s.path, err)
	}
	if s.d.Repos == nil {
		s.d.Repos = map[string]*RepoState{}
	}
	return s, nil
}

// repo returns the RepoState for repo, creating it if needed. Caller holds mu.
func (s *Store) repo(repo string) *RepoState {
	rs := s.d.Repos[repo]
	if rs == nil {
		rs = newRepoState()
		s.d.Repos[repo] = rs
	}
	// Guard against a partially-populated document loaded from disk.
	if rs.ProcessedIssues == nil {
		rs.ProcessedIssues = map[int64]bool{}
	}
	if rs.ProcessedPulls == nil {
		rs.ProcessedPulls = map[int64]bool{}
	}
	if rs.ProcessedComments == nil {
		rs.ProcessedComments = map[int64]bool{}
	}
	if rs.ActedIssues == nil {
		rs.ActedIssues = map[int64]bool{}
	}
	if rs.ActedPulls == nil {
		rs.ActedPulls = map[int64]bool{}
	}
	return rs
}

// IssueProcessed reports whether an issue index was already handled.
func (s *Store) IssueProcessed(repo string, index int64) bool {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.repo(repo).ProcessedIssues[index]
}

// PullProcessed reports whether a PR index was already handled.
func (s *Store) PullProcessed(repo string, index int64) bool {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.repo(repo).ProcessedPulls[index]
}

// CommentProcessed reports whether a comment ID was already handled.
func (s *Store) CommentProcessed(repo string, id int64) bool {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.repo(repo).ProcessedComments[id]
}

// ActedOnIssue reports whether teabot ran a session for an issue.
func (s *Store) ActedOnIssue(repo string, index int64) bool {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.repo(repo).ActedIssues[index]
}

// ActedOnPull reports whether teabot ran a session for a PR.
func (s *Store) ActedOnPull(repo string, index int64) bool {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.repo(repo).ActedPulls[index]
}

// MarkIssue records an issue index as processed and acted-on.
func (s *Store) MarkIssue(repo string, index int64) {
	s.mu.Lock()
	rs := s.repo(repo)
	rs.ProcessedIssues[index] = true
	rs.ActedIssues[index] = true
	s.mu.Unlock()
}

// MarkPull records a PR index as processed and acted-on.
func (s *Store) MarkPull(repo string, index int64) {
	s.mu.Lock()
	rs := s.repo(repo)
	rs.ProcessedPulls[index] = true
	rs.ActedPulls[index] = true
	s.mu.Unlock()
}

// MarkIssueProcessed records an issue index as processed WITHOUT marking it
// acted-on. Used for events teabot deliberately skips (bot-authored or
// non-allowlisted), so they never re-trigger yet are not treated as engaged
// threads for comment follow-ups.
func (s *Store) MarkIssueProcessed(repo string, index int64) {
	s.mu.Lock()
	s.repo(repo).ProcessedIssues[index] = true
	s.mu.Unlock()
}

// MarkPullProcessed records a PR index as processed WITHOUT marking it acted-on.
func (s *Store) MarkPullProcessed(repo string, index int64) {
	s.mu.Lock()
	s.repo(repo).ProcessedPulls[index] = true
	s.mu.Unlock()
}

// MarkComment records a comment ID as processed.
func (s *Store) MarkComment(repo string, id int64) {
	s.mu.Lock()
	s.repo(repo).ProcessedComments[id] = true
	s.mu.Unlock()
}

// Seeded reports whether a repo has completed its baseline seeding poll.
func (s *Store) Seeded(repo string) bool {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.repo(repo).Seeded
}

// MarkSeeded records that a repo has completed baseline seeding.
func (s *Store) MarkSeeded(repo string) {
	s.mu.Lock()
	s.repo(repo).Seeded = true
	s.mu.Unlock()
}

// LastPoll returns the time of the last completed poll for a repo.
func (s *Store) LastPoll(repo string) time.Time {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.repo(repo).LastPoll
}

// SetLastPoll records the time of the last completed poll for a repo.
func (s *Store) SetLastPoll(repo string, t time.Time) {
	s.mu.Lock()
	s.repo(repo).LastPoll = t
	s.mu.Unlock()
}

// Save atomically writes the state document to disk.
func (s *Store) Save() error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if err := os.MkdirAll(filepath.Dir(s.path), 0o755); err != nil {
		return fmt.Errorf("creating state dir: %w", err)
	}
	raw, err := json.MarshalIndent(s.d, "", "  ")
	if err != nil {
		return err
	}
	tmp := s.path + ".tmp"
	if err := os.WriteFile(tmp, raw, 0o644); err != nil {
		return fmt.Errorf("writing state: %w", err)
	}
	if err := os.Rename(tmp, s.path); err != nil {
		return fmt.Errorf("committing state: %w", err)
	}
	return nil
}