agent/store.go

package agent

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"path/filepath"
	"sort"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/hanzoai/orm"
	ormdb "github.com/hanzoai/orm/db"
	luxlog "github.com/luxfi/log"
)

// Conversation is one persisted chat thread. Org is the owning org — physical
// isolation (one SQLite file per org) already scopes it; Org is stored for clarity
// and defense-in-depth. User is the member who opened it: a thread is one
// person's inside a shared org, so it lists, opens and takes turns for them
// alone. A thread recorded before users were kept has none, and nothing names
// who opened it, so it is nobody's: no member lists, opens or writes it, and a
// SuperAdmin alone may read it. Pinned lists it first; Archived takes it out of
// the list, into the archived one. UpdatedAt is when it was last spoken in: a
// rename, a pin or an archive does not move it.
type Conversation struct {
	orm.Model[Conversation]
	Org      string `json:"org"`
	User     string `json:"user"`
	Title    string `json:"title"`
	Pinned   bool   `json:"pinned"`
	Archived bool   `json:"archived"`
}

// Message is one persisted turn. ConversationId is deliberately spelled with a
// lowercase-d so orm's PascalCase→camelCase filter (ToJSONFieldName lowercases
// only the first rune) maps Filter("ConversationId=") onto the stored
// "conversationId" JSON key. ToolCalls is the marshaled model tool_calls (nil for
// a plain user/assistant turn). Producer names the model whose completion this
// server received and stored as the turn — set only by the round; a turn written
// through the record endpoint carries none, whatever role it claims.
type Message struct {
	orm.Model[Message]
	ConversationId string          `json:"conversationId"`
	Org            string          `json:"org"`
	Role           string          `json:"role"`
	Content        string          `json:"content"`
	Producer       string          `json:"model,omitempty"`
	ToolCalls      json.RawMessage `json:"toolCalls,omitempty"`
}

func init() {
	orm.Register[Conversation]("agent-conversation")
	orm.Register[Message]("agent-message")
}

// idSeq disambiguates ids minted within the same nanosecond so the zero-padded
// id string sorts in creation order.
var idSeq atomic.Uint64

// newID mints a lexically-sortable unique id: zero-padded UnixNano + a rolling
// counter. Ordering messages by id is therefore chronological without a separate
// sequence column.
func newID() string {
	return fmt.Sprintf("%019d-%06d", time.Now().UnixNano(), idSeq.Add(1)%1000000)
}

// store is the lazily-opened, cached set of per-org orm.DBs. Each org's SQLite
// file is opened (and its schema auto-migrated) exactly once, at
// {dataDir}/orgs/{slug}/agent.db. Isolation is PHYSICAL: a distinct org resolves
// to a distinct file, so a query in one can never reach another's rows.
type store struct {
	dataDir string
	mu      sync.Mutex
	byOrg   map[string]orm.DB
	// keys is the share index, {dataDir}/shares.db (share.go), opened on first use.
	keys orm.DB
}

func newStore(dataDir string) *store {
	return &store{dataDir: dataDir, byOrg: map[string]orm.DB{}}
}

// dbFor returns the org's DB, opening + migrating it on first use.
func (s *store) dbFor(org string) (orm.DB, error) {
	slug, err := orgSlug(org)
	if err != nil {
		return nil, err
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	if db, ok := s.byOrg[slug]; ok {
		return db, nil
	}
	path := filepath.Join(s.dataDir, "orgs", slug, "agent.db")
	db, err := orm.OpenSQLite(&ormdb.SQLiteDBConfig{
		Path:   path,
		Config: ormdb.SQLiteConfig{BusyTimeout: 5000, JournalMode: "WAL"},
	})
	if err != nil {
		return nil, fmt.Errorf("agent: open org db: %w", err)
	}
	s.byOrg[slug] = db
	if n, err := memberless(db); err == nil && n > 0 {
		luxlog.Default().Info("agent: conversations recorded with no member open for nobody but a SuperAdmin", "org", slug, "count", n)
	}
	return db, nil
}

// memberless counts an org's conversations recorded with no member. A row
// written before users were kept has no `user` key at all, which a JSON filter
// on "" does not match, so the rows are read and counted here.
func memberless(db orm.DB) (int, error) {
	all, err := orm.TypedQuery[Conversation](db).GetAll(context.Background())
	if err != nil {
		return 0, err
	}
	n := 0
	for _, cv := range all {
		if nobodys(cv) {
			n++
		}
	}
	return n, nil
}

// closeAll closes every open per-org DB. Idempotent; returns the first error.
func (s *store) closeAll() error {
	s.mu.Lock()
	defer s.mu.Unlock()
	var first error
	for k, db := range s.byOrg {
		if err := db.Close(); err != nil && first == nil {
			first = err
		}
		delete(s.byOrg, k)
	}
	if s.keys != nil {
		if err := s.keys.Close(); err != nil && first == nil {
			first = err
		}
		s.keys = nil
	}
	return first
}

// orgSlug names the directory an org's history lives in: one org name, one
// directory, and never two. A lowercase name is its own slug, which is every
// org IAM holds and the name its files already have. Any other name is its
// lowercase form plus "~" and a SHA-256 prefix of its exact bytes: no lowercase
// name carries a "~", and two names that differ only in case differ in the hash,
// so neither can open the other's file. (Regression: the slug was the name
// lowercased, so `Acme` and `acme` shared one file.) A value that could leave
// the data dir, or carries whitespace a caller would trim, is refused. Org is
// the VALIDATED principal value (never a client-supplied field), but this fails
// closed on anything unsafe.
func orgSlug(org string) (string, error) {
	if org == "" {
		return "", fmt.Errorf("agent: empty org")
	}
	lower := strings.ToLower(org)
	for _, r := range lower {
		switch {
		case r >= 'a' && r <= 'z', r >= '0' && r <= '9', r == '-', r == '_', r == '.':
		default:
			return "", fmt.Errorf("agent: unsafe org %q", org)
		}
	}
	if lower == "." || lower == ".." || strings.Contains(lower, "..") {
		return "", fmt.Errorf("agent: unsafe org %q", org)
	}
	if lower == org {
		return org, nil
	}
	sum := sha256.Sum256([]byte(org))
	return lower + "~" + hex.EncodeToString(sum[:8]), nil
}

// loadOrCreateConversation returns the member's conversation for id (in the
// org's DB), and opens a fresh one of theirs when id is empty, not found, or not
// theirs. An id that is another member's, nobody's or another org's is answered
// exactly as one that does not exist, so a turn never lands in a thread that is
// not the caller's and the answer says nothing about whether the id was there.
func (s *store) loadOrCreateConversation(ctx context.Context, org, user, id, title string) (*Conversation, error) {
	if strings.TrimSpace(user) == "" {
		return nil, errNoMember
	}
	db, err := s.dbFor(org)
	if err != nil {
		return nil, err
	}
	if id = strings.TrimSpace(id); id != "" {
		conv, gerr := orm.Get[Conversation](db, id)
		if gerr == nil && conv.Org == org && owns(conv, user) {
			return conv, nil
		}
		if gerr != nil && !errors.Is(gerr, orm.ErrNotFound) {
			return nil, gerr
		}
	}
	conv := orm.New[Conversation](db)
	conv.SetId(newID())
	conv.Org = org
	conv.User = strings.TrimSpace(user)
	conv.Title = clampTitle(title)
	if err := conv.CreateCtx(ctx); err != nil {
		return nil, err
	}
	return conv, nil
}

// appendMessage persists one turn in a conversation. model is the model whose
// completion the turn is, or "" for a turn a caller wrote.
func (s *store) appendMessage(ctx context.Context, org, convID, role, content, model string, toolCalls json.RawMessage) (*Message, error) {
	db, err := s.dbFor(org)
	if err != nil {
		return nil, err
	}
	m := orm.New[Message](db)
	m.SetId(newID())
	m.ConversationId = convID
	m.Org = org
	m.Role = role
	m.Content = content
	m.Producer = model
	m.ToolCalls = toolCalls
	if err := m.CreateCtx(ctx); err != nil {
		return nil, err
	}
	return m, nil
}

// listConversations returns the conversations of the org that pass keep, those
// archived or those not: pinned first, then most recently spoken in.
func (s *store) listConversations(ctx context.Context, org string, archived bool, keep func(*Conversation) bool) ([]*Conversation, error) {
	db, err := s.dbFor(org)
	if err != nil {
		return nil, err
	}
	all, err := orm.TypedQuery[Conversation](db).GetAll(ctx)
	if err != nil {
		return nil, err
	}
	items := all[:0]
	for _, cv := range all {
		if cv.Org == org && keep(cv) && cv.Archived == archived {
			items = append(items, cv)
		}
	}
	sort.SliceStable(items, func(i, j int) bool {
		if items[i].Pinned != items[j].Pinned {
			return items[i].Pinned
		}
		return items[i].UpdatedAt.After(items[j].UpdatedAt)
	})
	return items, nil
}

// change is what a member may set on their own conversation; nil leaves it.
type change struct {
	Title    *string `json:"title"`
	Pinned   *bool   `json:"pinned"`
	Archived *bool   `json:"archived"`
}

// errEmptyTitle refuses a rename to nothing.
var errEmptyTitle = errors.New("agent: title must not be empty")

// updateConversation applies a change to a conversation the member owns. The
// row is written as it stands, so UpdatedAt keeps the last time it was spoken in.
func (s *store) updateConversation(ctx context.Context, org, user, convID string, ch change) (*Conversation, error) {
	if ch.Title != nil && strings.TrimSpace(*ch.Title) == "" {
		return nil, errEmptyTitle
	}
	db, conv, err := s.ownedConversation(org, user, convID)
	if err != nil {
		return nil, err
	}
	if ch.Title != nil {
		conv.Title = clampTitle(*ch.Title)
	}
	if ch.Pinned != nil {
		conv.Pinned = *ch.Pinned
	}
	if ch.Archived != nil {
		conv.Archived = *ch.Archived
	}
	if _, err := db.Put(ctx, conv.Key(), conv); err != nil {
		return nil, err
	}
	return conv, nil
}

// deleteConversation removes a conversation the member owns, for good: its links
// first, so none opens once the call returns, then its turns, then the thread.
// A failure part way leaves the thread listed, and deleting it again finishes.
func (s *store) deleteConversation(ctx context.Context, org, user, convID string) error {
	db, conv, err := s.ownedConversation(org, user, convID)
	if err != nil {
		return err
	}
	shares, err := orm.TypedQuery[Share](db).Filter("ConversationId=", conv.Id()).GetAll(ctx)
	if err != nil {
		return err
	}
	for _, sh := range shares {
		if err := s.endShare(ctx, sh); err != nil {
			return err
		}
		if err := sh.DeleteCtx(ctx); err != nil {
			return err
		}
	}
	turns, err := orm.TypedQuery[Message](db).Filter("ConversationId=", conv.Id()).GetAll(ctx)
	if err != nil {
		return err
	}
	for _, m := range turns {
		if err := m.DeleteCtx(ctx); err != nil {
			return err
		}
	}
	return conv.DeleteCtx(ctx)
}

// touch marks a conversation as spoken in now, and reports whether it is still
// there. It reads the row again rather than writing back a copy taken when the
// round began, so a rename, pin or archive made while the model answered
// stands, and a thread deleted meanwhile stays deleted.
func (s *store) touch(ctx context.Context, org, convID string) (bool, error) {
	db, err := s.dbFor(org)
	if err != nil {
		return false, err
	}
	conv, err := orm.Get[Conversation](db, convID)
	if errors.Is(err, orm.ErrNotFound) {
		return false, nil
	}
	if err != nil {
		return false, err
	}
	return true, conv.UpdateCtx(ctx)
}

// conversationMessages returns a conversation the caller reads (reads) and its
// messages in chronological order. Any other is not found.
func (s *store) conversationMessages(ctx context.Context, org, user string, super bool, convID string) (*Conversation, []*Message, error) {
	db, err := s.dbFor(org)
	if err != nil {
		return nil, nil, err
	}
	conv, err := orm.Get[Conversation](db, convID)
	if err != nil {
		return nil, nil, err
	}
	if conv.Org != org || !reads(conv, user, super) {
		return nil, nil, orm.ErrNotFound
	}
	items, err := orm.TypedQuery[Message](db).Filter("ConversationId=", convID).GetAll(ctx)
	if err != nil {
		return nil, nil, err
	}
	sort.SliceStable(items, func(i, j int) bool { return items[i].Id() < items[j].Id() })
	return conv, items, nil
}

// whole reports whether the stored transcript is everything a round's caller
// put in front of the model: no system text of the caller's, no tools of the
// caller's, and the messages sent equal, turn for turn, the user and assistant
// turns the conversation holds.
func (s *store) whole(ctx context.Context, org, user, convID string, body runRequest) (bool, error) {
	if strings.TrimSpace(body.System) != "" || len(body.Tools) > 0 {
		return false, nil
	}
	_, held, err := s.conversationMessages(ctx, org, user, false, convID)
	if errors.Is(err, orm.ErrNotFound) {
		// Deleted while the model answered: there is no transcript to be whole.
		return false, nil
	}
	if err != nil {
		return false, err
	}
	if len(held) != len(body.Messages) {
		return false, nil
	}
	for i, m := range held {
		sent := body.Messages[i]
		role := strings.TrimSpace(sent.Role)
		if role == "" {
			role = "user"
		}
		if (m.Role != "user" && m.Role != "assistant") || m.Role != role || m.Content != sent.Content {
			return false, nil
		}
	}
	return true, nil
}

// owns reports whether a conversation is this member's: it names them. One
// recorded with no member is nobody's, and a caller who names nobody owns
// nothing. (Regression: a thread with no member was every member's, so every
// customer of a shared org listed, read and wrote into the threads recorded
// before users were kept.)
func owns(cv *Conversation, user string) bool {
	u := strings.TrimSpace(user)
	return u != "" && strings.TrimSpace(cv.User) == u
}

// reads reports whether a caller may read a conversation: its owner, or a
// SuperAdmin for one recorded with no member (nobodys), which no member can be
// shown to have opened. Reading is all a SuperAdmin gets: a turn, a change or a
// share still needs the owner.
func reads(cv *Conversation, user string, super bool) bool {
	return owns(cv, user) || super && nobodys(cv)
}

// nobodys reports a conversation recorded with no member.
func nobodys(cv *Conversation) bool { return strings.TrimSpace(cv.User) == "" }

// errNoMember refuses a write by a caller who names nobody: the thread it opened
// would be nobody's.
var errNoMember = errors.New("agent: a member is required")

// clampTitle derives a short, single-line conversation title.
func clampTitle(s string) string {
	s = strings.TrimSpace(strings.ReplaceAll(s, "\n", " "))
	if s == "" {
		return "New conversation"
	}
	const max = 80
	if len(s) > max {
		return strings.TrimSpace(s[:max])
	}
	return s
}