velox-bot/internal/tldr/run.go

// Package tldr implements the /tldr pipeline: find what a user hasn't seen
// in the channel they ran the command in, turn it into a prompt-sized
// transcript, and have a local model summarize it. Messages are fetched on
// demand and never stored; the only thing persisted is a per-user,
// per-channel "read up to" cursor.
package tldr

import (
	"context"
	"errors"
	"fmt"
	"log"
	"net/http"
	"strconv"
	"sync"
	"time"

	"github.com/bwmarrin/discordgo"
)

const (
	defaultWindow = 24 * time.Hour     // first use in a channel, no cursor yet
	maxWindow     = 7 * 24 * time.Hour // a stale cursor is never trusted further back than this
	runTimeout    = 12 * time.Minute   // interaction tokens live 15m
	partialEvery  = 2 * time.Second    // minimum gap between edits while the answer streams in
)

type Discord interface {
	MessageFetcher
	UserChannelPermissions(userID, channelID string, fetchOptions ...discordgo.RequestOption) (int64, error)
}

type CursorStore interface {
	GetCursors(ctx context.Context, userID, guildID int64) (map[int64]int64, error)
	SetCursors(ctx context.Context, userID, guildID int64, cursors map[int64]int64) error
}

type Request struct {
	GuildID   string
	ChannelID string // the channel the command was run in, and the only one summarized
	UserID    string
	BotID     string
	// UserPerms is the caller's permission bitfield in ChannelID, as Discord
	// computes it and sends it with the interaction (this also covers threads).
	UserPerms int64
	// Since, when non-zero, overrides the stored cursors for this run and
	// leaves them untouched ("peek" mode).
	Since time.Duration
	Now   time.Time
}

type Result struct {
	Summary   string
	Messages  int
	Truncated bool
	Empty     bool
}

var (
	// ErrUserCannotRead means the caller lacks View Channel or Read Message
	// History in the channel: summarizing it would hand them messages they
	// cannot read themselves.
	ErrUserCannotRead = errors.New("the caller cannot read this channel's history")
	// ErrBotCannotRead means the bot is missing those permissions here.
	ErrBotCannotRead = errors.New("the bot cannot read this channel's history")
)

var (
	defaultSummarizer Summarizer

	// One model call at a time: the box has no GPU and a single Ollama slot,
	// so parallel runs would just fight each other (and Lavalink) for CPU.
	slot = make(chan struct{}, 1)

	inFlight sync.Map
)

func Init(s Summarizer) { defaultSummarizer = s }

func Default() Summarizer { return defaultSummarizer }

// TryStart reserves the user's single concurrent /tldr. Done releases it.
func TryStart(userID string) bool {
	_, loaded := inFlight.LoadOrStore(userID, struct{}{})
	return !loaded
}

func Done(userID string) { inFlight.Delete(userID) }

const needPerms = discordgo.PermissionViewChannel | discordgo.PermissionReadMessageHistory

func forbidden(err error) bool {
	var re *discordgo.RESTError
	return errors.As(err, &re) && re.Response != nil && re.Response.StatusCode == http.StatusForbidden
}

// Gather collects the messages in the request's channel that the user hasn't
// been summarized yet, as a single batch (nil when there are none).
func Gather(ctx context.Context, d Discord, store CursorStore, req Request) ([]ChannelBatch, error) {
	guildID, err := strconv.ParseInt(req.GuildID, 10, 64)
	if err != nil {
		return nil, fmt.Errorf("bad guild id: %w", err)
	}
	userID, err := strconv.ParseInt(req.UserID, 10, 64)
	if err != nil {
		return nil, fmt.Errorf("bad user id: %w", err)
	}
	channelID, err := strconv.ParseInt(req.ChannelID, 10, 64)
	if err != nil {
		return nil, fmt.Errorf("bad channel id: %w", err)
	}

	if req.UserPerms&needPerms != needPerms {
		return nil, ErrUserCannotRead
	}
	// A lookup error is not proof of missing access (the fetch below reports
	// a real 403), so only a successful lookup that lacks the bits refuses.
	if perms, err := d.UserChannelPermissions(req.BotID, req.ChannelID); err == nil && perms&needPerms != needPerms {
		return nil, ErrBotCannotRead
	}

	var after string
	if req.Since > 0 {
		after = SnowflakeFromTime(req.Now.Add(-req.Since))
	} else {
		cursors, err := store.GetCursors(ctx, userID, guildID)
		if err != nil {
			return nil, fmt.Errorf("loading cursors: %w", err)
		}
		if c, ok := cursors[channelID]; ok {
			after = strconv.FormatInt(c, 10)
			// A cursor from months ago would mean a huge backlog.
			if floor := SnowflakeFromTime(req.Now.Add(-maxWindow)); CompareIDs(after, floor) < 0 {
				after = floor
			}
		} else {
			after = SnowflakeFromTime(req.Now.Add(-defaultWindow))
		}
	}

	msgs, err := FetchAfter(d, req.ChannelID, after, maxMessages)
	if err != nil {
		if forbidden(err) {
			return nil, ErrBotCannotRead
		}
		return nil, fmt.Errorf("fetching messages: %w", err)
	}
	if len(msgs) == 0 {
		return nil, nil
	}
	return []ChannelBatch{{ID: req.ChannelID, Name: "", Messages: msgs}}, nil
}

// Commit advances the user's cursors to what this run covered.
func Commit(ctx context.Context, store CursorStore, req Request, lastProcessed map[string]string) error {
	guildID, err := strconv.ParseInt(req.GuildID, 10, 64)
	if err != nil {
		return err
	}
	userID, err := strconv.ParseInt(req.UserID, 10, 64)
	if err != nil {
		return err
	}

	cursors := make(map[int64]int64, len(lastProcessed))
	for ch, msg := range lastProcessed {
		chID, err1 := strconv.ParseInt(ch, 10, 64)
		msgID, err2 := strconv.ParseInt(msg, 10, 64)
		if err1 != nil || err2 != nil {
			continue
		}
		cursors[chID] = msgID
	}
	return store.SetCursors(ctx, userID, guildID, cursors)
}

// throttle lets the first call through at once and then at most one call per
// interval, so a streaming answer can update a Discord message without
// tripping its rate limit. Calls in between are dropped, which is fine
// because each carries the full text so far and a newer one always follows.
func throttle(every time.Duration, now func() time.Time, f func(string)) func(string) {
	if f == nil {
		return nil
	}
	var mu sync.Mutex
	var last time.Time
	return func(s string) {
		mu.Lock()
		defer mu.Unlock()
		if t := now(); last.IsZero() || t.Sub(last) >= every {
			last = t
			f(s)
		}
	}
}

// Run is the whole command minus Discord's interaction plumbing. status is
// called with short progress text for anything slow enough to look hung, and
// partial (optional) with the overview so far while it is being written.
func Run(ctx context.Context, d Discord, store CursorStore, sum Summarizer, req Request, status func(string), partial func(string)) (*Result, error) {
	if sum == nil {
		return nil, errors.New("the summarizer isn't configured")
	}
	if status == nil {
		status = func(string) {}
	}
	if req.Now.IsZero() {
		req.Now = time.Now()
	}
	peek := req.Since > 0

	batches, err := Gather(ctx, d, store, req)
	if err != nil {
		return nil, err
	}
	if len(batches) == 0 {
		return &Result{Empty: true}, nil
	}

	tr := BuildTranscript(batches)
	if tr.Messages == 0 {
		// Only bot/system/empty messages: still move the cursor past them.
		if !peek {
			if err := Commit(ctx, store, req, tr.LastProcessed); err != nil {
				log.Printf("tldr: committing cursors: %v", err)
			}
		}
		return &Result{Empty: true}, nil
	}

	select {
	case slot <- struct{}{}:
	default:
		status("Another summary is running, you're next in line…")
		select {
		case slot <- struct{}{}:
		case <-ctx.Done():
			return nil, ctx.Err()
		}
	}
	defer func() { <-slot }()

	status(fmt.Sprintf("Summarizing %d messages…", tr.Messages))
	summary, err := Summarize(ctx, sum, tr.Chunks, func(done, total int) {
		if total > 1 {
			status(fmt.Sprintf("Summarizing %d messages… (step %d of %d)", tr.Messages, done+1, total))
		}
	}, throttle(partialEvery, time.Now, partial))
	if err != nil {
		return nil, err
	}

	if !peek {
		if err := Commit(ctx, store, req, tr.LastProcessed); err != nil {
			log.Printf("tldr: committing cursors: %v", err)
		}
	}
	return &Result{Summary: summary, Messages: tr.Messages, Truncated: tr.Truncated}, nil
}

// RunTimeout is the context budget for one /tldr.
func RunTimeout() time.Duration { return runTimeout }