multica/server/internal/handler/runtime_update_redis_store.go

package handler

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"time"

	"github.com/redis/go-redis/v9"
)

// Redis-backed implementation of UpdateStore. CLI updates have the same
// pending-request shape as model-list and runtime-local-skill requests:
// frontend creates the request, daemon claims it on heartbeat, daemon reports
// a terminal result, and the UI polls by request ID. In multi-node deploys all
// four calls can hit different API replicas, so the lifecycle must live in
// shared storage.

const (
	updateKeyPrefix          = "mul:update:req:"
	updatePendingPrefix      = "mul:update:pending:"
	updateActivePrefix       = "mul:update:active:"
	updateRedisPopMaxRetries = 5
)

func updateKey(id string) string               { return updateKeyPrefix + id }
func updatePendingKey(runtimeID string) string { return updatePendingPrefix + runtimeID }
func updateActiveKey(runtimeID string) string  { return updateActivePrefix + runtimeID }

var deleteIfValueScript = redis.NewScript(`
if redis.call('GET', KEYS[1]) == ARGV[1] then
    return redis.call('DEL', KEYS[1])
end
return 0
`)

type RedisUpdateStore struct {
	rdb *redis.Client
}

func NewRedisUpdateStore(rdb *redis.Client) *RedisUpdateStore {
	return &RedisUpdateStore{rdb: rdb}
}

func (s *RedisUpdateStore) Create(ctx context.Context, runtimeID, targetVersion string) (*UpdateRequest, error) {
	now := time.Now()
	req := &UpdateRequest{
		ID:            randomID(),
		RuntimeID:     runtimeID,
		Status:        UpdatePending,
		TargetVersion: targetVersion,
		CreatedAt:     now,
		UpdatedAt:     now,
	}
	data, err := s.marshalRequest(req)
	if err != nil {
		return nil, err
	}

	activeKey := updateActiveKey(runtimeID)
	ok, err := s.rdb.SetNX(ctx, activeKey, req.ID, updateStoreRetention).Result()
	if err != nil {
		return nil, fmt.Errorf("reserve active update: %w", err)
	}
	if !ok {
		return nil, errUpdateInProgress
	}

	pipe := s.rdb.TxPipeline()
	pipe.Set(ctx, updateKey(req.ID), data, updateStoreRetention)
	pipe.ZAdd(ctx, updatePendingKey(runtimeID), redis.Z{
		Score:  float64(now.UnixNano()),
		Member: req.ID,
	})
	pipe.Expire(ctx, updatePendingKey(runtimeID), updateStoreRetention*2)
	if _, err := pipe.Exec(ctx); err != nil {
		_ = s.clearActiveIfMatches(ctx, runtimeID, req.ID)
		_ = s.rdb.Del(ctx, updateKey(req.ID)).Err()
		_ = s.rdb.ZRem(ctx, updatePendingKey(runtimeID), req.ID).Err()
		return nil, fmt.Errorf("persist update request: %w", err)
	}
	return req, nil
}

func (s *RedisUpdateStore) Get(ctx context.Context, id string) (*UpdateRequest, error) {
	return s.loadRequest(ctx, id)
}

func (s *RedisUpdateStore) loadRequest(ctx context.Context, id string) (*UpdateRequest, error) {
	raw, err := s.rdb.Get(ctx, updateKey(id)).Bytes()
	if errors.Is(err, redis.Nil) {
		return nil, nil
	}
	if err != nil {
		return nil, fmt.Errorf("get update request: %w", err)
	}
	req, err := s.unmarshalRequest(raw)
	if err != nil {
		return nil, err
	}
	if applyUpdateTimeout(req, time.Now()) {
		if err := s.persistRequest(ctx, req); err != nil {
			return nil, err
		}
		if err := s.clearActiveIfMatches(ctx, req.RuntimeID, req.ID); err != nil {
			return nil, err
		}
		s.rdb.ZRem(ctx, updatePendingKey(req.RuntimeID), req.ID)
	}
	return req, nil
}

func (s *RedisUpdateStore) persistRequest(ctx context.Context, req *UpdateRequest) error {
	data, err := s.marshalRequest(req)
	if err != nil {
		return err
	}
	if err := s.rdb.Set(ctx, updateKey(req.ID), data, updateStoreRetention).Err(); err != nil {
		return fmt.Errorf("persist update request: %w", err)
	}
	return nil
}

type redisUpdateEnvelope struct {
	Public       *UpdateRequest `json:"r"`
	RunStartedAt *time.Time     `json:"s,omitempty"`
}

func (s *RedisUpdateStore) marshalRequest(req *UpdateRequest) ([]byte, error) {
	env := redisUpdateEnvelope{Public: req, RunStartedAt: req.RunStartedAt}
	data, err := json.Marshal(env)
	if err != nil {
		return nil, fmt.Errorf("marshal update request: %w", err)
	}
	return data, nil
}

func (s *RedisUpdateStore) unmarshalRequest(raw []byte) (*UpdateRequest, error) {
	var env redisUpdateEnvelope
	if err := json.Unmarshal(raw, &env); err != nil {
		return nil, fmt.Errorf("decode update request: %w", err)
	}
	if env.Public == nil {
		return nil, fmt.Errorf("decode update request: missing payload")
	}
	env.Public.RunStartedAt = env.RunStartedAt
	return env.Public, nil
}

func (s *RedisUpdateStore) HasPending(ctx context.Context, runtimeID string) (bool, error) {
	cnt, err := s.rdb.ZCard(ctx, updatePendingKey(runtimeID)).Result()
	if err != nil {
		return false, fmt.Errorf("zcard pending updates: %w", err)
	}
	return cnt > 0, nil
}

func (s *RedisUpdateStore) PopPending(ctx context.Context, runtimeID string) (*UpdateRequest, error) {
	pendingKey := updatePendingKey(runtimeID)

	for attempt := 0; attempt < updateRedisPopMaxRetries; attempt++ {
		ids, err := s.rdb.ZRange(ctx, pendingKey, 0, 0).Result()
		if err != nil {
			return nil, fmt.Errorf("zrange pending updates: %w", err)
		}
		if len(ids) == 0 {
			return nil, nil
		}
		id := ids[0]

		req, err := s.loadRequest(ctx, id)
		if err != nil {
			return nil, err
		}
		if req == nil {
			s.rdb.ZRem(ctx, pendingKey, id)
			continue
		}
		if req.Status != UpdatePending {
			s.rdb.ZRem(ctx, pendingKey, id)
			continue
		}

		now := time.Now()
		req.Status = UpdateRunning
		req.RunStartedAt = &now
		req.UpdatedAt = now
		data, err := s.marshalRequest(req)
		if err != nil {
			return nil, err
		}

		result, err := claimPendingScript.Run(
			ctx, s.rdb,
			[]string{pendingKey, updateKey(id)},
			id, data, int(updateStoreRetention.Seconds()),
		).Int64()
		if err != nil {
			return nil, fmt.Errorf("claim pending update: %w", err)
		}
		if result == 0 {
			continue
		}
		return req, nil
	}
	return nil, nil
}

func (s *RedisUpdateStore) Complete(ctx context.Context, id string, output string) error {
	req, err := s.loadRequest(ctx, id)
	if err != nil {
		return err
	}
	if req == nil || updateRequestTerminal(req.Status) {
		return nil
	}
	req.Status = UpdateCompleted
	req.Output = output
	req.UpdatedAt = time.Now()
	if err := s.persistRequest(ctx, req); err != nil {
		return err
	}
	if err := s.clearActiveIfMatches(ctx, req.RuntimeID, req.ID); err != nil {
		return err
	}
	s.rdb.ZRem(ctx, updatePendingKey(req.RuntimeID), req.ID)
	return nil
}

func (s *RedisUpdateStore) Fail(ctx context.Context, id string, errMsg string) error {
	req, err := s.loadRequest(ctx, id)
	if err != nil {
		return err
	}
	if req == nil || updateRequestTerminal(req.Status) {
		return nil
	}
	req.Status = UpdateFailed
	req.Error = errMsg
	req.UpdatedAt = time.Now()
	if err := s.persistRequest(ctx, req); err != nil {
		return err
	}
	if err := s.clearActiveIfMatches(ctx, req.RuntimeID, req.ID); err != nil {
		return err
	}
	s.rdb.ZRem(ctx, updatePendingKey(req.RuntimeID), req.ID)
	return nil
}

func (s *RedisUpdateStore) clearActiveIfMatches(ctx context.Context, runtimeID, id string) error {
	if runtimeID == "" || id == "" {
		return nil
	}
	if err := deleteIfValueScript.Run(ctx, s.rdb, []string{updateActiveKey(runtimeID)}, id).Err(); err != nil {
		return fmt.Errorf("clear active update: %w", err)
	}
	return nil
}