diff --git a/src/CodeyBox.Orchestrator/CostUsageRecorder.cs b/src/CodeyBox.Orchestrator/CostUsageRecorder.cs new file mode 100644 index 00000000..9648d75c --- /dev/null +++ b/src/CodeyBox.Orchestrator/CostUsageRecorder.cs @@ -0,0 +1,346 @@ +using System.Text.Json; +using Microsoft.Extensions.Logging; +using CodeyBox.Agents; +using CodeyBox.Core; + +namespace CodeyBox.Orchestrator; + +/// +/// Cost and token-usage recording collaborator for . +/// Extracts token counts from agent output, calculates estimated USD spend, and persists +/// durable cost/usage rows. All persistence failures are swallowed with warnings so cost +/// capture never fails a pipeline phase. +/// Extracted mechanically from ; behavior is unchanged. +/// +internal sealed class CostUsageRecorder +{ + internal const string ElapsedFallbackMetadataSource = "elapsed_fallback"; + + private readonly IWorkItemCostStore? _costStore; + private readonly IAgentUsageStore? _usageStore; + private readonly AgentCostCalculator? _costCalculator; + private readonly IReadOnlyDictionary? _costExtractors; + private readonly ILogger _log; + + public CostUsageRecorder( + IWorkItemCostStore? costStore, + IAgentUsageStore? usageStore, + AgentCostCalculator? costCalculator, + IReadOnlyDictionary? costExtractors, + ILogger log) + { + _costStore = costStore; + _usageStore = usageStore; + _costCalculator = costCalculator; + _costExtractors = costExtractors; + _log = log ?? throw new ArgumentNullException(nameof(log)); + } + + /// + /// Best-effort cost summary lookup for webhook usage blocks. Returns null + /// when the cost store is absent, no rows exist for the work item, or the + /// read fails — usage is reported as absent in any of those cases. + /// + public async Task TryGetUsageSummaryAsync(WorkItemId id) + { + if (_costStore is null) return null; + try { return await _costStore.SummariseAsync(id.ToString(), CancellationToken.None); } + catch (OperationCanceledException) { throw; } + catch (Exception ex) + { + _log.LogDebug(ex, "Cost: failed to summarise usage for work item {Id}; webhook will omit usage", id); + return null; + } + } + + public async Task TryRecordCompletionCostAsync( + CheckAndActCompletionResult result, + WorkItem item, + string phase, + int? iteration, + DateTimeOffset startedAt, + DateTimeOffset endedAt) + { + if (_costStore is null && _usageStore is null) return; + + var snapshot = NormalizeCostSnapshot( + new AgentCostSnapshot( + result.Usage.InputTokens, + result.Usage.CachedInputTokens, + result.Usage.OutputTokens, + result.ModelId), + result.ModelId); + + var usd = 0m; + if (_costCalculator is not null) + { + try { usd = _costCalculator.Calculate(snapshot, result.AgentKind); } + catch (Exception ex) + { + _log.LogWarning(ex, + "Cost: calculator threw for check-and-act completion provider '{Provider}' phase '{Phase}'; recording tokens with zero estimated cost", + result.Provider, phase); + } + } + usd = Math.Max(0m, usd); + + if (_costStore is not null) + { + try + { + await _costStore.RecordAsync(new WorkItemCost + { + Id = Guid.NewGuid().ToString(), + WorkItemId = item.Id.ToString(), + Phase = phase, + Iteration = iteration, + AgentKind = result.AgentKind.Value, + AgentInstanceId = item.AgentInstanceId, + ModelId = snapshot.ModelId, + InputTokens = snapshot.InputTokens, + CachedInputTokens = snapshot.CachedInputTokens, + OutputTokens = snapshot.OutputTokens, + EstimatedUsd = (double)usd, + StartedAt = startedAt, + EndedAt = endedAt, + RawMetadataJson = JsonSerializer.Serialize(new + { + source = "check_and_act_completion", + provider = result.Provider, + cacheHit = result.Usage.CacheHit, + }), + HasExtractedTokenUsage = true, + }, CancellationToken.None); + + var model = snapshot.ModelId ?? "(default)"; + var agentTag = new KeyValuePair("agent.kind", result.AgentKind.Value); + var agentInstanceTag = new KeyValuePair("agent.instance", item.AgentInstanceId ?? result.AgentKind.Value); + var modelTag = new KeyValuePair("model", model); + CodeyBoxMeters.AgentTokens.Add(snapshot.InputTokens, agentTag, agentInstanceTag, modelTag, + new KeyValuePair("token_type", "input")); + CodeyBoxMeters.AgentTokens.Add(snapshot.CachedInputTokens, agentTag, agentInstanceTag, modelTag, + new KeyValuePair("token_type", "cached_input")); + CodeyBoxMeters.AgentTokens.Add(snapshot.OutputTokens, agentTag, agentInstanceTag, modelTag, + new KeyValuePair("token_type", "output")); + CodeyBoxMeters.AgentCostUsd.Add((double)usd, agentTag, agentInstanceTag, modelTag); + } + catch (Exception ex) + { + _log.LogWarning(ex, "Cost: failed to persist completion row for work item {Id} phase '{Phase}'", + item.Id, phase); + } + } + + if (_usageStore is not null) + { + try + { + await _usageStore.RecordAsync( + BuildUsageEvent(result.AgentKind, item.AgentInstanceId, result.ModelId, snapshot, usd, item.Id, endedAt, phase, startedAt), + CancellationToken.None); + } + catch (Exception ex) + { + _log.LogWarning(ex, "Usage: failed to persist completion event for work item {Id} phase '{Phase}'", + item.Id, phase); + } + } + } + + /// + /// Best-effort cost capture: extracts token counts from agent output, calculates + /// estimated USD, and persists a cost row. Any failure is swallowed with a warning + /// so cost capture never aborts a pipeline phase. + /// + public async Task TryRecordCostAsync( + string? stdout, + string? stderr, + AgentKind agentKind, + string? agentInstanceId, + WorkItemId workItemId, + string phase, + int? iteration, + DateTimeOffset startedAt, + DateTimeOffset endedAt, + string? dispatchModelId) + { + if (_costStore is null && _usageStore is null) return; + + AgentCostSnapshot? snapshot; + if (_costExtractors is not null && _costExtractors.TryGetValue(agentKind, out var extractor)) + { + try { snapshot = extractor.TryExtract(stdout, stderr); } + catch (Exception ex) + { + _log.LogWarning(ex, "Cost: extractor threw for agent '{Agent}' phase '{Phase}'; recording elapsed fallback", + agentKind.Value, phase); + snapshot = null; + } + } + else + { + snapshot = null; + } + + var usedElapsedFallback = snapshot is null; + snapshot ??= new AgentCostSnapshot( + InputTokens: 0, + CachedInputTokens: 0, + OutputTokens: 0, + ModelId: dispatchModelId); + snapshot = NormalizeCostSnapshot(snapshot, dispatchModelId); + + var usd = 0m; + if (!usedElapsedFallback && _costCalculator is not null) + { + try { usd = _costCalculator.Calculate(snapshot, agentKind); } + catch (Exception ex) + { + _log.LogWarning(ex, "Cost: calculator threw for agent '{Agent}' phase '{Phase}'; recording tokens with zero estimated cost", + agentKind.Value, phase); + } + } + usd = Math.Max(0m, usd); + + if (_costStore is not null) + { + try + { + await _costStore.RecordAsync(new WorkItemCost + { + Id = Guid.NewGuid().ToString(), + WorkItemId = workItemId.ToString(), + Phase = phase, + Iteration = iteration, + AgentKind = agentKind.Value, + AgentInstanceId = agentInstanceId, + ModelId = snapshot.ModelId, + InputTokens = snapshot.InputTokens, + CachedInputTokens = snapshot.CachedInputTokens, + OutputTokens = snapshot.OutputTokens, + EstimatedUsd = (double)usd, + StartedAt = startedAt, + EndedAt = endedAt, + RawMetadataJson = usedElapsedFallback + ? JsonSerializer.Serialize(new { source = ElapsedFallbackMetadataSource }) + : "{}", + HasExtractedTokenUsage = !usedElapsedFallback, + }, CancellationToken.None); + + // Emit the same accounting as OTel counters so dashboards align with + // the per-work-item cost rows (no double-counting — one emit per row). + var model = snapshot.ModelId ?? "(default)"; + var agentTag = new KeyValuePair("agent.kind", agentKind.Value); + var agentInstanceTag = new KeyValuePair("agent.instance", agentInstanceId ?? agentKind.Value); + var modelTag = new KeyValuePair("model", model); + CodeyBoxMeters.AgentTokens.Add(snapshot.InputTokens, agentTag, agentInstanceTag, modelTag, + new KeyValuePair("token_type", "input")); + CodeyBoxMeters.AgentTokens.Add(snapshot.CachedInputTokens, agentTag, agentInstanceTag, modelTag, + new KeyValuePair("token_type", "cached_input")); + CodeyBoxMeters.AgentTokens.Add(snapshot.OutputTokens, agentTag, agentInstanceTag, modelTag, + new KeyValuePair("token_type", "output")); + CodeyBoxMeters.AgentCostUsd.Add((double)usd, agentTag, agentInstanceTag, modelTag); + } + catch (Exception ex) + { + _log.LogWarning(ex, "Cost: failed to persist row for work item {Id} phase '{Phase}'", + workItemId, phase); + } + } + + if (_usageStore is not null) + { + try + { + await _usageStore.RecordAsync( + BuildUsageEvent(agentKind, agentInstanceId, dispatchModelId, snapshot, usd, workItemId, endedAt, phase, startedAt), + CancellationToken.None); + } + catch (Exception ex) + { + _log.LogWarning(ex, "Usage: failed to persist event for work item {Id} phase '{Phase}'", + workItemId, phase); + } + } + } + + internal static AgentCostSnapshot NormalizeCostSnapshot(AgentCostSnapshot snapshot, string? dispatchModelId) => new( + InputTokens: Math.Max(0, snapshot.InputTokens), + CachedInputTokens: Math.Max(0, snapshot.CachedInputTokens), + OutputTokens: Math.Max(0, snapshot.OutputTokens), + ModelId: ResolveCostRowModelId(snapshot.ModelId, dispatchModelId)); + + internal static AgentCostSnapshot ClampCostSnapshot(AgentCostSnapshot snapshot, string? dispatchModelId = null) => + NormalizeCostSnapshot(snapshot, dispatchModelId); + + internal static string? ResolveCostRowModelId(string? extractedModelId, string? dispatchModelId) + { + if (!string.IsNullOrWhiteSpace(extractedModelId)) + return extractedModelId; + + return string.IsNullOrWhiteSpace(dispatchModelId) ? null : dispatchModelId; + } + + /// + /// Builds the durable usage-accounting row for one agent invocation. + /// + /// The row is keyed by the DISPATCHED model id, never the model id parsed from + /// agent output (). The budget gate sums + /// spend filtered on the operator-configured member.ModelId — the same + /// value used to route/dispatch. Persisting under the parsed model id (which is + /// null on many human-readable footers and a provider-supplied string on JSON + /// paths) would store spend in a different or NULL bucket than the one being + /// gated, so the gate's SUM returns zero used and AvailablePct stays at 100% + /// while real cost accrues — a fail-open bypass of the operator spend cap. + /// == member.ModelId guarantees the + /// bucket the gate reads is the bucket spend lands in. + /// + /// + /// Token counts and cost come from parsing untrusted agent stdout/stderr. A + /// hostile or malformed CLI emission (e.g. completion_tokens:-999999999) + /// would otherwise persist a negative legacy cost unit, deflate the budget + /// window SUM, and keep AvailablePct artificially high — fail-open on the + /// spend cap. Every persisted component is clamped non-negative so a bad + /// emission can only ever over-report spend, never deflate it. + /// + /// + internal static AgentUsageEvent BuildUsageEvent( + AgentKind agentKind, + string? dispatchModelId, + AgentCostSnapshot snapshot, + decimal usd, + WorkItemId workItemId, + DateTimeOffset endedAt, + string? phase = null, + DateTimeOffset? startedAt = null) => + BuildUsageEvent(agentKind, null, dispatchModelId, snapshot, usd, workItemId, endedAt, phase, startedAt); + + internal static AgentUsageEvent BuildUsageEvent( + AgentKind agentKind, + string? agentInstanceId, + string? dispatchModelId, + AgentCostSnapshot snapshot, + decimal usd, + WorkItemId workItemId, + DateTimeOffset endedAt, + string? phase = null, + DateTimeOffset? startedAt = null) => new() + { + Id = Guid.NewGuid().ToString(), + TimeUtc = endedAt, + AgentKind = agentKind.Value, + AgentInstanceId = agentInstanceId, + ModelId = dispatchModelId, + Phase = phase, + StartedUtc = startedAt, + EndedUtc = endedAt, + ElapsedMs = startedAt is { } start + ? (long)Math.Max(0, (endedAt - start).TotalMilliseconds) + : 0, + InputTokens = Math.Max(0, snapshot.InputTokens), + CachedInputTokens = Math.Max(0, snapshot.CachedInputTokens), + OutputTokens = Math.Max(0, snapshot.OutputTokens), + CostMicroCents = Math.Max(0L, AgentUsageEvent.UsdToMicroCents(usd)), + WorkItemId = workItemId.ToString(), + }; +} diff --git a/src/CodeyBox.Orchestrator/PipelineRunner.cs b/src/CodeyBox.Orchestrator/PipelineRunner.cs index ee9d325b..81de9bc8 100644 --- a/src/CodeyBox.Orchestrator/PipelineRunner.cs +++ b/src/CodeyBox.Orchestrator/PipelineRunner.cs @@ -43,7 +43,6 @@ public sealed partial class PipelineRunner : IPipelineRunner // Synthetic quota probes only ask provider availability; router score is // irrelevant, but AgentMembership requires a valid score. private const int SyntheticQuotaProbeQualityScore = 100; - private const string ElapsedFallbackMetadataSource = "elapsed_fallback"; private const int CompletionReviewContextMaxChars = 64 * 1024; private const int CompletionReviewFileMaxChars = 8 * 1024; private const int CompletionReviewMaxFiles = 80; @@ -163,6 +162,10 @@ public sealed partial class PipelineRunner : IPipelineRunner // Pure prompt builders (extracted cold-tier cluster). Owns the Build*Prompt / // Build*EscalationMessage cluster; PipelineRunner delegates to it. private readonly PromptComposer _promptComposer; + // Cost and token-usage recording (extracted cold-tier cluster). Owns the + // TryRecordCostAsync / TryRecordCompletionCostAsync / TryGetUsageSummaryAsync / + // BuildUsageEvent cluster; PipelineRunner delegates to it. + private readonly CostUsageRecorder _costUsageRecorder; // Convergence-brief composer for the delegation phase. Null in minimal // compositions / tests that don't exercise delegation; a Delegating entry // with no composer parks to NeedsOperatorInput instead of running blind. @@ -546,6 +549,7 @@ public PipelineRunner( agentSupervision: _agentSupervision, authFailureClassifier: _authFailureClassifier); _promptComposer = new PromptComposer(); + _costUsageRecorder = new CostUsageRecorder(_costStore, _usageStore, _costCalculator, _costExtractors, _log); _disabledHostHooksPath = Path.Combine(Path.GetTempPath(), "codeybox-disabled-host-hooks-" + Guid.NewGuid().ToString("N")); Directory.CreateDirectory(_disabledHostHooksPath); _watchdogOptionsAccessor = watchdogOptionsAccessor; @@ -21657,17 +21661,8 @@ private static string MergeConflictFailureInvolvementOutcome(string? failureKind /// when the cost store is absent, no rows exist for the work item, or the /// read fails — usage is reported as absent in any of those cases. /// - private async Task TryGetUsageSummaryAsync(WorkItemId id) - { - if (_costStore is null) return null; - try { return await _costStore.SummariseAsync(id.ToString(), CancellationToken.None); } - catch (OperationCanceledException) { throw; } - catch (Exception ex) - { - _log.LogDebug(ex, "Cost: failed to summarise usage for work item {Id}; webhook will omit usage", id); - return null; - } - } + private Task TryGetUsageSummaryAsync(WorkItemId id) => + _costUsageRecorder.TryGetUsageSummaryAsync(id); private async Task TransitionFailed( WorkItem item, @@ -22380,108 +22375,27 @@ internal static string NormalizeQuotaRetryPhase(string phase) => _ => $"work_item.{state.ToString().ToLowerInvariant()}", }; - // ── Cost capture ──────────────────────────────────────────────────────── + // ── Cost capture (delegated) ───────────────────────────────────────────── + // + // Owned by CostUsageRecorder; the forwarders below keep the existing + // intra-pipeline and test call-sites unchanged. - private async Task TryRecordCompletionCostAsync( + private Task TryRecordCompletionCostAsync( CheckAndActCompletionResult result, WorkItem item, string phase, int? iteration, DateTimeOffset startedAt, - DateTimeOffset endedAt) - { - if (_costStore is null && _usageStore is null) return; - - var snapshot = NormalizeCostSnapshot( - new AgentCostSnapshot( - result.Usage.InputTokens, - result.Usage.CachedInputTokens, - result.Usage.OutputTokens, - result.ModelId), - result.ModelId); - - var usd = 0m; - if (_costCalculator is not null) - { - try { usd = _costCalculator.Calculate(snapshot, result.AgentKind); } - catch (Exception ex) - { - _log.LogWarning(ex, - "Cost: calculator threw for check-and-act completion provider '{Provider}' phase '{Phase}'; recording tokens with zero estimated cost", - result.Provider, phase); - } - } - usd = Math.Max(0m, usd); - - if (_costStore is not null) - { - try - { - await _costStore.RecordAsync(new WorkItemCost - { - Id = Guid.NewGuid().ToString(), - WorkItemId = item.Id.ToString(), - Phase = phase, - Iteration = iteration, - AgentKind = result.AgentKind.Value, - AgentInstanceId = item.AgentInstanceId, - ModelId = snapshot.ModelId, - InputTokens = snapshot.InputTokens, - CachedInputTokens = snapshot.CachedInputTokens, - OutputTokens = snapshot.OutputTokens, - EstimatedUsd = (double)usd, - StartedAt = startedAt, - EndedAt = endedAt, - RawMetadataJson = JsonSerializer.Serialize(new - { - source = "check_and_act_completion", - provider = result.Provider, - cacheHit = result.Usage.CacheHit, - }), - HasExtractedTokenUsage = true, - }, CancellationToken.None); - - var model = snapshot.ModelId ?? "(default)"; - var agentTag = new KeyValuePair("agent.kind", result.AgentKind.Value); - var agentInstanceTag = new KeyValuePair("agent.instance", item.AgentInstanceId ?? result.AgentKind.Value); - var modelTag = new KeyValuePair("model", model); - CodeyBoxMeters.AgentTokens.Add(snapshot.InputTokens, agentTag, agentInstanceTag, modelTag, - new KeyValuePair("token_type", "input")); - CodeyBoxMeters.AgentTokens.Add(snapshot.CachedInputTokens, agentTag, agentInstanceTag, modelTag, - new KeyValuePair("token_type", "cached_input")); - CodeyBoxMeters.AgentTokens.Add(snapshot.OutputTokens, agentTag, agentInstanceTag, modelTag, - new KeyValuePair("token_type", "output")); - CodeyBoxMeters.AgentCostUsd.Add((double)usd, agentTag, agentInstanceTag, modelTag); - } - catch (Exception ex) - { - _log.LogWarning(ex, "Cost: failed to persist completion row for work item {Id} phase '{Phase}'", - item.Id, phase); - } - } - - if (_usageStore is not null) - { - try - { - await _usageStore.RecordAsync( - BuildUsageEvent(result.AgentKind, item.AgentInstanceId, result.ModelId, snapshot, usd, item.Id, endedAt, phase, startedAt), - CancellationToken.None); - } - catch (Exception ex) - { - _log.LogWarning(ex, "Usage: failed to persist completion event for work item {Id} phase '{Phase}'", - item.Id, phase); - } - } - } + DateTimeOffset endedAt) => + _costUsageRecorder.TryRecordCompletionCostAsync( + result, item, phase, iteration, startedAt, endedAt); /// /// Best-effort cost capture: extracts token counts from agent output, calculates /// estimated USD, and persists a cost row. Any failure is swallowed with a warning /// so cost capture never aborts a pipeline phase. /// - private async Task TryRecordCostAsync( + private Task TryRecordCostAsync( string? stdout, string? stderr, AgentKind agentKind, @@ -22491,145 +22405,19 @@ private async Task TryRecordCostAsync( int? iteration, DateTimeOffset startedAt, DateTimeOffset endedAt, - string? dispatchModelId) - { - if (_costStore is null && _usageStore is null) return; - - AgentCostSnapshot? snapshot; - if (_costExtractors is not null && _costExtractors.TryGetValue(agentKind, out var extractor)) - { - try { snapshot = extractor.TryExtract(stdout, stderr); } - catch (Exception ex) - { - _log.LogWarning(ex, "Cost: extractor threw for agent '{Agent}' phase '{Phase}'; recording elapsed fallback", - agentKind.Value, phase); - snapshot = null; - } - } - else - { - snapshot = null; - } + string? dispatchModelId) => + _costUsageRecorder.TryRecordCostAsync( + stdout, stderr, agentKind, agentInstanceId, workItemId, phase, iteration, startedAt, endedAt, dispatchModelId); - var usedElapsedFallback = snapshot is null; - snapshot ??= new AgentCostSnapshot( - InputTokens: 0, - CachedInputTokens: 0, - OutputTokens: 0, - ModelId: dispatchModelId); - snapshot = NormalizeCostSnapshot(snapshot, dispatchModelId); + private static AgentCostSnapshot NormalizeCostSnapshot(AgentCostSnapshot snapshot, string? dispatchModelId) => + CostUsageRecorder.NormalizeCostSnapshot(snapshot, dispatchModelId); - var usd = 0m; - if (!usedElapsedFallback && _costCalculator is not null) - { - try { usd = _costCalculator.Calculate(snapshot, agentKind); } - catch (Exception ex) - { - _log.LogWarning(ex, "Cost: calculator threw for agent '{Agent}' phase '{Phase}'; recording tokens with zero estimated cost", - agentKind.Value, phase); - } - } - usd = Math.Max(0m, usd); + internal static AgentCostSnapshot ClampCostSnapshot(AgentCostSnapshot snapshot, string? dispatchModelId = null) => + CostUsageRecorder.ClampCostSnapshot(snapshot, dispatchModelId); - if (_costStore is not null) - { - try - { - await _costStore.RecordAsync(new WorkItemCost - { - Id = Guid.NewGuid().ToString(), - WorkItemId = workItemId.ToString(), - Phase = phase, - Iteration = iteration, - AgentKind = agentKind.Value, - AgentInstanceId = agentInstanceId, - ModelId = snapshot.ModelId, - InputTokens = snapshot.InputTokens, - CachedInputTokens = snapshot.CachedInputTokens, - OutputTokens = snapshot.OutputTokens, - EstimatedUsd = (double)usd, - StartedAt = startedAt, - EndedAt = endedAt, - RawMetadataJson = usedElapsedFallback - ? JsonSerializer.Serialize(new { source = ElapsedFallbackMetadataSource }) - : "{}", - HasExtractedTokenUsage = !usedElapsedFallback, - }, CancellationToken.None); + internal static string? ResolveCostRowModelId(string? extractedModelId, string? dispatchModelId) => + CostUsageRecorder.ResolveCostRowModelId(extractedModelId, dispatchModelId); - // Emit the same accounting as OTel counters so dashboards align with - // the per-work-item cost rows (no double-counting — one emit per row). - var model = snapshot.ModelId ?? "(default)"; - var agentTag = new KeyValuePair("agent.kind", agentKind.Value); - var agentInstanceTag = new KeyValuePair("agent.instance", agentInstanceId ?? agentKind.Value); - var modelTag = new KeyValuePair("model", model); - CodeyBoxMeters.AgentTokens.Add(snapshot.InputTokens, agentTag, agentInstanceTag, modelTag, - new KeyValuePair("token_type", "input")); - CodeyBoxMeters.AgentTokens.Add(snapshot.CachedInputTokens, agentTag, agentInstanceTag, modelTag, - new KeyValuePair("token_type", "cached_input")); - CodeyBoxMeters.AgentTokens.Add(snapshot.OutputTokens, agentTag, agentInstanceTag, modelTag, - new KeyValuePair("token_type", "output")); - CodeyBoxMeters.AgentCostUsd.Add((double)usd, agentTag, agentInstanceTag, modelTag); - } - catch (Exception ex) - { - _log.LogWarning(ex, "Cost: failed to persist row for work item {Id} phase '{Phase}'", - workItemId, phase); - } - } - - if (_usageStore is not null) - { - try - { - await _usageStore.RecordAsync( - BuildUsageEvent(agentKind, agentInstanceId, dispatchModelId, snapshot, usd, workItemId, endedAt, phase, startedAt), - CancellationToken.None); - } - catch (Exception ex) - { - _log.LogWarning(ex, "Usage: failed to persist event for work item {Id} phase '{Phase}'", - workItemId, phase); - } - } - } - - private static AgentCostSnapshot NormalizeCostSnapshot(AgentCostSnapshot snapshot, string? dispatchModelId) => new( - InputTokens: Math.Max(0, snapshot.InputTokens), - CachedInputTokens: Math.Max(0, snapshot.CachedInputTokens), - OutputTokens: Math.Max(0, snapshot.OutputTokens), - ModelId: ResolveCostRowModelId(snapshot.ModelId, dispatchModelId)); - - internal static string? ResolveCostRowModelId(string? extractedModelId, string? dispatchModelId) - { - if (!string.IsNullOrWhiteSpace(extractedModelId)) - return extractedModelId; - - return string.IsNullOrWhiteSpace(dispatchModelId) ? null : dispatchModelId; - } - - /// - /// Builds the durable usage-accounting row for one agent invocation. - /// - /// The row is keyed by the DISPATCHED model id, never the model id parsed from - /// agent output (). The budget gate sums - /// spend filtered on the operator-configured member.ModelId — the same - /// value used to route/dispatch. Persisting under the parsed model id (which is - /// null on many human-readable footers and a provider-supplied string on JSON - /// paths) would store spend in a different or NULL bucket than the one being - /// gated, so the gate's SUM returns zero used and AvailablePct stays at 100% - /// while real cost accrues — a fail-open bypass of the operator spend cap. - /// == member.ModelId guarantees the - /// bucket the gate reads is the bucket spend lands in. - /// - /// - /// Token counts and cost come from parsing untrusted agent stdout/stderr. A - /// hostile or malformed CLI emission (e.g. completion_tokens:-999999999) - /// would otherwise persist a negative legacy cost unit, deflate the budget - /// window SUM, and keep AvailablePct artificially high — fail-open on the - /// spend cap. Every persisted component is clamped non-negative so a bad - /// emission can only ever over-report spend, never deflate it. - /// - /// internal static AgentUsageEvent BuildUsageEvent( AgentKind agentKind, string? dispatchModelId, @@ -22639,7 +22427,7 @@ internal static AgentUsageEvent BuildUsageEvent( DateTimeOffset endedAt, string? phase = null, DateTimeOffset? startedAt = null) => - BuildUsageEvent(agentKind, null, dispatchModelId, snapshot, usd, workItemId, endedAt, phase, startedAt); + CostUsageRecorder.BuildUsageEvent(agentKind, dispatchModelId, snapshot, usd, workItemId, endedAt, phase, startedAt); internal static AgentUsageEvent BuildUsageEvent( AgentKind agentKind, @@ -22650,25 +22438,8 @@ internal static AgentUsageEvent BuildUsageEvent( WorkItemId workItemId, DateTimeOffset endedAt, string? phase = null, - DateTimeOffset? startedAt = null) => new() - { - Id = Guid.NewGuid().ToString(), - TimeUtc = endedAt, - AgentKind = agentKind.Value, - AgentInstanceId = agentInstanceId, - ModelId = dispatchModelId, - Phase = phase, - StartedUtc = startedAt, - EndedUtc = endedAt, - ElapsedMs = startedAt is { } start - ? (long)Math.Max(0, (endedAt - start).TotalMilliseconds) - : 0, - InputTokens = Math.Max(0, snapshot.InputTokens), - CachedInputTokens = Math.Max(0, snapshot.CachedInputTokens), - OutputTokens = Math.Max(0, snapshot.OutputTokens), - CostMicroCents = Math.Max(0L, AgentUsageEvent.UsdToMicroCents(usd)), - WorkItemId = workItemId.ToString(), - }; + DateTimeOffset? startedAt = null) => + CostUsageRecorder.BuildUsageEvent(agentKind, agentInstanceId, dispatchModelId, snapshot, usd, workItemId, endedAt, phase, startedAt); // ── Question parsing + NeedsOperatorInput parking ─────────────────────── diff --git a/tests/CodeyBox.Tests/CostUsageRecorderTests.cs b/tests/CodeyBox.Tests/CostUsageRecorderTests.cs new file mode 100644 index 00000000..a59aaa91 --- /dev/null +++ b/tests/CodeyBox.Tests/CostUsageRecorderTests.cs @@ -0,0 +1,325 @@ +using Microsoft.Extensions.Logging.Abstractions; +using CodeyBox.Agents; +using CodeyBox.Core; +using CodeyBox.Orchestrator; +using Xunit; + +namespace CodeyBox.Tests; + +public sealed class CostUsageRecorderTests +{ + private sealed class StubExtractor(AgentCostSnapshot? snapshot, Exception? toThrow = null) : IAgentCostExtractor + { + public AgentKind Kind => AgentKind.Claude; + public ModelRateConfig? DefaultPricing => null; + + public AgentCostSnapshot? TryExtract(string? stdout, string? stderr) + { + if (toThrow is not null) throw toThrow; + return snapshot; + } + } + + private sealed class SummarisingCostStore(WorkItemUsageSummary? summary, Exception? toThrow = null) : IWorkItemCostStore + { + public List Recorded { get; } = []; + + public Task RecordAsync(WorkItemCost cost, CancellationToken ct = default) + { + if (toThrow is not null) throw toThrow; + Recorded.Add(cost); + return Task.CompletedTask; + } + + public Task SummariseAsync(string workItemId, CancellationToken ct = default) + { + if (toThrow is not null) throw toThrow; + return Task.FromResult(summary); + } + + public Task> GetByWorkItemAsync(string workItemId, CancellationToken ct = default) + => Task.FromResult>(Recorded); + + public Task> GetByProjectAsync(string projectId, DateTimeOffset from, DateTimeOffset to, CancellationToken ct = default) + => Task.FromResult>([]); + + public Task> GetFleetCostSummaryAsync(DateTimeOffset from, DateTimeOffset to, CancellationToken ct = default) + => Task.FromResult>([]); + + public Task DeleteByWorkItemAsync(string workItemId, CancellationToken ct = default) + => Task.CompletedTask; + + public Task SumEstimatedUsdAsync(string projectId, DateTimeOffset from, DateTimeOffset to, CancellationToken ct = default) + => Task.FromResult(0m); + } + + [Fact] + public async Task TryRecordCostAsync_WhenStoresNull_ReturnsGracefully() + { + var recorder = new CostUsageRecorder( + costStore: null, + usageStore: null, + costCalculator: null, + costExtractors: null, + NullLogger.Instance); + + // Should return without throwing + await recorder.TryRecordCostAsync( + stdout: "output", + stderr: null, + agentKind: AgentKind.Claude, + agentInstanceId: "inst-1", + workItemId: new WorkItemId(Guid.NewGuid()), + phase: "work", + iteration: 1, + startedAt: DateTimeOffset.UtcNow.AddSeconds(-10), + endedAt: DateTimeOffset.UtcNow, + dispatchModelId: "claude-model"); + } + + [Fact] + public async Task TryRecordCostAsync_WithExtractor_RecordsCostAndUsage() + { + var costStore = new SummarisingCostStore(null); + var usageStore = new PipelineRunnerCostCaptureTests.RecordingUsageStore(); + var pricingOpts = new AgentPricingOptions + { + DefaultRates = new() + { + [AgentKind.Claude.Value] = new ModelRateConfig { InputPerMillion = 3.0, OutputPerMillion = 15.0 }, + } + }; + var calc = new AgentCostCalculator(pricingOpts); + + var extractors = new Dictionary + { + [AgentKind.Claude] = new StubExtractor(new AgentCostSnapshot(1000, 200, 500, "test-model")), + }; + + var recorder = new CostUsageRecorder(costStore, usageStore, calc, extractors, NullLogger.Instance); + + var itemId = new WorkItemId(Guid.NewGuid()); + var started = DateTimeOffset.UtcNow.AddSeconds(-5); + var ended = DateTimeOffset.UtcNow; + + await recorder.TryRecordCostAsync( + "out", "err", AgentKind.Claude, "inst-42", itemId, "work", 1, started, ended, "test-model"); + + Assert.Single(costStore.Recorded); + var cost = costStore.Recorded[0]; + Assert.Equal(itemId.ToString(), cost.WorkItemId); + Assert.Equal("work", cost.Phase); + Assert.Equal(1, cost.Iteration); + Assert.Equal(AgentKind.Claude.Value, cost.AgentKind); + Assert.Equal("inst-42", cost.AgentInstanceId); + Assert.Equal("test-model", cost.ModelId); + Assert.Equal(1000, cost.InputTokens); + Assert.Equal(200, cost.CachedInputTokens); + Assert.Equal(500, cost.OutputTokens); + Assert.True(cost.HasExtractedTokenUsage); + Assert.True(cost.EstimatedUsd > 0); + + Assert.Single(usageStore.Recorded); + var usage = usageStore.Recorded[0]; + Assert.Equal(itemId.ToString(), usage.WorkItemId); + Assert.Equal("test-model", usage.ModelId); + Assert.Equal(1000, usage.InputTokens); + Assert.Equal(200, usage.CachedInputTokens); + Assert.Equal(500, usage.OutputTokens); + Assert.True(usage.CostMicroCents > 0); + } + + [Fact] + public async Task TryRecordCostAsync_WhenExtractorThrows_RecordsElapsedFallback() + { + var costStore = new SummarisingCostStore(null); + var usageStore = new PipelineRunnerCostCaptureTests.RecordingUsageStore(); + var extractors = new Dictionary + { + [AgentKind.Claude] = new StubExtractor(null, new InvalidOperationException("boom")), + }; + + var recorder = new CostUsageRecorder(costStore, usageStore, null, extractors, NullLogger.Instance); + var itemId = new WorkItemId(Guid.NewGuid()); + + await recorder.TryRecordCostAsync( + "out", "err", AgentKind.Claude, null, itemId, "audit", null, DateTimeOffset.UtcNow, DateTimeOffset.UtcNow, "dispatch-model"); + + Assert.Single(costStore.Recorded); + var cost = costStore.Recorded[0]; + Assert.False(cost.HasExtractedTokenUsage); + Assert.Contains(CostUsageRecorder.ElapsedFallbackMetadataSource, cost.RawMetadataJson); + Assert.Equal(0, cost.InputTokens); + Assert.Equal(0, cost.OutputTokens); + + Assert.Single(usageStore.Recorded); + var usage = usageStore.Recorded[0]; + Assert.Equal(0, usage.InputTokens); + Assert.Equal("dispatch-model", usage.ModelId); + } + + [Fact] + public async Task TryRecordCostAsync_WhenStoreThrows_SwallowsException() + { + var costStore = new SummarisingCostStore(null, new InvalidOperationException("store db failure")); + var recorder = new CostUsageRecorder(costStore, null, null, null, NullLogger.Instance); + + // Does not throw + await recorder.TryRecordCostAsync( + "out", "err", AgentKind.Claude, null, new WorkItemId(Guid.NewGuid()), "work", 1, + DateTimeOffset.UtcNow, DateTimeOffset.UtcNow, null); + } + + [Fact] + public async Task TryRecordCompletionCostAsync_RecordsCostAndUsage() + { + var costStore = new SummarisingCostStore(null); + var usageStore = new PipelineRunnerCostCaptureTests.RecordingUsageStore(); + var pricingOpts = new AgentPricingOptions + { + DefaultRates = new() + { + [AgentKind.Codex.Value] = new ModelRateConfig { InputPerMillion = 2.5, OutputPerMillion = 10.0 }, + } + }; + var calc = new AgentCostCalculator(pricingOpts); + + var recorder = new CostUsageRecorder(costStore, usageStore, calc, null, NullLogger.Instance); + + var item = new WorkItem + { + Id = new WorkItemId(Guid.NewGuid()), + ProjectId = new ProjectId("test-proj"), + Title = "Test work item", + Prompt = "Run completion", + AgentInstanceId = "inst-completion", + }; + + var completionResult = new CheckAndActCompletionResult( + Provider: "openai", + AgentKind: AgentKind.Codex, + ModelId: "gpt-4o", + Output: "Completed successfully", + Usage: new CheckAndActCompletionUsage( + InputTokens: 500, + CachedInputTokens: 50, + OutputTokens: 120, + CacheHit: true)); + + var started = DateTimeOffset.UtcNow.AddSeconds(-2); + var ended = DateTimeOffset.UtcNow; + + await recorder.TryRecordCompletionCostAsync(completionResult, item, "work", 1, started, ended); + + Assert.Single(costStore.Recorded); + var cost = costStore.Recorded[0]; + Assert.Equal(item.Id.ToString(), cost.WorkItemId); + Assert.Equal("work", cost.Phase); + Assert.Equal("inst-completion", cost.AgentInstanceId); + Assert.Equal("gpt-4o", cost.ModelId); + Assert.Equal(500, cost.InputTokens); + Assert.Equal(50, cost.CachedInputTokens); + Assert.Equal(120, cost.OutputTokens); + Assert.True(cost.HasExtractedTokenUsage); + Assert.Contains("check_and_act_completion", cost.RawMetadataJson); + + Assert.Single(usageStore.Recorded); + var usage = usageStore.Recorded[0]; + Assert.Equal(item.Id.ToString(), usage.WorkItemId); + Assert.Equal("gpt-4o", usage.ModelId); + Assert.Equal(500, usage.InputTokens); + } + + [Fact] + public async Task TryGetUsageSummaryAsync_WhenStoreNull_ReturnsNull() + { + var recorder = new CostUsageRecorder(null, null, null, null, NullLogger.Instance); + var res = await recorder.TryGetUsageSummaryAsync(new WorkItemId(Guid.NewGuid())); + Assert.Null(res); + } + + [Fact] + public async Task TryGetUsageSummaryAsync_WhenStoreHasSummary_ReturnsSummary() + { + var summary = new WorkItemUsageSummary( + new WorkItemIterationUsage(1, 1000, 200, 0, 100, 0.05, 1000), + new WorkItemUsageTotal(1000, 200, 0, 100, 0.05, 1000)); + var costStore = new SummarisingCostStore(summary); + var recorder = new CostUsageRecorder(costStore, null, null, null, NullLogger.Instance); + + var res = await recorder.TryGetUsageSummaryAsync(new WorkItemId(Guid.NewGuid())); + Assert.NotNull(res); + Assert.Equal(1000, res.Total.TokensInput); + Assert.Equal(0.05, res.Total.CostUsd); + } + + [Fact] + public async Task TryGetUsageSummaryAsync_WhenStoreThrows_ReturnsNull() + { + var costStore = new SummarisingCostStore(null, new InvalidOperationException("sql error")); + var recorder = new CostUsageRecorder(costStore, null, null, null, NullLogger.Instance); + + var res = await recorder.TryGetUsageSummaryAsync(new WorkItemId(Guid.NewGuid())); + Assert.Null(res); + } + + [Fact] + public void BuildUsageEvent_ClampsNegativeValues() + { + var snapshot = new AgentCostSnapshot(-10, -5, -20, "parsed-id"); + var ev = CostUsageRecorder.BuildUsageEvent( + AgentKind.Claude, + "instance-1", + "dispatch-id", + snapshot, + usd: -10m, + new WorkItemId(Guid.NewGuid()), + endedAt: DateTimeOffset.UtcNow, + phase: "work", + startedAt: DateTimeOffset.UtcNow.AddSeconds(5)); // ended before started -> negative elapsed + + Assert.Equal(0, ev.InputTokens); + Assert.Equal(0, ev.CachedInputTokens); + Assert.Equal(0, ev.OutputTokens); + Assert.Equal(0L, ev.CostMicroCents); + Assert.Equal(0, ev.ElapsedMs); + Assert.Equal("dispatch-id", ev.ModelId); + } + + [Fact] + public void ClampCostSnapshot_NormalizesAndClamps() + { + var snapshot = new AgentCostSnapshot(-10, -2, -30, " "); + var clamped = CostUsageRecorder.ClampCostSnapshot(snapshot, "fallback-model"); + + Assert.Equal(0, clamped.InputTokens); + Assert.Equal(0, clamped.CachedInputTokens); + Assert.Equal(0, clamped.OutputTokens); + Assert.Equal("fallback-model", clamped.ModelId); + } + + [Fact] + public void PipelineRunner_Forwarders_DelegateCorrectly() + { + var snapshot = new AgentCostSnapshot(100, 20, 30, "parsed"); + var ev = PipelineRunner.BuildUsageEvent( + AgentKind.Codex, + "dispatch", + snapshot, + 1.0m, + new WorkItemId(Guid.NewGuid()), + DateTimeOffset.UtcNow); + + Assert.Equal(100, ev.InputTokens); + Assert.Equal("dispatch", ev.ModelId); + + var clamped = PipelineRunner.ClampCostSnapshot(new AgentCostSnapshot(-5, 10, -2, "model")); + Assert.Equal(0, clamped.InputTokens); + Assert.Equal(10, clamped.CachedInputTokens); + Assert.Equal(0, clamped.OutputTokens); + + Assert.Equal("extracted", PipelineRunner.ResolveCostRowModelId("extracted", "dispatch")); + Assert.Equal("dispatch", PipelineRunner.ResolveCostRowModelId("", "dispatch")); + Assert.Null(PipelineRunner.ResolveCostRowModelId(null, null)); + } +}