fix: improve agent startup (#3191)

* fix: improve agent startup

* fix initialization

agent/cmd/start.go

@@ -29,11 +29,14 @@ var StartCmd = cobra.Command{
 
 		log.Printf("starting agent [%s] connecting to %s", cfg.Name, cfg.ServerURL)
 
-		err = initialization.Start(ctx, cfg)
+		session, err := initialization.Start(ctx, cfg)
 		if err != nil {
 			fmt.Fprintln(os.Stderr, err)
 			ExitCLI(1)
 		}
+
+		session.WaitUntilDisconnected()
+		session.Close()
 	},
 }
 

agent/initialization/session.go

@@ -0,0 +1,16 @@
+package initialization
+
+import "github.com/kubeshop/tracetest/agent/client"
+
+type Session struct {
+	Token  string
+	client *client.Client
+}
+
+func (s *Session) Close() {
+	s.client.Close()
+}
+
+func (s *Session) WaitUntilDisconnected() {
+	s.client.WaitUntilDisconnected()
+}

agent/initialization/start.go

@@ -40,25 +40,27 @@ func NewClient(ctx context.Context, config config.Config, traceCache collector.T
 }
 
 // Start the agent with given configuration
-func Start(ctx context.Context, config config.Config) error {
+func Start(ctx context.Context, config config.Config) (*Session, error) {
 	traceCache := collector.NewTraceCache()
 	client, err := NewClient(ctx, config, traceCache)
 	if err != nil {
-		return err
+		return nil, err
 	}
 
 	err = client.Start(ctx)
 	if err != nil {
-		return err
+		return nil, err
 	}
 
 	err = StartCollector(ctx, config, traceCache)
 	if err != nil {
-		return err
+		return nil, err
 	}
 
-	client.WaitUntilDisconnected()
-	return nil
+	return &Session{
+		client: client,
+		Token:  client.SessionConfiguration().AgentIdentification.Token,
+	}, nil
 }
 
 func StartCollector(ctx context.Context, config config.Config, traceCache collector.TraceCache) error {

cli/cmd/start_cmd.go

@@ -3,6 +3,7 @@ package cmd
 import (
 	"context"
 
+	agentConfig "github.com/kubeshop/tracetest/agent/config"
 	"github.com/kubeshop/tracetest/cli/config"
 	"github.com/kubeshop/tracetest/cli/pkg/starter"
 	"github.com/spf13/cobra"
@@ -29,7 +30,16 @@ var startCmd = &cobra.Command{
 			AgentApiKey:    saveParams.agentApiKey,
 		}
 
-		err := start.Run(ctx, cliConfig, flags)
+		cfg, err := agentConfig.LoadConfig()
+		if err != nil {
+			return "", err
+		}
+
+		if cfg.APIKey != "" {
+			flags.AgentApiKey = cfg.APIKey
+		}
+
+		err = start.Run(ctx, cliConfig, flags)
 		return "", err
 	})),
 	PostRun: teardownCommand,

cli/pkg/starter/starter.go

@@ -6,7 +6,6 @@ import (
 	"fmt"
 
 	"github.com/golang-jwt/jwt/v4"
-	"github.com/kubeshop/tracetest/agent/collector"
 	agentConfig "github.com/kubeshop/tracetest/agent/config"
 	"github.com/kubeshop/tracetest/agent/initialization"
 
@@ -34,7 +33,7 @@ func (s *Starter) Run(ctx context.Context, cfg config.Config, flags config.Confi
 
 func (s *Starter) onStartAgent(ctx context.Context, cfg config.Config) {
 	if cfg.AgentApiKey != "" {
-		err := s.StartAgent(ctx, cfg.AgentEndpoint, "local", cfg.AgentApiKey, cfg.UIEndpoint)
+		err := s.StartAgent(ctx, cfg.AgentEndpoint, cfg.AgentApiKey, cfg.UIEndpoint)
 		if err != nil {
 			s.ui.Error(err.Error())
 		}
@@ -51,7 +50,7 @@ func (s *Starter) onStartAgent(ctx context.Context, cfg config.Config) {
 Connecting Agent with name %s to Organization %s and Environment %s
 `, "local", cfg.OrganizationID, env.Name))
 
-	err = s.StartAgent(ctx, cfg.AgentEndpoint, "local", env.AgentApiKey, cfg.UIEndpoint)
+	err = s.StartAgent(ctx, cfg.AgentEndpoint, env.AgentApiKey, cfg.UIEndpoint)
 	if err != nil {
 		s.ui.Error(err.Error())
 	}
@@ -95,35 +94,27 @@ func (s *Starter) getEnvironment(ctx context.Context, cfg config.Config) (enviro
 	return env, nil
 }
 
-func (s *Starter) StartAgent(ctx context.Context, endpoint, name, agentApiKey, uiEndpoint string) error {
-	cfg := agentConfig.Config{
-		ServerURL: endpoint,
-		APIKey:    agentApiKey,
-		Name:      name,
-		OTLPServer: agentConfig.OtlpServer{
-			GRPCPort: 4317,
-			HTTPPort: 4318,
-		},
-	}
-
-	s.ui.Info(fmt.Sprintf("Starting Agent with name %s...", name))
-	cache := collector.NewTraceCache()
-	client, err := initialization.NewClient(ctx, cfg, cache)
+func (s *Starter) StartAgent(ctx context.Context, endpoint, agentApiKey, uiEndpoint string) error {
+	cfg, err := agentConfig.LoadConfig()
 	if err != nil {
 		return err
 	}
 
-	err = client.Start(ctx)
-	if err != nil {
-		return err
+	if endpoint != "" {
+		cfg.ServerURL = endpoint
+	}
+
+	if agentApiKey != "" {
+		cfg.APIKey = agentApiKey
 	}
 
-	err = initialization.StartCollector(ctx, cfg, cache)
+	s.ui.Info(fmt.Sprintf("Starting Agent with name %s...", cfg.Name))
+	session, err := initialization.Start(ctx, cfg)
 	if err != nil {
 		return err
 	}
 
-	claims, err := s.getTokenClaims(client.SessionConfiguration().AgentIdentification.Token)
+	claims, err := s.getTokenClaims(session.Token)
 	if err != nil {
 		return err
 	}
@@ -140,7 +131,7 @@ func (s *Starter) StartAgent(ctx context.Context, endpoint, name, agentApiKey, u
 			Text: "Stop this agent",
 			Fn: func(_ ui.UI) {
 				isOpen = false
-				go client.Close()
+				session.Close()
 			},
 		}}