From 91a43e2f93a7218aa6cc4b99ccabdacd9f10f9ed Mon Sep 17 00:00:00 2001
From: Jacob Gillespie <jacobwgillespie@gmail.com>
Date: Fri, 24 Jan 2025 14:49:33 +0000
Subject: [PATCH] Fix handling of default build create path

---
 pkg/build/build.go   | 39 +++++++++++++++++++++++----------------
 pkg/helpers/build.go |  2 +-
 2 files changed, 24 insertions(+), 17 deletions(-)

diff --git a/pkg/build/build.go b/pkg/build/build.go
index 97f216f4..894a54d0 100644
--- a/pkg/build/build.go
+++ b/pkg/build/build.go
@@ -99,18 +99,15 @@ func NewBuild(ctx context.Context, req *cliv1.CreateBuildRequest, token string)
 		return Build{}, err
 	}
 
-	build, err := FromExistingBuild(ctx, res.Msg.BuildId, res.Msg.BuildToken)
+	build, err := FromExistingBuild(ctx, res.Msg.BuildId, res.Msg.BuildToken, res)
 	if err != nil {
 		return Build{}, err
 	}
 
-	build.Response = res
-	build.BuildURL = res.Msg.BuildUrl
-
 	return build, nil
 }
 
-func FromExistingBuild(ctx context.Context, buildID, token string) (Build, error) {
+func FromExistingBuild(ctx context.Context, buildID, token string, buildRes *connect.Response[cliv1.CreateBuildResponse]) (Build, error) {
 	client := depotapi.NewBuildClient()
 
 	finish := func(buildErr error) {
@@ -135,19 +132,29 @@ func FromExistingBuild(ctx context.Context, buildID, token string) (Build, error
 		}
 	}
 
-	req := cliv1.GetBuildRequest{BuildId: buildID}
-	res, err := client.GetBuild(ctx, depotapi.WithAuthentication(connect.NewRequest(&req), token))
-	if err != nil {
-		return Build{}, err
+	if buildRes == nil {
+		req := cliv1.GetBuildRequest{BuildId: buildID}
+		res, err := client.GetBuild(ctx, depotapi.WithAuthentication(connect.NewRequest(&req), token))
+		if err != nil {
+			return Build{}, err
+		}
+		return Build{
+			ID:        buildID,
+			Token:     token,
+			Finish:    finish,
+			BuildURL:  res.Msg.BuildUrl,
+			projectID: res.Msg.ProjectId,
+		}, nil
+	} else {
+		return Build{
+			ID:       buildID,
+			Token:    token,
+			Finish:   finish,
+			BuildURL: buildRes.Msg.BuildUrl,
+			Response: buildRes,
+		}, nil
 	}
 
-	return Build{
-		ID:        buildID,
-		Token:     token,
-		Finish:    finish,
-		BuildURL:  res.Msg.BuildUrl,
-		projectID: res.Msg.ProjectId,
-	}, nil
 }
 
 type authProvider struct {
diff --git a/pkg/helpers/build.go b/pkg/helpers/build.go
index 29d7e862..eb66970f 100644
--- a/pkg/helpers/build.go
+++ b/pkg/helpers/build.go
@@ -15,7 +15,7 @@ func BeginBuild(ctx context.Context, req *cliv1.CreateBuildRequest, token string
 	var build depotbuild.Build
 	var err error
 	if id := os.Getenv("DEPOT_BUILD_ID"); id != "" {
-		build, err = depotbuild.FromExistingBuild(ctx, id, token)
+		build, err = depotbuild.FromExistingBuild(ctx, id, token, nil)
 	} else {
 		build, err = depotbuild.NewBuild(ctx, req, token)
 	}