diff --git a/AuthorAgent.cs b/AuthorAgent.cs index b1b072f..073e896 100644 --- a/AuthorAgent.cs +++ b/AuthorAgent.cs @@ -64,6 +64,11 @@ public async Task InvokeAsync(ResearchState state) string content = response.Text; return !string.IsNullOrEmpty(content) ? content : "Draft in progress..."; } + catch (TokenCapExceededException) + { + // Budget breach is fatal — let it propagate so the app can shut down. + throw; + } catch (Exception e) { Console.WriteLine($"Author error: {e.Message}"); diff --git a/BlogWorkflow.cs b/BlogWorkflow.cs index bf25204..5f5e8ac 100644 --- a/BlogWorkflow.cs +++ b/BlogWorkflow.cs @@ -66,6 +66,14 @@ public async Task RunAsync(ResearchState state) case ExecutorFailedEvent failed: Console.WriteLine($"[workflow] ✗ {failed.ExecutorId} failed: {(failed.Data as Exception)?.Message}"); + + // A token-cap breach must abort the whole run, not just the + // node. Re-throw it so it unwinds to the application entry point. + if (failed.Data is Exception ex && FindTokenCap(ex) is { } capEx) + { + throw capEx; + } + break; case WorkflowOutputEvent { Data: ResearchState finalState }: @@ -78,4 +86,34 @@ public async Task RunAsync(ResearchState state) // Fall back to the input state only if no output event was ever produced. return result ?? state; } + + // Walks the exception chain (including AggregateException children) looking + // for a token-cap breach, which the workflow runtime may have wrapped. + private static TokenCapExceededException? FindTokenCap(Exception? exception) + { + while (exception is not null) + { + if (exception is TokenCapExceededException capEx) + { + return capEx; + } + + if (exception is AggregateException aggregate) + { + foreach (Exception inner in aggregate.InnerExceptions) + { + if (FindTokenCap(inner) is { } found) + { + return found; + } + } + + return null; + } + + exception = exception.InnerException; + } + + return null; + } } diff --git a/BloggerAgent.cs b/BloggerAgent.cs index cb2419f..07b1753 100644 --- a/BloggerAgent.cs +++ b/BloggerAgent.cs @@ -119,6 +119,11 @@ public async Task InvokeAsync(ResearchState state) return decision; } } + catch (TokenCapExceededException) + { + // Budget breach is fatal — let it propagate so the app can shut down. + throw; + } catch (Exception e) { Console.WriteLine($"LLM decision error: {e.Message}"); diff --git a/Program.cs b/Program.cs index 0c82fd0..367ec59 100644 --- a/Program.cs +++ b/Program.cs @@ -37,11 +37,17 @@ // Middleware is applied inner-to-outer, so function invocation wraps the raw // OpenAI client. (To add distributed tracing later, chain .UseOpenTelemetry() // here and register the source with a TracerProvider.) +// +// TokenCapChatClient is registered *after* function invocation, which makes it +// the innermost wrapper around the raw client — so it observes every individual +// model round-trip (including the extra calls tool invocation triggers) and +// enforces a hard 1000-token budget for the whole process. IChatClient llm = openAIClient .GetChatClient(modelName) .AsIChatClient() .AsBuilder() .UseFunctionInvocation() + .Use(inner => new TokenCapChatClient(inner, maxTotalTokens: 1000)) .Build(); var chatOptions = new ChatOptions @@ -108,11 +114,20 @@ }; ResearchState result; -using (Activity? runActivity = appActivitySource.StartActivity("BlogWriter.Run")) +try { + using Activity? runActivity = appActivitySource.StartActivity("BlogWriter.Run"); runActivity?.SetTag("blog.topic", topic); result = await app.RunAsync(initialState); } +catch (TokenCapExceededException ex) +{ + // Graceful shutdown: the exception unwinds the call stack so every `using` + // (logger factory, HTTP clients, etc.) is disposed before we exit. + Console.Error.WriteLine($"{ex.Message} Exiting application."); + Environment.ExitCode = 1; + return; +} Console.WriteLine("\n========== RESULTS =========="); Console.WriteLine($"Task: {result.MainTask}"); diff --git a/ResearcherAgent.cs b/ResearcherAgent.cs index 7f434cc..34fe6fb 100644 --- a/ResearcherAgent.cs +++ b/ResearcherAgent.cs @@ -66,6 +66,11 @@ public async Task InvokeAsync(string query) ? summary : $"Research completed on: {query}. Key information has been gathered from web sources."; } + catch (TokenCapExceededException) + { + // Budget breach is fatal — let it propagate so the app can shut down. + throw; + } catch (Exception e) { Console.WriteLine($"Research error: {e.Message}"); @@ -88,6 +93,11 @@ public async Task ResearchNodeAsync(ResearchState state) string preview = findings.Length > 100 ? findings[..100] : findings; Console.WriteLine($"Found: {preview}..."); } + catch (TokenCapExceededException) + { + // Budget breach is fatal — let it propagate so the app can shut down. + throw; + } catch (Exception e) { Console.WriteLine($"Research error: {e.Message}"); diff --git a/ReviewerAgent.cs b/ReviewerAgent.cs index 1c28bd2..20ff6ce 100644 --- a/ReviewerAgent.cs +++ b/ReviewerAgent.cs @@ -36,7 +36,7 @@ public ReviewerAgent(IChatClient llm, ChatOptions chatOptions, ILogger InvokeAsync(ResearchState state) @@ -66,6 +66,11 @@ public async Task InvokeAsync(ResearchState state) string content = response.Text; return !string.IsNullOrEmpty(content) ? content : "APPROVED"; } + catch (TokenCapExceededException) + { + // Budget breach is fatal — let it propagate so the app can shut down. + throw; + } catch (Exception e) { // Do NOT approve on failure — that would ship an unreviewed draft. diff --git a/TokenCapChatClient.cs b/TokenCapChatClient.cs new file mode 100644 index 0000000..dfebb06 --- /dev/null +++ b/TokenCapChatClient.cs @@ -0,0 +1,74 @@ +using System.Runtime.CompilerServices; +using Microsoft.Extensions.AI; + +namespace BlogWriter; + +/// +/// Chat-client middleware that enforces a hard cap on cumulative token usage +/// across every model round-trip in the process (including the extra calls made +/// during tool invocation). When the running total exceeds the cap, the +/// application is terminated with an explanatory message rather than continuing +/// to spend tokens. +/// +public sealed class TokenCapChatClient(IChatClient innerClient, long maxTotalTokens) + : DelegatingChatClient(innerClient) +{ + private long _totalTokens; + + public override async Task GetResponseAsync( + IEnumerable messages, + ChatOptions? options = null, + CancellationToken cancellationToken = default) + { + ChatResponse response = await base.GetResponseAsync(messages, options, cancellationToken); + Track(response.Usage); + return response; + } + + public override async IAsyncEnumerable GetStreamingResponseAsync( + IEnumerable messages, + ChatOptions? options = null, + [EnumeratorCancellation] CancellationToken cancellationToken = default) + { + await foreach (ChatResponseUpdate update in + base.GetStreamingResponseAsync(messages, options, cancellationToken)) + { + foreach (AIContent content in update.Contents) + { + if (content is UsageContent usageContent) + { + Track(usageContent.Details); + } + } + + yield return update; + } + } + + private void Track(UsageDetails? usage) + { + long used = usage?.TotalTokenCount ?? 0; + if (used == 0) + { + return; + } + + long total = Interlocked.Add(ref _totalTokens, used); + if (total > maxTotalTokens) + { + throw new TokenCapExceededException(total, maxTotalTokens); + } + } +} + +/// +/// Thrown when cumulative model token usage exceeds the configured cap. Callers +/// catch this to shut down gracefully instead of continuing to spend tokens. +/// +public sealed class TokenCapExceededException(long tokensUsed, long tokenLimit) + : Exception($"Token cap exceeded: consumed {tokensUsed} tokens, limit is {tokenLimit}.") +{ + public long TokensUsed { get; } = tokensUsed; + + public long TokenLimit { get; } = tokenLimit; +}