fix(protogen): map Python dependency imports

Python files generated for one module import dependency schemas through their
Go module paths. The dependency wheels install those schemas under
module-relative Python packages instead, so a consumer wheel can build
successfully and then fail during import.

Read module roots from vendor/modules.txt and replace the longest matching
prefix with the dependency package path. Require a dot boundary so similarly
named and undeclared modules remain unchanged.

Include the selected StarPC Python plugin and its frozen uv.lock in the cache
fingerprint. Record the new fingerprint only after generation succeeds, and
ignore Python lock changes when the Python plugin is not selected.

Install uv in CI before the Go suite runs its generated Python import check.
Generated cross-module imports now resolve without aliases, and plugin changes
cannot reuse stale service code.

Signed-off-by: Christian Stewart <christian@aperture.us>

Author: paralin

.github/workflows/tests.yml

@@ -19,6 +19,7 @@ jobs:
       - uses: actions/setup-go@v7
         with:
           go-version: '1.26'
+      - uses: astral-sh/setup-uv@v7
       - run: bun install
       - name: Cache tools
         uses: actions/cache@v6

protogen/cache.go

@@ -97,11 +97,15 @@ func (c *Cache) SetToolVersions(versions string) {
 // - The file content hash has changed
 // - The protoc flags have changed
 // - Force is true
-func (c *Cache) NeedsRegeneration(packageKey string, protoFiles []string, projectDir string, flagsHash string, force bool) (bool, error) {
+func (c *Cache) NeedsRegeneration(packageKey string, protoFiles []string, projectDir string, flagsHash string, toolVersions string, force bool) (bool, error) {
 	if force {
 		return true, nil
 	}
 
+	if c.ToolVersions != toolVersions {
+		return true, nil
+	}
+
 	// Check if flags changed
 	if c.ProtocFlagsHash != flagsHash {
 		return true, nil

protogen/cache_test.go

@@ -1,9 +1,14 @@
 package protogen
 
 import (
+	"context"
+	"crypto/sha256"
+	"encoding/hex"
+	"io"
 	"os"
 	"path/filepath"
 	"slices"
+	"strings"
 	"testing"
 )
 
@@ -194,3 +199,115 @@ func TestHashProtoFilesPortableAndContentSensitive(t *testing.T) {
 		t.Fatalf("content hash must change when a proto source changes")
 	}
 }
+
+func TestToolVersionChangeInvalidatesCachedPackage(t *testing.T) {
+	dir, files := writeProtoTree(t, map[string]string{
+		"foo.proto": "syntax = \"proto3\";\npackage foo;\n",
+	})
+	cache := NewCache()
+	cache.ProtocFlagsHash = "flags"
+	cache.ToolVersions = "starpc-python=old"
+	if err := cache.UpdatePackage("example/foo", files, nil, dir); err != nil {
+		t.Fatal(err)
+	}
+
+	stale, err := cache.NeedsRegeneration(
+		"example/foo", files, dir, "flags", "starpc-python=old", false,
+	)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if stale {
+		t.Fatal("matching tool state invalidated cached output")
+	}
+	stale, err = cache.NeedsRegeneration(
+		"example/foo", files, dir, "flags", "starpc-python=new", false,
+	)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if !stale {
+		t.Fatal("changed tool state reused cached output")
+	}
+}
+
+func TestGetToolVersionsIncludesUVLock(t *testing.T) {
+	dir := t.TempDir()
+	lock := []byte("version = 1\nrevision = 3\n")
+	if err := os.WriteFile(filepath.Join(dir, "uv.lock"), lock, 0o644); err != nil {
+		t.Fatal(err)
+	}
+	generator := &Generator{
+		ProjectDir: dir,
+		Config:     NewConfig(),
+		Plugins:    &Plugins{StarpcPython: &Plugin{}},
+	}
+	versions := generator.getToolVersions()
+	want := sha256.Sum256(lock)
+	if !strings.Contains(versions, "uv.lock="+hex.EncodeToString(want[:])) {
+		t.Fatalf("tool versions omit uv.lock digest: %q", versions)
+	}
+}
+
+func TestGetToolVersionsIgnoresUVLockWithoutStarpcPython(t *testing.T) {
+	dir := t.TempDir()
+	if err := os.WriteFile(filepath.Join(dir, "uv.lock"), []byte("changed"), 0o644); err != nil {
+		t.Fatal(err)
+	}
+	generator := &Generator{
+		ProjectDir: dir,
+		Config:     NewConfig(),
+		Plugins:    &Plugins{},
+	}
+	if versions := generator.getToolVersions(); strings.Contains(versions, "uv.lock=") {
+		t.Fatalf("unselected Python plugin invalidated tool state: %q", versions)
+	}
+}
+
+func TestGenerateFailureDoesNotPersistToolVersions(t *testing.T) {
+	dir := t.TempDir()
+	vendorDir := filepath.Join(dir, "vendor")
+	if err := os.MkdirAll(vendorDir, 0o755); err != nil {
+		t.Fatal(err)
+	}
+	if err := os.WriteFile(filepath.Join(dir, "bad.proto"), []byte("not protobuf"), 0o644); err != nil {
+		t.Fatal(err)
+	}
+	if err := os.WriteFile(filepath.Join(dir, "uv.lock"), []byte("changed plugin"), 0o644); err != nil {
+		t.Fatal(err)
+	}
+	config := NewConfig()
+	config.ProjectDir = dir
+	config.Targets = []string{"bad.proto"}
+	cache := NewCache()
+	cache.ToolVersions = "accepted"
+	plugins := &Plugins{
+		Languages:    Languages{LanguagePython: {}},
+		RPCLibraries: RPCLibraries{RPCLibraryStarpcPython: {}},
+		StarpcPython: &Plugin{
+			Name:       "starpc-python",
+			BinaryName: "protoc-gen-starpc-python",
+			Path:       "/bin/false",
+			Type:       PluginTypePython,
+			OutFlag:    "starpc-python_out",
+		},
+	}
+	generator := &Generator{
+		Config:     config,
+		Plugins:    plugins,
+		Cache:      cache,
+		ProjectDir: dir,
+		ModuleDir:  dir,
+		ModulePath: "example.com/project",
+		VendorDir:  vendorDir,
+		OutDir:     vendorDir,
+		Stdout:     io.Discard,
+		Stderr:     io.Discard,
+	}
+	if err := generator.Generate(context.Background()); err == nil {
+		t.Fatal("invalid source generation unexpectedly succeeded")
+	}
+	if cache.ToolVersions != "accepted" {
+		t.Fatalf("failed generation persisted tool state %q", cache.ToolVersions)
+	}
+}

protogen/generator.go

@@ -4,6 +4,8 @@ import (
 	"bufio"
 	"bytes"
 	"context"
+	"crypto/sha256"
+	"encoding/hex"
 	"fmt"
 	"io"
 	"os"
@@ -125,9 +127,8 @@ func (g *Generator) Generate(ctx context.Context) error {
 		fmt.Fprintf(g.Stdout, "Found %d proto files\n", len(protoFiles))
 	}
 
-	// Get tool versions for cache invalidation
+	// Get tool versions for cache invalidation.
 	toolVersions := g.getToolVersions()
-	g.Cache.SetToolVersions(toolVersions)
 
 	// Build protoc arguments
 	protocArgs := g.buildProtocArgs()
@@ -157,7 +158,7 @@ func (g *Generator) Generate(ctx context.Context) error {
 		currentPackages[packageKey] = struct{}{}
 
 		// Check if regeneration is needed
-		needsRegen, err := g.Cache.NeedsRegeneration(packageKey, files, g.ProjectDir, flagsHash, g.Config.Force)
+		needsRegen, err := g.Cache.NeedsRegeneration(packageKey, files, g.ProjectDir, flagsHash, toolVersions, g.Config.Force)
 		if err != nil {
 			return fmt.Errorf("failed to check cache for %s: %w", dir, err)
 		}
@@ -247,6 +248,7 @@ func (g *Generator) Generate(ctx context.Context) error {
 	g.Cache.CleanOrphanedPackages(currentPackages)
 
 	g.Cache.SetProtocFlags(protocArgs, g.ModuleDir)
+	g.Cache.SetToolVersions(toolVersions)
 	// Save cache
 	cacheFile, _ := g.Config.GetCacheFilePath()
 	if err := g.Cache.Save(cacheFile); err != nil {
@@ -453,6 +455,13 @@ func (g *Generator) getToolVersions() string {
 		}
 	}
 
+	if g.Plugins != nil && g.Plugins.StarpcPython != nil {
+		if data, err := os.ReadFile(filepath.Join(g.ProjectDir, "uv.lock")); err == nil {
+			digest := sha256.Sum256(data)
+			versions = append(versions, "uv.lock="+hex.EncodeToString(digest[:]))
+		}
+	}
+
 	return strings.Join(versions, ",")
 }
 

protogen/postprocess.go

@@ -6,6 +6,7 @@ import (
 	"os"
 	"path/filepath"
 	"regexp"
+	"slices"
 	"strings"
 )
 
@@ -221,13 +222,14 @@ func (p *PostProcessor) ProcessGoFile(filePath string) error {
 	return nil
 }
 
-// ProcessPythonFile rewrites local canonical module imports to adjacent imports.
+// ProcessPythonFile rewrites canonical Go module imports to the module-relative
+// packages installed by the current project and its vendored dependencies.
 func (p *PostProcessor) ProcessPythonFile(filePath string) error {
 	data, err := os.ReadFile(filePath)
 	if err != nil {
 		return err
 	}
-	prefix := strings.ReplaceAll(strings.ReplaceAll(p.ModulePath, "/", "."), "-", "_") + "."
+	prefixes := p.pythonImportPrefixes()
 	lines := strings.Split(string(data), "\n")
 	modified := false
 	for i, line := range lines {
@@ -236,24 +238,60 @@ func (p *PostProcessor) ProcessPythonFile(filePath string) error {
 			continue
 		}
 		from := strings.TrimPrefix(trimmed, "from ")
-		if !strings.HasPrefix(from, prefix) {
-			continue
-		}
-		rest := strings.TrimPrefix(from, prefix)
-		space := strings.IndexByte(rest, ' ')
-		if space <= 0 {
-			continue
+		for _, prefix := range prefixes {
+			if !strings.HasPrefix(from, prefix) {
+				continue
+			}
+			rest := strings.TrimPrefix(from, prefix)
+			if strings.IndexByte(rest, ' ') <= 0 {
+				continue
+			}
+			indent := line[:len(line)-len(trimmed)]
+			lines[i] = indent + "from " + rest
+			modified = true
+			break
 		}
-		indent := line[:len(line)-len(trimmed)]
-		lines[i] = indent + "from " + rest
-		modified = true
 	}
 	if !modified {
 		return nil
 	}
 	return os.WriteFile(filePath, []byte(strings.Join(lines, "\n")), 0o644) //nolint:gosec
 }
 
+func (p *PostProcessor) pythonImportPrefixes() []string {
+	modules := map[string]struct{}{p.ModulePath: {}}
+	modulesFile := filepath.Join(p.VendorDir, "modules.txt")
+	if data, err := os.ReadFile(modulesFile); err == nil {
+		scanner := bufio.NewScanner(bytes.NewReader(data))
+		for scanner.Scan() {
+			line := scanner.Text()
+			if !strings.HasPrefix(line, "# ") || strings.HasPrefix(line, "## ") {
+				continue
+			}
+			fields := strings.Fields(strings.TrimPrefix(line, "# "))
+			if len(fields) < 2 || (fields[1] != "=>" && !strings.HasPrefix(fields[1], "v")) {
+				continue
+			}
+			modules[fields[0]] = struct{}{}
+		}
+	}
+	prefixes := make([]string, 0, len(modules))
+	for module := range modules {
+		if module == "" {
+			continue
+		}
+		prefix := strings.ReplaceAll(strings.ReplaceAll(module, "/", "."), "-", "_")
+		prefixes = append(prefixes, prefix+".")
+	}
+	slices.SortFunc(prefixes, func(a, b string) int {
+		if len(a) != len(b) {
+			return len(b) - len(a)
+		}
+		return strings.Compare(a, b)
+	})
+	return prefixes
+}
+
 // ProcessTsFile processes a TypeScript file.
 // Rewrites relative import paths to @go/ format for vendor dependencies.
 //

protogen/postprocess_test.go

@@ -122,3 +122,53 @@ func TestProcessPythonFileRewritesPyi(t *testing.T) {
 		t.Fatalf("got %q", got)
 	}
 }
+
+func TestProcessPythonFileRewritesVendoredModuleImports(t *testing.T) {
+	projectDir := t.TempDir()
+	vendorDir := filepath.Join(projectDir, "vendor")
+	if err := os.MkdirAll(vendorDir, 0o755); err != nil {
+		t.Fatal(err)
+	}
+	modules := `# github.com/aperturerobotics/starpc v0.52.0
+## explicit; go 1.25.0
+# github.com/aperturerobotics/starpc/extensions v0.1.0
+## explicit; go 1.25.0
+# github.com/foo/bar v1.0.0
+## explicit; go 1.25.0
+github.com/foo/bar/pkg
+`
+	if err := os.WriteFile(filepath.Join(vendorDir, "modules.txt"), []byte(modules), 0o644); err != nil {
+		t.Fatal(err)
+	}
+	file := filepath.Join(projectDir, "resource_pb2.py")
+	content := `from github.com.s4wave.spacewave.local import local_pb2
+from github.com.aperturerobotics.starpc.rpcstream import rpcstream_pb2
+from github.com.aperturerobotics.starpc.extensions.echo import echo_pb2
+from github.com.foo.bar.pkg import package_pb2
+from github.com.foo.barista.foo import barista_pb2
+from google.protobuf import timestamp_pb2
+from undeclared.example import example_pb2
+`
+	if err := os.WriteFile(file, []byte(content), 0o644); err != nil {
+		t.Fatal(err)
+	}
+	pp := NewPostProcessor(projectDir, vendorDir, "github.com/s4wave/spacewave", nil, false)
+	if err := pp.ProcessPythonFile(file); err != nil {
+		t.Fatal(err)
+	}
+	got, err := os.ReadFile(file)
+	if err != nil {
+		t.Fatal(err)
+	}
+	want := `from local import local_pb2
+from rpcstream import rpcstream_pb2
+from echo import echo_pb2
+from pkg import package_pb2
+from github.com.foo.barista.foo import barista_pb2
+from google.protobuf import timestamp_pb2
+from undeclared.example import example_pb2
+`
+	if string(got) != want {
+		t.Fatalf("unexpected dependency import rewrite:\n%s", got)
+	}
+}