byteclaw/internal/store/store.go

// Package store persists sessions and messages in SQLite (pure Go driver, no CGO).
package store

import (
	"context"
	"crypto/rand"
	"database/sql"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"time"

	"github.com/crypto/byteclaw/internal/llm"
	_ "modernc.org/sqlite"
)

var ErrNotFound = errors.New("not found")

type Session struct {
	ID    string `json:"id"`
	Title string `json:"title"`
	// Provider/Model override the server-wide active provider for this session (empty = default).
	Provider  string    `json:"provider,omitempty"`
	Model     string    `json:"model,omitempty"`
	CreatedAt time.Time `json:"created_at"`
	UpdatedAt time.Time `json:"updated_at"`
}

type StoredMessage struct {
	llm.Message
	ID        int64     `json:"id"`
	CreatedAt time.Time `json:"created_at"`
}

type Store struct{ db *sql.DB }

const schema = `
PRAGMA journal_mode=WAL;
PRAGMA foreign_keys=ON;
CREATE TABLE IF NOT EXISTS sessions(
  id TEXT PRIMARY KEY, title TEXT NOT NULL DEFAULT '',
  created_at INTEGER NOT NULL, updated_at INTEGER NOT NULL);
CREATE TABLE IF NOT EXISTS messages(
  id INTEGER PRIMARY KEY AUTOINCREMENT,
  session_id TEXT NOT NULL REFERENCES sessions(id) ON DELETE CASCADE,
  role TEXT NOT NULL, content TEXT NOT NULL DEFAULT '',
  tool_calls TEXT NOT NULL DEFAULT '', tool_call_id TEXT NOT NULL DEFAULT '',
  name TEXT NOT NULL DEFAULT '', created_at INTEGER NOT NULL);
CREATE INDEX IF NOT EXISTS idx_messages_session ON messages(session_id, id);
CREATE TABLE IF NOT EXISTS summaries(
  id INTEGER PRIMARY KEY AUTOINCREMENT,
  session_id TEXT NOT NULL REFERENCES sessions(id) ON DELETE CASCADE,
  upto_message_id INTEGER NOT NULL, content TEXT NOT NULL, created_at INTEGER NOT NULL);
CREATE INDEX IF NOT EXISTS idx_summaries_session ON summaries(session_id, id);
CREATE TABLE IF NOT EXISTS settings(key TEXT PRIMARY KEY, value TEXT NOT NULL);
CREATE TABLE IF NOT EXISTS users(
  id INTEGER PRIMARY KEY AUTOINCREMENT,
  username TEXT NOT NULL UNIQUE COLLATE NOCASE,
  password_hash TEXT NOT NULL, created_at INTEGER NOT NULL);
CREATE TABLE IF NOT EXISTS auth_tokens(
  token_hash TEXT PRIMARY KEY,
  user_id INTEGER NOT NULL REFERENCES users(id) ON DELETE CASCADE,
  created_at INTEGER NOT NULL, expires_at INTEGER NOT NULL);
CREATE INDEX IF NOT EXISTS idx_auth_tokens_user ON auth_tokens(user_id);
`

// migrations add columns to databases created by older versions.
var migrations = []struct{ table, column, ddl string }{
	{"sessions", "provider", "ALTER TABLE sessions ADD COLUMN provider TEXT NOT NULL DEFAULT ''"},
	{"sessions", "model", "ALTER TABLE sessions ADD COLUMN model TEXT NOT NULL DEFAULT ''"},
	{"messages", "cancelled", "ALTER TABLE messages ADD COLUMN cancelled INTEGER NOT NULL DEFAULT 0"},
	{"sessions", "user_id", "ALTER TABLE sessions ADD COLUMN user_id INTEGER NOT NULL DEFAULT 0"},
}

func migrate(db *sql.DB) error {
	if _, err := db.Exec(schema); err != nil {
		return err
	}
	for _, m := range migrations {
		var n int
		if err := db.QueryRow(`SELECT COUNT(*) FROM pragma_table_info(?) WHERE name=?`, m.table, m.column).Scan(&n); err != nil {
			return err
		}
		if n == 0 {
			if _, err := db.Exec(m.ddl); err != nil {
				return err
			}
		}
	}
	return nil
}

func Open(path string) (*Store, error) {
	db, err := sql.Open("sqlite", "file:"+path+"?_pragma=busy_timeout(5000)")
	if err != nil {
		return nil, err
	}
	db.SetMaxOpenConns(1) // SQLite: serialize writers, avoids SQLITE_BUSY
	if err := migrate(db); err != nil {
		db.Close()
		return nil, fmt.Errorf("migrate: %w", err)
	}
	return &Store{db: db}, nil
}

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

func newID() string {
	b := make([]byte, 12)
	rand.Read(b)
	return hex.EncodeToString(b)
}

func (s *Store) CreateSession(ctx context.Context, title string) (Session, error) {
	return s.CreateSessionFor(ctx, 0, title)
}

// CreateSessionFor creates a session owned by userID (0 = legacy/admin).
func (s *Store) CreateSessionFor(ctx context.Context, userID int64, title string) (Session, error) {
	now := time.Now().UTC()
	ss := Session{ID: newID(), Title: title, CreatedAt: now, UpdatedAt: now}
	_, err := s.db.ExecContext(ctx, `INSERT INTO sessions(id,title,user_id,created_at,updated_at) VALUES(?,?,?,?,?)`,
		ss.ID, ss.Title, userID, now.UnixMilli(), now.UnixMilli())
	return ss, err
}

func (s *Store) GetSession(ctx context.Context, id string) (Session, error) {
	var ss Session
	var c, u int64
	err := s.db.QueryRowContext(ctx, `SELECT id,title,provider,model,created_at,updated_at FROM sessions WHERE id=?`, id).Scan(&ss.ID, &ss.Title, &ss.Provider, &ss.Model, &c, &u)
	if errors.Is(err, sql.ErrNoRows) {
		return ss, ErrNotFound
	}
	ss.CreatedAt, ss.UpdatedAt = time.UnixMilli(c).UTC(), time.UnixMilli(u).UTC()
	return ss, err
}

func (s *Store) ListSessions(ctx context.Context, limit int) ([]Session, error) {
	return s.ListSessionsFor(ctx, AllUsers, limit)
}

// AllUsers disables per-user scoping (admin / static token).
const AllUsers int64 = -1

// ListSessionsFor lists the sessions of userID (AllUsers = every session).
func (s *Store) ListSessionsFor(ctx context.Context, userID int64, limit int) ([]Session, error) {
	rows, err := s.db.QueryContext(ctx, `SELECT id,title,provider,model,created_at,updated_at FROM sessions
		WHERE (?<0 OR user_id=?) ORDER BY updated_at DESC LIMIT ?`, userID, userID, limit)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	out := []Session{}
	for rows.Next() {
		var ss Session
		var c, u int64
		if err := rows.Scan(&ss.ID, &ss.Title, &ss.Provider, &ss.Model, &c, &u); err != nil {
			return nil, err
		}
		ss.CreatedAt, ss.UpdatedAt = time.UnixMilli(c).UTC(), time.UnixMilli(u).UTC()
		out = append(out, ss)
	}
	return out, rows.Err()
}

func (s *Store) RenameSession(ctx context.Context, id, title string) error {
	res, err := s.db.ExecContext(ctx, `UPDATE sessions SET title=? WHERE id=?`, title, id)
	return affected(res, err)
}

func (s *Store) DeleteSession(ctx context.Context, id string) error {
	res, err := s.db.ExecContext(ctx, `DELETE FROM sessions WHERE id=?`, id)
	return affected(res, err)
}

func affected(res sql.Result, err error) error {
	if err != nil {
		return err
	}
	if n, _ := res.RowsAffected(); n == 0 {
		return ErrNotFound
	}
	return nil
}

func (s *Store) AppendMessage(ctx context.Context, sessionID string, m llm.Message) (int64, error) {
	tc := ""
	if len(m.ToolCalls) > 0 {
		b, _ := json.Marshal(m.ToolCalls)
		tc = string(b)
	}
	now := time.Now().UTC().UnixMilli()
	tx, err := s.db.BeginTx(ctx, nil)
	if err != nil {
		return 0, err
	}
	defer tx.Rollback()
	res, err := tx.ExecContext(ctx, `INSERT INTO messages(session_id,role,content,tool_calls,tool_call_id,name,cancelled,created_at) VALUES(?,?,?,?,?,?,?,?)`,
		sessionID, string(m.Role), m.Content, tc, m.ToolCallID, m.Name, m.Cancelled, now)
	if err != nil {
		return 0, err
	}
	if _, err := tx.ExecContext(ctx, `UPDATE sessions SET updated_at=? WHERE id=?`, now, sessionID); err != nil {
		return 0, err
	}
	id, _ := res.LastInsertId()
	return id, tx.Commit()
}

// Messages returns the last `limit` messages of a session in chronological order (limit<=0: all).
func (s *Store) Messages(ctx context.Context, sessionID string, limit int) ([]StoredMessage, error) {
	return s.MessagesAfter(ctx, sessionID, 0, limit)
}

// MessagesAfter returns the last `limit` messages with id > afterID (limit<=0: all), oldest first.
func (s *Store) MessagesAfter(ctx context.Context, sessionID string, afterID int64, limit int) ([]StoredMessage, error) {
	if limit <= 0 {
		limit = -1
	}
	rows, err := s.db.QueryContext(ctx, `SELECT * FROM (SELECT id,role,content,tool_calls,tool_call_id,name,cancelled,created_at
		FROM messages WHERE session_id=? AND id>? ORDER BY id DESC LIMIT ?) ORDER BY id ASC`, sessionID, afterID, limit)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	out := []StoredMessage{}
	for rows.Next() {
		var m StoredMessage
		var role, tc string
		var c int64
		if err := rows.Scan(&m.ID, &role, &m.Content, &tc, &m.ToolCallID, &m.Name, &m.Cancelled, &c); err != nil {
			return nil, err
		}
		m.Role, m.CreatedAt = llm.Role(role), time.UnixMilli(c).UTC()
		if tc != "" {
			if err := json.Unmarshal([]byte(tc), &m.ToolCalls); err != nil {
				return nil, err
			}
		}
		out = append(out, m)
	}
	return out, rows.Err()
}

func (s *Store) SetSessionProvider(ctx context.Context, id, provider, model string) error {
	res, err := s.db.ExecContext(ctx, `UPDATE sessions SET provider=?, model=? WHERE id=?`, provider, model, id)
	return affected(res, err)
}

// ---- summaries

type Summary struct {
	ID            int64     `json:"id"`
	SessionID     string    `json:"session_id"`
	UptoMessageID int64     `json:"upto_message_id"`
	Content       string    `json:"content"`
	CreatedAt     time.Time `json:"created_at"`
}

func (s *Store) AddSummary(ctx context.Context, sessionID string, upto int64, content string) (Summary, error) {
	now := time.Now().UTC()
	res, err := s.db.ExecContext(ctx, `INSERT INTO summaries(session_id,upto_message_id,content,created_at) VALUES(?,?,?,?)`,
		sessionID, upto, content, now.UnixMilli())
	if err != nil {
		return Summary{}, err
	}
	id, _ := res.LastInsertId()
	return Summary{ID: id, SessionID: sessionID, UptoMessageID: upto, Content: content, CreatedAt: now}, nil
}

// LatestSummary returns the newest summary of a session or ErrNotFound.
func (s *Store) LatestSummary(ctx context.Context, sessionID string) (Summary, error) {
	var sm Summary
	var c int64
	err := s.db.QueryRowContext(ctx, `SELECT id,session_id,upto_message_id,content,created_at FROM summaries
		WHERE session_id=? ORDER BY id DESC LIMIT 1`, sessionID).Scan(&sm.ID, &sm.SessionID, &sm.UptoMessageID, &sm.Content, &c)
	if errors.Is(err, sql.ErrNoRows) {
		return sm, ErrNotFound
	}
	sm.CreatedAt = time.UnixMilli(c).UTC()
	return sm, err
}

// ---- settings (key/value)

func (s *Store) GetSetting(ctx context.Context, key string) (string, error) {
	var v string
	err := s.db.QueryRowContext(ctx, `SELECT value FROM settings WHERE key=?`, key).Scan(&v)
	if errors.Is(err, sql.ErrNoRows) {
		return "", ErrNotFound
	}
	return v, err
}

func (s *Store) SetSetting(ctx context.Context, key, value string) error {
	_, err := s.db.ExecContext(ctx, `INSERT INTO settings(key,value) VALUES(?,?) ON CONFLICT(key) DO UPDATE SET value=excluded.value`, key, value)
	return err
}