workout-suggester/internal/session/store.go

// Package session keeps one plan and its follow-up conversation per browser,
// in memory.
package session

import (
	"crypto/rand"
	"encoding/hex"
	"sync"
	"time"

	"git.b0b.be/bdeb/workout-suggester/internal/coach"
	"git.b0b.be/bdeb/workout-suggester/internal/llm"
)

// State tracks an assistant reply through its lifecycle.
type State int

const (
	Pending   State = iota // created, nobody is generating it yet
	Streaming              // a stream handler owns it
	Done                   // finished (successfully or with Err set)
)

// Message is a chat turn.
type Message struct {
	ID      string
	Role    string // "user" or "assistant"
	Content string
	State   State
	Err     string
	Hidden  bool // the app's own question, not shown in the UI
}

// Conversation is one plan plus the chat about it.
type Conversation struct {
	mu       sync.Mutex
	plan     coach.Plan
	messages []*Message
	lastUsed time.Time
}

// Store maps session ids to conversations.
type Store struct {
	mu    sync.Mutex
	convs map[string]*Conversation
}

// NewStore returns an empty store.
func NewStore() *Store {
	return &Store{convs: make(map[string]*Conversation)}
}

// NewID returns a random hex id, suitable for sessions and messages.
func NewID() string {
	b := make([]byte, 16)
	rand.Read(b)
	return hex.EncodeToString(b)
}

// Get returns the conversation for a session, or nil if there is none.
func (s *Store) Get(session string) *Conversation {
	s.mu.Lock()
	defer s.mu.Unlock()
	c := s.convs[session]
	if c != nil {
		c.mu.Lock()
		c.lastUsed = time.Now()
		c.mu.Unlock()
	}
	return c
}

// Start replaces the session's conversation with a fresh plan and returns
// the pending first reply.
func (s *Store) Start(session string, p coach.Plan) (*Conversation, Message) {
	c := &Conversation{plan: p, lastUsed: time.Now()}
	q := &Message{ID: NewID(), Role: "user", Content: coach.FirstQuestion, State: Done, Hidden: true}
	r := &Message{ID: NewID(), Role: "assistant", State: Pending}
	c.messages = []*Message{q, r}
	s.mu.Lock()
	s.convs[session] = c
	s.mu.Unlock()
	return c, *r
}

// Prune drops conversations idle for longer than maxIdle.
func (s *Store) Prune(maxIdle time.Duration) {
	cutoff := time.Now().Add(-maxIdle)
	s.mu.Lock()
	defer s.mu.Unlock()
	for id, c := range s.convs {
		c.mu.Lock()
		idle := c.lastUsed.Before(cutoff)
		c.mu.Unlock()
		if idle {
			delete(s.convs, id)
		}
	}
}

// Plan returns the conversation's plan.
func (c *Conversation) Plan() coach.Plan {
	c.mu.Lock()
	defer c.mu.Unlock()
	return c.plan
}

// UpdatePlan changes the plan, e.g. once routes have been looked up.
func (c *Conversation) UpdatePlan(f func(*coach.Plan)) {
	c.mu.Lock()
	defer c.mu.Unlock()
	f(&c.plan)
}

// Messages returns a snapshot of the visible conversation.
func (c *Conversation) Messages() []Message {
	c.mu.Lock()
	defer c.mu.Unlock()
	var out []Message
	for _, m := range c.messages {
		if !m.Hidden {
			out = append(out, *m)
		}
	}
	return out
}

// Ask appends a user message plus a pending assistant reply and returns both.
func (c *Conversation) Ask(text string) (user, reply Message) {
	c.mu.Lock()
	defer c.mu.Unlock()
	u := &Message{ID: NewID(), Role: "user", Content: text, State: Done}
	r := &Message{ID: NewID(), Role: "assistant", State: Pending}
	c.messages = append(c.messages, u, r)
	return *u, *r
}

// maxHistory bounds how many earlier turns go back to the model; small
// models lose the thread (and speed) with long contexts.
const maxHistory = 12

// Claim marks a pending reply as streaming and returns the prompt to send:
// the plan's system prompt plus the finished messages before the reply. ok
// is false if the reply does not exist or someone else already claimed it;
// msg then holds its current state (if it exists).
func (c *Conversation) Claim(id string) (prompt []llm.Message, msg Message, ok bool) {
	c.mu.Lock()
	defer c.mu.Unlock()
	idx := c.index(id)
	if idx < 0 {
		return nil, Message{}, false
	}
	m := c.messages[idx]
	if m.State != Pending {
		return nil, *m, false
	}
	m.State = Streaming

	var history []llm.Message
	for _, h := range c.messages[:idx] {
		if h.State == Done && h.Err == "" && h.Content != "" {
			history = append(history, llm.Message{Role: h.Role, Content: h.Content})
		}
	}
	if len(history) > maxHistory { // keep the opening question and plan
		history = append(history[:2:2], history[len(history)-maxHistory+2:]...)
	}
	// Small models tend to repeat the whole plan template; nudge follow-ups.
	if n := len(history); n > 2 && history[n-1].Role == "user" {
		history[n-1].Content += "\n\n(" + coach.FollowUpHint + ")"
	}
	prompt = []llm.Message{{Role: "system", Content: c.plan.SystemPrompt()}}
	return append(prompt, history...), *m, true
}

// IsFirstReply reports whether id is the plan's opening reply.
func (c *Conversation) IsFirstReply(id string) bool {
	c.mu.Lock()
	defer c.mu.Unlock()
	return len(c.messages) > 1 && c.messages[1].ID == id
}

// Append adds streamed text to a reply.
func (c *Conversation) Append(id, delta string) {
	c.mu.Lock()
	defer c.mu.Unlock()
	if i := c.index(id); i >= 0 {
		c.messages[i].Content += delta
	}
}

// Finish marks a reply as done, recording err if generation failed, and
// returns its final state.
func (c *Conversation) Finish(id string, err error) Message {
	c.mu.Lock()
	defer c.mu.Unlock()
	i := c.index(id)
	if i < 0 {
		return Message{ID: id, Role: "assistant", State: Done}
	}
	m := c.messages[i]
	m.State = Done
	if err != nil {
		m.Err = err.Error()
	}
	return *m
}

func (c *Conversation) index(id string) int {
	for i, m := range c.messages {
		if m.ID == id {
			return i
		}
	}
	return -1
}