pikasTech-HWLAB/internal/workbenchruntime/realtime_sync.go

package workbenchruntime

import (
	"context"
	"database/sql"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"strings"
)

const (
	realtimeSyncContractVersion = "workbench-runtime-sync-v1"
	realtimeSyncDefaultLimit    = 100
	realtimeSyncMaxLimit        = 500
	realtimeSyncMaxIdentityLen  = 512
)

type realtimeSyncQuery struct {
	SessionID      string       `json:"sessionId"`
	TraceID        string       `json:"traceId"`
	AfterOutboxSeq int64        `json:"afterOutboxSeq"`
	Limit          int          `json:"limit"`
	SnapshotOnly   bool         `json:"snapshotOnly,omitempty"`
	DeltaOnly      bool         `json:"deltaOnly,omitempty"`
	Actor          sessionActor `json:"actor"`
}

type realtimeSyncSnapshot struct {
	Facts           map[string][]any    `json:"facts"`
	Events          []realtimeSyncEvent `json:"events"`
	CutoffOutboxSeq int64               `json:"cutoffOutboxSeq"`
	CursorOutboxSeq int64               `json:"cursorOutboxSeq"`
	HasMore         bool                `json:"hasMore"`
}

type realtimeSyncEvent struct {
	OutboxSeq          int64  `json:"outboxSeq"`
	OutboxEventID      string `json:"outboxEventId"`
	EntityFamily       string `json:"entityFamily"`
	EntityID           string `json:"entityId"`
	EventSeq           any    `json:"eventSeq"`
	AggregateID        any    `json:"aggregateId"`
	AggregateSeq       int64  `json:"aggregateSeq"`
	ProjectionRevision int64  `json:"projectionRevision"`
	TraceID            string `json:"traceId"`
	SessionID          any    `json:"sessionId"`
	TurnID             any    `json:"turnId"`
	MessageID          any    `json:"messageId"`
	ProjectedSeq       int64  `json:"projectedSeq"`
	SourceSeq          int64  `json:"sourceSeq"`
	SourceEventID      any    `json:"sourceEventId"`
	CommitType         string `json:"commitType"`
	Terminal           bool   `json:"terminal"`
	Sealed             bool   `json:"sealed"`
	Payload            any    `json:"payload"`
	CreatedAt          string `json:"createdAt"`
	ValuesRedacted     bool   `json:"valuesRedacted"`
}

type realtimeSyncFactSpec struct {
	Family     string
	Table      string
	JSONColumn string
	OrderBy    string
}

var realtimeSyncFactSpecs = []realtimeSyncFactSpec{
	{Family: "sessions", Table: "workbench_sessions", JSONColumn: "session_json", OrderBy: "updated_at ASC, session_id ASC"},
	{Family: "messages", Table: "workbench_messages", JSONColumn: "message_json", OrderBy: "updated_at ASC, message_id ASC"},
	{Family: "parts", Table: "workbench_parts", JSONColumn: "part_json", OrderBy: "updated_at ASC, part_index ASC, part_id ASC"},
	{Family: "turns", Table: "workbench_turns", JSONColumn: "turn_json", OrderBy: "updated_at ASC, turn_id ASC"},
	{Family: "traceEvents", Table: "workbench_trace_events", JSONColumn: "event_json", OrderBy: "projected_seq ASC, id ASC"},
	{Family: "checkpoints", Table: "workbench_projection_checkpoints", JSONColumn: "checkpoint_json", OrderBy: "updated_at ASC, trace_id ASC"},
}

var realtimeSyncSnapshotFactSpecs = []realtimeSyncFactSpec{
	realtimeSyncFactSpecs[0],
	realtimeSyncFactSpecs[1],
	realtimeSyncFactSpecs[3],
}

const realtimeSyncResolveScopeSQL = `
SELECT sessions.session_id
FROM workbench_sessions AS sessions
WHERE ($1 = '' OR sessions.session_id = $1)
  AND (
    $2 = ''
    OR sessions.last_trace_id = $2
    OR EXISTS (
      SELECT 1
      FROM workbench_projection_outbox AS scoped_outbox
      WHERE scoped_outbox.session_id = sessions.session_id
        AND scoped_outbox.trace_id = $2
    )
  )
  AND ($3 OR sessions.owner_user_id = $4)
ORDER BY
  CASE WHEN sessions.session_id = NULLIF($1, '') THEN 0 ELSE 1 END,
  sessions.updated_at DESC,
  sessions.session_id ASC
LIMIT 1`

const realtimeSyncCutoffSQL = `
SELECT COALESCE(MAX(outbox_seq), 0)::bigint
FROM workbench_projection_outbox
WHERE session_id = $1
  AND ($2 = '' OR trace_id = $2)`

const realtimeSyncEventsSQL = `
SELECT
  outbox_seq,
  outbox_event_id,
  entity_family,
  entity_id,
  event_seq,
  aggregate_id,
  aggregate_seq,
  projection_revision,
  trace_id,
  session_id,
  turn_id,
  message_id,
  projected_seq,
  source_seq,
  source_event_id,
  commit_type,
  terminal,
  sealed,
  payload_json,
  created_at
FROM workbench_projection_outbox
WHERE session_id = $1
  AND ($2 = '' OR trace_id = $2)
  AND outbox_seq > $3
  AND outbox_seq <= $4
ORDER BY outbox_seq ASC
LIMIT $5`

func (s *Server) handleRealtimeSync(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		methodNotAllowed(w, "POST")
		return
	}
	query, ok := decodeRealtimeSyncQuery(w, r)
	if !ok {
		return
	}
	query, validation := normalizeRealtimeSyncQuery(query)
	if validation != nil {
		writeAPIError(w, validation.Status, validation.Code, validation.Message)
		return
	}
	ctx, cancel := context.WithTimeout(r.Context(), s.config.QueryTimeout)
	defer cancel()
	snapshot, err := s.readRealtimeSync(ctx, query)
	if errors.Is(err, sql.ErrNoRows) {
		writeAPIError(w, http.StatusNotFound, "workbench_sync_scope_not_found", "workbench sync scope was not found")
		return
	}
	if err != nil {
		writeQueryError(w, err)
		return
	}
	writeJSON(w, http.StatusOK, map[string]any{
		"ok":                   true,
		"status":               "succeeded",
		"contractVersion":      realtimeSyncContractVersion,
		"facts":                snapshot.Facts,
		"events":               snapshot.Events,
		"cutoffOutboxSeq":      snapshot.CutoffOutboxSeq,
		"cursorOutboxSeq":      snapshot.CursorOutboxSeq,
		"hasMore":              snapshot.HasMore,
		"persistence":          s.persistenceSummary(),
		"servedBy":             serviceID,
		"valuesRedacted":       true,
		"secretMaterialStored": false,
	})
}

type realtimeSyncValidationError struct {
	Status  int
	Code    string
	Message string
}

func decodeRealtimeSyncQuery(w http.ResponseWriter, r *http.Request) (realtimeSyncQuery, bool) {
	defer r.Body.Close()
	decoder := json.NewDecoder(http.MaxBytesReader(w, r.Body, 1<<20))
	decoder.DisallowUnknownFields()
	var query realtimeSyncQuery
	if err := decoder.Decode(&query); err != nil {
		writeAPIError(w, http.StatusBadRequest, "invalid_json", "request body must be a closed workbench sync JSON object")
		return realtimeSyncQuery{}, false
	}
	var trailing any
	if err := decoder.Decode(&trailing); !errors.Is(err, io.EOF) {
		writeAPIError(w, http.StatusBadRequest, "invalid_json", "request body must contain exactly one JSON object")
		return realtimeSyncQuery{}, false
	}
	return query, true
}

func normalizeRealtimeSyncQuery(query realtimeSyncQuery) (realtimeSyncQuery, *realtimeSyncValidationError) {
	query.SessionID = strings.TrimSpace(query.SessionID)
	query.TraceID = strings.TrimSpace(query.TraceID)
	query.Actor.ID = strings.TrimSpace(query.Actor.ID)
	query.Actor.Role = strings.ToLower(strings.TrimSpace(query.Actor.Role))
	if query.SessionID == "" && query.TraceID == "" {
		return query, &realtimeSyncValidationError{Status: http.StatusBadRequest, Code: "workbench_sync_scope_required", Message: "sessionId or traceId is required"}
	}
	if len(query.SessionID) > realtimeSyncMaxIdentityLen || len(query.TraceID) > realtimeSyncMaxIdentityLen || len(query.Actor.ID) > realtimeSyncMaxIdentityLen {
		return query, &realtimeSyncValidationError{Status: http.StatusBadRequest, Code: "workbench_sync_identity_invalid", Message: "workbench sync identifiers are too long"}
	}
	if query.AfterOutboxSeq < 0 {
		return query, &realtimeSyncValidationError{Status: http.StatusBadRequest, Code: "workbench_sync_cursor_invalid", Message: "afterOutboxSeq must be non-negative"}
	}
	if query.Limit < 0 {
		return query, &realtimeSyncValidationError{Status: http.StatusBadRequest, Code: "workbench_sync_limit_invalid", Message: "limit must be non-negative"}
	}
	if query.SnapshotOnly && query.DeltaOnly {
		return query, &realtimeSyncValidationError{Status: http.StatusBadRequest, Code: "workbench_sync_mode_invalid", Message: "snapshotOnly and deltaOnly are mutually exclusive"}
	}
	query.Limit = boundedLimit(query.Limit, realtimeSyncDefaultLimit, realtimeSyncMaxLimit)
	if query.Actor.Role != "admin" && query.Actor.Role != "user" {
		return query, &realtimeSyncValidationError{Status: http.StatusBadRequest, Code: "workbench_sync_actor_invalid", Message: "actor.role must be user or admin"}
	}
	if query.Actor.Role != "admin" && query.Actor.ID == "" {
		return query, &realtimeSyncValidationError{Status: http.StatusForbidden, Code: "workbench_sync_actor_required", Message: "ordinary users require an actor id"}
	}
	return query, nil
}

func realtimeSyncTxOptions() *sql.TxOptions {
	return &sql.TxOptions{Isolation: sql.LevelRepeatableRead, ReadOnly: true}
}

func (s *Server) readRealtimeSync(ctx context.Context, query realtimeSyncQuery) (realtimeSyncSnapshot, error) {
	tx, err := s.db.BeginTx(ctx, realtimeSyncTxOptions())
	if err != nil {
		return realtimeSyncSnapshot{}, wrapQueryStage("workbench_runtime.db.sync.begin", err)
	}
	committed := false
	defer func() {
		if !committed {
			_ = tx.Rollback()
		}
	}()
	snapshot, err := readRealtimeSyncSnapshot(ctx, sqlRealtimeSyncQueryer{tx: tx}, query)
	if err != nil {
		return realtimeSyncSnapshot{}, err
	}
	if err := tx.Commit(); err != nil {
		return realtimeSyncSnapshot{}, wrapQueryStage("workbench_runtime.db.sync.commit", err)
	}
	committed = true
	return snapshot, nil
}

type realtimeSyncRow interface {
	Scan(...any) error
}

type realtimeSyncRows interface {
	Close() error
	Err() error
	Next() bool
	Scan(...any) error
}

type realtimeSyncQueryer interface {
	QueryContext(context.Context, string, ...any) (realtimeSyncRows, error)
	QueryRowContext(context.Context, string, ...any) realtimeSyncRow
}

type sqlRealtimeSyncQueryer struct {
	tx *sql.Tx
}

func (queryer sqlRealtimeSyncQueryer) QueryContext(ctx context.Context, query string, args ...any) (realtimeSyncRows, error) {
	return queryer.tx.QueryContext(ctx, query, args...)
}

func (queryer sqlRealtimeSyncQueryer) QueryRowContext(ctx context.Context, query string, args ...any) realtimeSyncRow {
	return queryer.tx.QueryRowContext(ctx, query, args...)
}

func readRealtimeSyncSnapshot(ctx context.Context, tx realtimeSyncQueryer, query realtimeSyncQuery) (realtimeSyncSnapshot, error) {
	sessionID, err := resolveRealtimeSyncSession(ctx, tx, query)
	if err != nil {
		return realtimeSyncSnapshot{}, err
	}
	cutoff, err := readRealtimeSyncCutoff(ctx, tx, sessionID, query.TraceID)
	if err != nil {
		return realtimeSyncSnapshot{}, err
	}
	events := make([]realtimeSyncEvent, 0)
	hasMore := false
	if !query.SnapshotOnly {
		allEvents, readErr := readRealtimeSyncEvents(ctx, tx, sessionID, query.TraceID, query.AfterOutboxSeq, cutoff, query.Limit+1)
		if readErr != nil {
			return realtimeSyncSnapshot{}, readErr
		}
		hasMore = len(allEvents) > query.Limit
		events = allEvents
		if hasMore {
			events = allEvents[:query.Limit]
		}
	}
	facts := make(map[string][]any)
	if !query.DeltaOnly {
		specs := realtimeSyncFactSpecs
		if query.SnapshotOnly {
			specs = realtimeSyncSnapshotFactSpecs
		}
		facts, err = readRealtimeSyncFacts(ctx, tx, sessionID, query.TraceID, specs)
		if err != nil {
			return realtimeSyncSnapshot{}, err
		}
	}
	cursor := query.AfterOutboxSeq
	if query.SnapshotOnly || cursor > cutoff {
		cursor = cutoff
	}
	if len(events) > 0 {
		cursor = events[len(events)-1].OutboxSeq
	}
	return realtimeSyncSnapshot{Facts: facts, Events: events, CutoffOutboxSeq: cutoff, CursorOutboxSeq: cursor, HasMore: hasMore}, nil
}

func resolveRealtimeSyncSession(ctx context.Context, tx realtimeSyncQueryer, query realtimeSyncQuery) (string, error) {
	var sessionID string
	err := tx.QueryRowContext(ctx, realtimeSyncResolveScopeSQL, query.SessionID, query.TraceID, query.Actor.Role == "admin", query.Actor.ID).Scan(&sessionID)
	if err != nil {
		return "", wrapQueryStage("workbench_runtime.db.sync.scope", err)
	}
	return sessionID, nil
}

func readRealtimeSyncCutoff(ctx context.Context, tx realtimeSyncQueryer, sessionID string, traceID string) (int64, error) {
	var cutoff int64
	if err := tx.QueryRowContext(ctx, realtimeSyncCutoffSQL, sessionID, traceID).Scan(&cutoff); err != nil {
		return 0, wrapQueryStage("workbench_runtime.db.sync.cutoff", err)
	}
	return cutoff, nil
}

func readRealtimeSyncEvents(ctx context.Context, tx realtimeSyncQueryer, sessionID string, traceID string, after int64, cutoff int64, limit int) ([]realtimeSyncEvent, error) {
	rows, err := tx.QueryContext(ctx, realtimeSyncEventsSQL, sessionID, traceID, after, cutoff, limit)
	if err != nil {
		return nil, wrapQueryStage("workbench_runtime.db.sync.events", err)
	}
	defer rows.Close()
	events := make([]realtimeSyncEvent, 0)
	for rows.Next() {
		event, err := scanRealtimeSyncEvent(rows)
		if err != nil {
			return nil, wrapQueryStage("workbench_runtime.db.sync.events", err)
		}
		events = append(events, event)
	}
	if err := rows.Err(); err != nil {
		return nil, wrapQueryStage("workbench_runtime.db.sync.events", err)
	}
	return events, nil
}

type realtimeSyncScanner interface {
	Scan(...any) error
}

func scanRealtimeSyncEvent(scanner realtimeSyncScanner) (realtimeSyncEvent, error) {
	var event realtimeSyncEvent
	var outboxEventID, entityFamily, entityID, traceID, commitType, payloadJSON, createdAt sql.NullString
	var eventSeq, aggregateSeq, projectionRevision, projectedSeq, sourceSeq sql.NullInt64
	var aggregateID, sessionID, turnID, messageID, sourceEventID sql.NullString
	var terminal, sealed sql.NullBool
	if err := scanner.Scan(
		&event.OutboxSeq,
		&outboxEventID,
		&entityFamily,
		&entityID,
		&eventSeq,
		&aggregateID,
		&aggregateSeq,
		&projectionRevision,
		&traceID,
		&sessionID,
		&turnID,
		&messageID,
		&projectedSeq,
		&sourceSeq,
		&sourceEventID,
		&commitType,
		&terminal,
		&sealed,
		&payloadJSON,
		&createdAt,
	); err != nil {
		return realtimeSyncEvent{}, err
	}
	payload := any(map[string]any{})
	if payloadJSON.Valid && strings.TrimSpace(payloadJSON.String) != "" {
		if err := json.Unmarshal([]byte(payloadJSON.String), &payload); err != nil {
			return realtimeSyncEvent{}, fmt.Errorf("decode projection outbox payload: %w", err)
		}
	}
	event.OutboxEventID = nullableStringValue(outboxEventID)
	event.EntityFamily = nullableStringValue(entityFamily)
	event.EntityID = nullableStringValue(entityID)
	event.EventSeq = realtimeSyncNullableInt(eventSeq)
	event.AggregateID = nullableString(aggregateID)
	event.AggregateSeq = realtimeSyncInt(aggregateSeq)
	event.ProjectionRevision = realtimeSyncInt(projectionRevision)
	event.TraceID = nullableStringValue(traceID)
	event.SessionID = nullableString(sessionID)
	event.TurnID = nullableString(turnID)
	event.MessageID = nullableString(messageID)
	event.ProjectedSeq = realtimeSyncInt(projectedSeq)
	event.SourceSeq = realtimeSyncInt(sourceSeq)
	event.SourceEventID = nullableString(sourceEventID)
	event.CommitType = nullableStringValue(commitType)
	event.Terminal = nullableBool(terminal)
	event.Sealed = nullableBool(sealed)
	event.Payload = payload
	event.CreatedAt = nullableTimeText(createdAt)
	event.ValuesRedacted = true
	return event, nil
}

func readRealtimeSyncFacts(ctx context.Context, tx realtimeSyncQueryer, sessionID string, traceID string, specs []realtimeSyncFactSpec) (map[string][]any, error) {
	facts := make(map[string][]any, len(specs))
	for _, spec := range specs {
		args := []any{sessionID}
		where := "session_id = $1"
		if spec.Family != "sessions" {
			args = append(args, traceID)
			where += " AND ($2 = '' OR trace_id = $2)"
		}
		sqlText := fmt.Sprintf("SELECT %s FROM %s WHERE %s ORDER BY %s", spec.JSONColumn, spec.Table, where, spec.OrderBy)
		rows, err := tx.QueryContext(ctx, sqlText, args...)
		if err != nil {
			return nil, wrapQueryStage("workbench_runtime.db.sync.facts."+spec.Family, err)
		}
		items, err := scanRealtimeSyncFacts(rows)
		if err != nil {
			return nil, wrapQueryStage("workbench_runtime.db.sync.facts."+spec.Family, err)
		}
		facts[spec.Family] = items
	}
	return facts, nil
}

func scanRealtimeSyncFacts(rows realtimeSyncRows) ([]any, error) {
	defer rows.Close()
	items := make([]any, 0)
	for rows.Next() {
		var raw sql.NullString
		if err := rows.Scan(&raw); err != nil {
			return nil, err
		}
		if !raw.Valid || strings.TrimSpace(raw.String) == "" {
			continue
		}
		var item any
		if err := json.Unmarshal([]byte(raw.String), &item); err != nil {
			return nil, fmt.Errorf("decode workbench fact: %w", err)
		}
		if item != nil {
			items = append(items, item)
		}
	}
	if err := rows.Err(); err != nil {
		return nil, err
	}
	return items, nil
}

func realtimeSyncNullableInt(value sql.NullInt64) any {
	if !value.Valid {
		return nil
	}
	return value.Int64
}

func realtimeSyncInt(value sql.NullInt64) int64 {
	if !value.Valid {
		return 0
	}
	return value.Int64
}