diff --git a/dotnet/server/aspnetcore/SmoothOperatorWebSocketExtensions.cs b/dotnet/server/aspnetcore/SmoothOperatorWebSocketExtensions.cs index e09d81c..7c9b127 100644 --- a/dotnet/server/aspnetcore/SmoothOperatorWebSocketExtensions.cs +++ b/dotnet/server/aspnetcore/SmoothOperatorWebSocketExtensions.cs @@ -81,7 +81,10 @@ private static FrameDispatcher BuildDispatcher(HttpContext context) // Tool-name patterns gated behind write-confirmation HITL (default none — the DI analog of // Python's ServerState.confirm_tools). Each connection gets its own ConfirmationRegistry // (a confirm_tool_action frame and the parked turn it resumes are always on the same one). - confirmTools: services.GetService()?.Patterns); + confirmTools: services.GetService()?.Patterns, + // Per-turn token/iteration limits + the resolved model's output ceiling (EPIC th-1cc9fa). + // Absent ⇒ the raised server defaults (max_tokens 8192, iterations 20) with no ceiling. + limits: services.GetService()); } private static async Task PumpAsync(WebSocket socket, FrameDispatcher dispatcher, CancellationToken cancellationToken) diff --git a/dotnet/server/host/Program.cs b/dotnet/server/host/Program.cs index a6dd712..811a50d 100644 --- a/dotnet/server/host/Program.cs +++ b/dotnet/server/host/Program.cs @@ -100,6 +100,20 @@ string Get(string key, string? fallback = null) => builder.Services.AddSingleton(new ConfirmTools(confirmTools)); } +// ── Per-turn limits: raised output-token budget + agentic-iteration cap (EPIC th-1cc9fa), plus the +// resolved model's HARD output ceiling clamped from the gateway's /model/info so a budget can +// never exceed what the model can physically emit (a reasoning model would otherwise burn the +// budget and return empty, or the upstream 400s). The ceiling fetch is best-effort at boot — this +// single-model server resolves it once rather than the Rust server's per-turn lazy cache; any +// gateway problem ⇒ null ⇒ unclamped. SMOOTH_MAX_TOKENS / SMOOTH_MAX_ITERATIONS override the +// defaults. ── +var maxTokens = int.TryParse(Get("SMOOTH_MAX_TOKENS"), out var mt) && mt > 0 ? mt : TurnLimits.DefaultMaxTokens; +var maxIterations = int.TryParse(Get("SMOOTH_MAX_ITERATIONS"), out var mi) && mi > 0 ? mi : TurnLimits.DefaultMaxIterations; +int? modelCeiling = string.IsNullOrEmpty(gatewayKey) + ? null + : await ModelInfo.FetchCeilingAsync(EmbeddingHttpClient(gatewayUrl, gatewayKey), model); +builder.Services.AddSingleton(new TurnLimits(maxTokens, maxIterations, modelCeiling)); + builder.Services.AddSmoothOperatorServer(); var app = builder.Build(); diff --git a/dotnet/server/src/FrameDispatcher.cs b/dotnet/server/src/FrameDispatcher.cs index 05d1e96..f2b8867 100644 --- a/dotnet/server/src/FrameDispatcher.cs +++ b/dotnet/server/src/FrameDispatcher.cs @@ -25,6 +25,7 @@ public sealed class FrameDispatcher private readonly IReadOnlyList _tools; private readonly IReadOnlyList _confirmTools; private readonly ConfirmationRegistry _confirmations; + private readonly TurnLimits _limits; // In-flight spawned send_message turns. A turn that calls a confirmation-gated tool parks // awaiting a later confirm_tool_action frame, so the turn runs as a background Task (not awaited @@ -41,7 +42,8 @@ public FrameDispatcher( IReranker? reranker = null, IReadOnlyList? tools = null, IReadOnlyList? confirmTools = null, - ConfirmationRegistry? confirmations = null) + ConfirmationRegistry? confirmations = null, + TurnLimits? limits = null) { _store = store ?? throw new ArgumentNullException(nameof(store)); _chatClient = chatClient ?? throw new ArgumentNullException(nameof(chatClient)); @@ -57,6 +59,9 @@ public FrameDispatcher( // Session-keyed pending-confirmation registry shared with each spawned turn so a // confirm_tool_action frame resolves the verdict a parked turn awaits. One per connection. _confirmations = confirmations ?? new ConfirmationRegistry(); + // Per-turn token/iteration limits + the resolved model's output ceiling (EPIC th-1cc9fa). + // Absent ⇒ the raised server defaults (max_tokens 8192, iterations 20) with no ceiling. + _limits = limits ?? TurnLimits.Default; } /// @@ -208,7 +213,7 @@ private async Task HandleSendMessageAsync(JsonObject frame, string? requestId, A // 2. Stream the turn, retrieving through knowledge SCOPED to this connection's access — so a // user only ever sees documents their groups grant (ACL enforced on the chat path). var scopedKnowledge = _knowledge?.ForAccess(_access); - var runner = new TurnRunner(_chatClient, _store, scopedKnowledge, _systemPrompt, _reranker, _tools, _confirmTools, _confirmations); + var runner = new TurnRunner(_chatClient, _store, scopedKnowledge, _systemPrompt, _reranker, _tools, _confirmTools, _confirmations, _limits); // Run the turn as a background task, NOT awaited inline. A turn that calls a // confirmation-gated tool PARKS awaiting a later confirm_tool_action frame; the connection's diff --git a/dotnet/server/src/ModelInfo.cs b/dotnet/server/src/ModelInfo.cs new file mode 100644 index 0000000..d88dc07 --- /dev/null +++ b/dotnet/server/src/ModelInfo.cs @@ -0,0 +1,84 @@ +using System.Text.Json.Nodes; + +namespace SmooAI.SmoothOperator.Server; + +/// +/// Reads a model's HARD output ceiling (max_output_tokens) from the LiteLLM gateway's +/// /model/info, so the chat path can clamp max_tokens to what the model can physically +/// emit (EPIC th-1cc9fa). The C# analog of the Rust server's admin::model_output_ceiling + +/// map_model_info — kept out of the engine so the published engine takes no LiteLLM-specific +/// HTTP dependency. Best-effort: any gateway error, an unknown model, or a model with no positive +/// ceiling ⇒ null ⇒ the engine leaves max_tokens unclamped (graceful, no behavior change). +/// +public static class ModelInfo +{ + /// + /// Parse the gateway's /model/info payload + /// ({ data: [{ model_name, model_info: { max_output_tokens } }] }) into a + /// model_name → ceiling map. Entries without a model_name or with a missing / + /// non-positive ceiling are dropped. Pure + network-free, so it's unit-testable on a sample payload. + /// + public static IReadOnlyDictionary ParseCeilings(JsonNode? payload) + { + var map = new Dictionary(StringComparer.Ordinal); + if (payload?["data"] is not JsonArray entries) + { + return map; + } + foreach (var entry in entries) + { + var name = (entry?["model_name"] as JsonValue)?.GetValue(); + if (string.IsNullOrEmpty(name)) + { + continue; + } + var ceiling = TryGetPositiveInt(entry?["model_info"]?["max_output_tokens"]); + if (ceiling is not null) + { + map[name] = ceiling.Value; + } + } + return map; + } + + /// + /// Fetch the output ceiling for from {gateway}/model/info via + /// (its is the gateway root — with a + /// trailing slash — and the auth header is already set). Returns null on ANY error, an + /// unknown model, or a model with no positive ceiling; the caller then leaves max_tokens + /// unclamped. Never throws. + /// + public static async Task FetchCeilingAsync(HttpClient http, string model, CancellationToken cancellationToken = default) + { + try + { + using var response = await http.GetAsync("model/info", cancellationToken).ConfigureAwait(false); + if (!response.IsSuccessStatusCode) + { + return null; + } + var body = await response.Content.ReadAsStringAsync(cancellationToken).ConfigureAwait(false); + var ceilings = ParseCeilings(JsonNode.Parse(body)); + return ceilings.TryGetValue(model, out var ceiling) ? ceiling : null; + } + catch + { + // Best-effort: a gateway blip must never fail a boot or a turn — just skip the clamp. + return null; + } + } + + /// The node as a positive (accepting int/long/double JSON numbers), or + /// null when it's absent, non-numeric, or ≤ 0 (a bogus ceiling must not clamp to nothing). + private static int? TryGetPositiveInt(JsonNode? node) + { + if (node is not JsonValue value) + { + return null; + } + long? number = value.TryGetValue(out var l) ? l + : value.TryGetValue(out var d) && d is >= long.MinValue and <= long.MaxValue ? (long)d + : null; + return number is > 0 and <= int.MaxValue ? (int)number.Value : null; + } +} diff --git a/dotnet/server/src/SmooAI.SmoothOperator.Server.csproj b/dotnet/server/src/SmooAI.SmoothOperator.Server.csproj index 5d7f99f..2cbc782 100644 --- a/dotnet/server/src/SmooAI.SmoothOperator.Server.csproj +++ b/dotnet/server/src/SmooAI.SmoothOperator.Server.csproj @@ -16,7 +16,7 @@ - + diff --git a/dotnet/server/src/TurnLimits.cs b/dotnet/server/src/TurnLimits.cs new file mode 100644 index 0000000..5470bf5 --- /dev/null +++ b/dotnet/server/src/TurnLimits.cs @@ -0,0 +1,26 @@ +namespace SmooAI.SmoothOperator.Server; + +/// +/// Per-turn token/iteration limits threaded into the agent config. Raised from the starvation-prone +/// legacy defaults (max_tokens 512, iterations 6) that made reasoning models return EMPTY — a +/// reasoning model needs headroom to think AND answer, and more than a handful of iterations to +/// actually use its tools (EPIC th-1cc9fa, matching the Rust server's DEFAULT_MAX_TOKENS / +/// DEFAULT_MAX_ITERATIONS raise). is the resolved model's +/// HARD output ceiling (from the gateway's /model/info); null ⇒ the engine leaves +/// max_tokens unclamped. A distinct carrier type (like ConfirmTools) so it can't collide +/// with other values in the DI container. +/// +public sealed record TurnLimits( + int MaxTokens = TurnLimits.DefaultMaxTokens, + int MaxIterations = TurnLimits.DefaultMaxIterations, + int? ModelMaxOutputTokens = null) +{ + /// Default per-turn output-token budget. Raised 512 → 8192 (EPIC th-1cc9fa). + public const int DefaultMaxTokens = 8_192; + + /// Default per-turn agentic-iteration cap. Raised 6 → 20 (EPIC th-1cc9fa). + public const int DefaultMaxIterations = 20; + + /// The "raised server defaults, no model ceiling resolved" instance. + public static readonly TurnLimits Default = new(); +} diff --git a/dotnet/server/src/TurnRunner.cs b/dotnet/server/src/TurnRunner.cs index 33bac80..be16132 100644 --- a/dotnet/server/src/TurnRunner.cs +++ b/dotnet/server/src/TurnRunner.cs @@ -30,8 +30,9 @@ public sealed class TurnRunner private readonly IReadOnlyList _tools; private readonly IReadOnlyList _confirmTools; private readonly ConfirmationRegistry? _confirmations; + private readonly TurnLimits _limits; - public TurnRunner(IChatClient chatClient, ISessionStore store, IKnowledgeBase? knowledge = null, string? systemPrompt = null, IReranker? reranker = null, IReadOnlyList? tools = null, IReadOnlyList? confirmTools = null, ConfirmationRegistry? confirmations = null) + public TurnRunner(IChatClient chatClient, ISessionStore store, IKnowledgeBase? knowledge = null, string? systemPrompt = null, IReranker? reranker = null, IReadOnlyList? tools = null, IReadOnlyList? confirmTools = null, ConfirmationRegistry? confirmations = null, TurnLimits? limits = null) { _chatClient = chatClient ?? throw new ArgumentNullException(nameof(chatClient)); _store = store ?? throw new ArgumentNullException(nameof(store)); @@ -45,6 +46,9 @@ public TurnRunner(IChatClient chatClient, ISessionStore store, IKnowledgeBase? k _confirmTools = confirmTools ?? Array.Empty(); // The session-keyed pending-confirmation registry the gate parks on (null → HITL off). _confirmations = confirmations; + // Per-turn output-token budget + agentic-iteration cap, plus the resolved model's hard output + // ceiling. Absent ⇒ the raised server defaults (max_tokens 8192, iterations 20; EPIC th-1cc9fa). + _limits = limits ?? TurnLimits.Default; } /// True when matches a confirmation-gated pattern (substring, @@ -89,7 +93,17 @@ public async Task RunAsync(string conversationId, string requestId, // Registered tools (default none) are passed straight to the engine's agentic loop; the // streaming block below already translates the resulting tool-call/result events into // stream_chunks, so enabling tools is purely a matter of supplying them here. - var options = new AgentOptions { Instructions = _systemPrompt, Knowledge = _knowledge }; + // MaxOutputTokens is clamped to the model's ModelMaxOutputTokens ceiling by the engine so a + // budget never exceeds what the model can physically emit (EPIC th-1cc9fa). The raised defaults + // (8192 / 20) give reasoning models room to think AND answer, and iterations to actually use tools. + var options = new AgentOptions + { + Instructions = _systemPrompt, + Knowledge = _knowledge, + MaxIterations = _limits.MaxIterations, + MaxOutputTokens = _limits.MaxTokens, + ModelMaxOutputTokens = _limits.ModelMaxOutputTokens, + }; foreach (var tool in _tools) { options.Tools.Add(tool); diff --git a/dotnet/server/tests/MaxTokensLimitsTests.cs b/dotnet/server/tests/MaxTokensLimitsTests.cs new file mode 100644 index 0000000..f50becb --- /dev/null +++ b/dotnet/server/tests/MaxTokensLimitsTests.cs @@ -0,0 +1,174 @@ +using System.Net; +using System.Runtime.CompilerServices; +using System.Text; +using System.Text.Json.Nodes; +using Microsoft.Extensions.AI; + +namespace SmooAI.SmoothOperator.Server.Tests; + +/// +/// Parity tests for the model-output-token ceiling clamp on the server side (EPIC th-1cc9fa): the +/// gateway /model/info ceiling parse/fetch, the raised starvation defaults, and TurnRunner +/// threading the budget into the model request (where the engine clamps it to the ceiling). +/// +public class MaxTokensLimitsTests +{ + private const string SamplePayload = """ + { + "data": [ + { "model_name": "groq-compound", "model_info": { "max_output_tokens": 8192, "input_cost_per_token": 0.0000001 } }, + { "model_name": "claude-haiku-4-5", "model_info": { "max_output_tokens": 65536 } }, + { "model_name": "no-ceiling-model", "model_info": { "input_cost_per_token": 0.0000002 } }, + { "model_name": "zero-ceiling-model", "model_info": { "max_output_tokens": 0 } }, + { "model_info": { "max_output_tokens": 1234 } } + ] + } + """; + + // ── ParseCeilings (pure, network-free) ─────────────────────────────────────────────────────── + + [Fact] + public void ParseCeilings_MapsPositiveCeilingsByModelName() + { + var map = ModelInfo.ParseCeilings(JsonNode.Parse(SamplePayload)); + + Assert.Equal(8192, map["groq-compound"]); + Assert.Equal(65536, map["claude-haiku-4-5"]); + } + + [Fact] + public void ParseCeilings_DropsMissingZeroAndNamelessEntries() + { + var map = ModelInfo.ParseCeilings(JsonNode.Parse(SamplePayload)); + + Assert.False(map.ContainsKey("no-ceiling-model")); // no max_output_tokens + Assert.False(map.ContainsKey("zero-ceiling-model")); // 0 is not a real ceiling + Assert.Equal(2, map.Count); // the nameless entry is skipped too + } + + [Fact] + public void ParseCeilings_ToleratesMalformedPayloads() + { + Assert.Empty(ModelInfo.ParseCeilings(null)); + Assert.Empty(ModelInfo.ParseCeilings(JsonNode.Parse("{}"))); + Assert.Empty(ModelInfo.ParseCeilings(JsonNode.Parse("{\"data\": \"nope\"}"))); + Assert.Empty(ModelInfo.ParseCeilings(JsonNode.Parse("{\"data\": []}"))); + } + + // ── FetchCeilingAsync (best-effort HTTP) ───────────────────────────────────────────────────── + + [Fact] + public async Task FetchCeiling_ReturnsCeilingForKnownModel() + { + var http = ClientReturning(HttpStatusCode.OK, SamplePayload); + + Assert.Equal(8192, await ModelInfo.FetchCeilingAsync(http, "groq-compound")); + } + + [Fact] + public async Task FetchCeiling_UnknownModel_IsNull() + { + var http = ClientReturning(HttpStatusCode.OK, SamplePayload); + + Assert.Null(await ModelInfo.FetchCeilingAsync(http, "some-other-model")); + } + + [Fact] + public async Task FetchCeiling_GatewayErrorOrGarbage_IsNull() + { + Assert.Null(await ModelInfo.FetchCeilingAsync(ClientReturning(HttpStatusCode.InternalServerError, "boom"), "groq-compound")); + Assert.Null(await ModelInfo.FetchCeilingAsync(ClientReturning(HttpStatusCode.OK, "not json"), "groq-compound")); + } + + // ── Raised starvation defaults ─────────────────────────────────────────────────────────────── + + [Fact] + public void TurnLimits_RaisesStarvationDefaults() + { + Assert.Equal(8192, TurnLimits.DefaultMaxTokens); + Assert.Equal(20, TurnLimits.DefaultMaxIterations); + Assert.Equal(8192, TurnLimits.Default.MaxTokens); + Assert.Equal(20, TurnLimits.Default.MaxIterations); + Assert.Null(TurnLimits.Default.ModelMaxOutputTokens); + } + + // ── TurnRunner threads the budget into the request (engine then clamps to the ceiling) ─────── + + [Fact] + public async Task TurnRunner_SendsConfiguredBudgetAsMaxTokens() + { + var chat = new MaxTokensCapturingClient(); + var store = new InMemorySessionStore(); + var session = await store.CreateSessionAsync("agent-1", null, null); + var runner = new TurnRunner(chat, store, limits: new TurnLimits(MaxTokens: 4096, MaxIterations: 20)); + + await runner.RunAsync(session.ConversationId, "r1", "hello", _ => { }); + + Assert.Equal(4096, chat.LastMaxOutputTokens); + } + + [Fact] + public async Task TurnRunner_DefaultLimits_SendRaisedBudget() + { + var chat = new MaxTokensCapturingClient(); + var store = new InMemorySessionStore(); + var session = await store.CreateSessionAsync("agent-1", null, null); + var runner = new TurnRunner(chat, store); // no limits ⇒ raised defaults + + await runner.RunAsync(session.ConversationId, "r1", "hello", _ => { }); + + Assert.Equal(TurnLimits.DefaultMaxTokens, chat.LastMaxOutputTokens); + } + + private static HttpClient ClientReturning(HttpStatusCode status, string body) => + new(new StubHandler(status, body)) { BaseAddress = new Uri("https://gateway.test/v1/") }; + + private sealed class StubHandler : HttpMessageHandler + { + private readonly HttpStatusCode _status; + private readonly string _body; + + public StubHandler(HttpStatusCode status, string body) + { + _status = status; + _body = body; + } + + protected override Task SendAsync(HttpRequestMessage request, CancellationToken cancellationToken) + { + // The fetch must hit {gateway}/model/info (relative to the BaseAddress). + Assert.Equal("https://gateway.test/v1/model/info", request.RequestUri!.ToString()); + return Task.FromResult(new HttpResponseMessage(_status) { Content = new StringContent(_body, Encoding.UTF8, "application/json") }); + } + } + + /// Records the streaming request's MaxOutputTokens (the C# analog of the wire + /// max_tokens) so a test can assert what the runner actually sent. + private sealed class MaxTokensCapturingClient : IChatClient + { + public int? LastMaxOutputTokens { get; private set; } + + public Task GetResponseAsync(IEnumerable messages, ChatOptions? options = null, CancellationToken cancellationToken = default) + { + LastMaxOutputTokens = options?.MaxOutputTokens; + return Task.FromResult(new ChatResponse(new ChatMessage(ChatRole.Assistant, "ok")) { ModelId = "capture" }); + } + + public async IAsyncEnumerable GetStreamingResponseAsync( + IEnumerable messages, ChatOptions? options = null, [EnumeratorCancellation] CancellationToken cancellationToken = default) + { + LastMaxOutputTokens = options?.MaxOutputTokens; + foreach (var update in new ChatResponse(new ChatMessage(ChatRole.Assistant, "ok")).ToChatResponseUpdates()) + { + await Task.Yield(); + yield return update; + } + } + + public object? GetService(Type serviceType, object? serviceKey = null) => null; + + public void Dispose() + { + } + } +}