AI-Studio/app/MindWork AI Studio/Program.cs

using AIStudio.Agents;
using AIStudio.Agents.AssistantAudit;
using AIStudio.Assistants.VisualBriefing;
using AIStudio.Settings;
using AIStudio.Tools.ToolCallingSystem;
using AIStudio.Tools.Databases;
using AIStudio.Tools.AIJobs;
using AIStudio.Tools.AssistantSessions;
using AIStudio.Tools.Media;
using AIStudio.Tools.PluginSystem;
using AIStudio.Tools.PluginSystem.Assistants;
using AIStudio.Tools.Rust;
using AIStudio.Tools.Security;
using AIStudio.Tools.Services;
using AIStudio.Tools.ToolCallingSystem.Harness;
using AIStudio.Tools.ToolCallingSystem.ToolCallingImplementations;
using AIStudio.Tools.ToolCallingSystem.ToolCallingImplementations.WebSearch;
using AIStudio.Tools.ToolCallingSystem.ToolCallingImplementations.WebSearch.SearXNG;
using AIStudio.Tools.ToolCallingSystem.ToolCallingImplementations.WebSearch.Staan;
using AIStudio.Tools.ToolCallingSystem.ToolCallingImplementations.WebSearch.Tavily;
using AIStudio.Tools.Web;

using Microsoft.AspNetCore.Components.Server.Circuits;
using Microsoft.AspNetCore.DataProtection;
using Microsoft.AspNetCore.Server.Kestrel.Core;
using Microsoft.Extensions.Logging.Console;

using MudBlazor.Services;

using MudExtensions.Services;

#if !DEBUG
using System.Reflection;
using Microsoft.Extensions.FileProviders;
#endif

namespace AIStudio;

internal sealed class Program
{
    public static RustService RUST_SERVICE = null!;
    public static Encryption ENCRYPTION = null!;
    public static string API_TOKEN = null!;
    public static IServiceProvider SERVICE_PROVIDER = null!;
    public static ILoggerFactory LOGGER_FACTORY = null!;
    public static DatabaseClientProvider DATABASE_CLIENT_PROVIDER = null!;
    
    public static async Task Main()
    {
        #if DEBUG
        // Read the environment variables from the .env file:
        var envFilePath = Path.Combine("..", "..", "startup.env");
        await EnvFile.Apply(envFilePath);
        #endif
        
        // Read the secret key for the IPC from the AI_STUDIO_SECRET_KEY environment variable:
        var secretPasswordEncoded = Environment.GetEnvironmentVariable("AI_STUDIO_SECRET_PASSWORD");
        if(string.IsNullOrWhiteSpace(secretPasswordEncoded))
        {
            Console.WriteLine("Error: The AI_STUDIO_SECRET_PASSWORD environment variable is not set.");
            return;
        }

        var secretPassword = Convert.FromBase64String(secretPasswordEncoded);
        var secretKeySaltEncoded = Environment.GetEnvironmentVariable("AI_STUDIO_SECRET_KEY_SALT");
        if(string.IsNullOrWhiteSpace(secretKeySaltEncoded))
        {
            Console.WriteLine("Error: The AI_STUDIO_SECRET_KEY_SALT environment variable is not set.");
            return;
        }

        var secretKeySalt = Convert.FromBase64String(secretKeySaltEncoded);
        
        var certificateFingerprint = Environment.GetEnvironmentVariable("AI_STUDIO_CERTIFICATE_FINGERPRINT");
        if(string.IsNullOrWhiteSpace(certificateFingerprint))
        {
            Console.WriteLine("Error: The AI_STUDIO_CERTIFICATE_FINGERPRINT environment variable is not set.");
            return;
        }
        
        var rustApiPort = Environment.GetEnvironmentVariable("AI_STUDIO_API_PORT");
        if(string.IsNullOrWhiteSpace(rustApiPort))
        {
            Console.WriteLine("Error: The AI_STUDIO_API_PORT environment variable is not set.");
            return;
        }
        
        var apiToken = Environment.GetEnvironmentVariable("AI_STUDIO_API_TOKEN");
        if(string.IsNullOrWhiteSpace(apiToken))
        {
            Console.WriteLine("Error: The AI_STUDIO_API_TOKEN environment variable is not set.");
            return;
        }
        
        API_TOKEN = apiToken;
        
        using var rust = new RustService(rustApiPort, certificateFingerprint);
        var appPort = await rust.GetAppPort();
        if(appPort == 0)
        {
            Console.WriteLine("Error: Failed to get the app port from Rust.");
            return;
        }
        
        var runtimeInfo = await rust.GetRuntimeInfo();
        var builder = WebApplication.CreateBuilder();
        builder.WebHost.ConfigureKestrel(kestrelServerOptions =>
        {
            kestrelServerOptions.ConfigureEndpointDefaults(listenOptions =>
            {
                listenOptions.Protocols = HttpProtocols.Http1AndHttp2AndHttp3;
            });
        });
        
        builder.Logging.ClearProviders();
        builder.Logging.SetMinimumLevel(LogLevel.Debug);
        builder.Logging.AddFilter("Microsoft", LogLevel.Information);
        builder.Logging.AddFilter("Microsoft.AspNetCore.Hosting.Diagnostics", LogLevel.Warning);
        builder.Logging.AddFilter("Microsoft.AspNetCore.Routing.EndpointMiddleware", LogLevel.Warning);
        builder.Logging.AddFilter("Microsoft.AspNetCore.StaticFiles", LogLevel.Warning);
        builder.Logging.AddFilter("MudBlazor", LogLevel.Information);
        builder.Logging.AddConsole(options =>
        {
            options.FormatterName = TerminalLogger.FORMATTER_NAME;
        }).AddConsoleFormatter<TerminalLogger, ConsoleFormatterOptions>();

        if(runtimeInfo.LinuxPackageType is LinuxPackageType.FLATPAK)
        {
            try
            {
                var tauriDataDirectory = await rust.GetDataDirectory();
                if(string.IsNullOrWhiteSpace(tauriDataDirectory))
                    throw new InvalidOperationException("Rust returned an empty Tauri data directory.");

                var dataProtectionKeysDirectory = Path.Combine(tauriDataDirectory, "data-protection-keys");
                Directory.CreateDirectory(dataProtectionKeysDirectory);
                var writeTestPath = Path.Combine(dataProtectionKeysDirectory, $".write-test-{Guid.NewGuid():N}");
                using (new FileStream(writeTestPath, FileMode.CreateNew, FileAccess.Write, FileShare.None, 1, FileOptions.DeleteOnClose))
                {
                }

                builder.Services.AddDataProtection()
                    .PersistKeysToFileSystem(new DirectoryInfo(dataProtectionKeysDirectory))
                    .SetApplicationName("org.mindworkai.AIStudio");
            }
            catch(Exception exception)
            {
                Console.WriteLine($"Error: Failed to configure Flatpak data-protection keys in the Tauri data directory: {exception.Message}");
                return;
            }
        }

        builder.Services.AddMudExtensions();
        builder.Services.AddMudServices(config =>
        {
            config.SnackbarConfiguration.PositionClass = Defaults.Classes.Position.BottomLeft;
            config.SnackbarConfiguration.PreventDuplicates = false;
            config.SnackbarConfiguration.NewestOnTop = false;
            config.SnackbarConfiguration.ShowCloseIcon = true;
            config.SnackbarConfiguration.VisibleStateDuration = 6_000; //milliseconds aka 6 seconds
            config.SnackbarConfiguration.HideTransitionDuration = 500;
            config.SnackbarConfiguration.ShowTransitionDuration = 500;
            config.SnackbarConfiguration.SnackbarVariant = Variant.Outlined;
        });

        builder.Services.AddMemoryCache(); // Needed for the Markdown library
        builder.Services.AddMudMarkdownServices();
        builder.Services.AddSingleton(new MudTheme());
        builder.Services.AddSingleton(MessageBus.INSTANCE);
        builder.Services.AddSingleton(rust);
        builder.Services.AddSingleton(typeof(RuntimeInfoResponse), runtimeInfo);
        builder.Services.AddMudMarkdownClipboardService<MarkdownClipboardService>();
        builder.Services.AddSingleton<SettingsManager>();
        builder.Services.AddSingleton<PromptInjectionGuardService>();
        builder.Services.AddSingleton<ToolSettingsService>();
        builder.Services.AddSingleton<WebPageRetrievalService>();
        builder.Services.AddSingleton<IToolImplementation, ReadWebPageTool>();
        builder.Services.AddSingleton<IWebSearchBackend, SearXNGSearchBackend>();
        builder.Services.AddSingleton<IWebSearchBackend, StaanSearchBackend>();
        builder.Services.AddSingleton<IWebSearchBackend, TavilySearchBackend>();
        builder.Services.AddSingleton<IToolImplementation, WebSearchTool>();
        builder.Services.AddSingleton<IToolDefinitionSource, CodeToolDefinitionSource>();
        builder.Services.AddSingleton<ToolRegistry>();
        builder.Services.AddSingleton<ToolExecutor>();
        builder.Services.AddSingleton<IToolCallingLoop, ToolCallingLoop>();
        builder.Services.AddSingleton<ThreadSafeRandom>();
        builder.Services.AddSingleton<AIJobService>();
        builder.Services.AddSingleton<AssistantSessionService>();
        builder.Services.AddSingleton<VoiceRecordingAvailabilityService>();
        builder.Services.AddSingleton<GlobalShortcutService>();
        builder.Services.AddSingleton<MediaTranscriptionService>();
        builder.Services.AddSingleton<VisualBriefingArtifactService>();
        builder.Services.AddSingleton<VisualBriefingStore>();
        builder.Services.AddSingleton<VisualBriefingBuildProgressService>();
        builder.Services.AddSingleton<VisualBriefingBuildOrchestrator>();
        builder.Services.AddSingleton<VisualBriefingPreviewTokenService>();
        builder.Services.AddSingleton<IMediaTranscriptStorage, VisualBriefingTranscriptStorage>();
        builder.Services.AddSingleton<PluginInstallService>();
        builder.Services.AddSingleton<UpdatePolicy>();
        builder.Services.AddSingleton<AssistantPluginGenerationService>();
        builder.Services.AddSingleton<DataSourceService>();
        builder.Services.AddSingleton<DataSourceEmbeddingService>();
        builder.Services.AddSingleton<DataSourceLocalRetrievalService>();
        builder.Services.AddSingleton<DirectChatService>();
        builder.Services.AddScoped<PandocAvailabilityService>();
        
        // Stateless: every method works on its arguments alone, so one instance serves everyone.
        builder.Services.AddSingleton<HTMLParser>();
        builder.Services.AddTransient<AgentDataSourceSelection>();
        builder.Services.AddTransient<AgentRetrievalContextValidation>();
        builder.Services.AddTransient<AgentTextContentCleaner>();
        builder.Services.AddTransient<AssistantAuditAgent>();
        builder.Services.AddTransient<AssistantPluginAuditService>();
        builder.Services.AddHostedService<UpdateService>();
        builder.Services.AddHostedService<TemporaryChatService>();
        builder.Services.AddHostedService<TranscriptStagingCleanupService>();
        builder.Services.AddHostedService<EnterpriseEnvironmentService>();
        builder.Services.AddHostedService(sp => sp.GetRequiredService<DataSourceEmbeddingService>());
        builder.Services.AddSingleton<DatabaseClientProvider>();
        builder.Services.AddHostedService<GlobalShortcutService>(serviceProvider => serviceProvider.GetRequiredService<GlobalShortcutService>());
        builder.Services.AddHostedService<RustAvailabilityMonitorService>();
        builder.Services.AddScoped<NativeShareService>();
        builder.Services.AddScoped<PluginShareService>();

        //
        // One circuit state per circuit, and the handler which keeps it up to date. Both are scoped,
        // because the circuit is the scope: every browser window gets its own pair.
        //
        builder.Services.AddScoped<CircuitStateService>();
        builder.Services.AddScoped<CircuitHandler, AIStudioCircuitHandler>();
        
        // ReSharper disable AccessToDisposedClosure
        builder.Services.AddHostedService<RustService>(_ => rust);
        // ReSharper restore AccessToDisposedClosure
        
        builder.Services.AddRazorComponents()
            .AddInteractiveServerComponents(options =>
            {
                //
                // We keep disconnected circuits for a long time on purpose: when the machine goes to
                // sleep, the WebView loses its connection. Without this retention period, the user would
                // return to a lost app state after waking up the machine (cf. issue #849). Since AI Studio
                // is a single-user desktop app, at most two circuits are retained, which bounds the memory
                // this costs us.
                //
                options.DisconnectedCircuitRetentionPeriod = TimeSpan.FromDays(30);
                options.DisconnectedCircuitMaxRetained = 2;
            })
            .AddHubOptions(options =>
            {
                options.MaximumReceiveMessageSize = null;
                options.ClientTimeoutInterval = TimeSpan.FromSeconds(120);
                options.HandshakeTimeout = TimeSpan.FromSeconds(30);
                options.KeepAliveInterval = TimeSpan.FromSeconds(30);
            });

        builder.Services.AddSingleton(new HttpClient
        {
            BaseAddress = new Uri($"http://localhost:{appPort}")
        });

        builder.WebHost.UseUrls($"http://localhost:{appPort}");

        #if DEBUG
        builder.WebHost.UseWebRoot("wwwroot");
        builder.WebHost.UseStaticWebAssets();
        #endif

        // Execute the builder to get the app:
        var app = builder.Build();

        // Get the logging factory for e.g., static classes:
        LOGGER_FACTORY = app.Services.GetRequiredService<ILoggerFactory>();
        MessageBus.INSTANCE.Initialize(LOGGER_FACTORY.CreateLogger<MessageBus>());
        
        // Get a program logger:
        var programLogger = app.Services.GetRequiredService<ILogger<Program>>();
        programLogger.LogInformation("Starting the AI Studio server.");

        //
        // Observe tasks whose exceptions nobody awaited. We register this before the server starts:
        // otherwise, everything the startup does — the plugin system, the first message bus traffic —
        // would fault outside of this handler. The sender of such a task says nothing about where it
        // came from, which is why we log each inner exception with its own stack trace.
        //
        TaskScheduler.UnobservedTaskException += (sender, taskArgs) =>
        {
            programLogger.LogError(taskArgs.Exception, $"Unobserved task exception by sender '{sender ?? "n/a"}'.");
            foreach (var innerException in taskArgs.Exception.Flatten().InnerExceptions)
                programLogger.LogError(innerException, $"Unobserved task exception detail: {innerException.GetType().FullName}.");

            taskArgs.SetObserved();
        };

        // Store the service provider (DI). We need it later for some classes,
        // which are not part of the request pipeline:
        SERVICE_PROVIDER = app.Services;
        
        // Initialize the encryption service:
        programLogger.LogInformation("Initializing the encryption service.");
        var encryptionLogger = app.Services.GetRequiredService<ILogger<Encryption>>();
        var encryption = new Encryption(encryptionLogger, secretPassword, secretKeySalt);
        var encryptionInitializer = encryption.Initialize();

        // Set the logger for the Rust service:
        programLogger.LogInformation("Initializing the Rust service.");
        var rustLogger = app.Services.GetRequiredService<ILogger<RustService>>();
        rust.SetLogger(rustLogger);
        rust.SetEncryptor(encryption);
        TerminalLogger.SetRustService(rust);

        RUST_SERVICE = rust;
        ENCRYPTION = encryption;
        
        DATABASE_CLIENT_PROVIDER = app.Services.GetRequiredService<DatabaseClientProvider>();

        programLogger.LogInformation("Initialize internal file system.");
        app.Use(Redirect.HandlerContentAsync);
        app.Use(FileHandler.HandlerAsync);

#if DEBUG
        app.UseStaticFiles();
        app.UseDeveloperExceptionPage();
#else
        var fileProvider = new ManifestEmbeddedFileProvider(Assembly.GetAssembly(type: typeof(Program))!, "wwwroot");
        app.UseStaticFiles(new StaticFileOptions
        {
            FileProvider = fileProvider,
            RequestPath = string.Empty,
        });
#endif

        app.UseAntiforgery();

        // Serves committed briefing revisions to the assistant's live preview iframe:
        app.MapVisualBriefingPreview();

        app.MapRazorComponents<App>()
            .AddInteractiveServerRenderMode();

        var serverTask = app.RunAsync();
        programLogger.LogInformation("Server was started successfully.");

        await encryptionInitializer;
        await rust.AppIsReady();
        programLogger.LogInformation("The AI Studio server is ready.");

        await serverTask;
        
        RUST_SERVICE.Dispose();
        DATABASE_CLIENT_PROVIDER.Dispose();
        PluginFactory.Dispose();
        programLogger.LogInformation("The AI Studio server was stopped.");
    }
}