core/server/store/postgres/message.go

package postgres

import (
	"context"
	"database/sql"
	"encoding/json"
	"time"

	"git.gobha.me/xcaliber/chat-switchboard/models"
	"git.gobha.me/xcaliber/chat-switchboard/store"
)

type MessageStore struct{}

func NewMessageStore() *MessageStore { return &MessageStore{} }

func (s *MessageStore) Create(ctx context.Context, m *models.Message) error {
	toolCallsJSON := ToJSON(m.ToolCalls)
	metadataJSON := ToJSON(m.Metadata)
	return DB.QueryRowContext(ctx, `
		INSERT INTO messages (channel_id, role, content, model, tokens_used, tool_calls,
			metadata, parent_id, sibling_index, participant_type, participant_id)
		VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11)
		RETURNING id, created_at`,
		m.ChannelID, m.Role, m.Content, m.Model, m.TokensUsed,
		toolCallsJSON, metadataJSON,
		models.NullString(m.ParentID), m.SiblingIndex,
		m.ParticipantType, m.ParticipantID,
	).Scan(&m.ID, &m.CreatedAt)
}

func (s *MessageStore) GetByID(ctx context.Context, id string) (*models.Message, error) {
	var m models.Message
	var parentID sql.NullString
	var toolCallsJSON, metadataJSON []byte
	var deletedAt sql.NullTime
	err := DB.QueryRowContext(ctx, `
		SELECT id, channel_id, role, content, model, tokens_used, tool_calls, metadata,
			parent_id, sibling_index, participant_type, participant_id, deleted_at, created_at
		FROM messages WHERE id = $1`, id).Scan(
		&m.ID, &m.ChannelID, &m.Role, &m.Content, &m.Model, &m.TokensUsed,
		&toolCallsJSON, &metadataJSON,
		&parentID, &m.SiblingIndex, &m.ParticipantType, &m.ParticipantID,
		&deletedAt, &m.CreatedAt,
	)
	if err != nil {
		return nil, err
	}
	m.ParentID = NullableStringPtr(parentID)
	json.Unmarshal(toolCallsJSON, &m.ToolCalls)
	json.Unmarshal(metadataJSON, &m.Metadata)
	if deletedAt.Valid {
		m.DeletedAt = &deletedAt.Time
	}
	return &m, nil
}

func (s *MessageStore) Update(ctx context.Context, id string, fields map[string]interface{}) error {
	b := NewUpdate("messages")
	for k, v := range fields {
		if k == "tool_calls" || k == "metadata" {
			b.SetJSON(k, v)
		} else {
			b.Set(k, v)
		}
	}
	if !b.HasSets() {
		return nil
	}
	b.Where("id", id)
	_, err := b.Exec(DB)
	return err
}

func (s *MessageStore) Delete(ctx context.Context, id string) error {
	now := time.Now()
	_, err := DB.ExecContext(ctx, "UPDATE messages SET deleted_at = $1 WHERE id = $2", now, id)
	return err
}

func (s *MessageStore) ListForChannel(ctx context.Context, channelID string, opts store.ListOptions) ([]models.Message, error) {
	b := NewSelect(
		"id, channel_id, role, content, model, tokens_used, tool_calls, metadata, parent_id, sibling_index, participant_type, participant_id, deleted_at, created_at",
		"messages",
	).Where("channel_id = ?", channelID).WhereRaw("deleted_at IS NULL")
	if opts.Sort == "" {
		b.OrderBy("created_at", "ASC")
	}
	b.Paginate(opts)

	q, args := b.Build()
	rows, err := DB.QueryContext(ctx, q, args...)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	return scanMessages(rows)
}

func (s *MessageStore) GetChildren(ctx context.Context, parentID string) ([]models.Message, error) {
	rows, err := DB.QueryContext(ctx, `
		SELECT id, channel_id, role, content, model, tokens_used, tool_calls, metadata,
			parent_id, sibling_index, participant_type, participant_id, deleted_at, created_at
		FROM messages WHERE parent_id = $1 AND deleted_at IS NULL
		ORDER BY sibling_index`, parentID)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	return scanMessages(rows)
}

func (s *MessageStore) GetSiblings(ctx context.Context, messageID string) ([]models.Message, error) {
	rows, err := DB.QueryContext(ctx, `
		SELECT id, channel_id, role, content, model, tokens_used, tool_calls, metadata,
			parent_id, sibling_index, participant_type, participant_id, deleted_at, created_at
		FROM messages
		WHERE parent_id = (SELECT parent_id FROM messages WHERE id = $1)
			AND deleted_at IS NULL
		ORDER BY sibling_index`, messageID)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	return scanMessages(rows)
}

func (s *MessageStore) GetPathToRoot(ctx context.Context, messageID string) ([]models.Message, error) {
	rows, err := DB.QueryContext(ctx, `
		WITH RECURSIVE path AS (
			SELECT id, channel_id, role, content, model, tokens_used, tool_calls, metadata,
				parent_id, sibling_index, participant_type, participant_id, deleted_at, created_at
			FROM messages WHERE id = $1
			UNION ALL
			SELECT m.id, m.channel_id, m.role, m.content, m.model, m.tokens_used, m.tool_calls, m.metadata,
				m.parent_id, m.sibling_index, m.participant_type, m.participant_id, m.deleted_at, m.created_at
			FROM messages m JOIN path p ON m.id = p.parent_id
		)
		SELECT * FROM path ORDER BY created_at ASC`, messageID)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	return scanMessages(rows)
}

func (s *MessageStore) GetNextSiblingIndex(ctx context.Context, parentID string) (int, error) {
	var maxIdx sql.NullInt64
	err := DB.QueryRowContext(ctx,
		"SELECT MAX(sibling_index) FROM messages WHERE parent_id = $1 AND deleted_at IS NULL",
		parentID).Scan(&maxIdx)
	if err != nil || !maxIdx.Valid {
		return 0, err
	}
	return int(maxIdx.Int64) + 1, nil
}

func (s *MessageStore) CountForChannel(ctx context.Context, channelID string) (int, error) {
	var count int
	err := DB.QueryRowContext(ctx,
		"SELECT COUNT(*) FROM messages WHERE channel_id = $1 AND deleted_at IS NULL",
		channelID).Scan(&count)
	return count, err
}

func scanMessages(rows *sql.Rows) ([]models.Message, error) {
	var result []models.Message
	for rows.Next() {
		var m models.Message
		var parentID sql.NullString
		var toolCallsJSON, metadataJSON []byte
		var deletedAt sql.NullTime
		err := rows.Scan(
			&m.ID, &m.ChannelID, &m.Role, &m.Content, &m.Model, &m.TokensUsed,
			&toolCallsJSON, &metadataJSON,
			&parentID, &m.SiblingIndex, &m.ParticipantType, &m.ParticipantID,
			&deletedAt, &m.CreatedAt,
		)
		if err != nil {
			return nil, err
		}
		m.ParentID = NullableStringPtr(parentID)
		json.Unmarshal(toolCallsJSON, &m.ToolCalls)
		json.Unmarshal(metadataJSON, &m.Metadata)
		if deletedAt.Valid {
			m.DeletedAt = &deletedAt.Time
		}
		result = append(result, m)
	}
	return result, rows.Err()
}