From c1c9955edfc3f0633b95872ef5d18063e1dee659 Mon Sep 17 00:00:00 2001
From: Adam Bull <adamjamesbull@googlemail.com>
Date: Wed, 28 Aug 2024 14:51:11 +0100
Subject: [PATCH] fix: update studio handlers to bootstrap overflow path from
 source  (#891)

---
 cmd/run.go                        | 21 +++++++++++++++++----
 internal/studio/studioHandlers.go | 25 ++++++++++++-------------
 2 files changed, 29 insertions(+), 17 deletions(-)

diff --git a/cmd/run.go b/cmd/run.go
index f4474a65..d9dacc69 100644
--- a/cmd/run.go
+++ b/cmd/run.go
@@ -59,7 +59,7 @@ var runCmd = &model.ExecutableCommand[RunFlags]{
 	Short:            "Run all the workflows defined in your workflow.yaml file. This can include multiple SDK generations from different OpenAPI sources",
 	Long:             utils.RenderMarkdown(runLong),
 	PreRun:           preRun,
-	Run:              runFunc,
+	Run:              runNonInteractive,
 	RunInteractive:   runInteractive,
 	RequiresAuth:     true,
 	UsesWorkflowFile: true,
@@ -284,13 +284,12 @@ func askForSource(sources []string) (string, error) {
 	return source, nil
 }
 
-func runFunc(ctx context.Context, flags RunFlags) error {
+func runNonInteractive(ctx context.Context, flags RunFlags) error {
 	if flags.GitHub {
 		return run.RunGitHub(ctx, flags.Target, flags.SetVersion, flags.Force)
 	}
 
-	workflow, err := run.NewWorkflow(
-		ctx,
+	opts := []run.Opt{
 		run.WithTarget(flags.Target),
 		run.WithSource(flags.Source),
 		run.WithRepo(flags.Repo),
@@ -304,7 +303,17 @@ func runFunc(ctx context.Context, flags RunFlags) error {
 		run.WithRegistryTags(flags.RegistryTags),
 		run.WithSetVersion(flags.SetVersion),
 		run.WithFrozenWorkflowLock(flags.FrozenWorkflowLock),
+	}
+
+	if flags.LaunchStudio {
+		opts = append(opts, run.WithSkipCleanup())
+	}
+
+	workflow, err := run.NewWorkflow(
+		ctx,
+		opts...,
 	)
+
 	if err != nil {
 		return err
 	}
@@ -315,6 +324,10 @@ func runFunc(ctx context.Context, flags RunFlags) error {
 
 	github.GenerateWorkflowSummary(ctx, workflow.RootStep)
 
+	if flags.LaunchStudio {
+		return studio.LaunchStudio(ctx, workflow)
+	}
+
 	return err
 }
 
diff --git a/internal/studio/studioHandlers.go b/internal/studio/studioHandlers.go
index aac913d7..8575ec1a 100644
--- a/internal/studio/studioHandlers.go
+++ b/internal/studio/studioHandlers.go
@@ -33,10 +33,10 @@ type StudioHandlers struct {
 	Ctx            context.Context
 }
 
-func NewStudioHandlers(ctx context.Context, workflow *run.Workflow) (*StudioHandlers, error) {
-	ret := &StudioHandlers{WorkflowRunner: *workflow, Ctx: ctx}
+func NewStudioHandlers(ctx context.Context, workflowRunner *run.Workflow) (*StudioHandlers, error) {
+	ret := &StudioHandlers{WorkflowRunner: *workflowRunner, Ctx: ctx}
 
-	sourceID, err := findWorkflowSourceIDBasedOnTarget(*workflow, workflow.Target)
+	sourceID, err := findWorkflowSourceIDBasedOnTarget(*workflowRunner, workflowRunner.Target)
 	if err != nil {
 		return ret, fmt.Errorf("error finding source: %w", err)
 	}
@@ -44,6 +44,15 @@ func NewStudioHandlers(ctx context.Context, workflow *run.Workflow) (*StudioHand
 		return ret, errors.New("unable to find source")
 	}
 	ret.SourceID = sourceID
+	sourceConfig := workflowRunner.GetWorkflowFile().Sources[sourceID]
+
+	for _, overlay := range sourceConfig.Overlays {
+		// If there are multiple modifications overlays - we take the last one
+		contents, _ := isStudioModificationsOverlay(overlay)
+		if contents != "" {
+			ret.OverlayPath = overlay.Document.Location
+		}
+	}
 
 	return ret, nil
 }
@@ -263,16 +272,6 @@ func (h *StudioHandlers) getSourceInner(ctx context.Context) (*run.SourceResult,
 		return nil, fmt.Errorf("error running source: %w", err)
 	}
 
-	sourceConfig := workflowRunner.GetWorkflowFile().Sources[sourceID]
-
-	for _, overlay := range sourceConfig.Overlays {
-		// If there are multiple modifications overlays - we take the last one
-		contents, _ := isStudioModificationsOverlay(overlay)
-		if contents != "" {
-			h.OverlayPath = overlay.Document.Location
-		}
-	}
-
 	return sourceResult, nil
 }