[client] UI refactor (#6069)

Refactor UI

---------

Co-authored-by: Eduard Gert <kontakt@eduardgert.de>
Co-authored-by: braginini <bangvalo@gmail.com>
Co-authored-by: Pascal Fischer <32096965+pascal-fischer@users.noreply.github.com>
Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: riccardom <riccardomanfrin@gmail.com>

client/server/server.go

@@ -19,6 +19,7 @@ import (
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/metadata"
 	gstatus "google.golang.org/grpc/status"
+	"google.golang.org/protobuf/types/known/timestamppb"
 
 	"github.com/netbirdio/netbird/client/internal/auth"
 	"github.com/netbirdio/netbird/client/internal/expose"
@@ -67,7 +68,19 @@ type Server struct {
 
 	logFile string
 
+	// uiLogPath is the desktop UI's absolute log path, reported via
+	// RegisterUILog. Guarded by mutex. Consumed by DebugBundle so the bundle
+	// can collect the GUI log even though the daemon runs as root and can't
+	// resolve the user's config dir. Last-writer-wins (one UI per socket).
+	uiLogPath string
+
 	oauthAuthFlow oauthAuthFlow
+	// extendAuthSessionFlow holds the pending PKCE flow created by
+	// RequestExtendAuthSession until WaitExtendAuthSession resolves it.
+	// Kept separate from oauthAuthFlow (which is reserved for the SSH
+	// JWT path) so a concurrent SSH auth doesn't clobber the session
+	// extend flow or vice versa.
+	extendAuthSessionFlow *auth.PendingFlow
 
 	mutex  sync.Mutex
 	config *profilemanager.Config
@@ -87,7 +100,7 @@ type Server struct {
 	statusRecorder *peer.Status
 	sessionWatcher *internal.SessionWatcher
 
-	lastProbe           time.Time
+	probeThrottle       *probeThrottle
 	persistSyncResponse bool
 	isSessionActive     atomic.Bool
 
@@ -135,6 +148,8 @@ func New(ctx context.Context, logFile string, configFile string, profilesDisable
 		captureEnabled:         captureEnabled,
 		networksDisabled:       networksDisabled,
 		jwtCache:               newJWTCache(),
+		extendAuthSessionFlow:  auth.NewPendingFlow(),
+		probeThrottle:          newProbeThrottle(probeThreshold),
 	}
 	agent := &serverAgent{s}
 	s.sleepHandler = sleephandler.New(agent)
@@ -152,6 +167,15 @@ func (s *Server) Start() error {
 	}
 
 	state := internal.CtxGetState(s.rootCtx)
+	// Every contextState.Set in the connect/login/server paths must push a
+	// SubscribeStatus snapshot, otherwise transitions that don't happen to
+	// be accompanied by a Mark{Management,Signal,...} call (e.g. plain
+	// StatusNeedsLogin after a PermissionDenied login, StatusLoginFailed
+	// after OAuth init failure, StatusIdle in the Login defer) leave the
+	// UI stuck on the previous status until the next unrelated peer event.
+	// Binding the recorder here means new state.Set callsites don't have
+	// to opt in individually.
+	state.SetOnChange(s.statusRecorder.NotifyStateChange)
 
 	if err := handlePanicLog(); err != nil {
 		log.Warnf("failed to redirect stderr: %v", err)
@@ -236,7 +260,7 @@ func (s *Server) Start() error {
 	s.clientRunningChan = make(chan struct{})
 	s.clientGiveUpChan = make(chan struct{})
 	go s.connectWithRetryRuns(ctx, config, s.statusRecorder, s.clientRunningChan, s.clientGiveUpChan)
-	s.publishConfigChangedEvent("startup")
+	s.publishConfigChangedEvent(proto.MetadataSourceStartup)
 	return nil
 }
 
@@ -253,6 +277,10 @@ func (s *Server) Start() error {
 // "intent" (clientRunning) is maintained by the RPC handlers, not by this
 // goroutine.
 func (s *Server) connectWithRetryRuns(ctx context.Context, profileConfig *profilemanager.Config, statusRecorder *peer.Status, runningChan chan struct{}, giveUpChan chan struct{}) {
+	// close(giveUpChan) MUST run on every exit path (DisableAutoConnect
+	// return, backoff.Retry return, panic) — Down() blocks for up to 5s
+	// waiting on this signal before flipping the state to Idle, and a
+	// missed close leaves Down() always hitting the timeout.
 	defer func() {
 		if giveUpChan != nil {
 			close(giveUpChan)
@@ -291,6 +319,15 @@ func (s *Server) connectWithRetryRuns(ctx context.Context, profileConfig *profil
 	runOperation := func() error {
 		err := s.connect(ctx, profileConfig, statusRecorder, runningChan)
 		if err != nil {
+			// PermissionDenied means the daemon transitioned to NeedsLogin
+			// inside connect(). Without backoff.Permanent the outer retry
+			// re-enters connect(), which resets the state to Connecting and
+			// makes the tray flicker between NeedsLogin and Connecting until
+			// the user logs in. Stop retrying and let the state stick.
+			if s, ok := gstatus.FromError(err); ok && s.Code() == codes.PermissionDenied {
+				log.Debugf("run client connection exited with PermissionDenied, waiting for login")
+				return backoff.Permanent(err)
+			}
 			log.Debugf("run client connection exited with error: %v. Will retry in the background", err)
 			return err
 		}
@@ -385,6 +422,16 @@ func (s *Server) SetConfig(callerCtx context.Context, msg *proto.SetConfigReques
 		return nil, fmt.Errorf("failed to update profile config: %w", err)
 	}
 
+	// Apply the lazy connection toggle to the running engine so it takes
+	// effect without a down/up. s.mutex is already held.
+	if msg.LazyConnectionEnabled != nil && s.connectClient != nil {
+		if engine := s.connectClient.Engine(); engine != nil {
+			if err := engine.SetLazyConnEnabled(msg.GetLazyConnectionEnabled()); err != nil {
+				log.Errorf("failed to apply lazy connection change at runtime: %v", err)
+			}
+		}
+	}
+
 	return &proto.SetConfigResponse{}, nil
 }
 
@@ -425,7 +472,7 @@ func (s *Server) setConfigInputFromRequest(msg *proto.SetConfigRequest) (profile
 		wgPort := int(*msg.WireguardPort)
 		config.WireguardPort = &wgPort
 	}
-	if msg.OptionalPreSharedKey != nil && *msg.OptionalPreSharedKey != "" {
+	if msg.OptionalPreSharedKey != nil {
 		config.PreSharedKey = msg.OptionalPreSharedKey
 	}
 
@@ -577,8 +624,6 @@ func (s *Server) Login(callerCtx context.Context, msg *proto.LoginRequest) (*pro
 		return &proto.LoginResponse{}, nil
 	}
 
-	state.Set(internal.StatusConnecting)
-
 	if msg.SetupKey == "" {
 		hint := ""
 		if msg.Hint != nil {
@@ -593,6 +638,7 @@ func (s *Server) Login(callerCtx context.Context, msg *proto.LoginRequest) (*pro
 		if s.oauthAuthFlow.flow != nil && s.oauthAuthFlow.flow.GetClientID(ctx) == oAuthFlow.GetClientID(ctx) {
 			if s.oauthAuthFlow.expiresAt.After(time.Now().Add(90 * time.Second)) {
 				log.Debugf("using previous oauth flow info")
+				state.Set(internal.StatusNeedsLogin)
 				return &proto.LoginResponse{
 					NeedsSSOLogin:           true,
 					VerificationURI:         s.oauthAuthFlow.info.VerificationURI,
@@ -629,6 +675,11 @@ func (s *Server) Login(callerCtx context.Context, msg *proto.LoginRequest) (*pro
 		}, nil
 	}
 
+	// Setup-key path: we are about to dial Management with the key, so the
+	// Connecting paint is meaningful here — unlike the SSO branch above,
+	// which returns NeedsLogin and parks on the browser leg.
+	state.Set(internal.StatusConnecting)
+
 	if loginStatus, err := s.loginAttempt(ctx, msg.SetupKey, ""); err != nil {
 		state.Set(loginStatus)
 		return nil, err
@@ -637,8 +688,43 @@ func (s *Server) Login(callerCtx context.Context, msg *proto.LoginRequest) (*pro
 	return &proto.LoginResponse{}, nil
 }
 
-// WaitSSOLogin uses the userCode to validate the TokenInfo and
-// waits for the user to continue with the login on a browser
+// WaitSSOLogin validates the supplied userCode against the in-flight OAuth
+// device/PKCE flow and blocks until the user finishes the browser leg.
+//
+// The daemon holds StatusNeedsLogin for the whole browser wait (set on
+// entry): the login is not done until the token returns, so a client that
+// (re)attaches mid-wait — a restarted UI, a second `netbird up` — reads
+// "login required" and offers the affordance, instead of a Connecting that
+// never resolves. The wait is also tied to the caller's context (see the
+// goroutine below), so a client that goes away cancels the wait instead of
+// orphaning it on rootCtx until the device-code window expires.
+//
+// State transitions on exit:
+//
+//	┌──────────────────────────────────────────┬──────────────────────────────────┐
+//	│ Outcome                                  │ contextState                     │
+//	├──────────────────────────────────────────┼──────────────────────────────────┤
+//	│ Success → loginAttempt ok                │ NeedsLogin held; the caller's Up │
+//	│                                          │   drives Connecting → Connected  │
+//	│ Success → loginAttempt → still-NeedsLogin│ StatusNeedsLogin (loginAttempt)  │
+//	│ Success → loginAttempt error             │ StatusLoginFailed (loginAttempt) │
+//	│ UserCode mismatch                        │ StatusLoginFailed                │
+//	│ WaitToken: context.Canceled              │ NeedsLogin held. Caller gone     │
+//	│   (caller went away — UI restart /       │   (UI/CLI) → a fresh client      │
+//	│   Ctrl+C — or internal abort: profile    │   shows the login affordance;    │
+//	│   switch / app quit / another            │   internal aborts are            │
+//	│   WaitSSOLogin via actCancel/waitCancel) │   overwritten by the next Up.    │
+//	│ WaitToken: context.DeadlineExceeded      │ StatusNeedsLogin                 │
+//	│   (OAuth device-code window expired      │   (retryable; the UI's "Connect" │
+//	│   while waiting on the browser leg)      │   re-enters the Login flow)      │
+//	│ WaitToken: any other error               │ StatusLoginFailed                │
+//	│   (access_denied, expired_token, HTTP    │   (genuine auth/IO failure;      │
+//	│   failure, token validation rejection)   │   surfaced verbatim to caller)   │
+//	└──────────────────────────────────────────┴──────────────────────────────────┘
+//
+// The defer still applies a StatusIdle fallback for the early
+// oauth-flow-not-initialized return (before the entry Set), so a half state
+// doesn't leak when there is nothing to wait on.
 func (s *Server) WaitSSOLogin(callerCtx context.Context, msg *proto.WaitSSOLoginRequest) (*proto.WaitSSOLoginResponse, error) {
 	s.mutex.Lock()
 	if s.actCancel != nil {
@@ -646,6 +732,21 @@ func (s *Server) WaitSSOLogin(callerCtx context.Context, msg *proto.WaitSSOLogin
 	}
 	ctx, cancel := context.WithCancel(s.rootCtx)
 
+	// Tie the in-flight browser wait to the caller. ctx stays rooted in
+	// rootCtx so CtxGetState resolves the daemon's contextState, but if the
+	// UI window or CLI that drove the login goes away mid-flow (restart,
+	// Ctrl+C) the gRPC callerCtx cancels and we cancel the wait instead of
+	// orphaning it on rootCtx until the OAuth device-code window expires.
+	// The goroutine exits as soon as either context completes, so it can't
+	// outlive the RPC.
+	go func() {
+		select {
+		case <-callerCtx.Done():
+			cancel()
+		case <-ctx.Done():
+		}
+	}()
+
 	md, ok := metadata.FromIncomingContext(callerCtx)
 	if ok {
 		ctx = metadata.NewOutgoingContext(ctx, md)
@@ -671,7 +772,11 @@ func (s *Server) WaitSSOLogin(callerCtx context.Context, msg *proto.WaitSSOLogin
 		}
 	}()
 
-	state.Set(internal.StatusConnecting)
+	// Hold NeedsLogin for the whole browser wait — the login is not done
+	// until the token returns, so a client that (re)attaches mid-wait
+	// (restarted UI, second `netbird up`) reads "login required" and offers
+	// the affordance instead of a Connecting that never resolves.
+	state.Set(internal.StatusNeedsLogin)
 
 	s.mutex.Lock()
 	flowInfo := s.oauthAuthFlow.info
@@ -698,7 +803,30 @@ func (s *Server) WaitSSOLogin(callerCtx context.Context, msg *proto.WaitSSOLogin
 		s.mutex.Lock()
 		s.oauthAuthFlow.expiresAt = time.Now()
 		s.mutex.Unlock()
-		state.Set(internal.StatusLoginFailed)
+		switch {
+		case errors.Is(err, context.Canceled):
+			// External abort. If our caller cancelled (the client closed
+			// the browser-login popup, or the UI went away — callerCtx is
+			// done), clear the abandoned OAuth flow so a fresh Login starts
+			// a new device code instead of reusing this one. The entry
+			// NeedsLogin stays in place, so a reattaching client shows the
+			// login affordance. An internal abort (actCancel from a new
+			// Login/WaitSSOLogin, callerCtx still live) leaves the flow for
+			// the new owner — don't clobber it.
+			if callerCtx.Err() != nil {
+				s.mutex.Lock()
+				s.oauthAuthFlow = oauthAuthFlow{}
+				s.mutex.Unlock()
+			}
+		case errors.Is(err, context.DeadlineExceeded):
+			// OAuth device-code window expired with no user action.
+			// Retryable — leave the daemon in NeedsLogin so the UI
+			// keeps the Login affordance instead of reading as a
+			// hard failure.
+			state.Set(internal.StatusNeedsLogin)
+		default:
+			state.Set(internal.StatusLoginFailed)
+		}
 		log.Errorf("waiting for browser login failed: %v", err)
 		return nil, err
 	}
@@ -755,6 +883,22 @@ func (s *Server) Up(callerCtx context.Context, msg *proto.UpRequest) (*proto.UpR
 		return nil, err
 	}
 
+	// StatusNeedsLogin is a legitimate fresh-start entry state: a successful
+	// WaitSSOLogin deliberately leaves the daemon in NeedsLogin (the login is
+	// done, the token is in hand, but the engine hasn't been brought up yet —
+	// see WaitSSOLogin's state-transition table). The same holds after a
+	// mid-session expiry tore the engine down (clientRunning == false) and the
+	// user re-authenticated. In both cases the caller's Up is expected to drive
+	// the connection; treat NeedsLogin like Idle and reset to Idle so the
+	// engine's own StatusConnecting → StatusConnected progression starts from a
+	// clean slate. Without this, the first Up after an SSO login fails with
+	// "up already in progress" and the user has to trigger Up a second time
+	// (CLI: re-run `netbird up`; GUI: click Connect again).
+	if status == internal.StatusNeedsLogin {
+		status = internal.StatusIdle
+		state.Set(internal.StatusIdle)
+	}
+
 	if status != internal.StatusIdle {
 		s.mutex.Unlock()
 		return nil, fmt.Errorf("up already in progress: current status %s", status)
@@ -817,9 +961,12 @@ func (s *Server) Up(callerCtx context.Context, msg *proto.UpRequest) (*proto.UpR
 	s.clientGiveUpChan = make(chan struct{})
 
 	go s.connectWithRetryRuns(ctx, s.config, s.statusRecorder, s.clientRunningChan, s.clientGiveUpChan)
-	s.publishConfigChangedEvent("up_rpc")
+	s.publishConfigChangedEvent(proto.MetadataSourceUpRPC)
 
 	s.mutex.Unlock()
+	if msg.GetAsync() {
+		return &proto.UpResponse{}, nil
+	}
 	return s.waitForUp(callerCtx)
 }
 
@@ -929,6 +1076,10 @@ func (s *Server) SwitchProfile(callerCtx context.Context, msg *proto.SwitchProfi
 
 	s.config = config
 
+	if msg != nil && msg.ProfileName != nil {
+		s.publishProfileListChanged(*msg.ProfileName)
+	}
+
 	return &proto.SwitchProfileResponse{Id: activeProf.ID.String()}, nil
 }
 
@@ -945,23 +1096,37 @@ func (s *Server) Down(ctx context.Context, _ *proto.DownRequest) (*proto.DownRes
 		return nil, err
 	}
 
-	state := internal.CtxGetState(s.rootCtx)
-	state.Set(internal.StatusIdle)
-
 	s.mutex.Unlock()
 
 	// Wait for the connectWithRetryRuns goroutine to finish with a short timeout.
 	// This prevents the goroutine from setting ErrResetConnection after Down() returns.
-	// The giveUpChan is closed at the end of connectWithRetryRuns.
+	// The giveUpChan is closed by the goroutine's deferred cleanup (see
+	// connectWithRetryRuns) on every exit path. A timeout here typically
+	// means the goroutine is still wedged inside a slow teardown step.
 	if giveUpChan != nil {
 		select {
 		case <-giveUpChan:
-			log.Debugf("client goroutine finished successfully")
+			log.Debugf("client goroutine finished, giveUpChan closed")
 		case <-time.After(5 * time.Second):
 			log.Warnf("timeout waiting for client goroutine to finish, proceeding anyway")
 		}
 	}
 
+	// Set Idle only after the retry goroutine has exited (or timed out).
+	// Setting it earlier races with the goroutine's own Set(StatusConnecting)
+	// at the top of each retry attempt, which would leave the snapshot
+	// stuck at Connecting long after the user asked to disconnect.
+	internal.CtxGetState(s.rootCtx).Set(internal.StatusIdle)
+
+	// Clear stale management/signal errors so the next Up() (typically for a
+	// different profile) starts with a clean status snapshot. Without this,
+	// a managementError left over from a LoginFailed cycle persists in the
+	// statusRecorder and appears in the new profile's initial
+	// SubscribeStatus snapshot, making the new profile look like it also
+	// failed to log in.
+	s.statusRecorder.MarkManagementDisconnected(nil)
+	s.statusRecorder.MarkSignalDisconnected(nil)
+
 	return &proto.DownResponse{}, nil
 }
 
@@ -1172,7 +1337,19 @@ func (s *Server) sendLogoutRequestWithConfig(ctx context.Context, config *profil
 		}
 	}()
 
-	return mgmClient.Logout()
+	if err := mgmClient.Logout(); err != nil {
+		// The peer is already gone from the management server (e.g. deleted
+		// from the dashboard). The logout's goal — deregistering this peer —
+		// is therefore already satisfied, so treat NotFound as success rather
+		// than blocking the logout/profile-removal flow.
+		if logoutPeerGone(err) {
+			log.Infof("peer already removed from management server, treating logout as successful")
+			return nil
+		}
+		return err
+	}
+
+	return nil
 }
 
 // Status returns the daemon status
@@ -1225,9 +1402,24 @@ func (s *Server) Status(
 		}
 	}
 
-	status, err := internal.CtxGetState(s.rootCtx).Status()
+	return s.buildStatusResponse(ctx, msg)
+}
+
+// buildStatusResponse composes a StatusResponse from the current daemon
+// state. Shared between the unary Status RPC and the SubscribeStatus
+// stream so both paths return identical snapshots. ctx scopes the health
+// probe runProbes may trigger — a caller that disconnects cancels it.
+func (s *Server) buildStatusResponse(ctx context.Context, msg *proto.StatusRequest) (*proto.StatusResponse, error) {
+	state := internal.CtxGetState(s.rootCtx)
+	status, err := state.Status()
 	if err != nil {
-		return nil, err
+		// state.Status() blanks the status when err is set (e.g. management
+		// retry loop wrapped a connection error). The underlying status is
+		// still meaningful and the failure is already surfaced via
+		// FullStatus.ManagementState.Error, so don't propagate err — that
+		// would tear down the SubscribeStatus stream and cause the UI to
+		// mark the daemon as unreachable on every retry.
+		status = state.CurrentStatus()
 	}
 
 	if status == internal.StatusNeedsLogin && s.isSessionActive.Load() {
@@ -1238,15 +1430,20 @@ func (s *Server) Status(
 
 	statusResponse := proto.StatusResponse{Status: string(status), DaemonVersion: version.NetbirdVersion()}
 
+	if deadline := s.statusRecorder.GetSessionExpiresAt(); !deadline.IsZero() {
+		statusResponse.SessionExpiresAt = timestamppb.New(deadline)
+	}
+
 	s.statusRecorder.UpdateManagementAddress(s.config.ManagementURL.String())
 	s.statusRecorder.UpdateRosenpass(s.config.RosenpassEnabled, s.config.RosenpassPermissive)
 
 	if msg.GetFullPeerStatus {
-		s.runProbes(msg.ShouldRunProbes)
+		s.runProbes(ctx, msg.ShouldRunProbes)
 		fullStatus := s.statusRecorder.GetFullStatus()
 		pbFullStatus := fullStatus.ToProto()
 		pbFullStatus.Events = s.statusRecorder.GetEventHistory()
 		pbFullStatus.SshServerState = s.getSSHServerState()
+		pbFullStatus.NetworksRevision = s.statusRecorder.GetNetworksRevision()
 		statusResponse.FullStatus = pbFullStatus
 	}
 
@@ -1467,6 +1664,151 @@ func (s *Server) WaitJWTToken(
 	}, nil
 }
 
+// RequestExtendAuthSession initiates the SSO session-extension flow and
+// returns the verification URI the UI should open. The flow state is held
+// in s.extendAuthSessionFlow until WaitExtendAuthSession resolves it.
+func (s *Server) RequestExtendAuthSession(
+	ctx context.Context,
+	msg *proto.RequestExtendAuthSessionRequest,
+) (*proto.RequestExtendAuthSessionResponse, error) {
+	if ctx.Err() != nil {
+		return nil, ctx.Err()
+	}
+
+	s.mutex.Lock()
+	config := s.config
+	connectClient := s.connectClient
+	s.mutex.Unlock()
+
+	if config == nil {
+		return nil, gstatus.Errorf(codes.FailedPrecondition, "client is not configured")
+	}
+	if connectClient == nil {
+		return nil, gstatus.Errorf(codes.FailedPrecondition, "client is not running")
+	}
+
+	hint := ""
+	if msg.Hint != nil {
+		hint = *msg.Hint
+	}
+	if hint == "" {
+		hint = profilemanager.GetLoginHint()
+	}
+
+	isDesktop := isUnixRunningDesktop()
+	oAuthFlow, err := auth.NewOAuthFlow(ctx, config, isDesktop, false, hint)
+	if err != nil {
+		return nil, gstatus.Errorf(codes.Internal, "failed to create OAuth flow: %v", err)
+	}
+
+	authInfo, err := oAuthFlow.RequestAuthInfo(ctx)
+	if err != nil {
+		return nil, gstatus.Errorf(codes.Internal, "failed to request auth info: %v", err)
+	}
+
+	s.extendAuthSessionFlow.Set(oAuthFlow, authInfo)
+
+	return &proto.RequestExtendAuthSessionResponse{
+		VerificationURI:         authInfo.VerificationURI,
+		VerificationURIComplete: authInfo.VerificationURIComplete,
+		UserCode:                authInfo.UserCode,
+		DeviceCode:              authInfo.DeviceCode,
+		ExpiresIn:               int64(authInfo.ExpiresIn),
+	}, nil
+}
+
+// WaitExtendAuthSession blocks until the user completes the SSO step
+// initiated by RequestExtendAuthSession, then forwards the resulting JWT
+// to the management server's ExtendAuthSession RPC. The returned deadline
+// is also applied locally via the engine so SubscribeStatus consumers see
+// the refreshed state.
+func (s *Server) WaitExtendAuthSession(
+	ctx context.Context,
+	req *proto.WaitExtendAuthSessionRequest,
+) (*proto.WaitExtendAuthSessionResponse, error) {
+	if ctx.Err() != nil {
+		return nil, ctx.Err()
+	}
+
+	oAuthFlow, authInfo, ok := s.extendAuthSessionFlow.Get()
+
+	s.mutex.Lock()
+	connectClient := s.connectClient
+	s.mutex.Unlock()
+
+	if !ok || authInfo.DeviceCode != req.DeviceCode {
+		return nil, gstatus.Errorf(codes.InvalidArgument, "invalid device code or no active extend-session flow")
+	}
+
+	// Preempt a previous WaitExtendAuthSession (e.g. when the tray
+	// notification and the about-to-expire dialog both start a flow on
+	// the same deadline). The older waiter exits via context.Canceled;
+	// the new one takes over the IdP poll.
+	s.extendAuthSessionFlow.CancelWait()
+
+	waitCtx, cancel := context.WithCancel(ctx)
+	defer cancel()
+	s.extendAuthSessionFlow.SetWaitCancel(cancel)
+
+	tokenInfo, err := oAuthFlow.WaitToken(waitCtx, authInfo)
+	if err != nil {
+		if errors.Is(err, context.Canceled) {
+			return nil, gstatus.Errorf(codes.Canceled, "extend-session flow preempted")
+		}
+		return nil, gstatus.Errorf(codes.Internal, "failed to obtain JWT token: %v", err)
+	}
+
+	// Clear pending flow before talking to mgm so a retry can re-initiate.
+	s.extendAuthSessionFlow.Clear()
+
+	if connectClient == nil {
+		return nil, gstatus.Errorf(codes.FailedPrecondition, "client is not running")
+	}
+	engine := connectClient.Engine()
+	if engine == nil {
+		return nil, gstatus.Errorf(codes.FailedPrecondition, "engine is not initialised")
+	}
+
+	deadline, err := engine.ExtendAuthSession(ctx, tokenInfo.GetTokenToUse())
+	if err != nil {
+		// Log the full wrapped chain, but return only the innermost gRPC
+		// status (code + clean desc) so the UI shows the root cause, not
+		// the daemon's wrapping layers.
+		log.Errorf("management ExtendAuthSession failed: %v", err)
+		if st := innermostStatus(err); st != nil {
+			return nil, gstatus.Error(st.Code(), st.Message())
+		}
+		return nil, gstatus.Errorf(codes.Internal, "%v", err)
+	}
+
+	resp := &proto.WaitExtendAuthSessionResponse{}
+	if !deadline.IsZero() {
+		resp.SessionExpiresAt = timestamppb.New(deadline)
+	}
+	return resp, nil
+}
+
+// DismissSessionWarning forwards the user's "Dismiss" click on the
+// T-WarningLead notification down to the engine's sessionWatcher so the
+// T-FinalWarningLead fallback is suppressed for the current deadline.
+// Best-effort: when the client/engine is not yet running the call is a
+// successful no-op (the watcher has no deadline to dismiss anyway).
+func (s *Server) DismissSessionWarning(
+	_ context.Context,
+	_ *proto.DismissSessionWarningRequest,
+) (*proto.DismissSessionWarningResponse, error) {
+	s.mutex.Lock()
+	connectClient := s.connectClient
+	s.mutex.Unlock()
+	if connectClient == nil {
+		return &proto.DismissSessionWarningResponse{}, nil
+	}
+	if engine := connectClient.Engine(); engine != nil {
+		engine.DismissSessionWarning()
+	}
+	return &proto.DismissSessionWarningResponse{}, nil
+}
+
 // ExposeService exposes a local port via the NetBird reverse proxy.
 func (s *Server) ExposeService(req *proto.ExposeServiceRequest, srv proto.DaemonService_ExposeServiceServer) error {
 	s.mutex.Lock()
@@ -1533,7 +1875,7 @@ func isUnixRunningDesktop() bool {
 	return os.Getenv("DESKTOP_SESSION") != "" || os.Getenv("XDG_CURRENT_DESKTOP") != ""
 }
 
-func (s *Server) runProbes(waitForProbeResult bool) {
+func (s *Server) runProbes(ctx context.Context, waitForProbeResult bool) {
 	if s.connectClient == nil {
 		return
 	}
@@ -1543,15 +1885,7 @@ func (s *Server) runProbes(waitForProbeResult bool) {
 		return
 	}
 
-	if time.Since(s.lastProbe) > probeThreshold {
-		if engine.RunHealthProbes(waitForProbeResult) {
-			s.lastProbe = time.Now()
-		}
-	} else {
-		if err := s.statusRecorder.RefreshWireGuardStats(); err != nil {
-			log.Debugf("failed to refresh WireGuard stats: %v", err)
-		}
-	}
+	s.probeThrottle.Run(ctx, engine, s.statusRecorder, waitForProbeResult)
 }
 
 // GetConfig of the daemon.
@@ -1681,6 +2015,8 @@ func (s *Server) AddProfile(ctx context.Context, msg *proto.AddProfileRequest) (
 		return nil, fmt.Errorf("failed to create profile: %w", err)
 	}
 
+	s.publishProfileListChanged(msg.ProfileName)
+
 	return &proto.AddProfileResponse{Id: created.ID.String()}, nil
 }
 
@@ -1707,6 +2043,8 @@ func (s *Server) RenameProfile(ctx context.Context, msg *proto.RenameProfileRequ
 		return nil, fmt.Errorf("failed to rename profile: %w", err)
 	}
 
+	s.publishProfileListChanged(msg.NewProfileName)
+
 	return &proto.RenameProfileResponse{OldProfileName: resolved.Name}, nil
 }
 
@@ -1737,9 +2075,51 @@ func (s *Server) RemoveProfile(ctx context.Context, msg *proto.RemoveProfileRequ
 		return nil, fmt.Errorf("failed to remove profile: %w", err)
 	}
 
+	s.publishProfileListChanged(msg.ProfileName)
+
 	return &proto.RemoveProfileResponse{Id: resolved.ID.String()}, nil
 }
 
+// publishProfileListChanged nudges the desktop UI to refresh its profile list
+// after a CLI-driven add/remove. The daemon exposes no dedicated
+// profile-changed RPC event, and a profile add/remove doesn't move the
+// connection status, so the UI's SubscribeStatus path never fires for it (and
+// the tray's status-string guard would swallow it anyway). Instead we publish
+// a marked INFO/SYSTEM event over SubscribeEvents: the UI's dispatchSystemEvent
+// recognises the metadata "kind" marker and translates it into its internal
+// profile-changed signal that both the tray menu and the React profile views
+// already subscribe to (see proto.MetadataKindProfileListChanged, recognised in
+// client/ui/services/daemon_feed.go). userMessage is intentionally empty so this
+// stays a silent refresh signal rather than a user-facing notification.
+func (s *Server) publishProfileListChanged(profileName string) {
+	s.statusRecorder.PublishEvent(
+		proto.SystemEvent_INFO,
+		proto.SystemEvent_SYSTEM,
+		"Profile list changed",
+		"",
+		map[string]string{proto.MetadataKindKey: proto.MetadataKindProfileListChanged, proto.MetadataProfileKey: profileName},
+	)
+}
+
+// publishLogLevelChanged signals the desktop UI that the daemon log level
+// changed, so it can attach/detach its rotated gui-client.log. Like
+// publishProfileListChanged, this rides the SubscribeEvents stream as a marked
+// INFO/SYSTEM event (kind "log-level-changed", level the lowercase logrus
+// name); the UI's dispatchSystemEvent recognises the marker and routes it to
+// the logging toggle instead of an OS toast (userMessage is empty so it stays
+// a silent control signal). The "level" value matches log.Level.String()
+// (e.g. "debug", "info") so the UI can parse it directly. See
+// proto.MetadataKindLogLevelChanged, recognised in client/ui/services/daemon_feed.go.
+func (s *Server) publishLogLevelChanged(level string) {
+	s.statusRecorder.PublishEvent(
+		proto.SystemEvent_INFO,
+		proto.SystemEvent_SYSTEM,
+		"Log level changed",
+		"",
+		map[string]string{proto.MetadataKindKey: proto.MetadataKindLogLevelChanged, proto.MetadataLevelKey: level},
+	)
+}
+
 // ListProfiles lists all profiles in the daemon.
 func (s *Server) ListProfiles(ctx context.Context, msg *proto.ListProfilesRequest) (*proto.ListProfilesResponse, error) {
 	s.mutex.Lock()
@@ -1811,11 +2191,33 @@ func (s *Server) GetFeatures(ctx context.Context, msg *proto.GetFeaturesRequest)
 		DisableProfiles:       s.checkProfilesDisabled(),
 		DisableUpdateSettings: s.checkUpdateSettingsDisabled(),
 		DisableNetworks:       s.checkNetworksDisabled(),
+		DisableAdvancedView:   s.checkDisableAdvancedView(),
 	}
 
 	return features, nil
 }
 
+// WailsUIReady is a no-op the Wails UI probes at startup; merely answering it
+// (rather than returning Unimplemented) tells the UI this daemon is new enough.
+func (s *Server) WailsUIReady(context.Context, *proto.WailsUIReadyRequest) (*proto.WailsUIReadyResponse, error) {
+	return &proto.WailsUIReadyResponse{}, nil
+}
+
+// checkDisableAdvancedView reports the MDM-policy directive for the
+// upcoming UI's advanced-view section. Tristate: returns nil when no
+// MDM directive is set so the UI applies its own default; returns
+// &true / &false when MDM explicitly enforces. No CLI flag backs
+// this feature — MDM is the sole source.
+func (s *Server) checkDisableAdvancedView() *bool {
+	if s.config == nil {
+		return nil
+	}
+	if v, ok := s.config.Policy().GetBool(mdm.KeyDisableAdvancedView); ok {
+		return &v
+	}
+	return nil
+}
+
 func (s *Server) connect(ctx context.Context, config *profilemanager.Config, statusRecorder *peer.Status, runningChan chan struct{}) error {
 	log.Tracef("running client connection")
 	client := internal.NewConnectClient(ctx, config, statusRecorder)
@@ -1985,3 +2387,28 @@ func persistLoginOverrides(activeProf *profilemanager.ActiveProfileState, manage
 	}
 	return nil
 }
+
+// logoutPeerGone reports whether a management Logout failed because the peer
+// no longer exists server-side (gRPC NotFound), walking the wrap chain since
+// the client wraps the gRPC status with fmt.Errorf.
+func logoutPeerGone(err error) bool {
+	for e := err; e != nil; e = errors.Unwrap(e) {
+		if s, ok := gstatus.FromError(e); ok && s.Code() == codes.NotFound {
+			return true
+		}
+	}
+	return false
+}
+
+// innermostStatus walks the wrap chain and returns the deepest gRPC status,
+// or nil when none is present. gstatus.FromError does not unwrap, so a status
+// wrapped with fmt.Errorf %w would otherwise be missed.
+func innermostStatus(err error) *gstatus.Status {
+	var found *gstatus.Status
+	for e := err; e != nil; e = errors.Unwrap(e) {
+		if s, ok := gstatus.FromError(e); ok {
+			found = s
+		}
+	}
+	return found
+}