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 } }