Improved Qdrant server startup & client initialization (#770)

app/MindWork AI Studio/Tools/Databases/DatabaseClient.cs

@@ -4,7 +4,11 @@ public abstract class DatabaseClient(string name, string path)
 {
     public string Name => name;
 
-    public virtual bool IsAvailable => true;
+    public virtual string CacheKey => name;
+
+    public virtual DatabaseClientStatus Status => DatabaseClientStatus.AVAILABLE;
+
+    public bool IsAvailable => this.Status is DatabaseClientStatus.AVAILABLE;
     
     private string Path => path;
 

app/MindWork AI Studio/Tools/Databases/DatabaseClientProvider.cs

@@ -0,0 +1,180 @@
+using AIStudio.Tools.Databases.Qdrant;
+using AIStudio.Tools.Rust;
+using AIStudio.Tools.Services;
+
+namespace AIStudio.Tools.Databases;
+
+public sealed class DatabaseClientProvider(RustService rustService, ILoggerFactory loggerFactory) : IDisposable
+{
+    private readonly Dictionary<DatabaseRole, DatabaseClient> clients = new();
+    private readonly Dictionary<DatabaseRole, SemaphoreSlim> locks = new();
+    private readonly Lock locksLock = new();
+    private readonly ILogger<DatabaseClientProvider> logger = loggerFactory.CreateLogger<DatabaseClientProvider>();
+    private readonly ILogger<DatabaseClient> databaseClientLogger = loggerFactory.CreateLogger<DatabaseClient>();
+
+    public async Task<DatabaseClient> GetClientAsync(DatabaseRole databaseRole, CancellationToken cancellationToken = default)
+    {
+        var databaseLock = this.GetLock(databaseRole);
+        await databaseLock.WaitAsync(cancellationToken);
+        try
+        {
+            if (this.clients.TryGetValue(databaseRole, out var cachedClient) && cachedClient.IsAvailable)
+                return cachedClient;
+
+            var client = await this.CreateClientAsync(databaseRole, cancellationToken);
+            return this.CacheIfAvailable(databaseRole, client);
+        }
+        finally
+        {
+            databaseLock.Release();
+        }
+    }
+
+    public async Task<DatabaseClient> RefreshClientAsync(DatabaseRole databaseRole, CancellationToken cancellationToken = default)
+    {
+        var databaseLock = this.GetLock(databaseRole);
+        await databaseLock.WaitAsync(cancellationToken);
+        try
+        {
+            var client = await this.CreateClientAsync(databaseRole, cancellationToken);
+            return this.CacheIfAvailable(databaseRole, client);
+        }
+        finally
+        {
+            databaseLock.Release();
+        }
+    }
+
+    private DatabaseClient CacheIfAvailable(DatabaseRole databaseRole, DatabaseClient client)
+    {
+        if (!client.IsAvailable)
+            return client;
+
+        if (this.clients.TryGetValue(databaseRole, out var cachedClient))
+        {
+            if (IsSameClient(cachedClient, client))
+            {
+                client.Dispose();
+                return cachedClient;
+            }
+
+            cachedClient.Dispose();
+        }
+
+        this.clients[databaseRole] = client;
+        return client;
+    }
+
+    private SemaphoreSlim GetLock(DatabaseRole databaseRole)
+    {
+        lock (this.locksLock)
+        {
+            if (this.locks.TryGetValue(databaseRole, out var databaseLock))
+                return databaseLock;
+
+            databaseLock = new SemaphoreSlim(1, 1);
+            this.locks[databaseRole] = databaseLock;
+            return databaseLock;
+        }
+    }
+
+    private async Task<DatabaseClient> CreateClientAsync(DatabaseRole databaseRole, CancellationToken cancellationToken) => databaseRole switch
+    {
+        DatabaseRole.VECTOR_STORE => await this.CreateQdrantClientAsync(cancellationToken),
+        _ => new NoDatabaseClient(databaseRole.ToString(), "The requested database role is not supported.")
+    };
+
+    private async Task<DatabaseClient> CreateQdrantClientAsync(CancellationToken cancellationToken)
+    {
+        var qdrantInfo = await rustService.GetQdrantInfo(cancellationToken);
+        if (qdrantInfo.Status is QdrantStatus.STARTING)
+        {
+            return this.CreateNoDatabaseClient(
+                "Qdrant",
+                "Qdrant is starting. Details will appear shortly.",
+                DatabaseClientStatus.STARTING);
+        }
+
+        if (!qdrantInfo.IsAvailable || qdrantInfo.Status is QdrantStatus.UNAVAILABLE)
+        {
+            var reason = qdrantInfo.UnavailableReason ?? "unknown";
+            this.logger.LogWarning("Qdrant is not available. Starting without vector database. Reason: '{Reason}'.", reason);
+            return this.CreateNoDatabaseClient("Qdrant", qdrantInfo.UnavailableReason, DatabaseClientStatus.UNAVAILABLE);
+        }
+
+        if (!HasValidQdrantConnectionInfo(qdrantInfo, out var invalidReason))
+            return this.CreateNoDatabaseClient("Qdrant", invalidReason, DatabaseClientStatus.UNAVAILABLE);
+
+        var client = new QdrantClientImplementation("Qdrant", qdrantInfo.Path, qdrantInfo.PortHttp, qdrantInfo.PortGrpc, qdrantInfo.Fingerprint, qdrantInfo.ApiToken);
+        client.SetLogger(this.databaseClientLogger);
+
+        try
+        {
+            await client.CheckAvailabilityAsync();
+            return client;
+        }
+        catch (Exception e)
+        {
+            client.Dispose();
+            this.logger.LogWarning(e, "Qdrant reported as available by Rust, but the health check failed.");
+            return this.CreateNoDatabaseClient("Qdrant", e.Message, DatabaseClientStatus.STARTING);
+        }
+    }
+
+    private static bool HasValidQdrantConnectionInfo(QdrantInfo qdrantInfo, out string invalidReason)
+    {
+        if (qdrantInfo.Path == string.Empty)
+        {
+            invalidReason = "Failed to get the Qdrant path from Rust.";
+            return false;
+        }
+
+        if (qdrantInfo.PortHttp == 0)
+        {
+            invalidReason = "Failed to get the Qdrant HTTP port from Rust.";
+            return false;
+        }
+
+        if (qdrantInfo.PortGrpc == 0)
+        {
+            invalidReason = "Failed to get the Qdrant gRPC port from Rust.";
+            return false;
+        }
+
+        if (qdrantInfo.Fingerprint == string.Empty)
+        {
+            invalidReason = "Failed to get the Qdrant fingerprint from Rust.";
+            return false;
+        }
+
+        if (qdrantInfo.ApiToken == string.Empty)
+        {
+            invalidReason = "Failed to get the Qdrant API token from Rust.";
+            return false;
+        }
+
+        invalidReason = string.Empty;
+        return true;
+    }
+
+    private NoDatabaseClient CreateNoDatabaseClient(string name, string? unavailableReason, DatabaseClientStatus status)
+    {
+        var client = new NoDatabaseClient(name, unavailableReason, status);
+        client.SetLogger(this.databaseClientLogger);
+        return client;
+    }
+
+    private static bool IsSameClient(DatabaseClient left, DatabaseClient right) =>
+        left.IsAvailable
+        && right.IsAvailable
+        && left.CacheKey == right.CacheKey;
+
+    public void Dispose()
+    {
+        foreach (var client in this.clients.Values)
+            client.Dispose();
+
+        foreach (var databaseLock in this.locks.Values)
+            databaseLock.Dispose();
+    }
+}

app/MindWork AI Studio/Tools/Databases/DatabaseClientStatus.cs

@@ -0,0 +1,8 @@
+namespace AIStudio.Tools.Databases;
+
+public enum DatabaseClientStatus
+{
+    STARTING,
+    AVAILABLE,
+    UNAVAILABLE,
+}

app/MindWork AI Studio/Tools/Databases/DatabaseRole.cs

@@ -0,0 +1,6 @@
+namespace AIStudio.Tools.Databases;
+
+public enum DatabaseRole
+{
+    VECTOR_STORE,
+}

app/MindWork AI Studio/Tools/Databases/NoDatabaseClient.cs

@@ -2,15 +2,19 @@ using AIStudio.Tools.PluginSystem;
 
 namespace AIStudio.Tools.Databases;
 
-public sealed class NoDatabaseClient(string name, string? unavailableReason) : DatabaseClient(name, string.Empty)
+public sealed class NoDatabaseClient(string name, string? unavailableReason, DatabaseClientStatus status = DatabaseClientStatus.UNAVAILABLE) : DatabaseClient(name, string.Empty)
 {
     private static string TB(string fallbackEN) => I18N.I.T(fallbackEN, typeof(NoDatabaseClient).Namespace, nameof(NoDatabaseClient));
     
-    public override bool IsAvailable => false;
+    public override DatabaseClientStatus Status => status;
 
     public override async IAsyncEnumerable<(string Label, string Value)> GetDisplayInfo()
     {
-        yield return (TB("Status"), TB("Unavailable"));
+        yield return (TB("Status"), status switch
+        {
+            DatabaseClientStatus.STARTING => TB("Starting"),
+            _ => TB("Unavailable")
+        });
 
         if (!string.IsNullOrWhiteSpace(unavailableReason))
             yield return (TB("Reason"), unavailableReason);

app/MindWork AI Studio/Tools/Databases/Qdrant/QdrantClientImplementation.cs

@@ -26,6 +26,8 @@ public class QdrantClientImplementation : DatabaseClient
         this.ApiToken = apiToken;
         this.GrpcClient = this.CreateQdrantClient();
     }
+
+    public override string CacheKey => $"{this.Name}:{this.HttpPort}:{this.GrpcPort}:{this.Fingerprint}";
     
     private const string IP_ADDRESS = "localhost";
 
@@ -47,6 +49,11 @@ public class QdrantClientImplementation : DatabaseClient
         return $"v{operation.Version}";
     }
 
+    public async Task CheckAvailabilityAsync()
+    {
+        await this.GrpcClient.HealthAsync();
+    }
+
     private async Task<string> GetCollectionsAmount()
     {
         var operation = await this.GrpcClient.ListCollectionsAsync();

app/MindWork AI Studio/Tools/Rust/QdrantInfo.cs

@@ -5,6 +5,8 @@
 /// </summary>
 public readonly record struct QdrantInfo
 {
+    public QdrantStatus Status { get; init; }
+
     public bool IsAvailable { get; init; }
 
     public string? UnavailableReason { get; init; }

app/MindWork AI Studio/Tools/Rust/QdrantStatus.cs

@@ -0,0 +1,8 @@
+namespace AIStudio.Tools.Rust;
+
+public enum QdrantStatus
+{
+    STARTING,
+    AVAILABLE,
+    UNAVAILABLE,
+}

app/MindWork AI Studio/Tools/Services/RustService.Databases.cs

@@ -4,13 +4,27 @@ namespace AIStudio.Tools.Services;
 
 public sealed partial class RustService
 {
-    public async Task<QdrantInfo> GetQdrantInfo()
+    public async Task<QdrantInfo> GetQdrantInfo(CancellationToken cancellationToken = default)
     {
         try
         {
-            var cts = new CancellationTokenSource(TimeSpan.FromSeconds(45));
-            var response = await this.http.GetFromJsonAsync<QdrantInfo>("/system/qdrant/info", this.jsonRustSerializerOptions, cts.Token);
-            return response;
+            using var cts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
+            cts.CancelAfter(TimeSpan.FromSeconds(45));
+            
+            return await this.http.GetFromJsonAsync<QdrantInfo>("/system/qdrant/info", this.jsonRustSerializerOptions, cts.Token);
+        }
+        catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested)
+        {
+            if(this.logger is not null)
+                this.logger.LogWarning("Fetching Qdrant info from Rust service was cancelled by caller.");
+            else
+                Console.WriteLine("Fetching Qdrant info from Rust service was cancelled by caller.");
+            
+            return new QdrantInfo
+            {
+                Status = QdrantStatus.UNAVAILABLE,
+                UnavailableReason = "Operation cancelled by caller."
+            };
         }
         catch (Exception e)
         {
@@ -19,7 +33,11 @@ public sealed partial class RustService
             else
                 Console.WriteLine($"Error while fetching Qdrant info from Rust service: '{e}'.");
             
-            return default;
+            return new QdrantInfo
+            {
+                Status = QdrantStatus.UNAVAILABLE,
+                UnavailableReason = e.Message
+            };
         }
     }
 }