Fix platform-specific image resolution paths

Author: cursoragent

cmd/test-prewarm/main.go

@@ -157,7 +157,11 @@ func ensureMirroredImage(ctx context.Context, inspector *images.OCIClient, regis
 		return manifestImage{}, err
 	}
 
-	if digest, err := inspector.InspectManifest(ctx, localRef); err == nil {
+	inspectDigest := func() (string, error) {
+		return inspector.InspectManifestForPlatform(ctx, localRef, img.Platform)
+	}
+
+	if digest, err := inspectDigest(); err == nil {
 		return manifestImage{Source: img.Source, LocalRef: localRef, Digest: digest, CacheHit: true}, nil
 	}
 
@@ -169,7 +173,7 @@ func ensureMirroredImage(ctx context.Context, inspector *images.OCIClient, regis
 		return manifestImage{}, err
 	}
 
-	digest, err := inspector.InspectManifest(ctx, localRef)
+	digest, err := inspectDigest()
 	if err != nil {
 		digest = res.Digest
 	}

lib/images/oci_public.go

@@ -3,6 +3,7 @@ package images
 import (
 	"context"
 	"fmt"
+	"strings"
 )
 
 // OCIClient is a public wrapper for system manager to use OCI operations
@@ -25,6 +26,19 @@ func (c *OCIClient) InspectManifest(ctx context.Context, imageRef string) (strin
 	return c.client.inspectManifest(ctx, imageRef)
 }
 
+// InspectManifestForPlatform inspects a remote image to get its digest for a
+// specific platform variant (or host platform when empty).
+func (c *OCIClient) InspectManifestForPlatform(ctx context.Context, imageRef, platform string) (string, error) {
+	if strings.TrimSpace(platform) == "" {
+		return c.client.inspectManifest(ctx, imageRef)
+	}
+	resolvedPlatform, err := ParsePlatform(platform)
+	if err != nil {
+		return "", err
+	}
+	return c.client.inspectManifestWithPlatform(ctx, imageRef, resolvedPlatform.ToGCR())
+}
+
 // InspectManifestForLinux is an alias for InspectManifest (all images target Linux)
 func (c *OCIClient) InspectManifestForLinux(ctx context.Context, imageRef string) (string, error) {
 	return c.InspectManifest(ctx, imageRef)

lib/instances/create.go

@@ -95,14 +95,24 @@ func (m *manager) createInstance(
 	imageCtx, imageSpanEnd := m.startLifecycleStep(ctx, "resolve_image",
 		attribute.String("operation", "resolve_image"),
 	)
+	waitNameForDigest := func(digest string) string {
+		if strings.TrimSpace(digest) == "" {
+			return req.Image
+		}
+		normalizedRef, parseErr := images.ParseNormalizedRef(req.Image)
+		if parseErr != nil {
+			return req.Image
+		}
+		return normalizedRef.Repository() + "@" + digest
+	}
 	imageInfo, err := m.imageManager.GetImage(imageCtx, req.Image)
 	if err != nil {
 		if err == images.ErrNotFound {
 			// Auto-pull: image not found locally, kick off the pull in the
 			// background and wait up to 5 seconds for it to complete. Thread the
 			// requested platform so boot-by-tag pulls the right architecture.
 			log.InfoContext(ctx, "image not found locally, auto-pulling", "image", req.Image, "platform", req.Platform)
-			_, pullErr := m.imageManager.CreateImage(imageCtx, images.CreateImageRequest{Name: req.Image, Platform: req.Platform})
+			pulledImage, pullErr := m.imageManager.CreateImage(imageCtx, images.CreateImageRequest{Name: req.Image, Platform: req.Platform})
 			if pullErr != nil {
 				imageSpanEnd(pullErr)
 				log.ErrorContext(ctx, "failed to auto-pull image", "image", req.Image, "error", pullErr)
@@ -112,7 +122,11 @@ func (m *manager) createInstance(
 			// we return an error but let it continue in the background.
 			pullCtx, pullCancel := context.WithTimeout(imageCtx, 5*time.Second)
 			defer pullCancel()
-			if waitErr := m.imageManager.WaitForReady(pullCtx, req.Image); waitErr != nil {
+			waitName := req.Image
+			if pulledImage != nil {
+				waitName = waitNameForDigest(pulledImage.Digest)
+			}
+			if waitErr := m.imageManager.WaitForReady(pullCtx, waitName); waitErr != nil {
 				imageSpanEnd(waitErr)
 				log.InfoContext(ctx, "image pull not ready within timeout, pull continues in background", "image", req.Image, "error", waitErr)
 				return nil, fmt.Errorf("%w: image %s is being pulled, please try again shortly", ErrImageNotReady, req.Image)
@@ -130,6 +144,52 @@ func (m *manager) createInstance(
 			return nil, fmt.Errorf("get image: %w", err)
 		}
 	}
+	if strings.TrimSpace(req.Platform) != "" {
+		requestedPlatform, parseErr := images.ParsePlatform(req.Platform)
+		if parseErr != nil {
+			imageSpanEnd(parseErr)
+			log.ErrorContext(ctx, "invalid requested image platform", "image", req.Image, "platform", req.Platform, "error", parseErr)
+			return nil, parseErr
+		}
+		resolvedPlatform, parseErr := images.ParsePlatform(imageInfo.Platform)
+		platformMismatch := parseErr != nil || resolvedPlatform.String() != requestedPlatform.String()
+		if platformMismatch {
+			log.InfoContext(ctx,
+				"resolved image platform differs from request, reconciling",
+				"image", req.Image,
+				"requested_platform", requestedPlatform.String(),
+				"resolved_platform", imageInfo.Platform,
+			)
+			pulledImage, pullErr := m.imageManager.CreateImage(imageCtx, images.CreateImageRequest{Name: req.Image, Platform: requestedPlatform.String()})
+			if pullErr != nil {
+				imageSpanEnd(pullErr)
+				log.ErrorContext(ctx, "failed to resolve requested image platform", "image", req.Image, "platform", requestedPlatform.String(), "error", pullErr)
+				return nil, fmt.Errorf("resolve image for platform %s: %w", requestedPlatform, pullErr)
+			}
+			pullCtx, pullCancel := context.WithTimeout(imageCtx, 5*time.Second)
+			defer pullCancel()
+			waitName := req.Image
+			if pulledImage != nil {
+				waitName = waitNameForDigest(pulledImage.Digest)
+			}
+			if waitErr := m.imageManager.WaitForReady(pullCtx, waitName); waitErr != nil {
+				imageSpanEnd(waitErr)
+				log.InfoContext(ctx,
+					"platform-specific image pull not ready within timeout, pull continues in background",
+					"image", req.Image,
+					"platform", requestedPlatform.String(),
+					"error", waitErr,
+				)
+				return nil, fmt.Errorf("%w: image %s for platform %s is being pulled, please try again shortly", ErrImageNotReady, req.Image, requestedPlatform)
+			}
+			imageInfo, err = m.imageManager.GetImage(imageCtx, req.Image)
+			if err != nil {
+				imageSpanEnd(err)
+				log.ErrorContext(ctx, "failed to get image after platform reconciliation", "image", req.Image, "platform", requestedPlatform.String(), "error", err)
+				return nil, fmt.Errorf("get image after platform reconciliation: %w", err)
+			}
+		}
+	}
 	imageSpanEnd(nil)
 
 	if imageInfo.Status != images.StatusReady {