Add support for pulling draft sources (#2395)

* Add support for pulling draft sources

* Draft sources can now be pulled using
draft:user/digest
* Extracts tarball logic into lower function, only
Significant change is the URL processing

* Remove console log

Author: 8W9aG

pkg/api/client.go

@@ -23,8 +23,12 @@ import (
 	"github.com/replicate/cog/pkg/web"
 )
 
+const DraftsPrefix = "draft:"
+
 var (
 	ErrorBadResponseNewVersionEndpoint = errors.New("Bad response from new version endpoint")
+	ErrorBadDraftFormat                = errors.New("Bad draft format")
+	ErrorBadDraftUsernameDigestFormat  = errors.New("Bad draft username/digest format")
 )
 
 type Client struct {
@@ -65,6 +69,14 @@ func (c *Client) PostNewPipeline(ctx context.Context, image string, tarball *byt
 }
 
 func (c *Client) PullSource(ctx context.Context, image string, tarFileProcess func(*tar.Header, *tar.Reader) error) error {
+	if strings.HasPrefix(image, DraftsPrefix) {
+		username, digest, err := decomposeDraftSlug(image)
+		if err != nil {
+			return err
+		}
+		return c.getDraftSource(ctx, username, digest, tarFileProcess)
+	}
+
 	_, entity, name, tag, err := decomposeImageName(image)
 	if err != nil {
 		return err
@@ -228,7 +240,21 @@ func (c *Client) getSource(ctx context.Context, entity string, name string, tag
 	}
 
 	sourceURL := newSourceURL(entity, name, tag)
-	req, err := http.NewRequestWithContext(ctx, http.MethodGet, sourceURL.String(), nil)
+	return c.downloadTarball(ctx, token, sourceURL, strings.Join([]string{entity, name}, "/"), tarFileProcess)
+}
+
+func (c *Client) getDraftSource(ctx context.Context, username string, digest string, tarFileProcess func(*tar.Header, *tar.Reader) error) error {
+	token, err := c.provideToken(ctx, username)
+	if err != nil {
+		return err
+	}
+
+	draftURL := newDraftSourceURL(digest)
+	return c.downloadTarball(ctx, token, draftURL, DraftsPrefix+strings.Join([]string{username, digest}, "/"), tarFileProcess)
+}
+
+func (c *Client) downloadTarball(ctx context.Context, token string, url url.URL, slug string, tarFileProcess func(*tar.Header, *tar.Reader) error) error {
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, url.String(), nil)
 	if err != nil {
 		return err
 	}
@@ -242,7 +268,7 @@ func (c *Client) getSource(ctx context.Context, entity string, name string, tag
 	defer resp.Body.Close()
 
 	if resp.StatusCode == http.StatusNotFound {
-		return fmt.Errorf("Model %s/%s does not have a source package associated with it.", entity, name)
+		return fmt.Errorf("Entity %s does not have a source package associated with it.", slug)
 	}
 
 	if resp.StatusCode >= 400 {
@@ -332,6 +358,12 @@ func newModelURL(entity string, name string) url.URL {
 	return newModelUrl
 }
 
+func newDraftSourceURL(digest string) url.URL {
+	newDraftSourceUrl := apiBaseURL()
+	newDraftSourceUrl.Path = strings.Join([]string{"", "v1", "drafts", digest, "source"}, "/")
+	return newDraftSourceUrl
+}
+
 func decomposeImageName(image string) (string, string, string, string, error) {
 	imageComponents := strings.Split(image, "/")
 
@@ -353,3 +385,17 @@ func decomposeImageName(image string) (string, string, string, string, error) {
 	}
 	return imageComponents[0], imageComponents[1], imageComponents[2], tag, nil
 }
+
+func decomposeDraftSlug(slug string) (string, string, error) {
+	slugComponents := strings.Split(slug, ":")
+	if len(slugComponents) != 2 {
+		return "", "", ErrorBadDraftFormat
+	}
+
+	draftComponents := strings.Split(slugComponents[1], "/")
+	if len(draftComponents) != 2 {
+		return "", "", ErrorBadDraftUsernameDigestFormat
+	}
+
+	return draftComponents[0], draftComponents[1], nil
+}

pkg/api/client_test.go

@@ -170,3 +170,84 @@ func TestPullSource(t *testing.T) {
 	})
 	require.NoError(t, err)
 }
+
+func TestPullDraftSource(t *testing.T) {
+	// Create file to pull
+	dir := t.TempDir()
+	predictPyPath := filepath.Join(dir, "predict.py")
+	handle, err := os.Create(predictPyPath)
+	require.NoError(t, err)
+	handle.WriteString("import cog")
+	err = handle.Close()
+	require.NoError(t, err)
+	info, err := os.Stat(predictPyPath)
+	require.NoError(t, err)
+
+	// Setup mock web server for cog.replicate.com (token exchange)
+	webServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		switch r.URL.Path {
+		case "/api/token/user":
+			// Mock token exchange response
+			//nolint:gosec
+			tokenResponse := `{
+				"keys": {
+					"cog": {
+						"key": "test-api-token",
+						"expires_at": "2024-12-31T23:59:59Z"
+					}
+				}
+			}`
+			w.WriteHeader(http.StatusOK)
+			w.Write([]byte(tokenResponse))
+		default:
+			w.WriteHeader(http.StatusNotFound)
+		}
+	}))
+	defer webServer.Close()
+
+	// Setup mock API server for api.replicate.com (model and source endpoints)
+	apiServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		switch r.URL.Path {
+		case "/v1/drafts/digest/source":
+			// Mock draft source pull endpoint
+			var buf bytes.Buffer
+			tw := tar.NewWriter(&buf)
+			header, err := tar.FileInfoHeader(info, info.Name())
+			require.NoError(t, err)
+			header.Name = "predict.py"
+			err = tw.WriteHeader(header)
+			require.NoError(t, err)
+			file, err := os.Open(predictPyPath)
+			require.NoError(t, err)
+			defer file.Close()
+			_, err = io.Copy(tw, file)
+			require.NoError(t, err)
+			err = tw.Close()
+			require.NoError(t, err)
+			w.WriteHeader(http.StatusOK)
+			w.Write(buf.Bytes())
+		default:
+			w.WriteHeader(http.StatusNotFound)
+		}
+	}))
+	defer apiServer.Close()
+
+	webURL, err := url.Parse(webServer.URL)
+	require.NoError(t, err)
+	apiURL, err := url.Parse(apiServer.URL)
+	require.NoError(t, err)
+
+	t.Setenv(env.SchemeEnvVarName, webURL.Scheme)
+	t.Setenv(env.WebHostEnvVarName, webURL.Host)
+	t.Setenv(env.APIHostEnvVarName, apiURL.Host)
+
+	// Setup mock command
+	command := dockertest.NewMockCommand()
+	webClient := web.NewClient(command, http.DefaultClient)
+
+	client := NewClient(command, http.DefaultClient, webClient)
+	err = client.PullSource(t.Context(), "draft:user/digest", func(header *tar.Header, tr *tar.Reader) error {
+		return nil
+	})
+	require.NoError(t, err)
+}