fix: sync stub references for each build group (#2039)

This PR modifies the build engine to generate stubs for all known
schemas at build time. After a build group is built, we then sync the
stub references to help users get auto completions and imports in their
IDEs or text editors.

It essentially changes the way we gather schemas for the module(s) we're
building.

In `go` syncing the stub references means that we update the `go.work`
file for each built module to include ALL shared module stubs. This
helps when a module is adding a shared module as a dependency.

buildengine/engine.go

@@ -574,17 +574,18 @@ func (e *Engine) buildWithCallback(ctx context.Context, callback buildCallback,
 	}
 	errCh := make(chan error, 1024)
 	for _, group := range topology {
-		groupSchemas := map[string]*schema.Module{}
-		metas, err := e.gatherGroupSchemas(builtModules, group, groupSchemas)
+		knownSchemas := map[string]*schema.Module{}
+		err := e.gatherSchemas(builtModules, knownSchemas)
 		if err != nil {
 			return err
 		}
 
+		metas := e.allModuleMetas()
 		moduleConfigs := make([]moduleconfig.ModuleConfig, len(metas))
 		for i, meta := range metas {
 			moduleConfigs[i] = meta.module.Config
 		}
-		err = GenerateStubs(ctx, e.projectRoot, maps.Values(groupSchemas), moduleConfigs)
+		err = GenerateStubs(ctx, e.projectRoot, maps.Values(knownSchemas), moduleConfigs)
 		if err != nil {
 			return err
 		}
@@ -616,6 +617,17 @@ func (e *Engine) buildWithCallback(ctx context.Context, callback buildCallback,
 		for sch := range schemas {
 			builtModules[sch.Name] = sch
 		}
+
+		moduleNames := []string{}
+		for _, module := range knownSchemas {
+			moduleNames = append(moduleNames, module.Name)
+		}
+
+		// Sync references to stubs if needed by the runtime
+		err = SyncStubReferences(ctx, e.projectRoot, moduleNames, moduleConfigs)
+		if err != nil {
+			return err
+		}
 	}
 
 	close(errCh)
@@ -679,14 +691,15 @@ func (e *Engine) build(ctx context.Context, moduleName string, builtModules map[
 	}
 
 	combined := map[string]*schema.Module{}
-	if err := e.gatherSchemas(builtModules, meta.module, combined); err != nil {
+	if err := e.gatherSchemas(builtModules, combined); err != nil {
 		return err
 	}
 	sch := &schema.Schema{Modules: maps.Values(combined)}
 
 	if e.listener != nil {
 		e.listener.OnBuildStarted(meta.module)
 	}
+
 	err := Build(ctx, e.projectRoot, sch, meta.module, e.watcher.GetTransaction(meta.module.Config.Dir))
 	if err != nil {
 		return err
@@ -700,54 +713,31 @@ func (e *Engine) build(ctx context.Context, moduleName string, builtModules map[
 	return nil
 }
 
-// Construct a combined schema for a group of modules and their transitive dependencies.
-func (e *Engine) gatherGroupSchemas(
-	moduleSchemas map[string]*schema.Module,
-	group []string,
-	out map[string]*schema.Module,
-) ([]moduleMeta, error) {
-	var metas []moduleMeta
-	for _, module := range group {
-		if module == "builtin" {
-			continue // Skip the builtin module
-		}
-
-		meta, ok := e.moduleMetas.Load(module)
-		if ok {
-			metas = append(metas, meta)
-			if err := e.gatherSchemas(moduleSchemas, meta.module, out); err != nil {
-				return nil, err
-			}
-		}
-	}
-	return metas, nil
+func (e *Engine) allModuleMetas() []moduleMeta {
+	var out []moduleMeta
+	e.moduleMetas.Range(func(name string, meta moduleMeta) bool {
+		out = append(out, meta)
+		return true
+	})
+	return out
 }
 
 // Construct a combined schema for a module and its transitive dependencies.
 func (e *Engine) gatherSchemas(
 	moduleSchemas map[string]*schema.Module,
-	module Module,
 	out map[string]*schema.Module,
 ) error {
-	latestModule, ok := e.moduleMetas.Load(module.Config.Module)
-	if !ok {
-		latestModule = moduleMeta{module: module}
-	}
-	for _, dep := range latestModule.module.Dependencies {
-		if moduleSchemas[dep] != nil {
-			out[dep] = moduleSchemas[dep]
-		}
+	e.controllerSchema.Range(func(name string, sch *schema.Module) bool {
+		out[name] = sch
+		return true
+	})
 
-		if dep != "builtin" {
-			depModule, ok := e.moduleMetas.Load(dep)
-			// TODO: should we be gathering schemas from dependencies without a module?
-			// This can happen if the schema is loaded from the controller
-			if ok {
-				if err := e.gatherSchemas(moduleSchemas, depModule.module, out); err != nil {
-					return err
-				}
-			}
+	e.moduleMetas.Range(func(name string, meta moduleMeta) bool {
+		if _, ok := moduleSchemas[name]; ok {
+			out[name] = moduleSchemas[name]
 		}
-	}
+		return true
+	})
+
 	return nil
 }

buildengine/stubs.go

@@ -21,6 +21,13 @@ func CleanStubs(ctx context.Context, projectRoot string) error {
 	return cleanGoStubs(ctx, projectRoot)
 }
 
+// SyncStubReferences syncs the references in the generated stubs.
+//
+// For Go, this means updating all the go.work files to include all known modules in the shared stubbed modules directory.
+func SyncStubReferences(ctx context.Context, projectRoot string, moduleNames []string, moduleConfigs []moduleconfig.ModuleConfig) error {
+	return syncGoStubReferences(ctx, projectRoot, moduleNames, moduleConfigs)
+}
+
 func generateGoStubs(ctx context.Context, projectRoot string, modules []*schema.Module, moduleConfigs []moduleconfig.ModuleConfig) error {
 	sch := &schema.Schema{Modules: modules}
 	err := compile.GenerateStubsForModules(ctx, projectRoot, moduleConfigs, sch)
@@ -37,3 +44,11 @@ func cleanGoStubs(ctx context.Context, projectRoot string) error {
 	}
 	return nil
 }
+
+func syncGoStubReferences(ctx context.Context, projectRoot string, moduleNames []string, moduleConfigs []moduleconfig.ModuleConfig) error {
+	err := compile.SyncGeneratedStubReferences(ctx, projectRoot, moduleNames, moduleConfigs)
+	if err != nil {
+		fmt.Printf("failed to sync go stub references: %v\n", err)
+	}
+	return nil
+}

go-runtime/compile/build.go

@@ -125,8 +125,8 @@ func Build(ctx context.Context, projectRootDir, moduleDir string, sch *schema.Sc
 	if err != nil {
 		return fmt.Errorf("failed to load module config: %w", err)
 	}
-	logger := log.FromContext(ctx)
 
+	logger := log.FromContext(ctx)
 	funcs := maps.Clone(scaffoldFuncs)
 
 	buildDir := buildDir(moduleDir)
@@ -319,6 +319,33 @@ func GenerateStubsForModules(ctx context.Context, projectRoot string, moduleConf
 	return nil
 }
 
+func SyncGeneratedStubReferences(ctx context.Context, projectRootDir string, stubbedModules []string, moduleConfigs []moduleconfig.ModuleConfig) error {
+	for _, moduleConfig := range moduleConfigs {
+		var sharedModulesPaths []string
+		for _, mod := range stubbedModules {
+			if mod == moduleConfig.Module {
+				continue
+			}
+			sharedModulesPaths = append(sharedModulesPaths, filepath.Join(projectRootDir, buildDirName, "go", "modules", mod))
+		}
+
+		_, goModVersion, err := updateGoModule(filepath.Join(moduleConfig.Dir, "go.mod"))
+		if err != nil {
+			return err
+		}
+
+		funcs := maps.Clone(scaffoldFuncs)
+		if err := internal.ScaffoldZip(mainWorkTemplateFiles(), moduleConfig.Dir, MainWorkContext{
+			GoVersion:          goModVersion,
+			SharedModulesPaths: sharedModulesPaths,
+		}, scaffolder.Exclude("^go.mod$"), scaffolder.Functions(funcs)); err != nil {
+			return fmt.Errorf("failed to scaffold zip: %w", err)
+		}
+	}
+
+	return nil
+}
+
 var scaffoldFuncs = scaffolder.FuncMap{
 	"comment": schema.EncodeComments,
 	"type":    genType,