| | | 1 | | using System.Collections.Concurrent; |
| | | 2 | | |
| | | 3 | | namespace Elsa.AI.Host.Streaming; |
| | | 4 | | |
| | | 5 | | public class AIStreamSessionManager |
| | | 6 | | { |
| | 13 | 7 | | private readonly ConcurrentDictionary<string, ReconnectState> _reconnectStates = new(StringComparer.OrdinalIgnoreCas |
| | | 8 | | |
| | | 9 | | public void MarkDisconnected(string conversationId, TimeSpan graceWindow) |
| | | 10 | | { |
| | 8 | 11 | | var now = DateTimeOffset.UtcNow; |
| | 8 | 12 | | _reconnectStates[conversationId] = new ReconnectState(now.Add(graceWindow), false); |
| | 8 | 13 | | PruneExpired(now); |
| | 8 | 14 | | } |
| | | 15 | | |
| | | 16 | | public bool CanReconnect(string conversationId) |
| | | 17 | | { |
| | | 18 | | while (true) |
| | | 19 | | { |
| | 20 | 20 | | var now = DateTimeOffset.UtcNow; |
| | 20 | 21 | | PruneExpired(now); |
| | | 22 | | |
| | 20 | 23 | | if (!_reconnectStates.TryGetValue(conversationId, out var state)) |
| | 10 | 24 | | return false; |
| | | 25 | | |
| | 10 | 26 | | if (state.Deadline < now) |
| | | 27 | | { |
| | 0 | 28 | | _reconnectStates.TryRemove(conversationId, out _); |
| | 0 | 29 | | return false; |
| | | 30 | | } |
| | | 31 | | |
| | 10 | 32 | | if (state.IsReserved) |
| | 1 | 33 | | return false; |
| | | 34 | | |
| | 9 | 35 | | if (_reconnectStates.TryUpdate(conversationId, state with { IsReserved = true }, state)) |
| | 9 | 36 | | return true; |
| | | 37 | | } |
| | | 38 | | } |
| | | 39 | | |
| | | 40 | | public void MarkConnected(string conversationId) |
| | | 41 | | { |
| | 2 | 42 | | _reconnectStates.TryRemove(conversationId, out _); |
| | 2 | 43 | | } |
| | | 44 | | |
| | | 45 | | public void ReleaseReconnect(string conversationId) |
| | | 46 | | { |
| | | 47 | | while (true) |
| | | 48 | | { |
| | 3 | 49 | | var now = DateTimeOffset.UtcNow; |
| | 3 | 50 | | PruneExpired(now); |
| | | 51 | | |
| | 3 | 52 | | if (!_reconnectStates.TryGetValue(conversationId, out var state)) |
| | 0 | 53 | | return; |
| | | 54 | | |
| | 3 | 55 | | if (state.Deadline < now) |
| | | 56 | | { |
| | 0 | 57 | | _reconnectStates.TryRemove(conversationId, out _); |
| | 0 | 58 | | return; |
| | | 59 | | } |
| | | 60 | | |
| | 3 | 61 | | if (!state.IsReserved) |
| | 0 | 62 | | return; |
| | | 63 | | |
| | 3 | 64 | | if (_reconnectStates.TryUpdate(conversationId, state with { IsReserved = false }, state)) |
| | 3 | 65 | | return; |
| | | 66 | | } |
| | | 67 | | } |
| | | 68 | | |
| | | 69 | | private void PruneExpired(DateTimeOffset now) |
| | | 70 | | { |
| | 112 | 71 | | foreach (var (conversationId, state) in _reconnectStates) |
| | | 72 | | { |
| | 25 | 73 | | if (state.Deadline < now) |
| | 1 | 74 | | _reconnectStates.TryRemove(conversationId, out _); |
| | | 75 | | } |
| | 31 | 76 | | } |
| | | 77 | | |
| | 63 | 78 | | private readonly record struct ReconnectState(DateTimeOffset Deadline, bool IsReserved); |
| | | 79 | | } |