feat(featureflags): support per-service context providers

packages/orchestrator/pkg/factories/featureflags_context.go

@@ -0,0 +1,27 @@
+//go:build linux
+
+package factories
+
+import (
+	"context"
+
+	"github.com/launchdarkly/go-sdk-common/v3/ldcontext"
+
+	"github.com/e2b-dev/infra/packages/shared/pkg/featureflags"
+)
+
+const (
+	orchestratorKind            ldcontext.Kind = "orchestrator"
+	orchestratorCommitAttribute string         = "commit"
+)
+
+func orchestratorContextProvider(nodeID, commit string) featureflags.ContextProvider {
+	versionContext := ldcontext.NewBuilder(nodeID).
+		Kind(orchestratorKind).
+		SetString(orchestratorCommitAttribute, commit).
+		Build()
+
+	return func(context.Context) ldcontext.Context {
+		return versionContext
+	}
+}

packages/orchestrator/pkg/factories/run.go

@@ -360,6 +360,7 @@ func run(config cfg.Config, opts Options) (success bool) {
 	closers = append(closers, closer{"feature flags", featureFlags.Close})
 
 	featureFlags.SetDeploymentName(config.DomainName)
+	featureFlags.RegisterContextProvider(orchestratorContextProvider(nodeID, commitSHA))
 
 	// gcp concurrent upload limiter
 	limiter, err := limit.New(ctx, featureFlags)

packages/orchestrator/pkg/server/sandboxes.go

@@ -117,7 +117,6 @@ func (s *Server) Create(ctx context.Context, req *orchestrator.SandboxCreateRequ
 		ldcontext.NewBuilder(req.GetSandbox().GetTeamId()).
 			Kind(featureflags.TeamKind).
 			Build(),
-		featureflags.VersionContext(s.info.ClientId, s.info.SourceCommit),
 	)
 
 	// BYOP egress proxy kill-switch; mirrors the API gate for direct gRPC

packages/shared/pkg/featureflags/client.go

@@ -25,11 +25,17 @@ var launchDarklyApiKey = os.Getenv("LAUNCH_DARKLY_API_KEY")
 const waitForInit = 5 * time.Second
 
 type Client struct {
-	ld             *ldclient.LDClient
-	deploymentName string
-	serviceName    string
+	ld               *ldclient.LDClient
+	deploymentName   string
+	serviceName      string
+	contextProviders []ContextProvider
 }
 
+// ContextProvider supplies an additional LD context on every flag evaluation.
+// Services register providers to inject specific contexts without leaking that
+// specificity into the shared client.
+type ContextProvider func(ctx context.Context) ldcontext.Context
+
 func NewClientWithDatasource(source *ldtestdata.TestDataSource) (*Client, error) {
 	ldClient, err := ldclient.MakeCustomClient(
 		"",
@@ -98,12 +104,18 @@ func (c *Client) SetServiceName(serviceName string) {
 	c.serviceName = serviceName
 }
 
+// RegisterContextProvider registers a provider whose contexts are appended to
+// every flag evaluation.
+func (c *Client) RegisterContextProvider(provider ContextProvider) {
+	c.contextProviders = append(c.contextProviders, provider)
+}
+
 func (c *Client) BoolFlag(ctx context.Context, flag BoolFlag, contexts ...ldcontext.Context) bool {
-	return getFlag(ctx, c.ld, c.ld.BoolVariationCtx, flag, c.allContexts(contexts))
+	return getFlag(ctx, c.ld, c.ld.BoolVariationCtx, flag, c.allContexts(ctx, contexts))
 }
 
 func (c *Client) JSONFlag(ctx context.Context, flag JSONFlag, contexts ...ldcontext.Context) ldvalue.Value {
-	return getFlag(ctx, c.ld, c.ld.JSONVariationCtx, flag, c.allContexts(contexts))
+	return getFlag(ctx, c.ld, c.ld.JSONVariationCtx, flag, c.allContexts(ctx, contexts))
 }
 
 func (c *Client) WatchJSONFlag(ctx context.Context, flag JSONFlag, contexts ...ldcontext.Context) (<-chan interfaces.FlagValueChangeEvent, func()) {
@@ -116,7 +128,7 @@ func (c *Client) WatchJSONFlag(ctx context.Context, flag JSONFlag, contexts ...l
 
 	listener := c.ld.GetFlagTracker().AddFlagValueChangeListener(
 		flag.Key(),
-		mergeContexts(ctx, c.allContexts(contexts)),
+		mergeContexts(ctx, c.allContexts(ctx, contexts)),
 		flag.Fallback(),
 	)
 
@@ -126,11 +138,11 @@ func (c *Client) WatchJSONFlag(ctx context.Context, flag JSONFlag, contexts ...l
 }
 
 func (c *Client) IntFlag(ctx context.Context, flag IntFlag, contexts ...ldcontext.Context) int {
-	return getFlag(ctx, c.ld, c.ld.IntVariationCtx, flag, c.allContexts(contexts))
+	return getFlag(ctx, c.ld, c.ld.IntVariationCtx, flag, c.allContexts(ctx, contexts))
 }
 
 func (c *Client) StringFlag(ctx context.Context, flag StringFlag, contexts ...ldcontext.Context) string {
-	return getFlag(ctx, c.ld, c.ld.StringVariationCtx, flag, c.allContexts(contexts))
+	return getFlag(ctx, c.ld, c.ld.StringVariationCtx, flag, c.allContexts(ctx, contexts))
 }
 
 type typedFlag[T any] interface {
@@ -174,13 +186,16 @@ func (c *Client) Close(ctx context.Context) error {
 	return nil
 }
 
-func (c *Client) allContexts(contexts []ldcontext.Context) []ldcontext.Context {
+func (c *Client) allContexts(ctx context.Context, contexts []ldcontext.Context) []ldcontext.Context {
 	if c.deploymentName != "" {
 		contexts = append(contexts, deploymentContext(c.deploymentName))
 	}
 	if c.serviceName != "" {
 		contexts = append(contexts, ServiceContext(c.serviceName))
 	}
+	for _, provider := range c.contextProviders {
+		contexts = append(contexts, provider(ctx))
+	}
 
 	return contexts
 }

packages/shared/pkg/featureflags/client_test.go

@@ -1,6 +1,7 @@
 package featureflags
 
 import (
+	"context"
 	"testing"
 
 	"github.com/launchdarkly/go-sdk-common/v3/ldcontext"
@@ -43,7 +44,7 @@ func TestAllContextsIncludesServiceAndDeployment(t *testing.T) {
 	client.SetDeploymentName("dev")
 	client.SetServiceName("orchestration-api")
 
-	merged := mergeContexts(t.Context(), client.allContexts(nil))
+	merged := mergeContexts(t.Context(), client.allContexts(t.Context(), nil))
 	contexts := merged.GetAllIndividualContexts(nil)
 
 	seen := map[ldcontext.Kind]string{}
@@ -54,3 +55,22 @@ func TestAllContextsIncludesServiceAndDeployment(t *testing.T) {
 	require.Equal(t, "dev", seen[deploymentKind])
 	require.Equal(t, "orchestration-api", seen[ServiceKind])
 }
+
+func TestAllContextsIncludesRegisteredProviders(t *testing.T) {
+	t.Parallel()
+
+	client := &Client{}
+	client.RegisterContextProvider(func(context.Context) ldcontext.Context {
+		return ldcontext.NewWithKind("node", "node-1")
+	})
+
+	merged := mergeContexts(t.Context(), client.allContexts(t.Context(), nil))
+	contexts := merged.GetAllIndividualContexts(nil)
+
+	seen := map[ldcontext.Kind]string{}
+	for _, item := range contexts {
+		seen[item.Kind()] = item.Key()
+	}
+
+	require.Equal(t, "node-1", seen["node"])
+}

packages/shared/pkg/featureflags/context.go

@@ -174,10 +174,3 @@ func CompressFileTypeContext(fileType string) ldcontext.Context {
 func CompressUseCaseContext(useCase string) ldcontext.Context {
 	return ldcontext.NewWithKind(CompressUseCaseKind, useCase)
 }
-
-func VersionContext(orchestratorID, commit string) ldcontext.Context {
-	return ldcontext.NewBuilder(orchestratorID).
-		Kind(OrchestratorKind).
-		SetString(OrchestratorCommitAttribute, commit).
-		Build()
-}

packages/shared/pkg/featureflags/flags.go

@@ -33,9 +33,6 @@ const (
 	VolumeKind           ldcontext.Kind = "volume"
 	CompressFileTypeKind ldcontext.Kind = "compress-file-type"
 	CompressUseCaseKind  ldcontext.Kind = "compress-use-case"
-
-	OrchestratorKind            ldcontext.Kind = "orchestrator"
-	OrchestratorCommitAttribute string         = "commit"
 )
 
 // All flags must be defined here: https://app.launchdarkly.com/projects/default/flags/