Skip to content

Fix LLM agent stability: conditional execution, error handling, tool registration #268

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ModerRAS merged 1 commit into from
Apr 17, 2026
Merged

Fix LLM agent stability: conditional execution, error handling, tool registration #268

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 8 additions & 1 deletion TelegramSearchBot.LLMAgent/LLMAgentProgram.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,8 +6,10 @@
using TelegramSearchBot.Common;
using TelegramSearchBot.Interface;
using TelegramSearchBot.Interface.AI.LLM;
using TelegramSearchBot.Interface.Tools;
using TelegramSearchBot.Model;
using TelegramSearchBot.Service.AI.LLM;
using TelegramSearchBot.Service.Tools;

namespace TelegramSearchBot.LLMAgent {
public static class LLMAgentProgram {
Expand All @@ -23,7 +25,10 @@ public static async Task RunAsync(string[] args) {

using var services = BuildServices(port);
var logger = services.GetRequiredService<ILoggerFactory>().CreateLogger("LLMAgent");
McpToolHelper.EnsureInitialized(typeof(Service.AgentToolService).Assembly, services, logger);
McpToolHelper.EnsureInitialized(
typeof(Service.AgentToolService).Assembly,
typeof(FileToolService).Assembly,
services, logger);

var loop = services.GetRequiredService<Service.AgentLoopService>();
using var shutdownCts = new CancellationTokenSource();
Expand Down Expand Up @@ -63,6 +68,8 @@ private static ServiceProvider BuildServices(int port) {
services.AddScoped<AnthropicService>();
services.AddScoped<Service.ToolExecutor>();
services.AddScoped<Service.AgentToolService>();
services.AddScoped<IFileToolService, FileToolService>();
services.AddScoped<IBashToolService, BashToolService>();
services.AddScoped<Service.IAgentTaskExecutor, Service.LlmServiceProxy>();
services.AddScoped<Service.LlmServiceProxy>();
services.AddSingleton<Service.GarnetClient>();
Expand Down
25 changes: 21 additions & 4 deletions TelegramSearchBot.LLMAgent/Service/AgentLoopService.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
using Microsoft.Extensions.DependencyInjection;
using Microsoft.Extensions.Logging;
using Newtonsoft.Json;
using StackExchange.Redis;
using TelegramSearchBot.Common;
using TelegramSearchBot.Model.AI;

Expand Down Expand Up @@ -42,7 +43,15 @@ public async Task RunAsync(long chatId, int port, CancellationToken cancellation
break;
}

var payload = await _garnetClient.BRPopAsync(LlmAgentRedisKeys.AgentTaskQueue, TimeSpan.FromSeconds(5));
string? payload;
try {
payload = await _garnetClient.BRPopAsync(LlmAgentRedisKeys.AgentTaskQueue, TimeSpan.FromSeconds(2));
} catch (RedisException ex) {
_logger.LogWarning(ex, "Redis error during BRPOP, retrying in 1s");
await Task.Delay(TimeSpan.FromSeconds(1), cancellationToken);
continue;
}

if (string.IsNullOrWhiteSpace(payload)) {
continue;
}
Expand Down Expand Up @@ -186,9 +195,17 @@ await _garnetClient.PublishChunkAsync(new AgentStreamChunk {

private async Task RunHeartbeatAsync(AgentSessionInfo session, CancellationToken cancellationToken) {
using var timer = new PeriodicTimer(TimeSpan.FromSeconds(Math.Max(1, Env.AgentHeartbeatIntervalSeconds)));
while (!cancellationToken.IsCancellationRequested && await timer.WaitForNextTickAsync(cancellationToken)) {
session.LastHeartbeatUtc = DateTime.UtcNow;
await _rpcClient.SaveSessionAsync(session);
try {
while (!cancellationToken.IsCancellationRequested && await timer.WaitForNextTickAsync(cancellationToken)) {
try {
session.LastHeartbeatUtc = DateTime.UtcNow;
await _rpcClient.SaveSessionAsync(session);
} catch (RedisException ex) {
_logger.LogWarning(ex, "Redis error during heartbeat, will retry next tick");
}
}
} catch (OperationCanceledException) {
// Normal shutdown – heartbeat stops
}
}

Expand Down
5 changes: 5 additions & 0 deletions TelegramSearchBot/Service/AI/LLM/AgentRegistryService.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -183,6 +183,11 @@ public async Task RunMaintenanceOnceAsync(CancellationToken cancellationToken =
}

protected override async Task ExecuteAsync(CancellationToken stoppingToken) {
if (!Env.EnableLLMAgentProcess) {
_logger.LogDebug("LLM agent process mode disabled – AgentRegistryService will not start");
return;
}
Comment on lines +186 to +189

@coderabbitai coderabbitai Bot Apr 17, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟡 Minor

Disabled-mode shutdown branch in RunMaintenanceOnceAsync becomes unreachable from the background loop.

With this early return, ExecuteAsync never calls RunMaintenanceOnceAsync, so the disabled-mode branch at lines 149–153 (which iterates _knownSessions and calls RequestShutdownAsync("agent mode disabled")) no longer fires from the hosted service. In practice this is mostly fine because EnsureAgentAsync throws when the flag is false so _knownSessions shouldn't grow, but if the flag is toggled from true→false at runtime, previously-running agents will no longer be gracefully asked to shut down from here.

Consider either:

  • Documenting that toggling requires a restart, or
  • Running one final RunMaintenanceOnceAsync pass (or just the shutdown-known-sessions block) before returning, so in-flight sessions get a graceful shutdown request.
🤖 Prompt for AI Agents 
Verify each finding against the current code and only fix it if needed.

In `@TelegramSearchBot/Service/AI/LLM/AgentRegistryService.cs` around lines 186 -
189, The early return when Env.EnableLLMAgentProcess is false prevents
ExecuteAsync from ever calling RunMaintenanceOnceAsync and leaves
previously-known agents running; update the AgentRegistryService start path so
that before returning when Env.EnableLLMAgentProcess is false you perform one
final shutdown pass: either call await
RunMaintenanceOnceAsync(CancellationToken.None) or iterate over _knownSessions
and call RequestShutdownAsync("agent mode disabled") for each (awaiting the
tasks) so existing sessions are asked to terminate gracefully; keep references
to Env.EnableLLMAgentProcess, ExecuteAsync, RunMaintenanceOnceAsync,
_knownSessions and RequestShutdownAsync when making the change.


while (!stoppingToken.IsCancellationRequested) {
try {
await RunMaintenanceOnceAsync(stoppingToken);
Expand Down
19 changes: 17 additions & 2 deletions TelegramSearchBot/Service/AI/LLM/ChunkPollingService.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -52,9 +52,24 @@ public async Task RunPollCycleAsync(CancellationToken cancellationToken = defaul
}

protected override async Task ExecuteAsync(CancellationToken stoppingToken) {
if (!Env.EnableLLMAgentProcess) {
return;
}

while (!stoppingToken.IsCancellationRequested) {
await RunPollCycleAsync(stoppingToken);
await Task.Delay(Math.Max(50, Env.AgentChunkPollingIntervalMilliseconds), stoppingToken);
try {
await RunPollCycleAsync(stoppingToken);
} catch (OperationCanceledException) {
break;
} catch (RedisException) {
// Transient Redis failure – wait before retrying
}

try {
await Task.Delay(Math.Max(50, Env.AgentChunkPollingIntervalMilliseconds), stoppingToken);
} catch (OperationCanceledException) {
break;
}
}
}

Expand Down
5 changes: 5 additions & 0 deletions TelegramSearchBot/Service/AI/LLM/TelegramTaskConsumer.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,11 @@ public TelegramTaskConsumer(
}

protected override async Task ExecuteAsync(CancellationToken stoppingToken) {
if (!Env.EnableLLMAgentProcess) {
_logger.LogDebug("LLM agent process mode disabled – TelegramTaskConsumer will not start");
return;
}

while (!stoppingToken.IsCancellationRequested) {
try {
// Use a 2-second block time so BRPOP returns well within SE.Redis's 5 s async timeout.
Expand Down
Loading