fix: sync stub references for each build group

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,