fix: guard the active runner environment behind accessors (#79) (#80)

Service.ActiveEnv was an exported interface field written by a
specialization's Init/Stop and read by the Code, Tooling and REPL
handlers on their own gRPC goroutines, with nothing ordering the two.
A specialization cannot close that gap: its lifecycle lock is not
reachable from these handlers.

The field is now unexported behind ActiveEnv/SetActiveEnv under an
RWMutex, and fixPython snapshots it once per request so a publish or
clear landing mid-fix cannot send ruff check and ruff format to two
different interpreters.

Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>

Author: antoinetoussaint-byte

agent.codefly.yaml

@@ -1,4 +1,4 @@
 publisher: codefly.dev
 kind: codefly:service
 name: python
-version: 0.0.59
+version: 0.0.60

pkg/code/callgraph.go

@@ -137,7 +137,7 @@ print(json.dumps({'calls': calls, 'implements': implements, 'module': workdir}))
 	// so AST analysis uses the same interpreter as the running service.
 	// Pre-Init falls through ResolveStandaloneEnvironment which picks Nix
 	// if declared, native otherwise — keeps mode consistency when possible.
-	env := c.Service.ActiveEnv
+	env := c.Service.ActiveEnv()
 	if env == nil {
 		var rctx *basev0.RuntimeContext
 		if c.Service.Base != nil && c.Service.Base.Runtime != nil {

pkg/code/code.go

@@ -75,6 +75,11 @@ func (c *Code) Execute(ctx context.Context, req *codev0.CodeRequest) (*codev0.Co
 // formatter. Both stages operate on stdin and return source on stdout, so the
 // core Code server remains the sole writer and can honor VFS/dry-run semantics.
 func (c *Code) fixPython(ctx context.Context, input corecode.FixInput) (corecode.FixResult, error) {
+	// One snapshot for the whole request: a specialization publishing or
+	// clearing its environment mid-fix must not send check and format to
+	// two different interpreters.
+	active := c.Service.ActiveEnv()
+
 	checkArgs := []string{"check", "--fix-only"}
 	if input.Mode == basev0.FixMode_FIX_MODE_AGGRESSIVE {
 		checkArgs = append(checkArgs, "--unsafe-fixes")
@@ -86,13 +91,14 @@ func (c *Code) fixPython(ctx context.Context, input corecode.FixInput) (corecode
 	// download the project's dependency graph. Factory Nix shells include Ruff,
 	// the Docker runtime exposes /venv/bin on PATH, and native users may supply
 	// either a project-local .venv binary or Ruff on the host PATH.
-	ruff := c.ruffCommand()
-	checked, checkLogBytes, err := runners.RunInput(ctx, c.runnerEnvironment(ctx), c.SourceDir(), input.Content, ruff, checkArgs...)
+	ruff := ruffCommand(active, c.SourceDir())
+	env := c.runnerEnvironment(ctx, active)
+	checked, checkLogBytes, err := runners.RunInput(ctx, env, c.SourceDir(), input.Content, ruff, checkArgs...)
 	checkLog := string(checkLogBytes)
 	if err != nil {
 		return corecode.FixResult{}, fmt.Errorf("ruff check --fix-only (Ruff must already be provisioned): %w: %s", err, strings.TrimSpace(checkLog))
 	}
-	formatted, formatLogBytes, err := runners.RunInput(ctx, c.runnerEnvironment(ctx), c.SourceDir(), checked, ruff, "format", "--stdin-filename", input.Path, "-")
+	formatted, formatLogBytes, err := runners.RunInput(ctx, env, c.SourceDir(), checked, ruff, "format", "--stdin-filename", input.Path, "-")
 	formatLog := string(formatLogBytes)
 	if err != nil {
 		return corecode.FixResult{}, fmt.Errorf("ruff format: %w: %s", err, strings.TrimSpace(formatLog))
@@ -105,20 +111,19 @@ func (c *Code) fixPython(ctx context.Context, input corecode.FixInput) (corecode
 	}, nil
 }
 
-func (c *Code) ruffCommand() string {
-	if c.Service.ActiveEnv == nil {
-		candidate := filepath.Join(c.SourceDir(), ".venv", "bin", "ruff")
+func ruffCommand(active runners.RunnerEnvironment, sourceDir string) string {
+	if active == nil {
+		candidate := filepath.Join(sourceDir, ".venv", "bin", "ruff")
 		if info, err := os.Stat(candidate); err == nil && info.Mode().IsRegular() && info.Mode().Perm()&0o111 != 0 {
 			return candidate
 		}
 	}
 	return "ruff"
 }
 
-func (c *Code) runnerEnvironment(ctx context.Context) runners.RunnerEnvironment {
-	env := c.Service.ActiveEnv
-	if env != nil {
-		return env
+func (c *Code) runnerEnvironment(ctx context.Context, active runners.RunnerEnvironment) runners.RunnerEnvironment {
+	if active != nil {
+		return active
 	}
 	var runtimeContext *basev0.RuntimeContext
 	if c.Service.Base != nil && c.Service.Base.Runtime != nil {

pkg/code/fix_test.go

@@ -6,10 +6,12 @@ import (
 	"os/exec"
 	"path/filepath"
 	"strings"
+	"sync"
 	"testing"
 
 	codev0 "github.com/codefly-dev/core/generated/go/codefly/services/code/v0"
 	"github.com/codefly-dev/core/resources"
+	runners "github.com/codefly-dev/core/runners/base"
 	pythoncode "github.com/codefly-dev/service-python/pkg/code"
 	pythonservice "github.com/codefly-dev/service-python/pkg/service"
 )
@@ -80,3 +82,63 @@ func TestApplyEditRunsRuffSafeFixerByDefault(t *testing.T) {
 		t.Fatalf("ApplyEdit did not commit returned content: err=%v content=%q", err, written)
 	}
 }
+
+// TestCodeIsSafeWhileTheActiveEnvironmentIsRepublished drives the two Code
+// entry points that resolve a runner environment — Fix (Ruff) and the call
+// graph (AST) — while a specialization publishes its environment from Init
+// and clears it from Stop. The RPCs run on one goroutine so the only thing
+// racing the publisher is the environment lookup itself.
+//
+// Both requests must keep answering: a Fix that lands between publish and
+// clear falls back to a standalone environment rather than failing.
+func TestCodeIsSafeWhileTheActiveEnvironmentIsRepublished(t *testing.T) {
+	if _, err := exec.LookPath("ruff"); err != nil {
+		t.Skip("ruff is not installed")
+	}
+	ctx := context.Background()
+	dir := t.TempDir()
+	if err := os.WriteFile(filepath.Join(dir, "sample.py"), []byte("def answer( ):\n return 42\n"), 0o644); err != nil {
+		t.Fatal(err)
+	}
+
+	svc := pythonservice.New(&resources.Agent{Kind: "codefly:service", Name: "python"})
+	svc.SourceLocation = dir
+	env, err := runners.NewNativeEnvironment(ctx, dir)
+	if err != nil {
+		t.Fatalf("native environment: %v", err)
+	}
+	server := pythoncode.New(svc)
+
+	done := make(chan struct{})
+	var publisher sync.WaitGroup
+	publisher.Add(1)
+	go func() {
+		defer publisher.Done()
+		for {
+			select {
+			case <-done:
+				return
+			default:
+				svc.SetActiveEnv(env)
+				svc.SetActiveEnv(nil)
+			}
+		}
+	}()
+
+	for i := 0; i < 10; i++ {
+		response, err := server.Execute(ctx, &codev0.CodeRequest{Operation: &codev0.CodeRequest_Fix{Fix: &codev0.FixRequest{
+			File: "sample.py", DryRun: true,
+		}}})
+		if err != nil {
+			t.Fatalf("fix %d: %v", i, err)
+		}
+		if fix := response.GetFix(); !fix.GetSuccess() {
+			t.Fatalf("fix %d did not run: %+v failure=%+v", i, fix, response.GetFailure())
+		}
+		if graph := server.ComputePythonCallGraph(ctx, dir); graph.Error != "" {
+			t.Fatalf("call graph %d: %s", i, graph.Error)
+		}
+	}
+	close(done)
+	publisher.Wait()
+}

pkg/runtime/commands.go

@@ -145,7 +145,7 @@ func (s *Runtime) getOrStartRepl(ctx context.Context) (*PythonRepl, error) {
 // resolveReplEnv returns the plugin's ActiveEnv or a standalone one
 // when Runtime.Init hasn't populated it. Same pattern as Code/Tooling.
 func (s *Runtime) resolveReplEnv(ctx context.Context) runners.RunnerEnvironment {
-	if env := s.Service.ActiveEnv; env != nil {
+	if env := s.Service.ActiveEnv(); env != nil {
 		return env
 	}
 	var rctx *basev0.RuntimeContext

pkg/runtime/runtime_internal_test.go

@@ -1,16 +1,19 @@
 package runtime
 
 import (
+	"context"
 	"errors"
 	"os"
 	"path/filepath"
 	"strings"
+	"sync"
 	"testing"
 	"time"
 	"unicode/utf8"
 
 	runtimev0 "github.com/codefly-dev/core/generated/go/codefly/services/runtime/v0"
 	"github.com/codefly-dev/core/resources"
+	runners "github.com/codefly-dev/core/runners/base"
 	pythonhelpers "github.com/codefly-dev/core/runners/python"
 	pythonservice "github.com/codefly-dev/service-python/pkg/service"
 )
@@ -421,3 +424,43 @@ func TestResolveTestFormulaOverlaysCommandlessHealOntoDerivedFormula(t *testing.
 		t.Fatalf("derived provisioning must survive the overlay, got %v", prov)
 	}
 }
+
+// TestResolveReplEnvIsSafeWhileTheActiveEnvironmentIsRepublished covers the
+// REPL's half of the same lookup: `exec` resolves an environment on its own
+// gRPC goroutine while a specialization's Init publishes one and Stop clears
+// it. Whichever side wins, the REPL must get an environment to boot into.
+func TestResolveReplEnvIsSafeWhileTheActiveEnvironmentIsRepublished(t *testing.T) {
+	ctx := context.Background()
+	service := pythonservice.New(&resources.Agent{Kind: "codefly:service", Name: "python"})
+	service.SourceLocation = t.TempDir()
+	rt := New(service)
+
+	env, err := runners.NewNativeEnvironment(ctx, service.SourceLocation)
+	if err != nil {
+		t.Fatalf("native environment: %v", err)
+	}
+
+	done := make(chan struct{})
+	var publisher sync.WaitGroup
+	publisher.Add(1)
+	go func() {
+		defer publisher.Done()
+		for {
+			select {
+			case <-done:
+				return
+			default:
+				service.SetActiveEnv(env)
+				service.SetActiveEnv(nil)
+			}
+		}
+	}()
+
+	for i := 0; i < 200; i++ {
+		if rt.resolveReplEnv(ctx) == nil {
+			t.Fatalf("resolveReplEnv returned nil on attempt %d", i)
+		}
+	}
+	close(done)
+	publisher.Wait()
+}

pkg/service/service.go

@@ -7,6 +7,7 @@ import (
 	"context"
 	"os"
 	"path/filepath"
+	"sync"
 
 	"github.com/codefly-dev/core/agents/services"
 	agentv0 "github.com/codefly-dev/core/generated/go/codefly/services/agent/v0"
@@ -45,11 +46,33 @@ type Service struct {
 	// Load). Specializations may override at Load time.
 	SourceLocation string
 
-	// ActiveEnv is a specialization's active RunnerEnvironment, consumed by
-	// Code / Tooling so every spawn routes through the same mode (native /
-	// docker / nix). The generic runtime leaves it nil; call sites fall back
-	// to a fresh NativeEnvironment.
-	ActiveEnv runners.RunnerEnvironment
+	// activeEnv is a specialization's active RunnerEnvironment, consumed by
+	// Code / Tooling / the REPL so every spawn routes through the same mode
+	// (native / docker / nix). The generic runtime leaves it nil; call sites
+	// fall back to a fresh NativeEnvironment.
+	//
+	// A specialization publishes it from Init and clears it from Stop, while
+	// Code, Tooling and the REPL read it on their own gRPC goroutines, so it
+	// is reachable only through ActiveEnv / SetActiveEnv.
+	activeEnvMu sync.RWMutex
+	activeEnv   runners.RunnerEnvironment
+}
+
+// ActiveEnv returns the environment a specialization published, or nil when
+// none is active — before Init, or after the specialization tore it down.
+func (s *Service) ActiveEnv() runners.RunnerEnvironment {
+	s.activeEnvMu.RLock()
+	defer s.activeEnvMu.RUnlock()
+	return s.activeEnv
+}
+
+// SetActiveEnv publishes the specialization's active RunnerEnvironment.
+// Pass nil once the environment is shut down: readers then fall back to a
+// standalone environment instead of spawning into a dead one.
+func (s *Service) SetActiveEnv(env runners.RunnerEnvironment) {
+	s.activeEnvMu.Lock()
+	defer s.activeEnvMu.Unlock()
+	s.activeEnv = env
 }
 
 // New builds a generic Python Service bound to the given agent manifest.

pkg/service/service_test.go

@@ -3,12 +3,14 @@ package service_test
 import (
 	"context"
 	"strings"
+	"sync"
 	"testing"
 
 	"gopkg.in/yaml.v3"
 
 	agentv0 "github.com/codefly-dev/core/generated/go/codefly/services/agent/v0"
 	"github.com/codefly-dev/core/resources"
+	runners "github.com/codefly-dev/core/runners/base"
 
 	pythonservice "github.com/codefly-dev/service-python/pkg/service"
 )
@@ -113,3 +115,61 @@ func TestGetAgentInformationGeneric(t *testing.T) {
 		}
 	}
 }
+
+// TestActiveEnvPublishAndClear locks in the contract specializations use:
+// what Init publishes is what Code, Tooling and the REPL read, and what Stop
+// clears sends them back to the standalone fallback.
+func TestActiveEnvPublishAndClear(t *testing.T) {
+	svc := pythonservice.New(&resources.Agent{Kind: "codefly:service", Name: "python"})
+	if svc.ActiveEnv() != nil {
+		t.Fatal("a Service with no specialization environment must read nil")
+	}
+
+	env, err := runners.NewNativeEnvironment(context.Background(), t.TempDir())
+	if err != nil {
+		t.Fatalf("native environment: %v", err)
+	}
+	svc.SetActiveEnv(env)
+	if svc.ActiveEnv() != env {
+		t.Fatal("published environment is not the one readers get")
+	}
+
+	svc.SetActiveEnv(nil)
+	if svc.ActiveEnv() != nil {
+		t.Fatal("cleared environment must read nil, not a shut-down one")
+	}
+}
+
+// TestActiveEnvUnderConcurrentPublishAndRead is the -race guard for the
+// field itself: a specialization publishes from Init and clears from Stop
+// while Code, Tooling and the REPL read on their own gRPC goroutines.
+func TestActiveEnvUnderConcurrentPublishAndRead(t *testing.T) {
+	svc := pythonservice.New(&resources.Agent{Kind: "codefly:service", Name: "python"})
+	env, err := runners.NewNativeEnvironment(context.Background(), t.TempDir())
+	if err != nil {
+		t.Fatalf("native environment: %v", err)
+	}
+
+	var wg sync.WaitGroup
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		for i := 0; i < 1000; i++ {
+			svc.SetActiveEnv(env)
+			svc.SetActiveEnv(nil)
+		}
+	}()
+	for readers := 0; readers < 3; readers++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			for i := 0; i < 1000; i++ {
+				if got := svc.ActiveEnv(); got != nil && got != env {
+					t.Errorf("read an environment that was never published: %#v", got)
+					return
+				}
+			}
+		}()
+	}
+	wg.Wait()
+}