feat: set up interruptible generation with progress updates

internal/log/logListener.go

@@ -12,4 +12,5 @@ var (
 	MsgWarn   MsgType = "warn"
 	MsgError  MsgType = "error"
 	MsgGithub MsgType = "github"
+	MsgStudio MsgType = "studio"
 )

internal/run/target.go

@@ -153,28 +153,36 @@ func (w *Workflow) runTarget(ctx context.Context, target string) (*SourceResult,
 	ctx = log.With(ctx, logger)
 	go genStep.ListenForSubsteps(logListener)
 
+	if w.GenerationProgress != nil {
+		cancelCtx, cancelFunc := context.WithCancel(ctx)
+		w.GenerationProgress.CancellableContext = cancelCtx
+		w.GenerationProgress.CancelGeneration = cancelFunc
+		w.GenerationProgress.LogListener = logListener
+	}
+
 	generationAccess, err := sdkgen.Generate(
 		ctx,
 		sdkgen.GenerateOptions{
-			CustomerID:      config.GetCustomerID(),
-			WorkspaceID:     config.GetWorkspaceID(),
-			Language:        t.Target,
-			SchemaPath:      sourcePath,
-			Header:          "",
-			Token:           "",
-			OutDir:          outDir,
-			CLIVersion:      events.GetSpeakeasyVersionFromContext(ctx),
-			InstallationURL: w.InstallationURLs[target],
-			Debug:           w.Debug,
-			AutoYes:         true,
-			Published:       published,
-			OutputTests:     false,
-			Repo:            w.Repo,
-			RepoSubDir:      w.RepoSubDirs[target],
-			Verbose:         w.Verbose,
-			Compile:         w.ShouldCompile,
-			TargetName:      target,
-			SkipVersioning:  w.SkipVersioning,
+			CustomerID:         config.GetCustomerID(),
+			WorkspaceID:        config.GetWorkspaceID(),
+			Language:           t.Target,
+			SchemaPath:         sourcePath,
+			Header:             "",
+			Token:              "",
+			OutDir:             outDir,
+			CLIVersion:         events.GetSpeakeasyVersionFromContext(ctx),
+			InstallationURL:    w.InstallationURLs[target],
+			Debug:              w.Debug,
+			AutoYes:            true,
+			Published:          published,
+			OutputTests:        false,
+			Repo:               w.Repo,
+			RepoSubDir:         w.RepoSubDirs[target],
+			Verbose:            w.Verbose,
+			Compile:            w.ShouldCompile,
+			TargetName:         target,
+			SkipVersioning:     w.SkipVersioning,
+			GenerationProgress: w.GenerationProgress,
 		},
 	)
 

internal/run/workflow.go

@@ -60,6 +60,9 @@ type Workflow struct {
 	Duration        time.Duration
 	criticalWarns   []string
 	Error           error
+
+	// Studio
+	GenerationProgress *sdkgen.GenerationProgress
 }
 
 type Opt func(w *Workflow)
@@ -88,21 +91,22 @@ func NewWorkflow(
 
 	// Default values
 	w := &Workflow{
-		workflowName:     "Workflow",
-		RepoSubDirs:      make(map[string]string),
-		InstallationURLs: make(map[string]string),
-		SDKOverviewURLs:  make(map[string]string),
-		Debug:            false,
-		ShouldCompile:    true,
-		workflow:         *wf,
-		ProjectDir:       projectDir,
-		ForceGeneration:  false,
-		SourceResults:    make(map[string]*SourceResult),
-		TargetResults:    make(map[string]*TargetResult),
-		OnSourceResult:   func(*SourceResult, string) {},
-		computedChanges:  make(map[string]bool),
-		lockfile:         lockfile,
-		lockfileOld:      lockfileOld,
+		workflowName:       "Workflow",
+		RepoSubDirs:        make(map[string]string),
+		InstallationURLs:   make(map[string]string),
+		SDKOverviewURLs:    make(map[string]string),
+		Debug:              false,
+		ShouldCompile:      true,
+		workflow:           *wf,
+		ProjectDir:         projectDir,
+		ForceGeneration:    false,
+		SourceResults:      make(map[string]*SourceResult),
+		TargetResults:      make(map[string]*TargetResult),
+		OnSourceResult:     func(*SourceResult, string) {},
+		computedChanges:    make(map[string]bool),
+		lockfile:           lockfile,
+		lockfileOld:        lockfileOld,
+		GenerationProgress: nil,
 	}
 
 	for _, opt := range opts {
@@ -252,6 +256,16 @@ func WithRegistryTags(registryTags []string) Opt {
 	}
 }
 
+// func WithGenerationProgress(onProgressUpdate func(sdkgen.ProgressUpdate), cancelGeneration context.CancelFunc) Opt {
+func WithGenerationProgress(onProgressUpdate func(sdkgen.ProgressUpdate)) Opt {
+	return func(w *Workflow) {
+		generationProgress := sdkgen.GenerationProgress{
+			OnProgressUpdate: onProgressUpdate,
+		}
+		w.GenerationProgress = &generationProgress
+	}
+}
+
 func (w *Workflow) CountDiagnostics() int {
 	count := 0
 	for _, sourceResult := range w.SourceResults {

internal/sdkgen/sdkgen.go

@@ -33,6 +33,18 @@
 	Level         *shared.Level
 }
 
+type ProgressUpdate = generate.ProgressUpdate
+type ProgressStep = generate.ProgressStep
+
+type GenerationProgress struct {
+	OnProgressUpdate   func(ProgressUpdate)
+	CancellableContext context.Context
+	CancelGeneration   context.CancelFunc
+	LogListener        chan log.Msg
+	SkipDiffs          bool
+	SkipSteps          bool
+}
+
 type GenerateOptions struct {
 	CustomerID      string
 	WorkspaceID     string
@@ -53,6 +65,10 @@
 	Compile         bool
 	TargetName      string
 	SkipVersioning  bool
+
+	// if GenerationProgress is not nil, the generation will be cancellable
+	// and progress updates may be sent to the OnProgressUpdate callback
+	GenerationProgress *GenerationProgress
 }
 
 func Generate(ctx context.Context, opts GenerateOptions) (*GenerationAccess, error) {
@@ -146,6 +162,19 @@
 		generatorOpts = append(generatorOpts, generate.WithSkipVersioning(opts.SkipVersioning))
 	}
 
+	var progressChan chan generate.ProgressUpdate
+	if opts.GenerationProgress != nil {
+		progressChan = make(chan generate.ProgressUpdate)
+		generatorOpts = append(
+			generatorOpts,
+			generate.WithProgressUpdates(
+				progressChan,
+				opts.GenerationProgress.SkipDiffs,
+				opts.GenerationProgress.SkipSteps,
+			),
+		)
+	}
+
 	g, err := generate.New(generatorOpts...)
 	if err != nil {
 		return &GenerationAccess{
@@ -157,15 +186,68 @@
 
 	err = events.Telemetry(ctx, shared.InteractionTypeTargetGenerate, func(ctx context.Context, event *shared.CliEvent) error {
 		event.GenerateTargetName = &opts.TargetName
-		if errs := g.Generate(ctx, schema, opts.SchemaPath, opts.Language, opts.OutDir, isRemote, opts.Compile); len(errs) > 0 {
-			for _, err := range errs {
+
+		var genErrs []error
+		genAborted := false
+
+		if opts.GenerationProgress != nil && opts.GenerationProgress.CancellableContext != nil {
+			genErrsChan := make(chan []error)
+			go func(cancellableContext context.Context) {
+				defer close(progressChan)
+				errs := g.Generate(cancellableContext, schema, opts.SchemaPath, opts.Language, opts.OutDir, isRemote, opts.Compile)
+				var nonHaltErrs []error
+				for _, err := range errs {
+					if haltErr, ok := generate.IsRunScriptHaltError(err); ok {
+						progressChan <- ProgressUpdate{
+							ID: "",
+							Step: &ProgressStep{
+								ID:      "cancel",
+								Message: fmt.Sprintf("Generation halted at %s", haltErr.Location),
+							},
+						}
+					} else {
+						nonHaltErrs = append(nonHaltErrs, err)
+					}
+				}
+
+				genErrsChan <- nonHaltErrs
+			}(opts.GenerationProgress.CancellableContext)
+
+		receiveProgressUpdates:
+			for {
+				select {
+				case <-opts.GenerationProgress.CancellableContext.Done():
+					genAborted = true
+					// waiting for the generator to exit
+				case progressUpdate := <-progressChan:
+					if opts.GenerationProgress.OnProgressUpdate != nil {
+						opts.GenerationProgress.OnProgressUpdate(progressUpdate)
+					}
+
+				case genErrs = <-genErrsChan:
+					close(genErrsChan)
+					break receiveProgressUpdates
+				}
+			}
+		} else {
+			genErrs = g.Generate(ctx, schema, opts.SchemaPath, opts.Language, opts.OutDir, isRemote, opts.Compile)
+		}
+
+		if len(genErrs) > 0 {
+			for _, err := range genErrs {
 				logger.Error("", zap.Error(err))
 			}
 
 			return fmt.Errorf("failed to generate SDKs for %s ✖", opts.Language)
 		}
+
+		if genAborted {
+			logger.Warn("Generation was aborted!")
+		}
+
 		return nil
 	})
+
 	if err != nil {
 		return &GenerationAccess{
 			AccessAllowed: generationAccess,

internal/studio/launchStudio.go

@@ -12,6 +12,7 @@ import (
 	"os"
 	"os/signal"
 	"slices"
+	"strconv"
 	"syscall"
 	"time"
 
@@ -64,10 +65,12 @@ func LaunchStudio(ctx context.Context, workflow *run.Workflow) error {
 	mux := http.NewServeMux()
 	mux.HandleFunc("/", handler(handlers.root))
 	mux.HandleFunc("/health", handler(handlers.health))
+	mux.HandleFunc("/cancel", handler(handlers.cancel))
 
 	mux.HandleFunc("/run", func(w http.ResponseWriter, r *http.Request) {
 		switch r.Method {
 		case http.MethodPost:
+			handlers.reRunOptions = parseReRunOptions(r)
 			handler(handlers.reRun)(w, r)
 		case http.MethodGet:
 			handler(handlers.getLastRunResult)(w, r)
@@ -102,7 +105,7 @@ func LaunchStudio(ctx context.Context, workflow *run.Workflow) error {
 		fmt.Println(listeningMessage+"Opening URL in your browser: ", handlers.StudioURL)
 	}
 
-	// After ten seconds, if the health check hasn't been seen then kill the server
+	// After 1 minute, if the health check hasn't been seen then kill the server
 	go func() {
 		time.Sleep(1 * time.Minute)
 		if !handlers.healthCheckSeen {
@@ -117,6 +120,29 @@ func LaunchStudio(ctx context.Context, workflow *run.Workflow) error {
 	return startServer(ctx, server, workflow)
 }
 
+func parseReRunOptions(r *http.Request) reRunOptions {
+	queryParams := r.URL.Query()
+	options := reRunOptions{}
+
+	compileVal := queryParams.Get("compile")
+	if compileVal != "" {
+		compile, err := strconv.ParseBool(compileVal)
+		if err != nil {
+			options.skipCompile = !compile
+		}
+	}
+
+	streamVal := queryParams.Get("stream")
+	if streamVal != "" {
+		stream, err := strconv.ParseBool(streamVal)
+		if err != nil {
+			options.skipProgress = !stream
+		}
+	}
+
+	return options
+}
+
 func corsMiddleware(next http.Handler) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		w.Header().Set("Access-Control-Allow-Origin", "*")
@@ -148,20 +174,23 @@ func authMiddleware(secret string, next http.Handler) http.Handler {
 
 func handler(h func(context.Context, http.ResponseWriter, *http.Request) error) http.HandlerFunc {
 	return func(w http.ResponseWriter, r *http.Request) {
+		ctx := r.Context()
+		logger := log.From(ctx)
+
 		start := time.Now()
 		id := generateRequestID()
 		method := fmt.Sprintf("%-6s", r.Method)  // Fixed width 6 characters
 		path := fmt.Sprintf("%-21s", r.URL.Path) // Fixed width 21 characters
 		base := fmt.Sprintf("%s %s %s", id, method, path)
-		log.From(r.Context()).Info(fmt.Sprintf("%s started", base))
-		ctx := r.Context()
+
+		logger.Info(fmt.Sprintf("%s started", base))
 		if err := h(ctx, w, r); err != nil {
 			log.From(ctx).Error(fmt.Sprintf("%s failed: %v", base, err))
 			respondJSONError(ctx, w, err)
 			return
 		}
 		duration := time.Since(start)
-		log.From(ctx).Info(fmt.Sprintf("%s completed in %s", base, duration))
+		logger.Info(fmt.Sprintf("%s completed in %s", base, duration))
 	}
 }