refactor inference polling to use a generic helper function based on context and ticker.

use duration flags for timeout and intervals.

Author: danmrichards

cmd/audio_inference.go

@@ -41,8 +41,8 @@ func init() {
 	f.Int("bitrate", 0, "Bitrate in kbps (32-320, compressed formats only)")
 	f.String("preset", "", "Named preset to apply")
 	f.Bool("dry-run", false, "Print the API request without executing")
-	f.Int("poll-interval", 5, "Polling interval in seconds for async results")
-	f.Int("timeout", 300, "Maximum wait time in seconds for audio generation")
+	f.Duration("poll-interval", defaultPollInterval, "Polling interval in seconds for async results")
+	f.Duration("timeout", defaultAudioGenerationTimeout, "Maximum wait time in seconds for audio generation")
 	f.Duration("download-timeout", defaultAudioDownloadTimeout, "timeout to use when downloading audio inference results")
 
 	audioInferenceCmd.RegisterFlagCompletionFunc("output-format", func(cmd *cobra.Command, args []string, toComplete string) ([]cobra.Completion, cobra.ShellCompDirective) { //nolint:errcheck,gosec
@@ -98,8 +98,8 @@ func runAudioInference(cmd *cobra.Command, args []string) error {
 	dryRun, _ := cmd.Flags().GetBool("dry-run")
 	sampleRate, _ := cmd.Flags().GetInt("sample-rate")
 	bitrate, _ := cmd.Flags().GetInt("bitrate")
-	pollInterval, _ := cmd.Flags().GetInt("poll-interval")
-	timeout, _ := cmd.Flags().GetInt("timeout")
+	pollInterval, _ := cmd.Flags().GetDuration("poll-interval")
+	timeout, _ := cmd.Flags().GetDuration("timeout")
 	downloadTimeout, _ := cmd.Flags().GetDuration("download-timeout")
 
 	// Validation
@@ -143,11 +143,12 @@ func runAudioInference(cmd *cobra.Command, args []string) error {
 	// Submit
 	s := output.NewSpinner(" Submitting audio generation task...")
 	s.Start()
-	defer s.Stop()
 
 	client := api.NewClient(key, config.GetBaseURL(), flagVerbose)
 	_, err := client.AudioInference(context.Background(), req)
 	if err != nil {
+		s.Stop()
+
 		if api.IsAuthError(err) {
 			output.Error("Authentication failed. Run 'runware auth login' to set your API key.")
 			return err
@@ -158,45 +159,40 @@ func runAudioInference(cmd *cobra.Command, args []string) error {
 	s.Suffix(" Generating audio (this may take a few minutes)...")
 
 	taskUUID := req.TaskUUID
-	deadline := time.Now().Add(time.Duration(timeout) * time.Second)
-	interval := time.Duration(pollInterval) * time.Second
-	var results []api.AudioInferenceResult
-
-	for time.Now().Before(deadline) {
-		rawData, err := client.GetResponse(context.Background(), taskUUID)
-		if err != nil {
-			if flagVerbose {
-				fmt.Fprintf(os.Stderr, "Poll: %s\n", err) //nolint:errcheck,gosec
-			}
-			time.Sleep(interval)
-			continue
-		}
 
-		for _, raw := range rawData {
+	pollCtx, cancel := context.WithTimeout(cmd.Context(), timeout)
+	defer cancel()
+
+	results, err := api.PollResults(
+		pollCtx,
+		client,
+		taskUUID,
+		pollInterval,
+		flagVerbose,
+		func(raw json.RawMessage) (api.AudioInferenceResult, bool) {
 			var r api.AudioInferenceResult
 			if err := json.Unmarshal(raw, &r); err != nil {
-				continue
-			}
-			if r.AudioURL != "" {
-				results = append(results, r)
+				return r, false
 			}
-		}
-
-		if len(results) > 0 {
-			break
-		}
-
-		time.Sleep(interval)
+			return r, r.AudioURL != ""
+		},
+	)
+	if err != nil {
+		s.Stop()
+		output.Error("Audo generation failed")
+		return err
 	}
 
 	if len(results) == 0 {
+		s.Stop()
 		output.Error("Audio generation timed out or returned no results")
-		return fmt.Errorf("no audio results after %ds", timeout)
+		return fmt.Errorf("no audio results after %v", timeout)
 	}
 
 	// JSON/YAML output
 	format := output.ParseFormat(getFormat())
 	if format != output.FormatTable {
+		s.Stop()
 		return output.Print(format, results, nil, nil)
 	}
 
@@ -215,9 +211,7 @@ func runAudioInference(cmd *cobra.Command, args []string) error {
 			rows = append(rows, row)
 		}
 
-		// Manually stop the spinner to ensure clean output of results.
 		s.Stop()
-
 		return output.Print(format, results, headers, rows)
 	}
 
@@ -254,8 +248,6 @@ func runAudioInference(cmd *cobra.Command, args []string) error {
 		rows = append(rows, row)
 	}
 
-	// Manually stop the spinner to ensure clean output of results.
 	s.Stop()
-
 	return output.Print(format, results, headers, rows)
 }

cmd/constants.go

@@ -14,6 +14,14 @@ const (
 )
 
 const (
+	// Poll intervals.
+	defaultPollInterval = 5 * time.Second
+
+	// Generation timeouts.
+	defaultAudioGenerationTimeout = 5 * time.Minute
+	defaultVideoGenerationTimeout = 10 * time.Minute
+
+	// Download timeouts.
 	defaultImageDownloadTimeout = 1 * time.Minute
 	defaultAudioDownloadTimeout = 5 * time.Minute
 	defaultVideoDownloadTimeout = 10 * time.Minute

cmd/image_inference.go

@@ -245,12 +245,12 @@ func runImageInference(cmd *cobra.Command, args []string) error {
 	// Start spinner
 	s := output.NewSpinner(" Generating image...")
 	s.Start()
-	defer s.Stop()
 
 	client := api.NewClient(key, config.GetBaseURL(), flagVerbose)
 
 	results, err := client.ImageInference(ctx, req)
 	if err != nil {
+		s.Stop()
 		if api.IsAuthError(err) {
 			output.Error("Authentication failed. Run 'runware auth login' to set your API key.")
 			return err
@@ -259,13 +259,15 @@ func runImageInference(cmd *cobra.Command, args []string) error {
 	}
 
 	if len(results) == 0 {
+		s.Stop()
 		output.Error("No images returned")
 		return fmt.Errorf("empty result")
 	}
 
 	// JSON/YAML output
 	format := output.ParseFormat(getFormat())
 	if format != output.FormatTable {
+		s.Stop()
 		return output.Print(format, results, nil, nil)
 	}
 
@@ -277,9 +279,7 @@ func runImageInference(cmd *cobra.Command, args []string) error {
 			rows = append(rows, []any{i + 1, r.ImageURL, r.Seed})
 		}
 
-		// Manually stop the spinner to ensure clean output of results.
 		s.Stop()
-
 		return output.Print(format, results, headers, rows)
 	}
 
@@ -309,9 +309,7 @@ func runImageInference(cmd *cobra.Command, args []string) error {
 		rows = append(rows, []any{i + 1, destPath, r.Seed})
 	}
 
-	// Manually stop the spinner to ensure clean output of results.
 	s.Stop()
-
 	return output.Print(format, results, headers, rows)
 }
 

cmd/text_inference.go

@@ -140,19 +140,18 @@ func runTextInference(cmd *cobra.Command, args []string) error {
 	// Submit
 	s := output.NewSpinner(" Generating text...")
 	s.Start()
-	defer s.Stop()
 
 	client := api.NewClient(key, config.GetBaseURL(), flagVerbose)
 	results, err := client.TextInference(context.Background(), req)
 	if err != nil {
+		s.Stop()
 		if api.IsAuthError(err) {
 			output.Error("Authentication failed. Run 'runware auth login' to set your API key.")
 			return err
 		}
 		return err
 	}
 
-	// Manually stop the spinner to ensure clean output of results.
 	s.Stop()
 
 	if len(results) == 0 {

cmd/video_inference.go

@@ -47,8 +47,8 @@ func init() {
 	f.Bool("include-cost", false, "Include cost info in response")
 	f.String("preset", "", "Named preset to apply")
 	f.Bool("dry-run", false, "Print the API request without executing")
-	f.Int("poll-interval", 5, "Polling interval in seconds for async results")
-	f.Int("timeout", 600, "Maximum wait time in seconds for video generation")
+	f.Duration("poll-interval", defaultPollInterval, "Polling interval for async results")
+	f.Duration("timeout", defaultVideoGenerationTimeout, "Maximum wait time for video generation")
 	f.Duration("download-timeout", defaultVideoDownloadTimeout, "timeout to use when downloading video inference results")
 
 	videoInferenceCmd.RegisterFlagCompletionFunc("output-format", func(cmd *cobra.Command, args []string, toComplete string) ([]cobra.Completion, cobra.ShellCompDirective) { //nolint:errcheck,gosec
@@ -119,8 +119,8 @@ func runVideoInference(cmd *cobra.Command, args []string) error {
 	dryRun, _ := cmd.Flags().GetBool("dry-run")
 	sourcePath, _ := cmd.Flags().GetString("source")
 	sourceLastPath, _ := cmd.Flags().GetString("source-last")
-	pollInterval, _ := cmd.Flags().GetInt("poll-interval")
-	timeout, _ := cmd.Flags().GetInt("timeout")
+	pollInterval, _ := cmd.Flags().GetDuration("poll-interval")
+	timeout, _ := cmd.Flags().GetDuration("timeout")
 	downloadTimeout, _ := cmd.Flags().GetDuration("download-timeout")
 
 	// Build request
@@ -192,12 +192,12 @@ func runVideoInference(cmd *cobra.Command, args []string) error {
 	// Submit the video generation task
 	s := output.NewSpinner(" Submitting video generation task...")
 	s.Start()
-	defer s.Stop()
 
 	client := api.NewClient(key, config.GetBaseURL(), flagVerbose)
 
 	submitResults, err := client.VideoInference(context.Background(), req)
 	if err != nil {
+		s.Stop()
 		if api.IsAuthError(err) {
 			output.Error("Authentication failed. Run 'runware auth login' to set your API key.")
 			return err
@@ -213,48 +213,39 @@ func runVideoInference(cmd *cobra.Command, args []string) error {
 		taskUUID = submitResults[0].TaskUUID
 	}
 
-	deadline := time.Now().Add(time.Duration(timeout) * time.Second)
-	interval := time.Duration(pollInterval) * time.Second
-	var results []api.VideoInferenceResult
+	pollCtx, cancel := context.WithTimeout(cmd.Context(), timeout)
+	defer cancel()
 
-	for time.Now().Before(deadline) {
-		rawData, err := client.GetResponse(context.Background(), taskUUID)
-		if err != nil {
-			// getResponse may return an error if results aren't ready yet — keep polling
-			if flagVerbose {
-				fmt.Fprintf(os.Stderr, "Poll: %s\n", err) //nolint:errcheck,gosec
-			}
-			time.Sleep(interval)
-			continue
-		}
-
-		// Parse results
-		for _, raw := range rawData {
+	results, err := api.PollResults(
+		pollCtx,
+		client,
+		taskUUID,
+		pollInterval,
+		flagVerbose,
+		func(raw json.RawMessage) (api.VideoInferenceResult, bool) {
 			var r api.VideoInferenceResult
 			if err := json.Unmarshal(raw, &r); err != nil {
-				continue
-			}
-			// Check if we have a video URL (indicates completion)
-			if r.VideoURL != "" || r.MediaURL != "" {
-				results = append(results, r)
+				return r, false
 			}
-		}
-
-		if len(results) > 0 {
-			break
-		}
-
-		time.Sleep(interval)
+			return r, r.VideoURL != "" || r.MediaURL != ""
+		},
+	)
+	if err != nil {
+		s.Stop()
+		output.Error("Video generation failed")
+		return err
 	}
 
 	if len(results) == 0 {
+		s.Stop()
 		output.Error("Video generation timed out or returned no results")
-		return fmt.Errorf("no video results after %ds", timeout)
+		return fmt.Errorf("no video results after %v", timeout)
 	}
 
 	// JSON/YAML output
 	format := output.ParseFormat(getFormat())
 	if format != output.FormatTable {
+		s.Stop()
 		return output.Print(format, results, nil, nil)
 	}
 
@@ -277,9 +268,7 @@ func runVideoInference(cmd *cobra.Command, args []string) error {
 			headers = append(headers, tableHeaderCost)
 		}
 
-		// Manually stop the spinner to ensure clean output of results.
 		s.Stop()
-
 		return output.Print(format, results, headers, rows)
 	}
 
@@ -328,7 +317,6 @@ func runVideoInference(cmd *cobra.Command, args []string) error {
 		return fmt.Errorf("all %d video downloads failed", len(results))
 	}
 
-	// Manually stop the spinner to ensure clean output of results.
 	s.Stop()
 
 	if err := output.Print(format, results, headers, rows); err != nil {

internal/api/poll.go

@@ -0,0 +1,45 @@
+package api
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"os"
+	"time"
+)
+
+// PollResults polls for generic results from a task.
+func PollResults[T any](ctx context.Context, client Client, taskUUID string, interval time.Duration, verbose bool, parse func(json.RawMessage) (T, bool)) ([]T, error) {
+	ticker := time.NewTicker(interval)
+	defer ticker.Stop()
+
+	var results []T
+	for {
+		rawData, err := client.GetResponse(ctx, taskUUID)
+		if err != nil {
+			var apiErr APIError
+			if IsAuthError(err) || errors.As(err, &apiErr) {
+				return nil, err
+			}
+			if verbose {
+				fmt.Fprintf(os.Stderr, "Poll: %s\n", err) //nolint:errcheck,gosec
+			}
+		} else {
+			for _, raw := range rawData {
+				if r, ok := parse(raw); ok {
+					results = append(results, r)
+				}
+			}
+			if len(results) > 0 {
+				return results, nil
+			}
+		}
+
+		select {
+		case <-ctx.Done():
+			return nil, nil
+		case <-ticker.C:
+		}
+	}
+}