diff --git a/ast-visual-studio-extension/CxPreferences/CxPreferencesUI.cs b/ast-visual-studio-extension/CxPreferences/CxPreferencesUI.cs index f0fe4e2b..e4268548 100644 --- a/ast-visual-studio-extension/CxPreferences/CxPreferencesUI.cs +++ b/ast-visual-studio-extension/CxPreferences/CxPreferencesUI.cs @@ -25,6 +25,8 @@ public partial class CxPreferencesUI : UserControl private static bool _isAuthenticated; internal static event Action AuthStateChanged; private static int _restoreAuthInProgress; + private static bool _isFreshLoginInProgress = false; + private static bool _welcomeDialogVisible = false; private bool _isValidationInProgress; private bool _isInitializing; private bool _hasLoaded; @@ -155,6 +157,8 @@ private async void OnValidateConnection(object sender, EventArgs e) if (_isValidationInProgress || string.IsNullOrWhiteSpace(tbApiKey.Text)) return; + _isFreshLoginInProgress = true; + UpdateAuthControlsState(); SetValidationMessage(CxConstants.AUTH_VALIDATE_IN_PROGRESS, isSuccess: true); await ValidateApiKeyAsync(showErrorOnFailure: true); } @@ -259,38 +263,54 @@ await Task.Run(() => // Success message shown inside CompleteAuthenticationSetupAsync after license check _ = CompleteAuthenticationSetupAsync(GetCxConfig(), showWelcomeDialog: showErrorOnFailure); + + // DO NOT update UI here - let CompleteAuthenticationSetupAsync handle it + // to ensure Welcome dialog flag is set before button is enabled } catch (CxException ex) when (showErrorOnFailure) { SetAuthState(false); + _isFreshLoginInProgress = false; + _welcomeDialogVisible = false; if (cxPreferencesModule != null) cxPreferencesModule.RestoreAuthenticatedSession = false; SetValidationMessage(string.Format(CxConstants.AUTH_VALIDATE_FAIL_TEMPLATE, SanitizeErrorMessage(ex.Message)), isSuccess: false); + _isValidationInProgress = false; + UpdateAuthControlsState(); } catch (Exception ex) when (showErrorOnFailure) { SetAuthState(false); + _isFreshLoginInProgress = false; + _welcomeDialogVisible = false; if (cxPreferencesModule != null) cxPreferencesModule.RestoreAuthenticatedSession = false; SetValidationMessage(CxConstants.AUTH_VALIDATE_ERROR, isSuccess: false); System.Diagnostics.Debug.WriteLine($"Authentication error: {LogForgingSanitizer.StripLineTermination(ex.Message)}"); + _isValidationInProgress = false; + UpdateAuthControlsState(); } catch { SetAuthState(false); + _isFreshLoginInProgress = false; + _welcomeDialogVisible = false; if (cxPreferencesModule != null) cxPreferencesModule.RestoreAuthenticatedSession = false; + _isValidationInProgress = false; + UpdateAuthControlsState(); } finally { _isValidationInProgress = false; - UpdateAuthControlsState(); } } private void ResetAuthState() { SetAuthState(false); + _isFreshLoginInProgress = false; + _welcomeDialogVisible = false; ClearValidationMessage(); } @@ -473,6 +493,8 @@ private async Task CompleteAuthenticationSetupAsync(CxConfig config, bool showWe if (config == null || string.IsNullOrWhiteSpace(config.ApiKey)) return; + _isFreshLoginInProgress = false; + try { var package = cxPreferencesModule?.GetOwnerPackage() as AsyncPackage; @@ -483,7 +505,7 @@ private async Task CompleteAuthenticationSetupAsync(CxConfig config, bool showWe bool previousMcpEnabled = oneAssistModule?.McpEnabled ?? false; bool mcpStatusPreviouslyChecked = oneAssistModule?.McpStatusChecked ?? false; - // Step 1: Check license + // Step 1: Check license (network call, can take time) bool hadAssist = await ApplyAssistTenantLicenseAndMcpFlagsAsync(package, config, GetType()); // Step 2: Show success message after license check @@ -496,24 +518,54 @@ private async Task CompleteAuthenticationSetupAsync(CxConfig config, bool showWe if (oneAssistModule == null) return; - // Step 3: Install MCP if enabled - if (oneAssistModule.McpEnabled) - { - var installService = new McpInstallService(); - await installService.InstallSilentlyAsync(config, GetType()); - } - if (showWelcomeDialog) { - // Step 4: Show welcome dialog — policy, persist and scan start ONLY after user closes it + // Step 3a: Show welcome dialog immediately (don't wait for MCP installation) ShowOneAssistWelcomeDialog(oneAssistModule, oneAssistModule.McpEnabled, previousMcpEnabled, mcpStatusPreviouslyChecked); + + // Step 3b: Install MCP in background while user is viewing the Welcome dialog + if (oneAssistModule.McpEnabled) + { + _ = Task.Run(async () => + { + try + { + var installService = new McpInstallService(); + await installService.InstallSilentlyAsync(config, GetType()); + } + catch (Exception ex) + { + System.Diagnostics.Debug.WriteLine($"Background MCP installation failed: {ex.Message}"); + } + }); + } } else { // No welcome dialog (background restore) — apply policy and persist immediately + _welcomeDialogVisible = true; + UpdateAuthControlsState(); + ApplyJetBrainsStyleRealtimeScannerPolicy(oneAssistModule, previousMcpEnabled, mcpStatusPreviouslyChecked); oneAssistModule.PersistSettings(); CxOneAssistSettingsUI.GetInstance()?.RefreshCheckboxesFromModule(); + + // Install MCP in background for restored sessions too + if (oneAssistModule.McpEnabled) + { + _ = Task.Run(async () => + { + try + { + var installService = new McpInstallService(); + await installService.InstallSilentlyAsync(config, GetType()); + } + catch (Exception ex) + { + System.Diagnostics.Debug.WriteLine($"Background MCP installation failed (restore): {ex.Message}"); + } + }); + } } } catch (Exception ex) @@ -544,6 +596,10 @@ private void ShowOneAssistWelcomeDialog(CxOneAssistSettingsModule module, bool m { using (var welcomeDialog = new CxOneAssistWelcomeDialog(module, mcpEnabled)) { + // Mark welcome dialog as visible before showing it + _welcomeDialogVisible = true; + UpdateAuthControlsState(); + welcomeDialog.ShowDialog(FindForm()); module.WelcomeShown = true; } @@ -579,7 +635,7 @@ private void UpdateAuthControlsState() : System.Drawing.SystemColors.Window; button1.Enabled = hasApiKey && !_isAuthenticated && !_isValidationInProgress; - btnLogout.Enabled = _isAuthenticated && !_isValidationInProgress; + btnLogout.Enabled = _isAuthenticated && !_isValidationInProgress && (!_isFreshLoginInProgress || _welcomeDialogVisible); // Temporary: hide assist navigation link from the Authentication page. goToAssistLink.Visible = false;