feat: set up interruptible generation in Studio (#1319)

GEN-1145

*WARNING* This PR relies on [openapi-generation
PR#2367](speakeasy-api/openapi-generation#2367)
being merged first and should be merged alongside [speakeasy-registry
PR#4134](speakeasy-api/speakeasy-registry#4134)

This PR sets up the ability to:
- Cancel generation before it is complete
- Receive progress udpates from the generator (new steps, file status
changes, main README content)
- Update the Studio UI based on received updates

---
**1. Cancelability**

To control cancelability, a `CancellableGeneration` field was added to
the `Workflow` struct to hold the *cancellable context* and the *cancel
function*.
When this field is non-nil, `g.GenerateWithCancel` will be used instead
of `g.Generate` (c.f. open-generation PR#2367)

---
**2. Streamability**

To control streamability, a `StreamableGeneneration` field was added to
the `Workflow` struct to hold the `OnProgressUpdate` callback function
that will get
called everytime the generator sends an update:
- if a new step has started then `progressUpdate.Step` will be non-nil
- if a file status has changed then the `progressUpdate.File` will
non-nil
- a `progressUpdate.File.IsMainReadme` boolean is also set to indicate
that the file being modified is the Readme.
If it is `true`, then the `progressUpdate.File.Content` will also be set

---
**3. Studio events**

At the moment, the Studio UI is being notified when:
- the `genSDK` and `compileSDK` steps started
- the main `README` content got received

Author: 2ynn

internal/log/logListener.go

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

internal/run/source.go

@@ -27,6 +27,26 @@ import (
 	"github.com/speakeasy-api/speakeasy/internal/workflowTracking"
 )
 
+type SourceResultCallback func(sourceRes *SourceResult, sourceStep SourceStepID) error
+
+type SourceStepID string
+
+const (
+	// CLI steps
+	SourceStepFetch     SourceStepID = "Fetching spec"
+	SourceStepOverlay   SourceStepID = "Overlaying"
+	SourceStepTransform SourceStepID = "Transforming"
+	SourceStepLint      SourceStepID = "Linting"
+	SourceStepUpload    SourceStepID = "Uploading spec"
+	// Generator steps
+	SourceStepStart    SourceStepID = "Started"
+	SourceStepGenerate SourceStepID = "Generating SDK"
+	SourceStepCompile  SourceStepID = "Compiling SDK"
+	SourceStepComplete SourceStepID = "Completed"
+	SourceStepCancel   SourceStepID = "Cancelling"
+	SourceStepExit     SourceStepID = "Exiting"
+)
+
 type SourceResult struct {
 	Source string
 	// The merged OAS spec that was input to the source contents as a string
@@ -55,7 +75,7 @@ func (e *LintingError) Error() string {
 	return fmt.Sprintf("linting failed: %s - %s", e.Document, errString)
 }
 
-func (w *Workflow) RunSource(ctx context.Context, parentStep *workflowTracking.WorkflowStep, sourceID, targetID string, targetLanguage string) (string, *SourceResult, error) {
+func (w *Workflow) RunSource(ctx context.Context, parentStep *workflowTracking.WorkflowStep, sourceID, targetID, targetLanguage string) (string, *SourceResult, error) {
 	rootStep := parentStep.NewSubstep(fmt.Sprintf("Source: %s", sourceID))
 	source := w.workflow.Sources[sourceID]
 	sourceRes := &SourceResult{
@@ -64,9 +84,9 @@ func (w *Workflow) RunSource(ctx context.Context, parentStep *workflowTracking.W
 	}
 	defer func() {
 		w.SourceResults[sourceID] = sourceRes
-		w.OnSourceResult(sourceRes, "")
+		w.OnSourceResult(sourceRes, SourceStepComplete)
 	}()
-	w.OnSourceResult(sourceRes, "Fetching spec")
+	w.OnSourceResult(sourceRes, SourceStepFetch)
 
 	rulesetToUse := "speakeasy-generation"
 	if source.Ruleset != nil {
@@ -122,29 +142,29 @@ func (w *Workflow) RunSource(ctx context.Context, parentStep *workflowTracking.W
 	}
 
 	if len(source.Overlays) > 0 && !w.FrozenWorkflowLock {
-		w.OnSourceResult(sourceRes, "Overlaying")
+		w.OnSourceResult(sourceRes, SourceStepOverlay)
 		currentDocument, err = NewOverlay(rootStep, source).Do(ctx, currentDocument)
 		if err != nil {
 			return "", nil, err
 		}
 	}
 
 	if len(source.Transformations) > 0 && !w.FrozenWorkflowLock {
-		w.OnSourceResult(sourceRes, "Transforming")
+		w.OnSourceResult(sourceRes, SourceStepTransform)
 		currentDocument, err = NewTransform(rootStep, source).Do(ctx, currentDocument)
 		if err != nil {
 			return "", nil, err
 		}
 	}
 
 	if err := writeToOutputLocation(ctx, currentDocument, outputLocation); err != nil {
-		return "", nil, fmt.Errorf("failed to write to output location: %w", err)
+		return "", nil, fmt.Errorf("failed to write to output location: %w %s?", err, outputLocation)
 	}
 	currentDocument = outputLocation
 	sourceRes.OutputPath = currentDocument
 
 	if !w.SkipLinting {
-		w.OnSourceResult(sourceRes, "Linting")
+		w.OnSourceResult(sourceRes, SourceStepLint)
 		sourceRes.LintResult, err = w.validateDocument(ctx, rootStep, sourceID, currentDocument, rulesetToUse, w.ProjectDir, targetLanguage)
 		if err != nil {
 			return "", sourceRes, &LintingError{Err: err, Document: currentDocument}
@@ -159,7 +179,7 @@ func (w *Workflow) RunSource(ctx context.Context, parentStep *workflowTracking.W
 	}
 	step.Succeed()
 
-	w.OnSourceResult(sourceRes, "Uploading spec")
+	w.OnSourceResult(sourceRes, SourceStepUpload)
 
 	if !w.SkipSnapshot {
 		err = w.snapshotSource(ctx, rootStep, sourceID, source, currentDocument)

internal/run/target.go

@@ -147,35 +147,57 @@ func (w *Workflow) runTarget(ctx context.Context, target string) (*SourceResult,
 		}
 	}
 
-	genStep := rootStep.NewSubstep(fmt.Sprintf("Generating %s SDK", utils.CapitalizeFirst(t.Target)))
-
 	logListener := make(chan log.Msg)
 	logger := log.From(ctx).WithListener(logListener)
 	ctx = log.With(ctx, logger)
+
+	if w.StreamableGeneration != nil && w.Debug {
+		w.StreamableGeneration.LogListener = logListener
+	}
+
+	if w.CancellableGeneration != nil {
+		cancelCtx, cancelFunc := context.WithCancel(ctx)
+		w.CancellableGeneration.CancellationMutex.Lock()
+		w.CancellableGeneration.CancellableContext = cancelCtx
+		w.CancellableGeneration.CancelGeneration = cancelFunc
+		w.CancellableGeneration.CancellationMutex.Unlock()
+
+		defer func() {
+			w.CancellableGeneration.CancellationMutex.Lock()
+			w.CancellableGeneration.CancelGeneration = nil
+			w.CancellableGeneration.CancellableContext = nil
+			w.CancellableGeneration.CancellationMutex.Unlock()
+			cancelFunc() // Ensure context is cleaned up
+		}()
+	}
+
+	genStep := rootStep.NewSubstep(fmt.Sprintf("Generating %s SDK", utils.CapitalizeFirst(t.Target)))
 	go genStep.ListenForSubsteps(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,
+			CancellableGeneration: w.CancellableGeneration,
+			StreamableGeneration:  w.StreamableGeneration,
 		},
 	)
 
@@ -388,3 +410,16 @@ func (w *Workflow) printTargetSuccessMessage(ctx context.Context) {
 	msg := fmt.Sprintf("%s\n%s\n", styles.Success.Render(heading), strings.Join(additionalLines, "\n"))
 	log.From(ctx).Println(msg)
 }
+
+func (w *Workflow) CancelGeneration() error {
+	if w.CancellableGeneration != nil {
+		w.CancellableGeneration.CancellationMutex.Lock()
+		defer w.CancellableGeneration.CancellationMutex.Unlock()
+		if w.CancellableGeneration.CancelGeneration != nil {
+			w.CancellableGeneration.CancelGeneration()
+			return nil
+		}
+	}
+
+	return fmt.Errorf("Generation is not cancellable")
+}

internal/run/workflow.go

@@ -3,10 +3,12 @@ package run
 import (
 	"context"
 	"fmt"
+	"sync"
 	"time"
 
 	"github.com/speakeasy-api/speakeasy/registry"
 
+	"github.com/speakeasy-api/openapi-generation/v2/pkg/generate"
 	"github.com/speakeasy-api/sdk-gen-config/workflow"
 	"github.com/speakeasy-api/speakeasy-core/events"
 	"github.com/speakeasy-api/speakeasy/internal/log"
@@ -57,10 +59,14 @@ type Workflow struct {
 	computedChanges map[string]bool
 	SourceResults   map[string]*SourceResult
 	TargetResults   map[string]*TargetResult
-	OnSourceResult  func(*SourceResult, string)
+	OnSourceResult  SourceResultCallback
 	Duration        time.Duration
 	criticalWarns   []string
 	Error           error
+
+	// Studio
+	CancellableGeneration *sdkgen.CancellableGeneration
+	StreamableGeneration  *sdkgen.StreamableGeneration
 }
 
 type Opt func(w *Workflow)
@@ -100,7 +106,7 @@ func NewWorkflow(
 		ForceGeneration:  false,
 		SourceResults:    make(map[string]*SourceResult),
 		TargetResults:    make(map[string]*TargetResult),
-		OnSourceResult:   func(*SourceResult, string) {},
+		OnSourceResult:   func(*SourceResult, SourceStepID) error { return nil },
 		computedChanges:  make(map[string]bool),
 		lockfile:         lockfile,
 		lockfileOld:      lockfileOld,
@@ -262,6 +268,42 @@ func WithRegistryTags(registryTags []string) Opt {
 	}
 }
 
+func WithSourceUpdates(onSourceResult SourceResultCallback) Opt {
+	if onSourceResult != nil {
+		return func(w *Workflow) {
+			w.OnSourceResult = onSourceResult
+		}
+	}
+
+	return func(w *Workflow) {
+		w.OnSourceResult = func(sourceRes *SourceResult, sourceStep SourceStepID) error { return nil }
+	}
+}
+
+func WithCancellableGeneration(cancellable bool) Opt {
+	return func(w *Workflow) {
+		if cancellable {
+			w.CancellableGeneration = &sdkgen.CancellableGeneration{
+				CancellationMutex: sync.Mutex{},
+				// CancelCtx and CancelFunc fields depend on the runTarget context
+				// and will be set right before generation starts.
+			}
+		} else {
+			w.CancellableGeneration = nil
+		}
+	}
+}
+
+func WithStreamableGeneration(onProgressUpdate func(generate.ProgressUpdate), genSteps, fileStatus bool) Opt {
+	return func(w *Workflow) {
+		w.StreamableGeneration = &sdkgen.StreamableGeneration{
+			OnProgressUpdate: onProgressUpdate,
+			GenSteps:         genSteps,
+			FileStatus:       fileStatus,
+		}
+	}
+}
+
 func (w *Workflow) CountDiagnostics() int {
 	count := 0
 	for _, sourceResult := range w.SourceResults {

internal/sdkgen/sdkgen.go

@@ -7,6 +7,7 @@ import (
 	"os"
 	"path/filepath"
 	"strings"
+	"sync"
 
 	"github.com/charmbracelet/lipgloss"
 	"github.com/speakeasy-api/speakeasy-core/auth"
@@ -33,6 +34,19 @@ type GenerationAccess struct {
 	Level         *shared.Level
 }
 
+type CancellableGeneration struct {
+	CancellationMutex  sync.Mutex         // protects both CancellableContext and CancelGeneration (exposed by w.CancelGeneration())
+	CancellableContext context.Context    // the context that can be cancelled to stop generation
+	CancelGeneration   context.CancelFunc // the function to call to cancel generation
+}
+
+type StreamableGeneration struct {
+	OnProgressUpdate func(generate.ProgressUpdate) // the callback function called on each progress update
+	GenSteps         bool                          // whether to receive an update before each generation step starts
+	FileStatus       bool                          // whether to receive updates on each file status change
+	LogListener      chan log.Msg                  // the channel to listen for log messages (Debug only)
+}
+
 type GenerateOptions struct {
 	CustomerID      string
 	WorkspaceID     string
@@ -53,6 +67,9 @@ type GenerateOptions struct {
 	Compile         bool
 	TargetName      string
 	SkipVersioning  bool
+
+	CancellableGeneration *CancellableGeneration
+	StreamableGeneration  *StreamableGeneration
 }
 
 func Generate(ctx context.Context, opts GenerateOptions) (*GenerationAccess, error) {
@@ -146,6 +163,18 @@ func Generate(ctx context.Context, opts GenerateOptions) (*GenerationAccess, err
 		generatorOpts = append(generatorOpts, generate.WithSkipVersioning(opts.SkipVersioning))
 	}
 
+	if opts.StreamableGeneration != nil {
+		generatorOpts = append(
+			generatorOpts,
+			generate.WithProgressUpdates(
+				opts.TargetName,
+				opts.StreamableGeneration.OnProgressUpdate,
+				opts.StreamableGeneration.GenSteps,
+				opts.StreamableGeneration.FileStatus,
+			),
+		)
+	}
+
 	g, err := generate.New(generatorOpts...)
 	if err != nil {
 		return &GenerationAccess{
@@ -157,13 +186,28 @@ func Generate(ctx context.Context, opts GenerateOptions) (*GenerationAccess, err
 
 	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 {
+
+		var errs []error
+		if opts.CancellableGeneration != nil && opts.CancellableGeneration.CancellableContext != nil {
+			cancelCtx := opts.CancellableGeneration.CancellableContext
+
+			var cancelled bool
+			cancelled, errs = g.GenerateWithCancel(cancelCtx, schema, opts.SchemaPath, opts.Language, opts.OutDir, isRemote, opts.Compile)
+			if cancelled {
+				return fmt.Errorf("Generation was aborted for %s ✖", opts.Language)
+			}
+		} else {
+			errs = g.Generate(ctx, schema, opts.SchemaPath, opts.Language, opts.OutDir, isRemote, opts.Compile)
+		}
+
+		if len(errs) > 0 {
 			for _, err := range errs {
 				logger.Error("", zap.Error(err))
 			}
 
 			return fmt.Errorf("failed to generate SDKs for %s ✖", opts.Language)
 		}
+
 		return nil
 	})
 	if err != nil {