Skip to content

Commit 82b831d

Browse files
author
Jicheng Lu
committed
add session reconnect
1 parent d91a552 commit 82b831d

File tree

18 files changed

+250
-141
lines changed

18 files changed

+250
-141
lines changed

‎src/Infrastructure/BotSharp.Abstraction/Infrastructures/Attributes/BotSharpAuthAttribute.cs

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,32 @@
1+
using BotSharp.Abstraction.Users;
2+
using Microsoft.AspNetCore.Mvc;
3+
using Microsoft.AspNetCore.Mvc.Filters;
4+
using Microsoft.Extensions.DependencyInjection;
5+
6+
namespace BotSharp.Abstraction.Infrastructures.Attributes;
7+
8+
/// <summary>
9+
/// BotSharp authorization: check whether the request user is admin or root role.
10+
/// </summary>
11+
[AttributeUsage(AttributeTargets.Class | AttributeTargets.Method)]
12+
public class BotSharpAuthAttribute : Attribute, IAsyncAuthorizationFilter
13+
{
14+
public BotSharpAuthAttribute()
15+
{
16+
17+
}
18+
19+
public async Task OnAuthorizationAsync(AuthorizationFilterContext context)
20+
{
21+
var services = context.HttpContext.RequestServices;
22+
23+
var userIdentity = services.GetRequiredService<IUserIdentity>();
24+
var userService = services.GetRequiredService<IUserService>();
25+
26+
var (isAdmin, user) = await userService.IsAdminUser(userIdentity.Id);
27+
if (!isAdmin || user == null)
28+
{
29+
context.Result = new BadRequestResult();
30+
}
31+
}
32+
}

‎src/Infrastructure/BotSharp.Abstraction/MLTasks/IRealTimeCompletion.cs

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,11 +14,13 @@ Task Connect(
1414
Func<string, string, Task> onModelAudioDeltaReceived,
1515
Func<Task> onModelAudioResponseDone,
1616
Func<string, Task> onModelAudioTranscriptDone,
17-
Func<List<RoleDialogModel>, Task> onModelResponseDone,
17+
Func<List<RoleDialogModel>, Task<bool>> onModelResponseDone,
1818
Func<string, Task> onConversationItemCreated,
1919
Func<RoleDialogModel, Task> onInputAudioTranscriptionDone,
2020
Func<Task> onInterruptionDetected);
2121

22+
Task Reconnect(RealtimeHubConnection conn);
23+
2224
Task AppenAudioBuffer(string message);
2325
Task AppenAudioBuffer(ArraySegment<byte> data, int length);
2426

‎src/Infrastructure/BotSharp.Abstraction/Realtime/IRealtimeHook.cs

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
using BotSharp.Abstraction.Hooks;
22
using BotSharp.Abstraction.MLTasks;
3+
using BotSharp.Abstraction.Realtime.Models;
34

45
namespace BotSharp.Abstraction.Realtime;
56

@@ -8,4 +9,5 @@ public interface IRealtimeHook : IHookBase
89
Task OnModelReady(Agent agent, IRealTimeCompletion completer);
910
string[] OnModelTranscriptPrompt(Agent agent);
1011
Task OnTranscribeCompleted(RoleDialogModel message, TranscriptionData data);
12+
Task<bool> ShouldReconnect(RealtimeHubConnection conn) => Task.FromResult(false);
1113
}

‎src/Infrastructure/BotSharp.Abstraction/Utilities/StringExtensions.cs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@ namespace BotSharp.Abstraction.Utilities;
55

66
public static class StringExtensions
77
{
8-
public static string IfNullOrEmptyAs(this string str, string defaultValue)
8+
public static string IfNullOrEmptyAs(this string? str, string defaultValue)
99
=> string.IsNullOrEmpty(str) ? defaultValue : str;
1010

1111
public static string SubstringMax(this string str, int maxLength)

‎src/Infrastructure/BotSharp.Core.Realtime/Services/RealtimeHub.cs

Lines changed: 12 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -98,17 +98,15 @@ await HookEmitter.Emit<IRoutingHook>(_services, async hook => await hook.OnRouti
9898
}
9999

100100
await routing.InvokeFunction(message.FunctionName, message);
101-
dialogs.Add(message);
102-
storage.Append(_conn.ConversationId, message);
103101
}
104102
else
105103
{
106104
// append output audio transcript to conversation
107105
dialogs.Add(message);
108106
storage.Append(_conn.ConversationId, message);
109107

110-
var hooks = _services.GetHooksOrderByPriority<IConversationHook>(_conn.CurrentAgentId);
111-
foreach (var hook in hooks)
108+
var convHooks = _services.GetHooksOrderByPriority<IConversationHook>(_conn.CurrentAgentId);
109+
foreach (var hook in convHooks)
112110
{
113111
hook.SetAgent(agent)
114112
.SetConversation(conversation);
@@ -117,6 +115,16 @@ await HookEmitter.Emit<IRoutingHook>(_services, async hook => await hook.OnRouti
117115
}
118116
}
119117
}
118+
119+
var isReconnect = false;
120+
var realtimeHooks = _services.GetHooks<IRealtimeHook>(_conn.CurrentAgentId);
121+
foreach (var hook in realtimeHooks)
122+
{
123+
isReconnect = await hook.ShouldReconnect(_conn);
124+
if (isReconnect) break;
125+
}
126+
127+
return isReconnect;
120128
},
121129
onConversationItemCreated: async response =>
122130
{

‎src/Infrastructure/BotSharp.Core/Agents/Services/AgentService.RefreshAgents.cs

Lines changed: 0 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -16,14 +16,6 @@ public async Task<string> RefreshAgents()
1616
return refreshResult;
1717
}
1818

19-
var userIdentity = _services.GetRequiredService<IUserIdentity>();
20-
var userService = _services.GetRequiredService<IUserService>();
21-
var (isValid, _) = await userService.IsAdminUser(userIdentity.Id);
22-
if (!isValid)
23-
{
24-
return "Unauthorized user.";
25-
}
26-
2719
var agentDir = Path.Combine(AppDomain.CurrentDomain.BaseDirectory,
2820
dbSettings.FileRepository,
2921
_agentSettings.DataDir);

‎src/Infrastructure/BotSharp.Core/Session/BotSharpRealtimeSession.cs

Lines changed: 13 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ public class BotSharpRealtimeSession : IDisposable
1111
private readonly ChatSessionOptions? _sessionOptions;
1212
private readonly object _singleReceiveLock = new();
1313
private AsyncWebsocketDataCollectionResult _receivedCollectionResult;
14+
private bool _disposed = false;
1415

1516
public BotSharpRealtimeSession(
1617
IServiceProvider services,
@@ -57,23 +58,30 @@ private ChatSessionUpdate HandleSessionResult(ClientResult result)
5758

5859
public async Task SendEventAsync(string message)
5960
{
60-
if (_websocket.State == WebSocketState.Open)
61+
if (_disposed || _websocket.State != WebSocketState.Open)
6162
{
62-
var buffer = Encoding.UTF8.GetBytes(message);
63-
await _websocket.SendAsync(new ArraySegment<byte>(buffer), WebSocketMessageType.Text, true, CancellationToken.None);
63+
return;
6464
}
65+
66+
var buffer = Encoding.UTF8.GetBytes(message);
67+
await _websocket.SendAsync(new ArraySegment<byte>(buffer), WebSocketMessageType.Text, true, CancellationToken.None);
6568
}
6669

6770
public async Task DisconnectAsync()
6871
{
69-
if (_websocket.State == WebSocketState.Open)
72+
if (_disposed || _websocket.State != WebSocketState.Open)
7073
{
71-
await _websocket.CloseAsync(WebSocketCloseStatus.Empty, null, CancellationToken.None);
74+
return;
7275
}
76+
77+
await _websocket.CloseAsync(WebSocketCloseStatus.Empty, null, CancellationToken.None);
7378
}
7479

7580
public void Dispose()
7681
{
82+
if (_disposed) return;
83+
84+
_disposed = true;
7785
_websocket.Dispose();
7886
}
7987
}

‎src/Infrastructure/BotSharp.Core/Session/LlmRealtimeSession.cs

Lines changed: 24 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ public class LlmRealtimeSession : IDisposable
1313
private readonly object _singleReceiveLock = new();
1414
private readonly SemaphoreSlim _clientEventSemaphore = new(initialCount: 1, maxCount: 1);
1515
private AsyncWebsocketDataCollectionResult _receivedCollectionResult;
16+
private bool _disposed = false;
1617

1718
public LlmRealtimeSession(
1819
IServiceProvider services,
@@ -24,6 +25,7 @@ public LlmRealtimeSession(
2425

2526
public async Task ConnectAsync(Uri uri, Dictionary<string, string>? headers = null, CancellationToken cancellationToken = default)
2627
{
28+
_disposed = false;
2729
_webSocket?.Dispose();
2830
_webSocket = new ClientWebSocket();
2931

@@ -73,31 +75,43 @@ private ChatSessionUpdate HandleSessionResult(ClientResult result)
7375

7476
public async Task SendEventToModelAsync(object message)
7577
{
76-
if (_webSocket.State != WebSocketState.Open)
78+
try
7779
{
78-
return;
79-
}
80+
if (_disposed)
81+
{
82+
return;
83+
}
8084

81-
await _clientEventSemaphore.WaitAsync();
85+
await _clientEventSemaphore.WaitAsync();
86+
87+
if (_webSocket.State != WebSocketState.Open)
88+
{
89+
return;
90+
}
8291

83-
try
84-
{
8592
if (message is not string data)
8693
{
8794
data = JsonSerializer.Serialize(message, _sessionOptions?.JsonOptions);
8895
}
8996

97+
//Console.WriteLine($"Sending event to model {data.Substring(0, 20)}");
98+
9099
var buffer = Encoding.UTF8.GetBytes(data);
91100
await _webSocket.SendAsync(new ArraySegment<byte>(buffer), WebSocketMessageType.Text, true, CancellationToken.None);
92101
}
93102
finally
94103
{
95-
_clientEventSemaphore.Release();
104+
if (!_disposed)
105+
{
106+
_clientEventSemaphore.Release();
107+
}
96108
}
97109
}
98110

99111
public async Task DisconnectAsync()
100112
{
113+
if (_disposed) return;
114+
101115
if (_webSocket.State == WebSocketState.Open)
102116
{
103117
await _webSocket.CloseAsync(WebSocketCloseStatus.Empty, null, CancellationToken.None);
@@ -106,6 +120,9 @@ public async Task DisconnectAsync()
106120

107121
public void Dispose()
108122
{
123+
if (_disposed) return;
124+
125+
_disposed = true;
109126
_clientEventSemaphore?.Dispose();
110127
_webSocket?.Dispose();
111128
}

‎src/Infrastructure/BotSharp.OpenAPI/Controllers/AgentController.cs

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
using BotSharp.Abstraction.Agents.Models;
2+
using BotSharp.Abstraction.Infrastructures.Attributes;
23

34
namespace BotSharp.OpenAPI.Controllers;
45

@@ -108,6 +109,7 @@ public async Task<AgentViewModel> CreateAgent(AgentCreationModel agent)
108109
return AgentViewModel.FromAgent(createdAgent);
109110
}
110111

112+
[BotSharpAuth]
111113
[HttpPost("/refresh-agents")]
112114
public async Task<string> RefreshAgents()
113115
{

‎src/Infrastructure/BotSharp.OpenAPI/Controllers/PluginController.cs

Lines changed: 4 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
using BotSharp.Abstraction.Infrastructures.Attributes;
12
using BotSharp.Abstraction.Plugins.Models;
23
using BotSharp.Abstraction.Users.Enums;
34
using BotSharp.Core.Plugins;
@@ -10,15 +11,10 @@ public class PluginController(IServiceProvider services, IUserIdentity user, Plu
1011
{
1112
private readonly IUserIdentity _user = user;
1213

14+
[BotSharpAuth]
1315
[HttpGet("/plugins")]
1416
public async Task<PagedItems<PluginDef>> GetPlugins([FromQuery] PluginFilter filter)
1517
{
16-
var isValid = await IsValidUser();
17-
if (!isValid)
18-
{
19-
return new PagedItems<PluginDef>();
20-
}
21-
2218
var loader = services.GetRequiredService<PluginLoader>();
2319
return loader.GetPagedPlugins(services, filter);
2420
}
@@ -72,24 +68,19 @@ public async Task<List<PluginMenuDef>> GetPluginMenu()
7268
return menu;
7369
}
7470

71+
[BotSharpAuth]
7572
[HttpPost("/plugin/{id}/install")]
7673
public PluginDef InstallPlugin([FromRoute] string id)
7774
{
7875
var loader = services.GetRequiredService<PluginLoader>();
7976
return loader.UpdatePluginStatus(services, id, true);
8077
}
8178

79+
[BotSharpAuth]
8280
[HttpPost("/plugin/{id}/remove")]
8381
public PluginDef RemovePluginStats([FromRoute] string id)
8482
{
8583
var loader = services.GetRequiredService<PluginLoader>();
8684
return loader.UpdatePluginStatus(services, id, false);
8785
}
88-
89-
private async Task<bool> IsValidUser()
90-
{
91-
var userService = services.GetRequiredService<IUserService>();
92-
var (isAdmin, _) = await userService.IsAdminUser(_user.Id);
93-
return isAdmin;
94-
}
9586
}

‎src/Infrastructure/BotSharp.OpenAPI/Controllers/RoleController.cs

Lines changed: 4 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
using BotSharp.Abstraction.Infrastructures.Attributes;
12
using BotSharp.Abstraction.Roles;
23

34
namespace BotSharp.OpenAPI.Controllers;
@@ -20,15 +21,10 @@ public RoleController(
2021
_user = user;
2122
}
2223

24+
[BotSharpAuth]
2325
[HttpPost("/role/refresh")]
2426
public async Task<bool> RefreshRoles()
2527
{
26-
var isValid = await IsValidUser();
27-
if (!isValid)
28-
{
29-
return false;
30-
}
31-
3228
return await _roleService.RefreshRoles();
3329
}
3430

@@ -39,6 +35,7 @@ public async Task<IEnumerable<string>> GetRoleOptions()
3935
return await _roleService.GetRoleOptions();
4036
}
4137

38+
[BotSharpAuth]
4239
[HttpPost("/roles")]
4340
public async Task<IEnumerable<RoleViewModel>> GetRoles([FromBody] RoleFilter? filter = null)
4441
{
@@ -47,12 +44,6 @@ public async Task<IEnumerable<RoleViewModel>> GetRoles([FromBody] RoleFilter? fi
4744
filter = RoleFilter.Empty();
4845
}
4946

50-
var isValid = await IsValidUser();
51-
if (!isValid)
52-
{
53-
return Enumerable.Empty<RoleViewModel>();
54-
}
55-
5647
var roles = await _roleService.GetRoles(filter);
5748
return roles.Select(x => RoleViewModel.FromRole(x)).ToList();
5849
}
@@ -64,25 +55,13 @@ public async Task<RoleViewModel> GetRoleDetails([FromRoute] string id)
6455
return RoleViewModel.FromRole(role);
6556
}
6657

58+
[BotSharpAuth]
6759
[HttpPut("/role")]
6860
public async Task<bool> UpdateRole([FromBody] RoleUpdateModel model)
6961
{
7062
if (model == null) return false;
7163

72-
var isValid = await IsValidUser();
73-
if (!isValid)
74-
{
75-
return false;
76-
}
77-
7864
var role = RoleUpdateModel.ToRole(model);
7965
return await _roleService.UpdateRole(role, isUpdateRoleAgents: true);
8066
}
81-
82-
private async Task<bool> IsValidUser()
83-
{
84-
var userService = _services.GetRequiredService<IUserService>();
85-
var (isAdmin, _) = await userService.IsAdminUser(_user.Id);
86-
return isAdmin;
87-
}
8867
}

0 commit comments

Comments
 (0)