feat(agent): allow user to define environment id using `--environment…

…` flag (#3840)

* agent: update connection request to support env id

* agent: send environment id to authentication endpoint

agent/client/client.go

@@ -26,9 +26,10 @@ const (
 )
 
 type Config struct {
-	APIKey     string
-	AgentName  string
-	PingPeriod time.Duration
+	APIKey        string
+	AgentName     string
+	EnvironmentID string
+	PingPeriod    time.Duration
 }
 
 type SessionConfig struct {
@@ -170,8 +171,9 @@ func (c *Client) getConnectionRequest() (*proto.ConnectRequest, error) {
 	}
 
 	request := proto.ConnectRequest{
-		ApiKey: c.config.APIKey,
-		Name:   name,
+		ApiKey:        c.config.APIKey,
+		EnvironmentID: c.config.EnvironmentID,
+		Name:          name,
 	}
 
 	return &request, nil

agent/client/options.go

@@ -21,6 +21,12 @@ func WithAgentName(name string) Option {
 	}
 }
 
+func WithEnvironmentID(id string) Option {
+	return func(c *Client) {
+		c.config.EnvironmentID = id
+	}
+}
+
 func WithPingPeriod(period time.Duration) Option {
 	return func(c *Client) {
 		c.config.PingPeriod = period

agent/config/config.go

@@ -13,6 +13,7 @@ import (
 type Config struct {
 	APIKey            string `mapstructure:"api_key"`
 	Name              string `mapstructure:"agent_name"`
+	EnvironmentID     string `mapstructure:"environment_id"`
 	ServerURL         string `mapstructure:"server_url"`
 	CollectorEndpoint string `mapstructure:"collector_endpoint"`
 	Mode              string `mapstructure:"mode"`
@@ -46,6 +47,7 @@ func LoadConfig() (Config, error) {
 
 	vp.SetDefault("AGENT_NAME", getHostname())
 	vp.SetDefault("API_KEY", "")
+	vp.SetDefault("ENVIRONMENT_ID", "")
 	vp.SetDefault("SERVER_URL", "https://app.tracetest.io")
 	vp.SetDefault("COLLECTOR_ENDPOINT", "")
 	vp.SetDefault("MODE", "")

agent/config/config_test.go

@@ -18,6 +18,7 @@ func TestConfigDefaults(t *testing.T) {
 
 	assert.Equal(t, "", cfg.APIKey)
 	assert.Equal(t, hostname, cfg.Name)
+	assert.Empty(t, cfg.EnvironmentID)
 	assert.Equal(t, "https://app.tracetest.io", cfg.ServerURL)
 	assert.Equal(t, 4317, cfg.OTLPServer.GRPCPort)
 	assert.Equal(t, 4318, cfg.OTLPServer.HTTPPort)
@@ -31,10 +32,12 @@ func TestConfigWithEnvs(t *testing.T) {
 		os.Unsetenv("TRACETEST_SERVER_URL")
 		os.Unsetenv("TRACETEST_OTLP_SERVER_GRPC_PORT")
 		os.Unsetenv("TRACETEST_OTLP_SERVER_HTTP_PORT")
+		os.Unsetenv("TRACETEST_ENVIRONMENT_ID")
 	})
 
 	os.Setenv("TRACETEST_AGENT_NAME", "my-agent-name")
 	os.Setenv("TRACETEST_API_KEY", "my-agent-api-key")
+	os.Setenv("TRACETEST_ENVIRONMENT_ID", "123456")
 	os.Setenv("TRACETEST_DEV_MODE", "true")
 	os.Setenv("TRACETEST_SERVER_URL", "https://custom.server.com")
 	os.Setenv("TRACETEST_OTLP_SERVER_GRPC_PORT", "1234")
@@ -46,6 +49,7 @@ func TestConfigWithEnvs(t *testing.T) {
 
 	assert.Equal(t, "my-agent-api-key", cfg.APIKey)
 	assert.Equal(t, "my-agent-name", cfg.Name)
+	assert.Equal(t, "123456", cfg.EnvironmentID)
 	assert.Equal(t, "https://custom.server.com", cfg.ServerURL)
 	assert.Equal(t, 1234, cfg.OTLPServer.GRPCPort)
 	assert.Equal(t, 1235, cfg.OTLPServer.HTTPPort)

agent/proto/orchestrator.proto

@@ -64,6 +64,7 @@ message StopRequest {
 message ConnectRequest {
     string apiKey = 1;
     string name = 2;
+    string environmentID = 3;
 }
 
 // ConnectResponse is the Orchestrator response to the ConnectRequest. It contains

agent/runner/runner.go

@@ -88,7 +88,7 @@ Once started, Tracetest Agent exposes OTLP ports 4317 and 4318 to ingest traces
 
 func (s *Runner) onStartAgent(ctx context.Context, cfg config.Config) {
 	if cfg.AgentApiKey != "" {
-		err := s.StartAgent(ctx, cfg.AgentEndpoint, cfg.AgentApiKey, cfg.UIEndpoint)
+		err := s.StartAgent(ctx, cfg.AgentEndpoint, cfg.AgentApiKey, cfg.UIEndpoint, cfg.EnvironmentID)
 		if err != nil {
 			s.ui.Error(err.Error())
 		}
@@ -106,13 +106,13 @@ func (s *Runner) onStartAgent(ctx context.Context, cfg config.Config) {
 		return
 	}
 
-	err = s.StartAgent(ctx, cfg.AgentEndpoint, env.AgentApiKey, cfg.UIEndpoint)
+	err = s.StartAgent(ctx, cfg.AgentEndpoint, env.AgentApiKey, cfg.UIEndpoint, cfg.EnvironmentID)
 	if err != nil {
 		s.ui.Error(err.Error())
 	}
 }
 
-func (s *Runner) StartAgent(ctx context.Context, endpoint, agentApiKey, uiEndpoint string) error {
+func (s *Runner) StartAgent(ctx context.Context, endpoint, agentApiKey, uiEndpoint, environmentID string) error {
 	cfg, err := agentConfig.LoadConfig()
 	s.logger.Debug("Loaded agent config", zap.Any("config", cfg))
 	if err != nil {
@@ -132,6 +132,12 @@ func (s *Runner) StartAgent(ctx context.Context, endpoint, agentApiKey, uiEndpoi
 	}
 	s.logger.Debug("Agent api key", zap.String("apiKey", cfg.APIKey))
 
+	if environmentID != "" {
+		s.logger.Debug("Overriding agent environment id", zap.String("environment", environmentID))
+		cfg.EnvironmentID = environmentID
+	}
+	s.logger.Debug("Agent environment id", zap.String("environmentID", cfg.EnvironmentID))
+
 	if s.mode == agentConfig.Mode_Desktop {
 		s.logger.Debug("Starting agent in desktop mode")
 		return s.RunDesktopStrategy(ctx, cfg, uiEndpoint)

agent/runner/session.go

@@ -120,6 +120,7 @@ func newControlPlaneClient(ctx context.Context, config config.Config, traceCache
 	controlPlaneClient, err := client.Connect(ctx, config.ServerURL,
 		client.WithAPIKey(config.APIKey),
 		client.WithAgentName(config.Name),
+		client.WithEnvironmentID(config.EnvironmentID),
 		client.WithLogger(logger),
 	)
 	if err != nil {

cli/cmd/start_cmd.go

@@ -70,6 +70,10 @@ var startCmd = &cobra.Command{
 			flags.Mode = agentConfig.Mode(cfg.Mode)
 		}
 
+		if flags.EnvironmentID != "" {
+			cfg.EnvironmentID = flags.EnvironmentID
+		}
+
 		err = agentRunner.Run(ctx, cliConfig, flags, verbose)
 		return "", err
 	})),