From c2272209bf61c84b55b4a3415fab3c15b51eb59c Mon Sep 17 00:00:00 2001
From: Joshua Rich <joshua.rich@gmail.com>
Date: Tue, 26 Sep 2023 22:43:30 +0000
Subject: [PATCH] fix(agent): broken registration flow after recent changes

---
 internal/agent/agent.go     | 37 +++++++++++++++++++++++--------------
 internal/agent/ui/fyneUI.go |  3 +++
 internal/agent/ui/ui.go     |  1 +
 3 files changed, 27 insertions(+), 14 deletions(-)

diff --git a/internal/agent/agent.go b/internal/agent/agent.go
index 62a9f27ad..a00a08813 100644
--- a/internal/agent/agent.go
+++ b/internal/agent/agent.go
@@ -38,13 +38,14 @@ const (
 // This includes the data structure for the UI elements and tray and some
 // strings such as app name and version.
 type Agent struct {
-	ui      AgentUI
-	config  AgentConfig
-	sensors *tracker.SensorTracker
-	done    chan struct{}
-	name    string
-	id      string
-	version string
+	ui       AgentUI
+	config   AgentConfig
+	sensors  *tracker.SensorTracker
+	done     chan struct{}
+	name     string
+	id       string
+	version  string
+	headless bool
 }
 
 //go:generate moq -out mock_AgentUI_test.go . AgentUI
@@ -64,9 +65,10 @@ type AgentOptions struct {
 
 func newAgent(appID string, headless bool) *Agent {
 	a := &Agent{
-		id:      appID,
-		version: version,
-		done:    make(chan struct{}),
+		id:       appID,
+		version:  version,
+		done:     make(chan struct{}),
+		headless: headless,
 	}
 	a.ui = ui.NewFyneUI(a, headless)
 	a.config = config.NewFyneConfig()
@@ -109,6 +111,7 @@ func Run(options AgentOptions) {
 
 		workerWg.Add(1)
 		go func() {
+			defer workerWg.Done()
 			device.StartWorkers(ctx, agent.sensors, sensorWorkers...)
 		}()
 		workerWg.Add(1)
@@ -118,13 +121,15 @@ func Run(options AgentOptions) {
 		}()
 	}()
 
-	<-configDone
-	agent.handleSignals(cancelFunc)
-	agent.handleShutdown(ctx)
+	go func() {
+		<-configDone
+		agent.handleSignals(cancelFunc)
+		agent.handleShutdown(ctx)
+		agent.ui.DisplayTrayIcon(ctx, agent)
+	}()
 
 	// If we are not running in headless mode, show a tray icon
 	if !options.Headless {
-		agent.ui.DisplayTrayIcon(ctx, agent)
 		agent.ui.Run()
 	}
 	workerWg.Wait()
@@ -241,6 +246,10 @@ func (agent *Agent) handleShutdown(ctx context.Context) {
 
 // Agent satisfies ui.Agent, tracker.Agent and api.Agent interfaces
 
+func (agent *Agent) IsHeadless() bool {
+	return agent.headless
+}
+
 func (agent *Agent) AppName() string {
 	return agent.name
 }
diff --git a/internal/agent/ui/fyneUI.go b/internal/agent/ui/fyneUI.go
index 19861af43..8ce2a8b8d 100644
--- a/internal/agent/ui/fyneUI.go
+++ b/internal/agent/ui/fyneUI.go
@@ -80,6 +80,9 @@ func NewFyneUI(agent Agent, headless bool) *fyneUI {
 // DisplayTrayIcon displays an icon in the desktop tray with a menu for
 // controlling the agent and showing other informational windows.
 func (ui *fyneUI) DisplayTrayIcon(ctx context.Context, agent Agent) {
+	if agent.IsHeadless() {
+		return
+	}
 	if desk, ok := ui.app.(desktop.App); ok {
 		menuItemQuit := fyne.NewMenuItem(ui.text.Translate("Quit"), func() {
 			agent.Stop()
diff --git a/internal/agent/ui/ui.go b/internal/agent/ui/ui.go
index 667602b6e..907e391b6 100644
--- a/internal/agent/ui/ui.go
+++ b/internal/agent/ui/ui.go
@@ -13,6 +13,7 @@ import (
 
 //go:generate moq -out mock_Agent_test.go . Agent
 type Agent interface {
+	IsHeadless() bool
 	AppVersion() string
 	AppName() string
 	AppID() string