using AIStudio.Tools.PluginSystem; using AIStudio.Settings; using System.Security.Cryptography; using System.Text; namespace AIStudio.Tools.Services; public sealed class EnterpriseEnvironmentService(ILogger logger, RustService rustService) : BackgroundService { public static List CURRENT_ENVIRONMENTS = []; public static bool HasValidEnterpriseSnapshot { get; private set; } private static EnterpriseSecretSnapshot CURRENT_SECRET_SNAPSHOT; private readonly record struct EnterpriseEnvironmentSnapshot(Guid ConfigurationId, string ConfigurationServerUrl, string? ETag); private readonly record struct EnterpriseSecretSnapshot(bool HasSecret, string Fingerprint); private readonly record struct EnterpriseSecretTarget(string SecretId, string SecretName, SecretStoreType StoreType) : ISecretId; #if DEBUG private static readonly TimeSpan CHECK_INTERVAL = TimeSpan.FromMinutes(6); #else private static readonly TimeSpan CHECK_INTERVAL = TimeSpan.FromMinutes(16); #endif #region Overrides of BackgroundService protected override async Task ExecuteAsync(CancellationToken stoppingToken) { logger.LogInformation("The enterprise environment service was initialized."); await this.StartUpdating(isFirstRun: true); while (!stoppingToken.IsCancellationRequested) { await Task.Delay(CHECK_INTERVAL, stoppingToken); await this.StartUpdating(); } } #endregion private async Task StartUpdating(bool isFirstRun = false) { try { logger.LogInformation("Start updating of the enterprise environment."); HasValidEnterpriseSnapshot = false; var previousSnapshot = BuildNormalizedSnapshot(CURRENT_ENVIRONMENTS); var previousSecretSnapshot = CURRENT_SECRET_SNAPSHOT; // // Step 1: Fetch all active configurations. // List fetchedConfigs; try { fetchedConfigs = await rustService.EnterpriseEnvConfigs(); } catch (Exception e) { logger.LogError(e, "Failed to fetch the enterprise configurations from the Rust service."); await MessageBus.INSTANCE.SendMessage(null, Event.RUST_SERVICE_UNAVAILABLE, "EnterpriseEnvConfigs failed"); return; } string enterpriseEncryptionSecret; try { enterpriseEncryptionSecret = await rustService.EnterpriseEnvConfigEncryptionSecret(); } catch (Exception e) { logger.LogError(e, "Failed to fetch the enterprise encryption secret from the Rust service."); await MessageBus.INSTANCE.SendMessage(null, Event.RUST_SERVICE_UNAVAILABLE, "EnterpriseEnvConfigEncryptionSecret failed"); return; } var nextSecretSnapshot = await BuildSecretSnapshot(enterpriseEncryptionSecret); var wasSecretChanged = previousSecretSnapshot != nextSecretSnapshot; // // Step 2: Determine ETags and build the list of reachable configurations. // IMPORTANT: when one config server fails, we continue with the others. // var reachableEnvironments = new List(); var failedConfigIds = new HashSet(); var currentEnvironmentsById = CURRENT_ENVIRONMENTS .GroupBy(env => env.ConfigurationId) .ToDictionary(group => group.Key, group => group.Last()); var activeFetchedEnvironmentsById = fetchedConfigs .Where(config => config.IsActive) .GroupBy(config => config.ConfigurationId) .ToDictionary(group => group.Key, group => group.Last()); foreach (var config in fetchedConfigs) { if (!config.IsActive) { logger.LogWarning("Skipping inactive enterprise configuration with ID '{ConfigId}'. There is either no valid server URL or config ID set.", config.ConfigurationId); continue; } var etagResponse = await PluginFactory.DetermineConfigPluginETagAsync(config.ConfigurationId, config.ConfigurationServerUrl); if (!etagResponse.Success) { failedConfigIds.Add(config.ConfigurationId); logger.LogWarning("Failed to read enterprise config metadata for '{ConfigId}' from '{ServerUrl}': {Issue}. Keeping the current plugin state for this configuration.", config.ConfigurationId, config.ConfigurationServerUrl, etagResponse.Issue ?? "Unknown issue"); continue; } reachableEnvironments.Add(config with { ETag = etagResponse.ETag }); } // // Step 3: Compare with current environments and process changes. // Download per configuration. A single failure must not block others. // var shouldDeferStartupDownloads = isFirstRun && !PluginFactory.IsInitialized; var effectiveEnvironmentsById = new Dictionary(); // Process new or changed configs: foreach (var nextEnv in reachableEnvironments) { var hasCurrentEnvironment = currentEnvironmentsById.TryGetValue(nextEnv.ConfigurationId, out var currentEnv); if (hasCurrentEnvironment && currentEnv == nextEnv) // Hint: This relies on the record equality to check if anything relevant has changed (e.g. server URL or ETag). { logger.LogInformation("Enterprise configuration '{ConfigId}' has not changed. No update required.", nextEnv.ConfigurationId); effectiveEnvironmentsById[nextEnv.ConfigurationId] = nextEnv; continue; } if(!hasCurrentEnvironment) logger.LogInformation("Detected new enterprise configuration with ID '{ConfigId}' and server URL '{ServerUrl}'.", nextEnv.ConfigurationId, nextEnv.ConfigurationServerUrl); else logger.LogInformation("Detected change in enterprise configuration with ID '{ConfigId}'. Server URL or ETag has changed.", nextEnv.ConfigurationId); if (shouldDeferStartupDownloads) { MessageBus.INSTANCE.DeferMessage(null, Event.STARTUP_ENTERPRISE_ENVIRONMENT, nextEnv); effectiveEnvironmentsById[nextEnv.ConfigurationId] = nextEnv; } else { var wasDownloadSuccessful = await PluginFactory.TryDownloadingConfigPluginAsync(nextEnv.ConfigurationId, nextEnv.ConfigurationServerUrl); if (!wasDownloadSuccessful) { failedConfigIds.Add(nextEnv.ConfigurationId); if (hasCurrentEnvironment) { logger.LogWarning("Failed to update enterprise configuration '{ConfigId}'. Keeping the previously active version.", nextEnv.ConfigurationId); effectiveEnvironmentsById[nextEnv.ConfigurationId] = currentEnv; } else logger.LogWarning("Failed to download the new enterprise configuration '{ConfigId}'. Skipping activation for now.", nextEnv.ConfigurationId); continue; } effectiveEnvironmentsById[nextEnv.ConfigurationId] = nextEnv; } } // Retain configurations for all failed IDs. On cold start there might be no // previous in-memory snapshot yet, so we also keep the current fetched entry // to protect it from cleanup while the server is unreachable. foreach (var failedConfigId in failedConfigIds) { if (effectiveEnvironmentsById.ContainsKey(failedConfigId)) continue; if (!currentEnvironmentsById.TryGetValue(failedConfigId, out var retainedEnvironment)) { if (!activeFetchedEnvironmentsById.TryGetValue(failedConfigId, out retainedEnvironment)) continue; logger.LogWarning("Could not refresh enterprise configuration '{ConfigId}'. Protecting it from cleanup until connectivity is restored.", failedConfigId); } else logger.LogWarning("Could not refresh enterprise configuration '{ConfigId}'. Keeping the previously active version.", failedConfigId); effectiveEnvironmentsById[failedConfigId] = retainedEnvironment; } var effectiveEnvironments = effectiveEnvironmentsById.Values.ToList(); // Cleanup is only allowed after a successful sync cycle: if (PluginFactory.IsInitialized && !shouldDeferStartupDownloads) PluginFactory.RemoveUnreferencedManagedConfigurationPlugins(effectiveEnvironmentsById.Keys.ToHashSet()); if (effectiveEnvironments.Count == 0) logger.LogInformation("AI Studio runs without any enterprise configurations."); var effectiveSnapshot = BuildNormalizedSnapshot(effectiveEnvironments); if (PluginFactory.IsInitialized && wasSecretChanged) { logger.LogInformation("The enterprise encryption secret changed. Refreshing the enterprise encryption service and reloading plugins."); PluginFactory.InitializeEnterpriseEncryption(enterpriseEncryptionSecret); await this.RemoveEnterpriseManagedApiKeysAsync(); await PluginFactory.LoadAll(); } CURRENT_ENVIRONMENTS = effectiveEnvironments; CURRENT_SECRET_SNAPSHOT = nextSecretSnapshot; HasValidEnterpriseSnapshot = true; if (!previousSnapshot.SequenceEqual(effectiveSnapshot) || wasSecretChanged) await MessageBus.INSTANCE.SendMessage(null, Event.ENTERPRISE_ENVIRONMENTS_CHANGED); } catch (Exception e) { logger.LogError(e, "An error occurred while updating the enterprise environment."); } } private static List BuildNormalizedSnapshot(IEnumerable environments) { return environments .Where(environment => environment.IsActive) .Select(environment => new EnterpriseEnvironmentSnapshot( environment.ConfigurationId, NormalizeServerUrl(environment.ConfigurationServerUrl), environment.ETag?.ToString())) .OrderBy(environment => environment.ConfigurationId) .ToList(); } private static async Task BuildSecretSnapshot(string secret) { if (string.IsNullOrWhiteSpace(secret)) return new EnterpriseSecretSnapshot(false, string.Empty); return new EnterpriseSecretSnapshot(true, await ComputeSecretFingerprint(secret)); } private static async Task ComputeSecretFingerprint(string secret) { using var secretStream = new MemoryStream(Encoding.UTF8.GetBytes(secret)); var hash = await SHA256.HashDataAsync(secretStream); return Convert.ToHexString(hash); } private static string NormalizeServerUrl(string serverUrl) { return serverUrl.Trim().TrimEnd('/'); } private async Task RemoveEnterpriseManagedApiKeysAsync() { var secretTargets = GetEnterpriseManagedSecretTargets(); if (secretTargets.Count == 0) { logger.LogInformation("No enterprise-managed API keys are currently known in the settings. No keyring cleanup is required."); return; } logger.LogInformation("Removing {SecretCount} enterprise-managed API key(s) from the OS keyring after an enterprise encryption secret change.", secretTargets.Count); foreach (var target in secretTargets) { try { var deleteResult = await rustService.DeleteAPIKey(target, target.StoreType); if (deleteResult.Success) { if (deleteResult.WasEntryFound) logger.LogInformation("Successfully deleted enterprise-managed API key '{SecretName}' from the OS keyring.", target.SecretName); else logger.LogInformation("Enterprise-managed API key '{SecretName}' was already absent from the OS keyring.", target.SecretName); } else logger.LogWarning("Failed to delete enterprise-managed API key '{SecretName}' from the OS keyring: {Issue}", target.SecretName, deleteResult.Issue); } catch (Exception e) { logger.LogWarning(e, "Failed to delete enterprise-managed API key '{SecretName}' from the OS keyring.", target.SecretName); } } } private static List GetEnterpriseManagedSecretTargets() { var configurationData = Program.SERVICE_PROVIDER.GetRequiredService().ConfigurationData; var secretTargets = new HashSet(); AddEnterpriseManagedSecretTargets(configurationData.Providers, SecretStoreType.LLM_PROVIDER, secretTargets); AddEnterpriseManagedSecretTargets(configurationData.EmbeddingProviders, SecretStoreType.EMBEDDING_PROVIDER, secretTargets); AddEnterpriseManagedSecretTargets(configurationData.TranscriptionProviders, SecretStoreType.TRANSCRIPTION_PROVIDER, secretTargets); return secretTargets.ToList(); } private static void AddEnterpriseManagedSecretTargets( IEnumerable secrets, SecretStoreType storeType, ISet secretTargets) where TSecret : ISecretId, IConfigurationObject { foreach (var secret in secrets) { if (!secret.IsEnterpriseConfiguration || secret.EnterpriseConfigurationPluginId == Guid.Empty) continue; secretTargets.Add(new EnterpriseSecretTarget(secret.SecretId, secret.SecretName, storeType)); } } }