fix

Author: BenjaminMichaelis

EssentialCSharp.Chat.Shared/Extensions/ServiceCollectionExtensions.cs

@@ -101,7 +101,10 @@ public static IServiceCollection AddAzureOpenAIServices(
 
     /// <summary>
     /// Adds PostgreSQL vector store with managed identity authentication support.
-    /// Uses periodic token refresh to ensure tokens are renewed before expiry.
+    /// Uses per-connection token refresh via <c>UsePasswordProvider</c>, which calls
+    /// <see cref="TokenCredential.GetTokenAsync"/> on every new physical connection.
+    /// <see cref="DefaultAzureCredential"/> caches tokens internally and auto-refreshes
+    /// ~5 minutes before expiry, so this does not add Azure AD overhead.
     /// </summary>
     /// <param name="services">The service collection to add services to</param>
     /// <param name="connectionString">The PostgreSQL connection string (without password)</param>
@@ -133,18 +136,30 @@ private static IServiceCollection AddPostgresVectorStoreWithManagedIdentity(
                     dataSourceBuilder.ConnectionStringBuilder.SslMode = SslMode.Require;
                 }
 
-                // Use periodic token refresh instead of a one-shot token at startup.
-                // Azure AD tokens expire after ~1 hour; refreshing every 50 minutes
-                // ensures uninterrupted connectivity for long-running applications.
-                dataSourceBuilder.UsePeriodicPasswordProvider(
-                    async (_, ct) =>
-                    {
-                        var tokenRequestContext = new TokenRequestContext(_PostgresScopes);
-                        var accessToken = await credential.GetTokenAsync(tokenRequestContext, ct);
-                        return accessToken.Token;
-                    },
-                    TimeSpan.FromMinutes(50),
-                    TimeSpan.FromSeconds(10));
+                var tokenRequestContext = new TokenRequestContext(_PostgresScopes);
+
+                // UsePasswordProvider is called for every new physical connection.
+                // DefaultAzureCredential caches tokens internally and auto-refreshes ~5 min before
+                // expiry — no extra Azure AD load. This is the approach recommended by the Npgsql
+                // docs for cloud providers that implement their own caching (Azure MI does).
+                // UsePeriodicPasswordProvider is only for token sources without built-in caching.
+                // See: https://www.npgsql.org/doc/security.html
+                // See: https://github.com/npgsql/npgsql/issues/5186
+                //
+                // Note: The username is expected to be set in the connection string already
+                // (Aspire sets it during deployment for Azure PostgreSQL Flexible Server).
+                // If a standalone username-extraction fallback is ever needed, use the
+                // Microsoft.Azure.PostgreSQL.Auth package (UseEntraAuthentication extension)
+                // once it ships on NuGet.
+                dataSourceBuilder.UsePasswordProvider(
+                    passwordProvider: _ => credential.GetToken(tokenRequestContext, default).Token,
+                    passwordProviderAsync: async (_, ct) =>
+                        (await credential.GetTokenAsync(tokenRequestContext, ct)).Token);
+
+                // Recycle pooled connections after 50 min, well before the 60-min JWT token TTL.
+                // Combined with UsePasswordProvider (called on every new physical connection),
+                // this ensures no pooled connection ever holds an expired token.
+                dataSourceBuilder.ConnectionStringBuilder.ConnectionLifetime = 3000;
             }
 
             return dataSourceBuilder.Build();

EssentialCSharp.Chat.Shared/Services/AIChatService.cs

@@ -118,14 +118,14 @@ private async Task<string> EnrichPromptWithContext(string prompt, bool enableCon
             return prompt;
         }
 
-        var searchResults = await _SearchService.ExecuteVectorSearch(prompt);
+        var searchResults = await _SearchService.ExecuteVectorSearch(prompt, cancellationToken: cancellationToken);
         var contextualInfo = new System.Text.StringBuilder();
 
         contextualInfo.AppendLine("## Contextual Information");
         contextualInfo.AppendLine("The following information might be relevant to your question:");
         contextualInfo.AppendLine();
 
-        await foreach (var result in searchResults)
+        foreach (var result in searchResults)
         {
             contextualInfo.AppendLine(System.Globalization.CultureInfo.InvariantCulture, $"**From: {result.Record.Heading}**");
             contextualInfo.AppendLine(result.Record.ChunkText);

EssentialCSharp.Chat.Shared/Services/AISearchService.cs

@@ -1,27 +1,54 @@
-using EssentialCSharp.Chat.Common.Models;
+using System.Diagnostics;
+using EssentialCSharp.Chat.Common.Models;
+using Microsoft.Extensions.Logging;
 using Microsoft.Extensions.VectorData;
+using Npgsql;
 
 namespace EssentialCSharp.Chat.Common.Services;
 
-public class AISearchService(VectorStore vectorStore, EmbeddingService embeddingService)
+public class AISearchService(
+    VectorStore vectorStore,
+    EmbeddingService embeddingService,
+    ILogger<AISearchService> logger)
 {
     // TODO: Implement Hybrid Search functionality, may need to switch db providers to support full text search?
 
-    public async Task<IAsyncEnumerable<VectorSearchResult<BookContentChunk>>> ExecuteVectorSearch(string query, string? collectionName = null)
+    public async Task<IReadOnlyList<VectorSearchResult<BookContentChunk>>> ExecuteVectorSearch(
+        string query, string? collectionName = null, CancellationToken cancellationToken = default)
     {
         collectionName ??= EmbeddingService.CollectionName;
 
         VectorStoreCollection<string, BookContentChunk> collection = vectorStore.GetCollection<string, BookContentChunk>(collectionName);
 
-        ReadOnlyMemory<float> searchVector = await embeddingService.GenerateEmbeddingAsync(query);
+        ReadOnlyMemory<float> searchVector = await embeddingService.GenerateEmbeddingAsync(query, cancellationToken);
 
         var vectorSearchOptions = new VectorSearchOptions<BookContentChunk>
         {
             VectorProperty = x => x.TextEmbedding,
         };
 
-        var searchResults = collection.SearchAsync(searchVector, options: vectorSearchOptions, top: 3);
-
-        return searchResults;
+        for (int attempt = 0; attempt <= 1; attempt++)
+        {
+            try
+            {
+                var results = new List<VectorSearchResult<BookContentChunk>>();
+                await foreach (var result in collection.SearchAsync(searchVector, options: vectorSearchOptions, top: 3, cancellationToken: cancellationToken))
+                {
+                    results.Add(result);
+                }
+                return results;
+            }
+            catch (PostgresException ex) when (ex.SqlState == "28000" && attempt == 0)
+            {
+                // The pooled connection held an expired Entra ID token. Npgsql automatically
+                // removes the broken connection from the pool on error — no manual pool clearing
+                // needed (clearing would evict all healthy connections, hurting concurrent users).
+                // The retry opens a fresh physical connection, which calls UsePasswordProvider
+                // and gets a new token from DefaultAzureCredential.
+                logger.LogWarning(ex, "Entra ID token expired on pooled connection (SqlState 28000); retrying once.");
+            }
+        }
+
+        throw new UnreachableException("Retry loop exited without returning or throwing.");
     }
 }