feat(pack): add optional OCI subject linkage

Author: rishi-jat

pkg/cmd/pack/pack.go

@@ -18,21 +18,27 @@ package pack
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"io"
 	"os"
+	"strings"
 
 	"github.com/kitops-ml/kitops/pkg/artifact"
+	cmdoptions "github.com/kitops-ml/kitops/pkg/cmd/options"
 	"github.com/kitops-ml/kitops/pkg/lib/constants"
 	"github.com/kitops-ml/kitops/pkg/lib/constants/mediatype"
 	"github.com/kitops-ml/kitops/pkg/lib/filesystem"
 	"github.com/kitops-ml/kitops/pkg/lib/filesystem/ignore"
 	kfutils "github.com/kitops-ml/kitops/pkg/lib/kitfile"
 	"github.com/kitops-ml/kitops/pkg/lib/repo/local"
+	"github.com/kitops-ml/kitops/pkg/lib/repo/remote"
 	"github.com/kitops-ml/kitops/pkg/lib/repo/util"
 	"github.com/kitops-ml/kitops/pkg/output"
 
+	"github.com/opencontainers/go-digest"
 	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
+	"oras.land/oras-go/v2/errdef"
 )
 
 // runPack compresses and stores a modelkit based on a Kitfile. Returns an error if packing
@@ -79,13 +85,21 @@ func runPack(ctx context.Context, options *packOptions) error {
 
 func pack(ctx context.Context, opts *packOptions, kitfile *artifact.KitFile, localRepo local.LocalRepo) (*ocispec.Descriptor, error) {
 	var extraLayerPaths []string
+	var subject *ocispec.Descriptor
 	if kitfile.Model != nil && artifact.IsModelKitReference(kitfile.Model.Path) {
 		baseRef := artifact.FormatRepositoryForDisplay(opts.modelRef.String())
 		parentKitfile, err := kfutils.ResolveKitfile(ctx, opts.configHome, kitfile.Model.Path, baseRef)
 		if err != nil {
 			return nil, fmt.Errorf("Failed to resolve referenced modelkit %s: %w", kitfile.Model.Path, err)
 		}
 		extraLayerPaths = util.LayerPathsFromKitfile(parentKitfile)
+
+		baseDesc, err := resolveBaseManifestDescriptor(ctx, opts.configHome, kitfile.Model.Path)
+		if err != nil {
+			output.Logf(output.LogLevelWarn, "failed to resolve base manifest descriptor for OCI subject: %v", err)
+		} else if baseDesc != nil {
+			subject = baseDesc
+		}
 	}
 
 	ignore, err := ignore.NewFromContext(opts.contextDir, kitfile, extraLayerPaths...)
@@ -105,13 +119,61 @@ func pack(ctx context.Context, opts *packOptions, kitfile *artifact.KitFile, loc
 		ModelFormat: modelFormat,
 		Compression: compression,
 		LayerFormat: mediatype.TarFormat,
+		Subject:     subject,
 	})
 	if err != nil {
 		return nil, err
 	}
 	return manifestDesc, nil
 }
 
+func resolveBaseManifestDescriptor(ctx context.Context, configHome, baseRef string) (*ocispec.Descriptor, error) {
+	if strings.HasPrefix(baseRef, "sha256:") {
+		parsedDigest, err := digest.Parse(baseRef)
+		if err != nil {
+			return nil, fmt.Errorf("failed to parse digest-only base reference: %w", err)
+		}
+		return &ocispec.Descriptor{MediaType: ocispec.MediaTypeImageManifest, Digest: parsedDigest}, nil
+	}
+
+	ref, _, err := artifact.ParseReference(baseRef)
+	if err != nil {
+		return nil, fmt.Errorf("failed to parse base reference: %w", err)
+	}
+
+	localRepo, err := local.NewLocalRepo(constants.StoragePath(configHome), ref)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create local repository: %w", err)
+	}
+
+	desc, _, _, err := util.ResolveManifestAndConfig(ctx, localRepo, ref.Reference)
+	if err == nil {
+		return &desc, nil
+	}
+	if errors.Is(err, util.ErrNoKitfile) || errors.Is(err, util.ErrNotAModelKit) {
+		return nil, nil
+	}
+	if !errors.Is(err, errdef.ErrNotFound) && !errors.Is(err, errdef.ErrMissingReference) {
+		return nil, fmt.Errorf("failed to resolve base descriptor locally: %w", err)
+	}
+
+	repository, err := remote.NewRepository(ctx, ref.Registry, ref.Repository, cmdoptions.DefaultNetworkOptions(configHome))
+	if err != nil {
+		return nil, fmt.Errorf("failed to create remote repository: %w", err)
+	}
+
+	desc, _, _, err = util.ResolveManifestAndConfig(ctx, repository, ref.Reference)
+	if err != nil {
+		if errors.Is(err, util.ErrNoKitfile) || errors.Is(err, util.ErrNotAModelKit) {
+			return nil, nil
+		}
+		output.Logf(output.LogLevelWarn,
+			"failed to resolve base descriptor, skipping subject linkage: %v", err)
+		return nil, nil
+	}
+	return &desc, nil
+}
+
 func readKitfile(modelFile string) (*artifact.KitFile, error) {
 	// 1. Read the model file
 	kitfile := &artifact.KitFile{}

pkg/lib/filesystem/local-storage.go

@@ -43,6 +43,7 @@ type SaveModelOptions struct {
 	ModelFormat mediatype.ModelFormat
 	Compression mediatype.CompressionType
 	LayerFormat mediatype.Format
+	Subject     *ocispec.Descriptor
 }
 
 // SaveModel saves an *artifact.Model to the provided oras.Target, compressing layers. It attempts to block
@@ -63,6 +64,9 @@ func SaveModel(ctx context.Context, localRepo local.LocalRepo, kitfile *artifact
 	if err != nil {
 		return nil, fmt.Errorf("error creating manifest: %w", err)
 	}
+	if opts.Subject != nil {
+		manifest.Subject = opts.Subject
+	}
 
 	// If not storing a Kitfile, save the Kitfile to an annotation since it will be lost otherwise
 	if opts.ModelFormat == mediatype.ModelPackFormat {