From 38d9a2bc1975b5e982cbfa426a5edc366ae512bd Mon Sep 17 00:00:00 2001 From: Aaron Stannard Date: Fri, 20 Mar 2026 14:15:25 +0000 Subject: [PATCH] fix(stats): parse token usage in OpenAI-compatible provider and relax stats guard OpenAiCompatibleChatClient never parsed the `usage` JSON field from API responses, so token counts were always null. This caused `netclaw stats` to report zero tokens in/out. - Add ParseUsage() to extract prompt_tokens/completion_tokens from both non-streaming and streaming responses - Request stream_options.include_usage for streaming calls so providers return usage in the final SSE chunk - Relax SessionCatalogService guard from requiring InputTokens to accepting either InputTokens or OutputTokens, preventing token loss when only one field is populated --- .../OpenAiCompatibleChatClientTests.cs | 97 ++++++++++++++++++- .../Gateway/SessionCatalogServiceTests.cs | 94 +++++++++++++++++- .../Gateway/SessionCatalogService.cs | 26 ++--- .../SelfHosted/OpenAiCompatibleChatClient.cs | 40 +++++++- 4 files changed, 241 insertions(+), 16 deletions(-) diff --git a/src/Netclaw.Daemon.Tests/Configuration/OpenAiCompatibleChatClientTests.cs b/src/Netclaw.Daemon.Tests/Configuration/OpenAiCompatibleChatClientTests.cs index 00173a391..31cb7df77 100644 --- a/src/Netclaw.Daemon.Tests/Configuration/OpenAiCompatibleChatClientTests.cs +++ b/src/Netclaw.Daemon.Tests/Configuration/OpenAiCompatibleChatClientTests.cs @@ -479,6 +479,96 @@ public async Task ThrowsProviderException_OnHttpError() Assert.Equal(404, ex.StatusCode); } + [Fact] + public async Task GetResponseAsync_ParsesUsageFromResponse() + { + const string json = """ + {"id":"1","model":"test","choices":[{"finish_reason":"stop","message":{"role":"assistant","content":"hi"}}], + "usage":{"prompt_tokens":100,"completion_tokens":25,"total_tokens":125}} + """; + + using var handler = new RecordingHandler(_ => new HttpResponseMessage(HttpStatusCode.OK) + { + Content = new StringContent(json, Encoding.UTF8, "application/json") + }); + using var httpClient = new HttpClient(handler) { BaseAddress = new Uri("http://localhost:8000") }; + var endpoint = OpenAiCompatibleEndpoint.FromBaseUrl("http://localhost:8000"); + var client = new OpenAiCompatibleChatClient(httpClient, endpoint, "test-model"); + + var response = await client.GetResponseAsync([new ChatMessage(ChatRole.User, "hello")]); + + Assert.NotNull(response.Usage); + Assert.Equal(100, response.Usage!.InputTokenCount); + Assert.Equal(25, response.Usage.OutputTokenCount); + Assert.Equal(125, response.Usage.TotalTokenCount); + } + + [Fact] + public async Task StreamingResponse_EmitsUsageContent_WhenPresent() + { + const string sse = """ + data: {"id":"abc","model":"test","choices":[{"index":0,"delta":{"role":"assistant","content":"hi"}}]} + + data: {"id":"abc","model":"test","choices":[{"index":0,"delta":{},"finish_reason":"stop"}],"usage":{"prompt_tokens":50,"completion_tokens":10,"total_tokens":60}} + + data: [DONE] + + """; + + using var handler = new RecordingHandler(_ => new HttpResponseMessage(HttpStatusCode.OK) + { + Content = new StringContent(sse, Encoding.UTF8, "text/event-stream") + }); + using var httpClient = new HttpClient(handler) { BaseAddress = new Uri("http://localhost:8000") }; + var endpoint = OpenAiCompatibleEndpoint.FromBaseUrl("http://localhost:8000"); + var client = new OpenAiCompatibleChatClient(httpClient, endpoint, "test-model"); + + var updates = new List(); + await foreach (var update in client.GetStreamingResponseAsync([new ChatMessage(ChatRole.User, "hello")])) + updates.Add(update); + + var usageContents = updates.SelectMany(u => u.Contents.OfType()).ToList(); + Assert.Single(usageContents); + Assert.Equal(50, usageContents[0].Details.InputTokenCount); + Assert.Equal(10, usageContents[0].Details.OutputTokenCount); + } + + [Fact] + public void ParseUsage_ReturnsNull_WhenUsageFieldMissing() + { + using var doc = JsonDocument.Parse("""{"id":"1","model":"test","choices":[]}"""); + Assert.Null(OpenAiCompatibleChatClient.ParseUsage(doc.RootElement)); + } + + [Fact] + public async Task StreamingRequest_IncludesStreamOptions() + { + const string sse = """ + data: {"id":"abc","model":"test","choices":[{"index":0,"delta":{"content":"ok"},"finish_reason":"stop"}]} + + data: [DONE] + + """; + + using var handler = new RecordingHandler(_ => new HttpResponseMessage(HttpStatusCode.OK) + { + Content = new StringContent(sse, Encoding.UTF8, "text/event-stream") + }); + using var httpClient = new HttpClient(handler) { BaseAddress = new Uri("http://localhost:8000") }; + var endpoint = OpenAiCompatibleEndpoint.FromBaseUrl("http://localhost:8000"); + var client = new OpenAiCompatibleChatClient(httpClient, endpoint, "test-model"); + + await foreach (var _ in client.GetStreamingResponseAsync([new ChatMessage(ChatRole.User, "hello")])) + { + // consume + } + + using var doc = JsonDocument.Parse(handler.RequestBodies.Single()); + var root = doc.RootElement; + Assert.True(root.GetProperty("stream").GetBoolean()); + Assert.True(root.GetProperty("stream_options").GetProperty("include_usage").GetBoolean()); + } + private sealed class RecordingHandler : HttpMessageHandler { private readonly Func _handler; @@ -489,11 +579,14 @@ public RecordingHandler(Func handler) } public List Requests { get; } = []; + public List RequestBodies { get; } = []; - protected override Task SendAsync(HttpRequestMessage request, CancellationToken cancellationToken) + protected override async Task SendAsync(HttpRequestMessage request, CancellationToken cancellationToken) { Requests.Add(request); - return Task.FromResult(_handler(request)); + if (request.Content is not null) + RequestBodies.Add(await request.Content.ReadAsStringAsync(cancellationToken)); + return _handler(request); } } } diff --git a/src/Netclaw.Daemon.Tests/Gateway/SessionCatalogServiceTests.cs b/src/Netclaw.Daemon.Tests/Gateway/SessionCatalogServiceTests.cs index 3cfbebe7a..72339cf93 100644 --- a/src/Netclaw.Daemon.Tests/Gateway/SessionCatalogServiceTests.cs +++ b/src/Netclaw.Daemon.Tests/Gateway/SessionCatalogServiceTests.cs @@ -1,6 +1,8 @@ using Microsoft.Data.Sqlite; using Microsoft.Extensions.Logging.Abstractions; using Netclaw.Actors.Channels; +using Netclaw.Actors.Protocol; +using Netclaw.Actors.Telemetry; using Netclaw.Configuration; using Netclaw.Daemon.Gateway; using Xunit; @@ -18,8 +20,8 @@ private NetclawPaths CreatePaths() return paths; } - private SessionCatalogService CreateService(NetclawPaths paths) - => new(paths, TimeProvider.System, NullLogger.Instance); + private SessionCatalogService CreateService(NetclawPaths paths, ISessionMetrics? metrics = null) + => new(paths, TimeProvider.System, NullLogger.Instance, metrics); public void Dispose() { @@ -259,6 +261,94 @@ public void GetStats_ReturnsZeros_WhenNoSessions() Assert.Equal(0, stats.TotalTurns); } + [Fact] + public void OnOutput_RecordsTokenUsage_WhenOnlyOutputTokensPresent() + { + var paths = CreatePaths(); + var metrics = new FakeMetrics(); + var service = CreateService(paths, metrics); + var sessionId = new SessionId("signalr/test-usage-1"); + + service.OnSessionCreated(sessionId, ChannelType.SignalR); + + // Simulate streaming response: InputTokens is null (Anthropic SDK bug), + // but OutputTokens has a value. + service.OnOutput(new UsageOutput + { + SessionId = sessionId, + InputTokens = null, + OutputTokens = 500 + }); + + Assert.Single(metrics.TokenUsageCalls); + Assert.Equal(0, metrics.TokenUsageCalls[0].Input); + Assert.Equal(500, metrics.TokenUsageCalls[0].Output); + } + + [Fact] + public void OnOutput_RecordsTokenUsage_WhenBothTokenFieldsPresent() + { + var paths = CreatePaths(); + var metrics = new FakeMetrics(); + var service = CreateService(paths, metrics); + var sessionId = new SessionId("signalr/test-usage-2"); + + service.OnSessionCreated(sessionId, ChannelType.SignalR); + + service.OnOutput(new UsageOutput + { + SessionId = sessionId, + InputTokens = 1000, + OutputTokens = 250 + }); + + Assert.Single(metrics.TokenUsageCalls); + Assert.Equal(1000, metrics.TokenUsageCalls[0].Input); + Assert.Equal(250, metrics.TokenUsageCalls[0].Output); + + // Verify last_input_tokens was updated in SQLite + using var conn = OpenConn(paths); + using var cmd = conn.CreateCommand(); + cmd.CommandText = "SELECT last_input_tokens FROM sessions WHERE persistence_id = $pid"; + cmd.Parameters.AddWithValue("$pid", $"session-{sessionId.Value}"); + var result = cmd.ExecuteScalar(); + Assert.Equal(1000L, result); + } + + [Fact] + public void OnOutput_DoesNotRecord_WhenBothTokenFieldsNull() + { + var paths = CreatePaths(); + var metrics = new FakeMetrics(); + var service = CreateService(paths, metrics); + var sessionId = new SessionId("signalr/test-usage-3"); + + service.OnSessionCreated(sessionId, ChannelType.SignalR); + + service.OnOutput(new UsageOutput + { + SessionId = sessionId, + InputTokens = null, + OutputTokens = null + }); + + Assert.Empty(metrics.TokenUsageCalls); + } + + private sealed class FakeMetrics : ISessionMetrics + { + public List<(long Input, long Output)> TokenUsageCalls { get; } = []; + + public void RecordTokenUsage(long inputTokens, long outputTokens) + => TokenUsageCalls.Add((inputTokens, outputTokens)); + + public void RecordTurnCompleted() { } + public void RecordSessionCreated() { } + public void RecordMemoriesFormed(int count) { } + public void RecordMemoriesRecalled(int count) { } + public void RecordSkillsLoaded(int count) { } + } + private static SqliteConnection OpenConn(NetclawPaths paths) { var conn = new SqliteConnection(new SqliteConnectionStringBuilder diff --git a/src/Netclaw.Daemon/Gateway/SessionCatalogService.cs b/src/Netclaw.Daemon/Gateway/SessionCatalogService.cs index d747512c8..487944342 100644 --- a/src/Netclaw.Daemon/Gateway/SessionCatalogService.cs +++ b/src/Netclaw.Daemon/Gateway/SessionCatalogService.cs @@ -117,18 +117,22 @@ UPDATE sessions SET _metrics?.RecordTurnCompleted(); break; - case UsageOutput usage when usage.InputTokens.HasValue: - UpdateSession(conn, persistenceId, cmd => + case UsageOutput usage + when usage.InputTokens.HasValue || usage.OutputTokens.HasValue: + if (usage.InputTokens.HasValue) { - cmd.CommandText = - """ - UPDATE sessions SET - last_input_tokens = $tokens, - last_activity = $now - WHERE persistence_id = $pid - """; - cmd.Parameters.AddWithValue("$tokens", usage.InputTokens.Value); - }); + UpdateSession(conn, persistenceId, cmd => + { + cmd.CommandText = + """ + UPDATE sessions SET + last_input_tokens = $tokens, + last_activity = $now + WHERE persistence_id = $pid + """; + cmd.Parameters.AddWithValue("$tokens", usage.InputTokens.Value); + }); + } _metrics?.RecordTokenUsage( usage.InputTokens ?? 0, usage.OutputTokens ?? 0); diff --git a/src/Netclaw.Providers/SelfHosted/OpenAiCompatibleChatClient.cs b/src/Netclaw.Providers/SelfHosted/OpenAiCompatibleChatClient.cs index 53d4de0f0..d5b877aaf 100644 --- a/src/Netclaw.Providers/SelfHosted/OpenAiCompatibleChatClient.cs +++ b/src/Netclaw.Providers/SelfHosted/OpenAiCompatibleChatClient.cs @@ -152,6 +152,11 @@ private JsonObject BuildPayload(IEnumerable messages, ChatOptions? ["stream"] = stream }; + if (stream) + { + body["stream_options"] = new JsonObject { ["include_usage"] = true }; + } + if (options?.Temperature is { } temperature) body["temperature"] = temperature; if (options?.TopP is { } topP) @@ -454,7 +459,8 @@ private static ChatResponse ParseChatResponse(JsonElement root) { ModelId = root.TryGetProperty("model", out var model) ? model.GetString() : null, ResponseId = root.TryGetProperty("id", out var id) ? id.GetString() : null, - FinishReason = finishReason + FinishReason = finishReason, + Usage = ParseUsage(root) }; } @@ -539,6 +545,11 @@ private static IEnumerable ParseStreamingUpdates(JsonElement pendingToolCalls.Clear(); } + // Usage may appear in the final streaming chunk (when stream_options.include_usage is set) + var usage = ParseUsage(root); + if (usage is not null) + contents.Add(new UsageContent(usage)); + if (contents.Count == 0 && finishReason is null) yield break; @@ -565,6 +576,33 @@ private static IEnumerable ParseStreamingUpdates(JsonElement } } + /// + /// Parses the usage object from an OpenAI-compatible response or streaming chunk. + /// Returns null when the field is absent or not an object. + /// + internal static UsageDetails? ParseUsage(JsonElement root) + { + if (!root.TryGetProperty("usage", out var usage) || usage.ValueKind != JsonValueKind.Object) + return null; + + long? promptTokens = usage.TryGetProperty("prompt_tokens", out var pt) && pt.ValueKind == JsonValueKind.Number + ? pt.GetInt64() : null; + long? completionTokens = usage.TryGetProperty("completion_tokens", out var ct) && ct.ValueKind == JsonValueKind.Number + ? ct.GetInt64() : null; + long? totalTokens = usage.TryGetProperty("total_tokens", out var tt) && tt.ValueKind == JsonValueKind.Number + ? tt.GetInt64() : null; + + if (promptTokens is null && completionTokens is null && totalTokens is null) + return null; + + return new UsageDetails + { + InputTokenCount = promptTokens, + OutputTokenCount = completionTokens, + TotalTokenCount = totalTokens ?? (promptTokens ?? 0) + (completionTokens ?? 0) + }; + } + private static ChatFinishReason? ParseFinishReason(JsonElement choice) { if (!choice.TryGetProperty("finish_reason", out var finishReason)