fix: prune stale container pins and move UpdateContainerPins to after compilation

pkg/cli/update_command.go

@@ -175,13 +175,6 @@ func RunUpdateWorkflows(ctx context.Context, opts UpdateWorkflowsOptions) error
 		fmt.Fprintln(os.Stderr, console.FormatWarningMessage(fmt.Sprintf("Warning: Failed to update actions-lock.json: %v", err)))
 	}
 
-	// Resolve and store SHA-256 digest pins for container images referenced in lock files.
-	updateLog.Print("Updating container image digest pins")
-	if err := UpdateContainerPins(ctx, opts.WorkflowsDir, opts.Verbose); err != nil {
-		// Non-fatal: Docker may not be available in all environments.
-		fmt.Fprintln(os.Stderr, console.FormatWarningMessage(fmt.Sprintf("Warning: Failed to update container pins: %v", err)))
-	}
-
 	// Update action references in user-provided steps within workflow .md files.
 	// By default all org/repo@version references are updated to the latest major version.
 	updateLog.Print("Updating action references in workflow .md files")
@@ -190,6 +183,16 @@ func RunUpdateWorkflows(ctx context.Context, opts UpdateWorkflowsOptions) error
 		fmt.Fprintln(os.Stderr, console.FormatWarningMessage(fmt.Sprintf("Warning: Failed to update action references in workflow files: %v", err)))
 	}
 
+	// Resolve and store SHA-256 digest pins for container images referenced in lock files.
+	// This runs after compilation (via UpdateActionsInWorkflowFiles) so that the lock files
+	// already reflect the current AWF version; stale pins from superseded versions are pruned
+	// and new versions are resolved in a single pass.
+	updateLog.Print("Updating container image digest pins")
+	if err := UpdateContainerPins(ctx, opts.WorkflowsDir, opts.Verbose); err != nil {
+		// Non-fatal: Docker may not be available in all environments.
+		fmt.Fprintln(os.Stderr, console.FormatWarningMessage(fmt.Sprintf("Warning: Failed to update container pins: %v", err)))
+	}
+
 	updateLog.Printf("Update process complete: had_error=%v", firstErr != nil)
 	return firstErr
 }

pkg/cli/update_container_pins.go

@@ -88,6 +88,21 @@ func UpdateContainerPins(ctx context.Context, workflowDir string, verbose bool)
 		}
 	}
 
+	// Build a set of images currently referenced in the compiled lock files so
+	// that stale entries (e.g. superseded AWF versions) can be pruned.
+	imageSet := make(map[string]bool, len(images))
+	for _, img := range images {
+		imageSet[img] = true
+	}
+
+	// Remove any container pin entries that are no longer referenced by the
+	// compiled lock files.  This keeps actions-lock.json consistent with what
+	// compile actually emits and prevents stale version accumulation.
+	prunedCount := actionCache.PruneStaleContainerPins(imageSet)
+	if prunedCount > 0 {
+		containerPinsLog.Printf("Pruned %d stale container pin(s) from actions-lock.json", prunedCount)
+	}
+
 	// Resolve digests for images that are not yet pinned.
 	type pinnedEntry struct {
 		image       string
@@ -141,6 +156,11 @@ func UpdateContainerPins(ctx context.Context, workflowDir string, verbose bool)
 		fmt.Fprintln(os.Stderr, "")
 	}
 
+	if prunedCount > 0 {
+		fmt.Fprintln(os.Stderr, console.FormatInfoMessage(fmt.Sprintf("Pruned %d stale container pin(s) from actions-lock.json", prunedCount)))
+		fmt.Fprintln(os.Stderr, "")
+	}
+
 	if len(skippedImages) > 0 && verbose {
 		fmt.Fprintln(os.Stderr, console.FormatInfoMessage(fmt.Sprintf("%d container image(s) already up to date", len(skippedImages))))
 		fmt.Fprintln(os.Stderr, "")
@@ -154,7 +174,7 @@ func UpdateContainerPins(ctx context.Context, workflowDir string, verbose bool)
 		fmt.Fprintln(os.Stderr, "")
 	}
 
-	if len(updatedImages) > 0 {
+	if len(updatedImages) > 0 || prunedCount > 0 {
 		if err := actionCache.Save(); err != nil {
 			return fmt.Errorf("failed to save actions-lock.json: %w", err)
 		}

pkg/cli/update_container_pins_test.go

@@ -7,6 +7,7 @@ import (
 	"path/filepath"
 	"testing"
 
+	"github.com/github/gh-aw/pkg/workflow"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
@@ -168,3 +169,75 @@ Manifests:
 		})
 	}
 }
+
+// TestUpdateContainerPins_PrunesStaleEntries verifies that UpdateContainerPins
+// removes container pin entries from actions-lock.json that are no longer
+// referenced in the compiled lock files (e.g. superseded AWF versions).
+func TestUpdateContainerPins_PrunesStaleEntries(t *testing.T) {
+	// Create a temp directory acting as the repo root.
+	tmpDir := t.TempDir()
+
+	// Write an actions-lock.json with a stale container pin and a live one.
+	// The live pin (0.27.2) should be kept; the stale one (0.27.0) should be pruned.
+	actionsLockDir := filepath.Join(tmpDir, ".github", "aw")
+	require.NoError(t, os.MkdirAll(actionsLockDir, 0755))
+	actionsLockContent := `{
+  "entries": {},
+  "containers": {
+    "ghcr.io/github/gh-aw-firewall/agent:0.27.0": {
+      "image": "ghcr.io/github/gh-aw-firewall/agent:0.27.0",
+      "digest": "sha256:olddigest0000000000000000000000000000000000000000000000000000000",
+      "pinned_image": "ghcr.io/github/gh-aw-firewall/agent:0.27.0@sha256:olddigest0000000000000000000000000000000000000000000000000000000"
+    },
+    "ghcr.io/github/gh-aw-firewall/agent:0.27.2": {
+      "image": "ghcr.io/github/gh-aw-firewall/agent:0.27.2",
+      "digest": "sha256:newdigest0000000000000000000000000000000000000000000000000000000",
+      "pinned_image": "ghcr.io/github/gh-aw-firewall/agent:0.27.2@sha256:newdigest0000000000000000000000000000000000000000000000000000000"
+    }
+  }
+}
+`
+	actionsLockPath := filepath.Join(actionsLockDir, "actions-lock.json")
+	require.NoError(t, os.WriteFile(actionsLockPath, []byte(actionsLockContent), 0644))
+
+	// Write a lock file referencing the NEW AWF version (0.27.2), not the old one.
+	workflowsDir := filepath.Join(tmpDir, ".github", "workflows")
+	require.NoError(t, os.MkdirAll(workflowsDir, 0755))
+	lockFileContent := `name: test
+jobs:
+  setup:
+    steps:
+      - name: Download container images
+        run: bash "${RUNNER_TEMP}/gh-aw/actions/download_docker_images.sh" ghcr.io/github/gh-aw-firewall/agent:0.27.2
+`
+	require.NoError(t, os.WriteFile(filepath.Join(workflowsDir, "my-workflow.lock.yml"), []byte(lockFileContent), 0644))
+
+	// collectImagesFromLockFiles should find the new version only.
+	images, err := collectImagesFromLockFiles(workflowsDir)
+	require.NoError(t, err)
+	assert.Equal(t, []string{"ghcr.io/github/gh-aw-firewall/agent:0.27.2"}, images)
+
+	// Load the cache and exercise PruneStaleContainerPins directly (Docker is not
+	// available in unit tests so we can't call the full UpdateContainerPins function).
+	cache := workflow.NewActionCache(tmpDir)
+	require.NoError(t, cache.Load())
+
+	imageSet := map[string]bool{"ghcr.io/github/gh-aw-firewall/agent:0.27.2": true}
+	pruned := cache.PruneStaleContainerPins(imageSet)
+	assert.Equal(t, 1, pruned, "stale 0.27.0 entry should be pruned")
+
+	_, ok := cache.GetContainerPin("ghcr.io/github/gh-aw-firewall/agent:0.27.0")
+	assert.False(t, ok, "old-version pin should not be in cache after prune")
+
+	pin, ok := cache.GetContainerPin("ghcr.io/github/gh-aw-firewall/agent:0.27.2")
+	require.True(t, ok, "current-version pin should still be in cache")
+	assert.Equal(t, "sha256:newdigest0000000000000000000000000000000000000000000000000000000", pin.Digest)
+
+	// Save and verify the stale entry is gone from disk.
+	require.NoError(t, cache.Save())
+
+	data, err := os.ReadFile(actionsLockPath)
+	require.NoError(t, err)
+	assert.NotContains(t, string(data), "0.27.0", "stale version should not appear in saved file")
+	assert.Contains(t, string(data), "0.27.2", "current version should be in saved file")
+}

pkg/cli/upgrade_command.go

@@ -282,21 +282,6 @@ func runUpgradeCommand(opts upgradeOptions) error {
 		}
 	}
 
-	// Step 3b: Update container image digest pins (unless --no-fix or --no-actions is specified)
-	// Container pins are stored alongside action pins in .github/aw/actions-lock.json.
-	// Running this before compilation means the next compile step will embed the
-	// pinned @sha256: references in the generated lock files.
-	if !opts.noFix && !opts.noActions {
-		upgradeLog.Print("Updating container image digest pins")
-		if err := UpdateContainerPins(opts.ctx, opts.workflowDir, opts.verbose); err != nil {
-			upgradeLog.Printf("Failed to update container pins: %v", err)
-			// Non-critical — Docker may not be available in all environments.
-			fmt.Fprintln(os.Stderr, console.FormatWarningMessage(fmt.Sprintf("Warning: Failed to update container pins: %v", err)))
-		} else if opts.verbose {
-			fmt.Fprintln(os.Stderr, console.FormatSuccessMessage("✓ Updated container image pins"))
-		}
-	}
-
 	// Step 4: Compile all workflows (unless --no-fix or --no-compile is specified)
 	if !opts.noFix && !opts.noCompile {
 		fmt.Fprintln(os.Stderr, console.FormatInfoMessage("Compiling all workflows..."))
@@ -344,6 +329,24 @@ func runUpgradeCommand(opts upgradeOptions) error {
 		}
 	}
 
+	// Step 4b: Update container image digest pins (unless --no-fix or --no-actions is specified)
+	// Container pins are stored alongside action pins in .github/aw/actions-lock.json.
+	// This runs AFTER compilation so that the compiled lock files already reflect the
+	// current AWF version; stale pins from superseded versions are pruned and new
+	// versions are resolved in a single pass.  When --no-compile is set, the existing
+	// lock files are used as-is — pins are still pruned and refreshed against whatever
+	// lock files are currently on disk.
+	if !opts.noFix && !opts.noActions {
+		upgradeLog.Print("Updating container image digest pins")
+		if err := UpdateContainerPins(opts.ctx, opts.workflowDir, opts.verbose); err != nil {
+			upgradeLog.Printf("Failed to update container pins: %v", err)
+			// Non-critical — Docker may not be available in all environments.
+			fmt.Fprintln(os.Stderr, console.FormatWarningMessage(fmt.Sprintf("Warning: Failed to update container pins: %v", err)))
+		} else if opts.verbose {
+			fmt.Fprintln(os.Stderr, console.FormatSuccessMessage("✓ Updated container image pins"))
+		}
+	}
+
 	// Print success message
 	fmt.Fprintln(os.Stderr, "")
 	fmt.Fprintln(os.Stderr, console.FormatSuccessMessage("Upgrade complete"))

pkg/workflow/action_cache.go

@@ -96,6 +96,26 @@ func (c *ActionCache) DeleteContainerPin(image string) {
 	}
 }
 
+// PruneStaleContainerPins removes container pin entries whose keys are not present
+// in knownImages. It returns the number of entries that were removed.
+// This is used to keep actions-lock.json consistent with the set of images
+// actually referenced by the compiled lock files.
+func (c *ActionCache) PruneStaleContainerPins(knownImages map[string]bool) int {
+	if c.ContainerPins == nil {
+		return 0
+	}
+	pruned := 0
+	for image := range c.ContainerPins {
+		if !knownImages[image] {
+			delete(c.ContainerPins, image)
+			c.dirty = true
+			pruned++
+			actionCacheLog.Printf("Pruned stale container pin for image=%s", image)
+		}
+	}
+	return pruned
+}
+
 // Load loads the cache from disk
 func (c *ActionCache) Load() error {
 	actionCacheLog.Printf("Loading action cache from: %s", c.path)