Cyrene_v2/backend/services/gateway/store/chat_store.go

// Package store provides PostgreSQL persistence for chat sessions and messages.
package store

import (
	"context"
	"database/sql"
	"fmt"
	"strings"
	"time"
)

// Session represents a chat session.
type Session struct {
	ID        string    `json:"id"`
	UserID    string    `json:"user_id"`
	Title     string    `json:"title"`
	CreatedAt time.Time `json:"created_at"`
	UpdatedAt time.Time `json:"updated_at"`
}

// Message represents a single chat message.
type Message struct {
	ID        int64     `json:"id"`
	SessionID string    `json:"session_id"`
	Role      string    `json:"role"` // "user" | "assistant"
	Content   string    `json:"content"`
	CreatedAt time.Time `json:"created_at"`
}

// ChatStore persists sessions and messages in PostgreSQL.
type ChatStore struct {
	db *sql.DB
}

// NewChatStore creates a store backed by the given database.
func NewChatStore(db *sql.DB) *ChatStore {
	return &ChatStore{db: db}
}

// CreateTableSQL returns the schema for session/message persistence.
func CreateTableSQL() string {
	return `
CREATE TABLE IF NOT EXISTS chat_sessions (
    id         TEXT PRIMARY KEY,
    user_id    TEXT NOT NULL,
    title      TEXT NOT NULL DEFAULT '',
    created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
    updated_at TIMESTAMPTZ NOT NULL DEFAULT now()
);
CREATE INDEX IF NOT EXISTS idx_chat_sessions_user ON chat_sessions(user_id, updated_at DESC);

CREATE TABLE IF NOT EXISTS chat_messages (
    id         BIGSERIAL PRIMARY KEY,
    session_id TEXT NOT NULL REFERENCES chat_sessions(id) ON DELETE CASCADE,
    user_id    TEXT NOT NULL,
    role       TEXT NOT NULL,
    content    TEXT NOT NULL,
    created_at TIMESTAMPTZ NOT NULL DEFAULT now()
);
CREATE INDEX IF NOT EXISTS idx_chat_messages_session ON chat_messages(session_id, id);
CREATE INDEX IF NOT EXISTS idx_chat_messages_user ON chat_messages(user_id, id);
`
}

// --- Sessions ---

// CreateSession creates a new session for the user.
func (s *ChatStore) CreateSession(ctx context.Context, id, userID, title string) (*Session, error) {
	sess := &Session{ID: id, UserID: userID, Title: title}
	err := s.db.QueryRowContext(ctx, `
		INSERT INTO chat_sessions (id, user_id, title) VALUES ($1, $2, $3)
		RETURNING created_at, updated_at
	`, id, userID, title).Scan(&sess.CreatedAt, &sess.UpdatedAt)
	if err != nil {
		return nil, fmt.Errorf("create session: %w", err)
	}
	return sess, nil
}

// ListSessions returns the user's sessions, most recently active first.
func (s *ChatStore) ListSessions(ctx context.Context, userID string) ([]Session, error) {
	rows, err := s.db.QueryContext(ctx, `
		SELECT id, user_id, title, created_at, updated_at
		FROM chat_sessions WHERE user_id = $1
		ORDER BY updated_at DESC LIMIT 100
	`, userID)
	if err != nil {
		return nil, fmt.Errorf("list sessions: %w", err)
	}
	defer rows.Close()

	sessions := []Session{}
	for rows.Next() {
		var sess Session
		if err := rows.Scan(&sess.ID, &sess.UserID, &sess.Title, &sess.CreatedAt, &sess.UpdatedAt); err != nil {
			return nil, err
		}
		sessions = append(sessions, sess)
	}
	return sessions, rows.Err()
}

// DeleteSession removes a session and its messages (CASCADE).
func (s *ChatStore) DeleteSession(ctx context.Context, id, userID string) error {
	res, err := s.db.ExecContext(ctx, `DELETE FROM chat_sessions WHERE id = $1 AND user_id = $2`, id, userID)
	if err != nil {
		return fmt.Errorf("delete session: %w", err)
	}
	if n, _ := res.RowsAffected(); n == 0 {
		return fmt.Errorf("session not found")
	}
	return nil
}

// --- Messages ---

// AddMessage appends a message, creating the session row if needed and
// bumping updated_at. The first user message becomes the session title.
func (s *ChatStore) AddMessage(ctx context.Context, sessionID, userID, role, content string) error {
	tx, err := s.db.BeginTx(ctx, nil)
	if err != nil {
		return fmt.Errorf("begin tx: %w", err)
	}
	defer tx.Rollback()

	if _, err := tx.ExecContext(ctx, `
		INSERT INTO chat_sessions (id, user_id) VALUES ($1, $2) ON CONFLICT (id) DO NOTHING
	`, sessionID, userID); err != nil {
		return fmt.Errorf("ensure session: %w", err)
	}

	if _, err := tx.ExecContext(ctx, `
		INSERT INTO chat_messages (session_id, user_id, role, content) VALUES ($1, $2, $3, $4)
	`, sessionID, userID, role, content); err != nil {
		return fmt.Errorf("insert message: %w", err)
	}

	title := ""
	if role == "user" {
		title = truncateRunes(strings.TrimSpace(content), 24)
	}
	if _, err := tx.ExecContext(ctx, `
		UPDATE chat_sessions SET
			updated_at = now(),
			title = CASE WHEN title = '' AND $3 <> '' THEN $3 ELSE title END
		WHERE id = $1 AND user_id = $2
	`, sessionID, userID, title); err != nil {
		return fmt.Errorf("touch session: %w", err)
	}

	return tx.Commit()
}

// ListMessages returns the most recent messages of a session (ascending order).
func (s *ChatStore) ListMessages(ctx context.Context, sessionID string, limit int) ([]Message, error) {
	if limit <= 0 || limit > 200 {
		limit = 100
	}
	rows, err := s.db.QueryContext(ctx, `
		SELECT id, session_id, role, content, created_at FROM (
			SELECT id, session_id, role, content, created_at
			FROM chat_messages WHERE session_id = $1
			ORDER BY id DESC LIMIT $2
		) recent ORDER BY id ASC
	`, sessionID, limit)
	if err != nil {
		return nil, fmt.Errorf("list messages: %w", err)
	}
	defer rows.Close()

	messages := []Message{}
	for rows.Next() {
		var m Message
		if err := rows.Scan(&m.ID, &m.SessionID, &m.Role, &m.Content, &m.CreatedAt); err != nil {
			return nil, err
		}
		messages = append(messages, m)
	}
	return messages, rows.Err()
}

// ClearMessages removes all messages of a session.
func (s *ChatStore) ClearMessages(ctx context.Context, sessionID string) error {
	_, err := s.db.ExecContext(ctx, `DELETE FROM chat_messages WHERE session_id = $1`, sessionID)
	return err
}

// SearchMessages does a substring search across the user's messages.
func (s *ChatStore) SearchMessages(ctx context.Context, userID, query string, limit int) ([]Message, error) {
	if limit <= 0 || limit > 100 {
		limit = 20
	}
	rows, err := s.db.QueryContext(ctx, `
		SELECT id, session_id, role, content, created_at
		FROM chat_messages
		WHERE user_id = $1 AND content ILIKE '%' || $2 || '%'
		ORDER BY id DESC LIMIT $3
	`, userID, query, limit)
	if err != nil {
		return nil, fmt.Errorf("search messages: %w", err)
	}
	defer rows.Close()

	messages := []Message{}
	for rows.Next() {
		var m Message
		if err := rows.Scan(&m.ID, &m.SessionID, &m.Role, &m.Content, &m.CreatedAt); err != nil {
			return nil, err
		}
		messages = append(messages, m)
	}
	return messages, rows.Err()
}

// --- Helpers ---

// truncateRunes cuts s to at most n runes, appending an ellipsis when cut.
func truncateRunes(s string, n int) string {
	runes := []rune(s)
	if len(runes) <= n {
		return s
	}
	return string(runes[:n]) + "…"
}