agentmon/internal/store/postgres/sessions.go

package postgres

import (
	"context"
	"fmt"
	"strings"
	"time"
)

type SessionRow struct {
	SessionID string     `json:"session_id"`
	StartedAt time.Time  `json:"started_at"`
	EndedAt   *time.Time `json:"ended_at,omitempty"`
	Framework string     `json:"framework"`
	ClientID  string     `json:"client_id,omitempty"`
	Host      string     `json:"host"`
	RunCount  int        `json:"run_count"`
}

type SessionsFilter struct {
	From      *time.Time
	To        *time.Time
	Framework string
	Host      string
	Limit     int
	Cursor    *time.Time // cursor is the last started_at seen
}

func (d *DB) ListSessions(ctx context.Context, f SessionsFilter) ([]SessionRow, *time.Time, error) {
	if f.Limit <= 0 {
		f.Limit = 50
	}
	if f.Limit > 200 {
		f.Limit = 200
	}

	// Build query dynamically using a CTE so cursor compares against
	// the grouped started_at rather than individual event timestamps.
	var innerConditions []string
	var outerConditions []string
	var args []any
	argN := 1

	// Default time range: last 24h
	if f.From == nil {
		t := time.Now().Add(-24 * time.Hour)
		f.From = &t
	}
	innerConditions = append(innerConditions, fmt.Sprintf("ts >= $%d", argN))
	args = append(args, *f.From)
	argN++

	if f.To != nil {
		innerConditions = append(innerConditions, fmt.Sprintf("ts <= $%d", argN))
		args = append(args, *f.To)
		argN++
	}

	if f.Framework != "" {
		innerConditions = append(innerConditions, fmt.Sprintf("source_framework = $%d", argN))
		args = append(args, f.Framework)
		argN++
	}

	// Host filter applies to an aggregate, so it goes in the outer WHERE
	if f.Host != "" {
		outerConditions = append(outerConditions, fmt.Sprintf("host = $%d", argN))
		args = append(args, f.Host)
		argN++
	}

	// Cursor compares against grouped started_at, so it goes in the outer WHERE
	if f.Cursor != nil {
		outerConditions = append(outerConditions, fmt.Sprintf("started_at < $%d", argN))
		args = append(args, *f.Cursor)
		argN++
	}

	innerWhere := strings.Join(innerConditions, " AND ")

	outerWhere := ""
	if len(outerConditions) > 0 {
		outerWhere = "WHERE " + strings.Join(outerConditions, " AND ")
	}

	query := fmt.Sprintf(`
		WITH session_groups AS (
			SELECT
				session_id,
				MIN(ts) as started_at,
				MAX(CASE WHEN type = 'session.end' THEN ts END) as ended_at,
				COALESCE((ARRAY_AGG(source_framework ORDER BY CASE WHEN type = 'session.start' THEN 0 ELSE 1 END, ts) FILTER (WHERE source_framework IS NOT NULL))[1], 'unknown') as framework,
				MAX(client_id) as client_id,
				MAX(payload->'event'->'source'->>'host') as host,
				COUNT(DISTINCT run_id) as run_count
			FROM events
			WHERE session_id IS NOT NULL AND %s
			GROUP BY session_id
		)
		SELECT session_id, started_at, ended_at, framework, client_id, host, run_count
		FROM session_groups
		%s
		ORDER BY started_at DESC
		LIMIT $%d
	`, innerWhere, outerWhere, argN)
	args = append(args, f.Limit+1) // fetch one extra to detect next page

	rows, err := d.sql.QueryContext(ctx, query, args...)
	if err != nil {
		return nil, nil, err
	}
	defer rows.Close()

	var out []SessionRow
	for rows.Next() {
		var r SessionRow
		var clientID *string
		var host *string
		if err := rows.Scan(&r.SessionID, &r.StartedAt, &r.EndedAt, &r.Framework, &clientID, &host, &r.RunCount); err != nil {
			return nil, nil, err
		}
		if clientID != nil {
			r.ClientID = *clientID
		}
		if host != nil {
			r.Host = *host
		}
		out = append(out, r)
	}

	var nextCursor *time.Time
	if len(out) > f.Limit {
		out = out[:f.Limit]
		nextCursor = &out[len(out)-1].StartedAt
	}

	return out, nextCursor, rows.Err()
}

// CountSessions returns the total number of sessions matching the filter (without limit or cursor).
func (d *DB) CountSessions(ctx context.Context, f SessionsFilter) (int, error) {
	var innerConditions []string
	var outerConditions []string
	var args []any
	argN := 1

	if f.From == nil {
		t := time.Now().Add(-24 * time.Hour)
		f.From = &t
	}
	innerConditions = append(innerConditions, fmt.Sprintf("ts >= $%d", argN))
	args = append(args, *f.From)
	argN++

	if f.To != nil {
		innerConditions = append(innerConditions, fmt.Sprintf("ts <= $%d", argN))
		args = append(args, *f.To)
		argN++
	}

	if f.Framework != "" {
		innerConditions = append(innerConditions, fmt.Sprintf("source_framework = $%d", argN))
		args = append(args, f.Framework)
		argN++
	}

	if f.Host != "" {
		outerConditions = append(outerConditions, fmt.Sprintf("host = $%d", argN))
		args = append(args, f.Host)
		argN++
	}

	innerWhere := strings.Join(innerConditions, " AND ")
	outerWhere := ""
	if len(outerConditions) > 0 {
		outerWhere = "WHERE " + strings.Join(outerConditions, " AND ")
	}

	query := fmt.Sprintf(`
		WITH session_groups AS (
			SELECT
				session_id,
				MAX(payload->'event'->'source'->>'host') as host
			FROM events
			WHERE session_id IS NOT NULL AND %s
			GROUP BY session_id
		)
		SELECT COUNT(*) FROM session_groups %s
	`, innerWhere, outerWhere)

	var count int
	err := d.sql.QueryRowContext(ctx, query, args...).Scan(&count)
	return count, err
}