Add middleware to cap token usage

AuthorAgent.cs

@@ -64,6 +64,11 @@ public async Task<string> 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}");

BlogWorkflow.cs

@@ -66,6 +66,14 @@ public async Task<ResearchState> 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<ResearchState> 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;
+    }
 }

BloggerAgent.cs

@@ -119,6 +119,11 @@ public async Task<BloggerDecision> 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}");

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}");

ResearcherAgent.cs

@@ -66,6 +66,11 @@ public async Task<string> 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<ResearchState> 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}");

ReviewerAgent.cs

@@ -36,7 +36,7 @@ public ReviewerAgent(IChatClient llm, ChatOptions chatOptions, ILogger<ReviewerA
                 MaxOutputTokens = chatOptions.MaxOutputTokens,
             },
         });
-        _logger.LogInformation("ReviewerAgent initialized.");   
+        _logger.LogInformation("ReviewerAgent initialized.");
     }
 
     public async Task<string> InvokeAsync(ResearchState state)
@@ -66,6 +66,11 @@ public async Task<string> 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.

TokenCapChatClient.cs

@@ -0,0 +1,74 @@
+using System.Runtime.CompilerServices;
+using Microsoft.Extensions.AI;
+
+namespace BlogWriter;
+
+/// <summary>
+/// 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.
+/// </summary>
+public sealed class TokenCapChatClient(IChatClient innerClient, long maxTotalTokens)
+    : DelegatingChatClient(innerClient)
+{
+    private long _totalTokens;
+
+    public override async Task<ChatResponse> GetResponseAsync(
+        IEnumerable<ChatMessage> messages,
+        ChatOptions? options = null,
+        CancellationToken cancellationToken = default)
+    {
+        ChatResponse response = await base.GetResponseAsync(messages, options, cancellationToken);
+        Track(response.Usage);
+        return response;
+    }
+
+    public override async IAsyncEnumerable<ChatResponseUpdate> GetStreamingResponseAsync(
+        IEnumerable<ChatMessage> 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);
+        }
+    }
+}
+
+/// <summary>
+/// Thrown when cumulative model token usage exceeds the configured cap. Callers
+/// catch this to shut down gracefully instead of continuing to spend tokens.
+/// </summary>
+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;
+}