Skip to content

Commit a0f8435

Browse files
sfmskywalkersfmskywalker
committed
Harden AI orchestration lifetimes
1 parent 35c4982 commit a0f8435

6 files changed

Lines changed: 178 additions & 35 deletions

File tree

‎src/modules/Elsa.AI.Host/Context/AiContextResolver.cs‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,11 @@
11
using System.Text.RegularExpressions;
22
using Elsa.AI.Abstractions.Contracts;
33
using Elsa.AI.Abstractions.Models;
4+
using Microsoft.Extensions.DependencyInjection;
45

56
namespace Elsa.AI.Host.Context;
67

7-
public class AiContextResolver(IEnumerable<IAiContextProvider> providers)
8+
public class AiContextResolver(IServiceScopeFactory scopeFactory)
89
{
910
private static readonly string[] SensitiveKeyFragments = ["secret", "token", "password", "apikey", "api_key", "api-key", "authorization", "credential", "bearer"];
1011
private static readonly Regex[] SensitiveValuePatterns =
@@ -13,16 +14,17 @@ public class AiContextResolver(IEnumerable<IAiContextProvider> providers)
1314
new(@"\b(?:api[_-]?key|token|secret|password)\s*[:=]\s*['""]?[A-Za-z0-9._~+/\-=]{8,}['""]?", RegexOptions.IgnoreCase | RegexOptions.CultureInvariant)
1415
];
1516
private const string Redacted = "[redacted]";
16-
private readonly Dictionary<string, IAiContextProvider> _providers = BuildProviders(providers);
1717

1818
public async ValueTask<IReadOnlyCollection<AiResolvedContext>> ResolveAsync(AiChatRequest request, CancellationToken cancellationToken = default)
1919
{
20+
using var scope = scopeFactory.CreateScope();
21+
var providers = BuildProviders(scope.ServiceProvider.GetServices<IAiContextProvider>());
2022
var resolved = new List<AiResolvedContext>();
2123
var attachmentsWithProviders = request.Attachments
2224
.Select(attachment => new
2325
{
2426
Attachment = attachment,
25-
Provider = _providers.GetValueOrDefault(attachment.Kind)
27+
Provider = providers.GetValueOrDefault(attachment.Kind)
2628
})
2729
.Where(x => x.Provider != null);
2830

‎src/modules/Elsa.AI.Host/Services/AiOrchestrator.cs‎

Lines changed: 27 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -70,7 +70,7 @@ public async IAsyncEnumerable<AiStreamEvent> ExecuteChatAsync(AiChatRequest requ
7070
var knownToolCallIds = RestoreToolResults(messages).Select(x => x.ToolCallId).ToHashSet(StringComparer.OrdinalIgnoreCase);
7171
var pendingToolResults = isDuplicateReconnectMessage ? RestorePendingToolResults(messages) : new List<AiToolTurnResult>();
7272

73-
await SaveConversationAsync(conversationId, request, AiConversationStatus.Active, messages, conversation, providerSessionId, cancellationToken);
73+
await TrySaveConversationAsync(conversationId, request, AiConversationStatus.Active, messages, conversation, providerSessionId, cancellationToken);
7474
await RecordChatAuditAsync("chat.started", request, conversationId, provider?.Name, cancellationToken);
7575

7676
IReadOnlyCollection<AiResolvedContext> context = [];
@@ -102,7 +102,7 @@ public async IAsyncEnumerable<AiStreamEvent> ExecuteChatAsync(AiChatRequest requ
102102
["content"] = content
103103
});
104104
messages.Add(CreateMessage(conversationId, AiMessageRole.Assistant, content, sequence - 1));
105-
await SaveConversationAsync(conversationId, request, AiConversationStatus.Failed, messages, conversation, providerSessionId, cancellationToken);
105+
await TrySaveConversationAsync(conversationId, request, AiConversationStatus.Failed, messages, conversation, providerSessionId, cancellationToken);
106106
await RecordChatAuditAsync("chat.failed", request, conversationId, provider?.Name, cancellationToken);
107107
yield return CreateEvent("conversation.completed", conversationId, sequence);
108108
yield break;
@@ -191,7 +191,7 @@ public async IAsyncEnumerable<AiStreamEvent> ExecuteChatAsync(AiChatRequest requ
191191
providerHistory.Add(userMessage);
192192

193193
providerHistory.AddRange(currentTurnMessages);
194-
await SaveConversationAsync(conversationId, request, AiConversationStatus.Active, messages, conversation, providerSessionId, cancellationToken);
194+
await TrySaveConversationAsync(conversationId, request, AiConversationStatus.Active, messages, conversation, providerSessionId, cancellationToken);
195195

196196
if (turn == MaxProviderTurns - 1)
197197
{
@@ -205,7 +205,7 @@ public async IAsyncEnumerable<AiStreamEvent> ExecuteChatAsync(AiChatRequest requ
205205
}
206206
}
207207

208-
await SaveConversationAsync(conversationId, request, AiConversationStatus.Completed, messages, conversation, providerSessionId, cancellationToken);
208+
await TrySaveConversationAsync(conversationId, request, AiConversationStatus.Completed, messages, conversation, providerSessionId, cancellationToken);
209209
await RecordChatAuditAsync("chat.completed", request, conversationId, provider?.Name, cancellationToken);
210210
yield return CreateEvent("conversation.completed", conversationId, sequence);
211211
}
@@ -571,25 +571,32 @@ private static bool BelongsToUser(AiConversation conversation, string userId) =>
571571
private static string NormalizeTenantId(string? tenantId) =>
572572
string.IsNullOrWhiteSpace(tenantId) ? "" : tenantId;
573573

574-
private async ValueTask SaveConversationAsync(string conversationId, AiChatRequest request, AiConversationStatus status, IReadOnlyCollection<AiMessage> messages, AiConversation? conversation, string? providerSessionId, CancellationToken cancellationToken)
574+
private async ValueTask TrySaveConversationAsync(string conversationId, AiChatRequest request, AiConversationStatus status, IReadOnlyCollection<AiMessage> messages, AiConversation? conversation, string? providerSessionId, CancellationToken cancellationToken)
575575
{
576-
var now = DateTimeOffset.UtcNow;
577-
var retentionMode = conversation?.RetentionMode ?? AiRetentionMode.Configured;
578-
var retentionExpiresAt = conversation?.RetentionExpiresAt ?? (retentionMode == AiRetentionMode.Configured ? now.Add(options.Value.ConversationRetention) : null);
576+
try
577+
{
578+
var now = DateTimeOffset.UtcNow;
579+
var retentionMode = conversation?.RetentionMode ?? AiRetentionMode.Configured;
580+
var retentionExpiresAt = conversation?.RetentionExpiresAt ?? (retentionMode == AiRetentionMode.Configured ? now.Add(options.Value.ConversationRetention) : null);
579581

580-
await conversationStore.SaveAsync(new AiConversation
582+
await conversationStore.SaveAsync(new AiConversation
583+
{
584+
Id = conversationId,
585+
TenantId = request.TenantId,
586+
UserId = request.UserId,
587+
Status = status,
588+
CreatedAt = conversation is null || conversation.CreatedAt == default ? now : conversation.CreatedAt,
589+
UpdatedAt = now,
590+
ProviderSessionId = providerSessionId ?? conversation?.ProviderSessionId,
591+
RetentionMode = retentionMode,
592+
RetentionExpiresAt = retentionExpiresAt,
593+
Messages = messages.ToList()
594+
}, cancellationToken);
595+
}
596+
catch (Exception e) when (e is not OperationCanceledException)
581597
{
582-
Id = conversationId,
583-
TenantId = request.TenantId,
584-
UserId = request.UserId,
585-
Status = status,
586-
CreatedAt = conversation is null || conversation.CreatedAt == default ? now : conversation.CreatedAt,
587-
UpdatedAt = now,
588-
ProviderSessionId = providerSessionId ?? conversation?.ProviderSessionId,
589-
RetentionMode = retentionMode,
590-
RetentionExpiresAt = retentionExpiresAt,
591-
Messages = messages.ToList()
592-
}, cancellationToken);
598+
logger.LogWarning(e, "Failed to persist AI conversation {ConversationId} with status {ConversationStatus}.", conversationId, status);
599+
}
593600
}
594601

595602
private readonly record struct ToolCall(string Id, string Name, JsonObject Arguments);

‎src/modules/Elsa.AI.Host/Services/AiToolRegistry.cs‎

Lines changed: 12 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -21,18 +21,22 @@ public ValueTask<IReadOnlyCollection<AiToolDefinition>> ListAsync(AiToolQuery qu
2121
public ValueTask<IAiTool?> FindAsync(string name, AiToolQuery query, CancellationToken cancellationToken = default)
2222
{
2323
var scope = scopeFactory.CreateScope();
24-
var tool = scope.ServiceProvider.GetServices<IAiTool>().FirstOrDefault(x => string.Equals(x.Definition.Name, name, StringComparison.OrdinalIgnoreCase));
25-
if (tool == null || !IsVisible(tool.Definition, query) || !enablementService.IsEnabled(tool.Definition))
24+
try
2625
{
27-
scope.Dispose();
28-
tool = null;
26+
var tool = scope.ServiceProvider.GetServices<IAiTool>().FirstOrDefault(x => string.Equals(x.Definition.Name, name, StringComparison.OrdinalIgnoreCase));
27+
if (tool == null || !IsVisible(tool.Definition, query) || !enablementService.IsEnabled(tool.Definition))
28+
{
29+
scope.Dispose();
30+
return ValueTask.FromResult<IAiTool?>(null);
31+
}
32+
33+
return ValueTask.FromResult<IAiTool?>(new ScopedAiTool(scope, tool));
2934
}
30-
else
35+
catch
3136
{
32-
tool = new ScopedAiTool(scope, tool);
37+
scope.Dispose();
38+
throw;
3339
}
34-
35-
return ValueTask.FromResult(tool);
3640
}
3741

3842
private static bool IsVisible(AiToolDefinition definition, AiToolQuery query) =>

‎test/integration/Elsa.AI.IntegrationTests/AiChatEndpointTests.cs‎

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -241,6 +241,30 @@ public async Task ChatOrchestrationContinuesWhenAuditSinkFails()
241241
Assert.Contains(events, x => x.Type == "conversation.completed");
242242
}
243243

244+
[Fact(DisplayName = "Chat orchestration continues when conversation persistence fails")]
245+
public async Task ChatOrchestrationContinuesWhenConversationPersistenceFails()
246+
{
247+
var services = new ServiceCollection();
248+
services.AddAiHostServices();
249+
services.RemoveAll<IAiConversationStore>();
250+
services.AddSingleton<IAiConversationStore, ThrowingConversationStore>();
251+
using var provider = services.BuildServiceProvider();
252+
var orchestrator = provider.GetRequiredService<IAiOrchestrator>();
253+
var events = new List<AiStreamEvent>();
254+
255+
await foreach (var streamEvent in orchestrator.ExecuteChatAsync(new AiChatRequest
256+
{
257+
ConversationId = "conversation-1",
258+
UserId = "user-1",
259+
Message = "Explain this workflow"
260+
}))
261+
events.Add(streamEvent);
262+
263+
Assert.Contains(events, x => x.Type == "conversation.started");
264+
Assert.Contains(events, x => x.Type == "assistant.delta");
265+
Assert.Contains(events, x => x.Type == "conversation.completed");
266+
}
267+
244268
[Fact(DisplayName = "Chat orchestration emits terminal events when context resolution fails")]
245269
public async Task ChatOrchestrationEmitsTerminalEventsWhenContextResolutionFails()
246270
{
@@ -1288,4 +1312,13 @@ private class ThrowingAuditSink : IAiAuditSink
12881312
public ValueTask RecordAsync(AiAuditEvent auditEvent, CancellationToken cancellationToken = default) =>
12891313
throw new InvalidOperationException("Audit sink unavailable.");
12901314
}
1315+
1316+
private class ThrowingConversationStore : IAiConversationStore
1317+
{
1318+
public ValueTask<AiConversation?> FindAsync(string id, CancellationToken cancellationToken = default) =>
1319+
ValueTask.FromResult<AiConversation?>(null);
1320+
1321+
public ValueTask SaveAsync(AiConversation conversation, CancellationToken cancellationToken = default) =>
1322+
throw new InvalidOperationException("Conversation store unavailable.");
1323+
}
12911324
}

‎test/unit/Elsa.AI.Host.UnitTests/AiToolRegistryTests.cs‎

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -171,6 +171,22 @@ public async Task ToolRegistryHonorsTenantAndActorAllowlists()
171171
Assert.Equal("matching", tool.Name);
172172
}
173173

174+
[Fact(DisplayName = "Tool registry disposes find scope when tool resolution throws")]
175+
public async Task ToolRegistryDisposesFindScopeWhenToolResolutionThrows()
176+
{
177+
var tracker = new ScopeDisposalTracker();
178+
var services = new ServiceCollection();
179+
services.AddSingleton(tracker);
180+
services.AddScoped<ScopedDependency>();
181+
services.AddScoped<IAiTool, ThrowingDefinitionTool>();
182+
using var provider = services.BuildServiceProvider(new ServiceProviderOptions { ValidateScopes = true });
183+
var registry = new AiToolRegistry(provider.GetRequiredService<IServiceScopeFactory>(), new AiToolEnablementService());
184+
185+
await Assert.ThrowsAsync<InvalidOperationException>(async () => await registry.FindAsync("throwing", new AiToolQuery()));
186+
187+
Assert.Equal(1, tracker.DisposeCount);
188+
}
189+
174190
private static AiToolRegistry CreateRegistry(IReadOnlyCollection<IAiTool> tools)
175191
{
176192
var services = new ServiceCollection();
@@ -188,4 +204,34 @@ private class TestTool(AiToolDefinition definition) : IAiTool
188204
public ValueTask<AiToolResult> ExecuteAsync(AiToolExecutionContext context, CancellationToken cancellationToken = default) =>
189205
ValueTask.FromResult(new AiToolResult());
190206
}
207+
208+
private class ThrowingDefinitionTool(ScopedDependency dependency) : IAiTool
209+
{
210+
private readonly ScopedDependency _dependency = dependency;
211+
212+
public AiToolDefinition Definition
213+
{
214+
get
215+
{
216+
_ = _dependency;
217+
throw new InvalidOperationException("Definition unavailable.");
218+
}
219+
}
220+
221+
public ValueTask<AiToolResult> ExecuteAsync(AiToolExecutionContext context, CancellationToken cancellationToken = default) =>
222+
ValueTask.FromResult(new AiToolResult());
223+
}
224+
225+
private class ScopedDependency(ScopeDisposalTracker tracker) : IDisposable
226+
{
227+
public void Dispose()
228+
{
229+
tracker.DisposeCount++;
230+
}
231+
}
232+
233+
private class ScopeDisposalTracker
234+
{
235+
public int DisposeCount { get; set; }
236+
}
191237
}

‎test/unit/Elsa.AI.Host.UnitTests/Context/AiContextResolverTests.cs‎

Lines changed: 55 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
using Elsa.AI.Abstractions.Contracts;
22
using Elsa.AI.Abstractions.Models;
33
using Elsa.AI.Host.Context;
4+
using Microsoft.Extensions.DependencyInjection;
45
using System.Text.Json.Nodes;
56

67
namespace Elsa.AI.Host.UnitTests.Context;
@@ -10,7 +11,8 @@ public class AiContextResolverTests
1011
[Fact(DisplayName = "Context resolver resolves workflow definition references server-side")]
1112
public async Task ContextResolverResolvesWorkflowDefinitionReferences()
1213
{
13-
var resolver = new AiContextResolver([new WorkflowDefinitionContextProvider()]);
14+
using var provider = CreateProvider(services => services.AddSingleton<IAiContextProvider, WorkflowDefinitionContextProvider>());
15+
var resolver = provider.GetRequiredService<AiContextResolver>();
1416

1517
var result = await resolver.ResolveAsync(new AiChatRequest
1618
{
@@ -34,7 +36,8 @@ public async Task ContextResolverResolvesWorkflowDefinitionReferences()
3436
[Fact(DisplayName = "Context resolver redacts sensitive data and metadata keys")]
3537
public async Task ContextResolverRedactsSensitiveDataAndMetadataKeys()
3638
{
37-
var resolver = new AiContextResolver([new SensitiveContextProvider()]);
39+
using var provider = CreateProvider(services => services.AddSingleton<IAiContextProvider, SensitiveContextProvider>());
40+
var resolver = provider.GetRequiredService<AiContextResolver>();
3841

3942
var result = await resolver.ResolveAsync(new AiChatRequest
4043
{
@@ -59,7 +62,12 @@ public async Task ContextResolverRedactsSensitiveDataAndMetadataKeys()
5962
[Fact(DisplayName = "Context resolver uses the last provider for duplicate provider kinds")]
6063
public async Task ContextResolverUsesTheLastProviderForDuplicateProviderKinds()
6164
{
62-
var resolver = new AiContextResolver([new DuplicateContextProvider("first"), new DuplicateContextProvider("second")]);
65+
using var provider = CreateProvider(services =>
66+
{
67+
services.AddSingleton<IAiContextProvider>(new DuplicateContextProvider("first"));
68+
services.AddSingleton<IAiContextProvider>(new DuplicateContextProvider("second"));
69+
});
70+
var resolver = provider.GetRequiredService<AiContextResolver>();
6371

6472
var result = await resolver.ResolveAsync(new AiChatRequest
6573
{
@@ -74,7 +82,12 @@ public async Task ContextResolverUsesTheLastProviderForDuplicateProviderKinds()
7482
[Fact(DisplayName = "Context resolver prefers real providers over placeholders")]
7583
public async Task ContextResolverPrefersRealProvidersOverPlaceholders()
7684
{
77-
var resolver = new AiContextResolver([new DuplicateContextProvider("real", "WorkflowDefinition"), new WorkflowDefinitionContextProvider()]);
85+
using var provider = CreateProvider(services =>
86+
{
87+
services.AddSingleton<IAiContextProvider>(new DuplicateContextProvider("real", "WorkflowDefinition"));
88+
services.AddSingleton<IAiContextProvider, WorkflowDefinitionContextProvider>();
89+
});
90+
var resolver = provider.GetRequiredService<AiContextResolver>();
7891

7992
var result = await resolver.ResolveAsync(new AiChatRequest
8093
{
@@ -86,6 +99,30 @@ public async Task ContextResolverPrefersRealProvidersOverPlaceholders()
8699
Assert.Equal("real", context.Summary);
87100
}
88101

102+
[Fact(DisplayName = "Context resolver resolves scoped providers from request scopes")]
103+
public async Task ContextResolverResolvesScopedProvidersFromRequestScopes()
104+
{
105+
using var provider = CreateProvider(services => services.AddScoped<IAiContextProvider, ScopedContextProvider>());
106+
var resolver = provider.GetRequiredService<AiContextResolver>();
107+
108+
var result = await resolver.ResolveAsync(new AiChatRequest
109+
{
110+
UserId = "user-1",
111+
Attachments = [new AiContextAttachment { Kind = "Scoped" }]
112+
});
113+
114+
var context = Assert.Single(result);
115+
Assert.Equal("scoped", context.Summary);
116+
}
117+
118+
private static ServiceProvider CreateProvider(Action<IServiceCollection> configure)
119+
{
120+
var services = new ServiceCollection();
121+
services.AddSingleton<AiContextResolver>();
122+
configure(services);
123+
return services.BuildServiceProvider(new ServiceProviderOptions { ValidateScopes = true });
124+
}
125+
89126
private class SensitiveContextProvider : IAiContextProvider
90127
{
91128
public string Kind => "Sensitive";
@@ -123,4 +160,18 @@ public ValueTask<AiResolvedContext> ResolveAsync(AiContextResolutionRequest requ
123160
});
124161
}
125162
}
163+
164+
private class ScopedContextProvider : IAiContextProvider
165+
{
166+
public string Kind => "Scoped";
167+
168+
public ValueTask<AiResolvedContext> ResolveAsync(AiContextResolutionRequest request, CancellationToken cancellationToken = default)
169+
{
170+
return ValueTask.FromResult(new AiResolvedContext
171+
{
172+
Kind = Kind,
173+
Summary = "scoped"
174+
});
175+
}
176+
}
126177
}

0 commit comments

Comments
 (0)