browser/plugins/examples/scum-server-plugin/companion/orchestration.go

package companion

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

const (
	SmokeClaimLimit                 = 1
	SmokeIsolationEnvironment       = "SCUM_COMPANION_SMOKE_SCOPE"
	SmokeIsolationValue             = "isolated-non-production"
	companionDiagnosticsCommandType = "companion.diagnostics"
	companionHealthSnapshotType     = "companion.health"
	companionHealthSchemaVersion    = "1"
	companionHealthKeepForSeconds   = 7 * 24 * 60 * 60
	companionHealthMaxRecords       = 1000
)

type SmokeOptions struct {
	IsolatedNonProduction bool
	StreamSuffix          func() (string, error)
}

type SmokeResult struct {
	ClaimedCount          int      `json:"claimedCount"`
	CompletedCount        int      `json:"completedCount"`
	UnsupportedCount      int      `json:"unsupportedCount"`
	CompletedCommandIDs   []string `json:"completedCommandIds"`
	UnsupportedCommandIDs []string `json:"unsupportedCommandIds"`
	SnapshotID            string   `json:"snapshotId,omitempty"`
}

func RunOneShotSmoke(ctx context.Context, client *Client, options SmokeOptions) (SmokeResult, error) {
	result := SmokeResult{CompletedCommandIDs: []string{}, UnsupportedCommandIDs: []string{}}
	if client == nil {
		return result, fmt.Errorf("companion client is required")
	}
	if !options.IsolatedNonProduction {
		return result, fmt.Errorf("one-shot smoke requires an isolated non-production queue")
	}
	streamSuffix := options.StreamSuffix
	if streamSuffix == nil {
		streamSuffix = randomNonce
	}
	suffix, err := streamSuffix()
	if err != nil || !validNonce(suffix) {
		return result, fmt.Errorf("create isolated smoke stream")
	}
	streamKey := "smoke-" + suffix
	if len(streamKey) > 80 {
		return result, fmt.Errorf("isolated smoke stream is too long")
	}
	if _, err := client.Register(ctx); err != nil {
		return result, err
	}
	observedAt := client.now().UTC()
	if _, err := client.Heartbeat(ctx, HealthReport{Status: "healthy", Reason: "one-shot smoke ready"}); err != nil {
		return result, err
	}
	commands, err := client.ClaimCommands(ctx, SmokeClaimLimit)
	if err != nil {
		return result, err
	}
	result.ClaimedCount = len(commands)
	if len(commands) != 1 {
		return result, fmt.Errorf("one-shot smoke requires exactly one isolated command")
	}
	command := commands[0]
	if err := validateSmokeDiagnosticsCommand(command, observedAt); err != nil {
		result.UnsupportedCount = 1
		if command.ID != "" {
			result.UnsupportedCommandIDs = append(result.UnsupportedCommandIDs, command.ID)
		}
		return result, err
	}
	ack, err := client.AckCommand(ctx, command.ID, command.FencingToken)
	if err != nil {
		return result, err
	}
	if ack.CommandID != command.ID || ack.FencingToken != command.FencingToken || ack.State != "claimed" {
		return result, fmt.Errorf("platform returned an invalid smoke acknowledgement")
	}
	diagnostics := map[string]any{
		"status":          "online",
		"version":         client.config.Component.Version,
		"lastHeartbeatAt": observedAt.Format(time.RFC3339Nano),
	}
	completed, err := client.CompleteCommand(ctx, command.ID, command.FencingToken, CommandResult{Status: "succeeded", Summary: "companion diagnostics available", Payload: diagnostics})
	if err != nil {
		return result, err
	}
	if completed.CommandID != command.ID || completed.State != "succeeded" {
		return result, fmt.Errorf("platform returned an invalid smoke command result")
	}
	result.CompletedCount = 1
	result.CompletedCommandIDs = append(result.CompletedCommandIDs, command.ID)

	health := map[string]any{
		"status":       "online",
		"version":      client.config.Component.Version,
		"observedAt":   observedAt.Format(time.RFC3339Nano),
		"capabilities": append([]string(nil), client.config.Capabilities...),
	}
	accepted, err := client.UploadSnapshot(ctx, Snapshot{
		Type:           companionHealthSnapshotType,
		SchemaVersion:  companionHealthSchemaVersion,
		StreamKey:      streamKey,
		Sequence:       1,
		ObservedAt:     observedAt,
		Payload:        health,
		KeepForSeconds: companionHealthKeepForSeconds,
		MaxRecords:     companionHealthMaxRecords,
	})
	if err != nil {
		return result, err
	}
	if accepted.SnapshotID == "" || accepted.ProfileKey != ProfileKey || accepted.Type != companionHealthSnapshotType || accepted.SchemaVersion != companionHealthSchemaVersion || accepted.StreamKey != streamKey || accepted.Sequence != 1 {
		return result, fmt.Errorf("platform returned an invalid smoke snapshot acknowledgement")
	}
	result.SnapshotID = accepted.SnapshotID
	return result, nil
}

func validateSmokeDiagnosticsCommand(command ClaimedCommand, stamp time.Time) error {
	if command.CommandType != companionDiagnosticsCommandType {
		return fmt.Errorf("isolated smoke claimed an unsupported command")
	}
	if command.ID == "" || command.ProfileKey != ProfileKey || command.FencingToken == 0 {
		return fmt.Errorf("isolated smoke command identity is invalid")
	}
	if command.LeaseExpiresAt.IsZero() || command.ExpiresAt.IsZero() || !stamp.Before(command.LeaseExpiresAt) || !stamp.Before(command.ExpiresAt) {
		return fmt.Errorf("isolated smoke command lease is not live")
	}
	if command.Payload == nil {
		return fmt.Errorf("isolated smoke diagnostics payload is invalid")
	}
	for key, value := range command.Payload {
		switch key {
		case "includeWindowState":
			if _, ok := value.(bool); !ok {
				return fmt.Errorf("isolated smoke diagnostics payload is invalid")
			}
		case "maxEntries":
			if !boundedDiagnosticsEntries(value) {
				return fmt.Errorf("isolated smoke diagnostics payload is invalid")
			}
		default:
			return fmt.Errorf("isolated smoke diagnostics payload is invalid")
		}
	}
	return nil
}

func boundedDiagnosticsEntries(value any) bool {
	switch typed := value.(type) {
	case int:
		return typed >= 1 && typed <= 20
	case int64:
		return typed >= 1 && typed <= 20
	case float64:
		return typed >= 1 && typed <= 20 && typed == float64(int(typed))
	default:
		return false
	}
}