diff --git a/pkg/cli/logs_orchestrator_download.go b/pkg/cli/logs_orchestrator_download.go index 683a0908cab..62304dec6d3 100644 --- a/pkg/cli/logs_orchestrator_download.go +++ b/pkg/cli/logs_orchestrator_download.go @@ -422,7 +422,7 @@ func appendProcessedWorkflowRuns( if shouldSkipProcessedWorkflowRun(result, opts.verbose) || applyRunFilters(activeCtx, result, opts.filters, opts.verbose) { continue } - processedRun := buildProcessedRun(result, opts.verbose, true) + processedRun := buildProcessedRun(activeCtx, result, opts.verbose, true) parseWorkflowRunArtifacts(result, processedRun, opts.parse, opts.verbose) processedRuns = append(processedRuns, processedRun) batchProcessed++ diff --git a/pkg/cli/logs_orchestrator_filters.go b/pkg/cli/logs_orchestrator_filters.go index 62c285cf030..3fa97adb698 100644 --- a/pkg/cli/logs_orchestrator_filters.go +++ b/pkg/cli/logs_orchestrator_filters.go @@ -149,7 +149,7 @@ func applyRunFilters(ctx context.Context, result DownloadResult, opts runFilterO // buildProcessedRun constructs a ProcessedRun from a DownloadResult, computing // duration, action minutes, effective tokens, and job-failure counts. -func buildProcessedRun(result DownloadResult, verbose, logFailedJobs bool) ProcessedRun { +func buildProcessedRun(ctx context.Context, result DownloadResult, verbose, logFailedJobs bool) ProcessedRun { run := result.Run run.TokenUsage = result.Metrics.TokenUsage applyMetricsTurnsToRun(&run, result.Metrics) @@ -164,7 +164,7 @@ func buildProcessedRun(result DownloadResult, verbose, logFailedJobs bool) Proce } // Add failed jobs to error count. - if failedJobCount, err := fetchJobStatusesForProcessedRun(context.Background(), run.DatabaseID, verbose); err == nil { + if failedJobCount, err := fetchJobStatusesForProcessedRun(ctx, run.DatabaseID, verbose); err == nil { run.ErrorCount += failedJobCount if verbose && logFailedJobs && failedJobCount > 0 { fmt.Fprintln(os.Stderr, console.FormatInfoMessage(fmt.Sprintf("Added %d failed jobs to error count for run %d", failedJobCount, run.DatabaseID))) diff --git a/pkg/cli/logs_orchestrator_filters_test.go b/pkg/cli/logs_orchestrator_filters_test.go index 285eee2e652..42c5b941b1f 100644 --- a/pkg/cli/logs_orchestrator_filters_test.go +++ b/pkg/cli/logs_orchestrator_filters_test.go @@ -225,7 +225,7 @@ func TestBuildProcessedRun(t *testing.T) { AwContext: awCtx, } - pr := buildProcessedRun(result, false, false) + pr := buildProcessedRun(context.Background(), result, false, false) assert.Equal(t, int64(1234), pr.Run.DatabaseID) assert.Equal(t, tmpDir, pr.Run.LogsPath) @@ -246,7 +246,7 @@ func TestBuildProcessedRun(t *testing.T) { LogsPath: t.TempDir(), } - pr := buildProcessedRun(result, false, false) + pr := buildProcessedRun(context.Background(), result, false, false) assert.Equal(t, 90*time.Second, pr.Run.Duration) assert.InDelta(t, 2.0, pr.Run.ActionMinutes, 0.001) // ceil(1.5) = 2 @@ -258,7 +258,7 @@ func TestBuildProcessedRun(t *testing.T) { Run: WorkflowRun{DatabaseID: 7}, LogsPath: t.TempDir(), } - pr := buildProcessedRun(result, false, false) + pr := buildProcessedRun(context.Background(), result, false, false) assert.Equal(t, time.Duration(0), pr.Run.Duration) assert.InDelta(t, 0.0, pr.Run.ActionMinutes, 0.001) }) @@ -271,7 +271,7 @@ func TestBuildProcessedRun(t *testing.T) { LogsPath: t.TempDir(), TokenUsage: usage, } - pr := buildProcessedRun(result, false, false) + pr := buildProcessedRun(context.Background(), result, false, false) assert.Equal(t, 5000, pr.Run.EffectiveTokens) }) @@ -283,14 +283,19 @@ func TestBuildProcessedRun(t *testing.T) { LogsPath: t.TempDir(), TokenUsage: usage, } - pr := buildProcessedRun(result, false, false) + pr := buildProcessedRun(context.Background(), result, false, false) assert.Equal(t, 0, pr.Run.EffectiveTokens) }) t.Run("failed job count is added via test seam", func(t *testing.T) { - stubFetchJobStatusesForProcessedRun(t, func(_ context.Context, runID int64, verbose bool) (int, error) { + type ctxKey string + const key ctxKey = "request-id" + ctx := context.WithValue(context.Background(), key, "abc123") + + stubFetchJobStatusesForProcessedRun(t, func(fetchCtx context.Context, runID int64, verbose bool) (int, error) { assert.Equal(t, int64(88), runID) assert.False(t, verbose) + assert.Equal(t, "abc123", fetchCtx.Value(key)) return 2, nil }) result := DownloadResult{ @@ -298,7 +303,7 @@ func TestBuildProcessedRun(t *testing.T) { LogsPath: t.TempDir(), } - pr := buildProcessedRun(result, false, false) + pr := buildProcessedRun(ctx, result, false, false) assert.Equal(t, 2, pr.Run.ErrorCount) }) diff --git a/pkg/cli/logs_orchestrator_stdin.go b/pkg/cli/logs_orchestrator_stdin.go index 9c907502a53..a9877019d0c 100644 --- a/pkg/cli/logs_orchestrator_stdin.go +++ b/pkg/cli/logs_orchestrator_stdin.go @@ -178,7 +178,7 @@ func DownloadWorkflowLogsFromStdin(ctx context.Context, opts StdinLogsOptions) e continue } - processedRun := buildProcessedRun(result, opts.Verbose, false) + processedRun := buildProcessedRun(ctx, result, opts.Verbose, false) if opts.Parse { awInfoPath := filepath.Join(result.LogsPath, "aw_info.json") diff --git a/pkg/cli/update_extension_check.go b/pkg/cli/update_extension_check.go index 3725a75d514..cbb522d8922 100644 --- a/pkg/cli/update_extension_check.go +++ b/pkg/cli/update_extension_check.go @@ -46,7 +46,7 @@ const backupCleanupRetryDelay = 300 * time.Millisecond // baked in. The caller should re-launch the freshly-installed binary (at // installPath) so that subsequent work (e.g. lock-file compilation) uses the // correct new version string. -func upgradeExtensionIfOutdated(verbose bool, includePrereleases bool) (bool, string, error) { +func upgradeExtensionIfOutdated(ctx context.Context, verbose bool, includePrereleases bool) (bool, string, error) { currentVersion := GetVersion() updateExtensionCheckLog.Printf("Checking if extension needs upgrade (current: %s)", currentVersion) @@ -60,7 +60,7 @@ func upgradeExtensionIfOutdated(verbose bool, includePrereleases bool) (bool, st } // Query GitHub API for latest release - latestVersion, err := getLatestRelease(context.Background(), includePrereleases) + latestVersion, err := getLatestRelease(ctx, includePrereleases) if err != nil { // Fail silently - don't block the upgrade command if we can't reach GitHub updateExtensionCheckLog.Printf("Failed to check for latest release (silently ignoring): %v", err) diff --git a/pkg/cli/update_extension_check_test.go b/pkg/cli/update_extension_check_test.go index 250893672fb..45f3cb9fd11 100644 --- a/pkg/cli/update_extension_check_test.go +++ b/pkg/cli/update_extension_check_test.go @@ -4,6 +4,7 @@ package cli import ( "bytes" + "context" "errors" "os" "path/filepath" @@ -27,7 +28,7 @@ func TestUpgradeExtensionIfOutdated_DevBuild(t *testing.T) { // Verify the function exits before making any API calls. // If it did make API calls we'd see a network error in test environments, // but the function must return (false, "", nil) immediately. - upgraded, installPath, err := upgradeExtensionIfOutdated(false, false) + upgraded, installPath, err := upgradeExtensionIfOutdated(context.Background(), false, false) require.NoError(t, err, "Should not return error for dev builds") assert.False(t, upgraded, "Should not report upgrade for dev builds") assert.Empty(t, installPath, "installPath should be empty for dev builds") @@ -44,7 +45,7 @@ func TestUpgradeExtensionIfOutdated_SilentFailureOnAPIError(t *testing.T) { // Use a release version so the API call is attempted SetVersionInfo("v0.1.0") - upgraded, installPath, err := upgradeExtensionIfOutdated(false, false) + upgraded, installPath, err := upgradeExtensionIfOutdated(context.Background(), false, false) require.NoError(t, err, "Should fail silently on API errors") assert.False(t, upgraded, "Should not report upgrade when API is unreachable") assert.Empty(t, installPath, "installPath should be empty when API is unreachable") diff --git a/pkg/cli/upgrade_command.go b/pkg/cli/upgrade_command.go index bb9d968f4be..e3093cfc9b7 100644 --- a/pkg/cli/upgrade_command.go +++ b/pkg/cli/upgrade_command.go @@ -242,7 +242,7 @@ func runUpgradeCommand(opts upgradeOptions) error { // prevents the re-launched process from entering this branch again. if !opts.skipExtensionUpgrade { fmt.Fprintln(os.Stderr, console.FormatInfoMessage("Checking gh-aw extension version...")) - upgraded, installPath, err := upgradeExtensionIfOutdated(opts.verbose, opts.preReleases) + upgraded, installPath, err := upgradeExtensionIfOutdated(opts.ctx, opts.verbose, opts.preReleases) if err != nil { upgradeLog.Printf("Extension upgrade failed: %v", err) return err