feat: parallelize forecast usage-artifact downloads to fix 45-min timeout

.github/workflows/daily-spending-forecast.md

@@ -21,6 +21,21 @@ steps:
       set -euo pipefail
       make build
       "$GITHUB_WORKSPACE/gh-aw" --version
+  - name: Prefetch forecast usage artifacts
+    continue-on-error: true
+    env:
+      REPOSITORY: ${{ github.repository }}
+    run: |
+      # Download usage artifacts for the last 30 days in parallel so the main
+      # forecast step reads from the local cache and produces output quickly.
+      DEBUG='*' "$GITHUB_WORKSPACE/gh-aw" forecast \
+        --repo "$REPOSITORY" \
+        --days 30 \
+        --sample 100 \
+        --concurrency 8 \
+        --timeout 25 \
+        --verbose \
+        > /dev/null 2>&1 || true
   - name: Run spending forecast
     id: spending_forecast
     continue-on-error: true
@@ -37,7 +52,8 @@ steps:
         --days 30 \
         --period month \
         --sample 100 \
-        --timeout 30 \
+        --concurrency 8 \
+        --timeout 10 \
         --verbose \
         --json \
         > >(tee "$output_dir/forecast.json") \

pkg/cli/forecast_command.go

@@ -29,6 +29,9 @@ type ForecastConfig struct {
 	// TimeoutMinutes gracefully cancels forecast computation after the configured
 	// number of minutes. Zero disables timeout.
 	TimeoutMinutes int
+	// DownloadConcurrency is the maximum number of usage-artifact downloads to run in
+	// parallel. Zero or negative uses the default (defaultForecastDownloadConcurrency).
+	DownloadConcurrency int
 }
 
 // NewForecastCommand creates the forecast command.
@@ -84,20 +87,22 @@ Backtesting (--eval):
 			sampleSize, _ := cmd.Flags().GetInt("sample")
 			evalMode, _ := cmd.Flags().GetBool("eval")
 			timeoutMinutes, _ := cmd.Flags().GetInt("timeout")
+			downloadConcurrency, _ := cmd.Flags().GetInt("concurrency")
 
 			forecastRunLog.Printf("Forecast command invoked: workflow_count=%d, days=%d, period=%s, sample_size=%d, eval=%v, timeout_minutes=%d, json=%v, repo=%q",
 				len(args), days, period, sampleSize, evalMode, timeoutMinutes, jsonOutput, repoOverride)
 
 			config := ForecastConfig{
-				WorkflowIDs:    args,
-				Days:           days,
-				Period:         period,
-				JSONOutput:     jsonOutput,
-				Verbose:        verbose,
-				RepoOverride:   repoOverride,
-				SampleSize:     sampleSize,
-				EvalMode:       evalMode,
-				TimeoutMinutes: timeoutMinutes,
+				WorkflowIDs:         args,
+				Days:                days,
+				Period:              period,
+				JSONOutput:          jsonOutput,
+				Verbose:             verbose,
+				RepoOverride:        repoOverride,
+				SampleSize:          sampleSize,
+				EvalMode:            evalMode,
+				TimeoutMinutes:      timeoutMinutes,
+				DownloadConcurrency: downloadConcurrency,
 			}
 
 			return RunForecast(config)
@@ -109,6 +114,7 @@ Backtesting (--eval):
 	cmd.Flags().Int("sample", 100, "Maximum number of completed runs to sample per workflow")
 	cmd.Flags().Bool("eval", false, "Evaluate forecast quality against past data (backtesting mode)")
 	cmd.Flags().Int("timeout", 0, "Gracefully stop forecast computation after this many minutes (0 = no timeout)")
+	cmd.Flags().Int("concurrency", 0, "Maximum number of concurrent usage-artifact downloads (0 = use default)")
 	addRepoFlag(cmd)
 	addJSONFlag(cmd)
 

pkg/cli/forecast_compute.go

@@ -10,6 +10,7 @@ import (
 	"path/filepath"
 	"slices"
 	"sort"
+	"sync"
 	"time"
 
 	"github.com/github/gh-aw/pkg/console"
@@ -19,6 +20,13 @@ import (
 	"github.com/github/gh-aw/pkg/workflow"
 )
 
+// defaultForecastDownloadConcurrency is the number of usage-artifact downloads that
+// run in parallel when DownloadConcurrency is not set (or is <= 0).
+// Increasing this value reduces wall-clock time for the download phase at the cost of
+// more simultaneous GitHub API requests; 8 is chosen as a balance between throughput
+// and rate-limit headroom.
+const defaultForecastDownloadConcurrency = 8
+
 var (
 	forecastLoadCachedRunAIC = loadCachedRunAIC
 	// forecastDownloadRunArtifacts uses a forecast-specific implementation that downloads
@@ -93,8 +101,11 @@ func forecastWorkflow(ctx context.Context, workflowName, startDate string, confi
 	aicObservations := make([]int, 0, len(completed))
 	samples := make([]ForecastRunSample, 0, len(completed))
 
+	// Download usage artifacts in parallel, then process results in run order.
+	aicMap := parallelLoadRunAICs(ctx, completed, config)
+
 	for _, r := range completed {
-		runAIC := forecastLoadCachedRunAIC(ctx, r.DatabaseID, config.Verbose)
+		runAIC := aicMap[r.DatabaseID]
 		if runAIC <= 0 {
 			forecastRunLog.Printf("Skipping run %d for %s: AIC=%.3f treated as missing data", r.DatabaseID, workflowName, runAIC)
 			continue
@@ -178,6 +189,65 @@ func forecastWorkflow(ctx context.Context, workflowName, startDate string, confi
 	return result, nil
 }
 
+// parallelLoadRunAICs fetches AIC data for all completed runs concurrently and returns
+// a map from DatabaseID to AIC value. Downloads run at most concurrency goroutines at a
+// time; when config.DownloadConcurrency is <= 0, defaultForecastDownloadConcurrency is
+// used. Runs whose AIC cannot be determined (download failure, no artifact, context
+// cancelled) are omitted from the map so the caller's existing zero-AIC guard takes effect.
+func parallelLoadRunAICs(ctx context.Context, runs []WorkflowRun, config ForecastConfig) map[int64]float64 {
+	n := len(runs)
+	if n == 0 {
+		return make(map[int64]float64)
+	}
+	concurrency := config.DownloadConcurrency
+	if concurrency <= 0 {
+		concurrency = defaultForecastDownloadConcurrency
+	}
+	if concurrency > n {
+		concurrency = n
+	}
+
+	type aicResult struct {
+		runID int64
+		aic   float64
+	}
+
+	sem := make(chan struct{}, concurrency)
+	resultsCh := make(chan aicResult, n)
+	var wg sync.WaitGroup
+
+	for _, r := range runs {
+		if ctx.Err() != nil {
+			break
+		}
+		wg.Add(1)
+		runID := r.DatabaseID
+		go func() {
+			defer wg.Done()
+			// Acquire semaphore slot; abort if context is cancelled while waiting.
+			select {
+			case sem <- struct{}{}:
+			case <-ctx.Done():
+				return
+			}
+			defer func() { <-sem }()
+			aic := forecastLoadCachedRunAIC(ctx, runID, config.Verbose)
+			resultsCh <- aicResult{runID: runID, aic: aic}
+		}()
+	}
+
+	go func() {
+		wg.Wait()
+		close(resultsCh)
+	}()
+
+	aicMap := make(map[int64]float64, n)
+	for r := range resultsCh {
+		aicMap[r.runID] = r.aic
+	}
+	return aicMap
+}
+
 // loadCachedRunAIC looks up a locally-cached RunSummary for the given
 // run ID and returns the TotalAIC from its TokenUsage summary.
 // Returns 0 when no cache exists or the cache does not contain AIC data.

pkg/cli/forecast_test.go

@@ -463,3 +463,75 @@ func TestLoadCachedRunAIC_MissingUsageReturnsZero(t *testing.T) {
 	require.False(t, analyzeCalled)
 	require.Equal(t, []string{"usage"}, downloaded)
 }
+
+// ── parallelLoadRunAICs ───────────────────────────────────────────────────────
+
+// TestParallelLoadRunAICs_ReturnsAllAICValues verifies that parallelLoadRunAICs collects
+// AIC values for every run, regardless of concurrency level.
+func TestParallelLoadRunAICs_ReturnsAllAICValues(t *testing.T) {
+	originalLoadAIC := forecastLoadCachedRunAIC
+	t.Cleanup(func() { forecastLoadCachedRunAIC = originalLoadAIC })
+
+	wantAIC := map[int64]float64{
+		1: 1.0,
+		2: 2.0,
+		3: 3.0,
+		4: 4.0,
+		5: 5.0,
+	}
+	forecastLoadCachedRunAIC = func(_ context.Context, runID int64, _ bool) float64 {
+		return wantAIC[runID]
+	}
+
+	runs := []WorkflowRun{
+		{DatabaseID: 1, Status: "completed", Conclusion: "success"},
+		{DatabaseID: 2, Status: "completed", Conclusion: "success"},
+		{DatabaseID: 3, Status: "completed", Conclusion: "failure"},
+		{DatabaseID: 4, Status: "completed", Conclusion: "success"},
+		{DatabaseID: 5, Status: "completed", Conclusion: "success"},
+	}
+
+	got := parallelLoadRunAICs(context.Background(), runs, ForecastConfig{DownloadConcurrency: 3})
+	assert.Equal(t, wantAIC, got)
+}
+
+// TestParallelLoadRunAICs_RespectsContextCancellation verifies that parallelLoadRunAICs
+// stops issuing new downloads and returns promptly when the context is cancelled.
+func TestParallelLoadRunAICs_RespectsContextCancellation(t *testing.T) {
+	originalLoadAIC := forecastLoadCachedRunAIC
+	t.Cleanup(func() { forecastLoadCachedRunAIC = originalLoadAIC })
+
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel() // cancel immediately
+
+	forecastLoadCachedRunAIC = func(_ context.Context, runID int64, _ bool) float64 {
+		return float64(runID)
+	}
+
+	runs := []WorkflowRun{
+		{DatabaseID: 10, Status: "completed", Conclusion: "success"},
+		{DatabaseID: 11, Status: "completed", Conclusion: "success"},
+	}
+
+	// Should return quickly without panicking, regardless of which goroutines completed.
+	got := parallelLoadRunAICs(ctx, runs, ForecastConfig{DownloadConcurrency: 1})
+	assert.NotNil(t, got)
+}
+
+// TestParallelLoadRunAICs_EmptyRunsReturnsEmptyMap verifies the empty-input edge case.
+func TestParallelLoadRunAICs_EmptyRunsReturnsEmptyMap(t *testing.T) {
+	got := parallelLoadRunAICs(context.Background(), nil, ForecastConfig{})
+	assert.Empty(t, got)
+}
+
+// TestNewForecastCommand_ConcurrencyFlag verifies that the --concurrency flag is
+// registered with the expected default and usage text.
+func TestNewForecastCommand_ConcurrencyFlag(t *testing.T) {
+	cmd := NewForecastCommand()
+	require.NotNil(t, cmd)
+
+	flag := cmd.Flags().Lookup("concurrency")
+	require.NotNil(t, flag, "forecast command should register --concurrency")
+	assert.Equal(t, "0", flag.DefValue)
+	assert.Contains(t, flag.Usage, "concurrent")
+}