diff --git a/src/CodexSdk.MAF/CodexAIAgent.cs b/src/CodexSdk.MAF/CodexAIAgent.cs index 229375d..0b7b28f 100644 --- a/src/CodexSdk.MAF/CodexAIAgent.cs +++ b/src/CodexSdk.MAF/CodexAIAgent.cs @@ -100,7 +100,7 @@ protected override async Task RunCoreAsync( switch (threadEvent) { case TurnCompletedEvent turnCompleted: - usage = CreateUsageDetails(turnCompleted.Usage); + usage = turnCompleted.Usage.ToUsageDetails(); break; case TurnFailedEvent turnFailed: @@ -200,16 +200,6 @@ await NotifyThreadStartedIfNeededAsync( await SaveNewMessagesAsync(safeSession, mergedMessages, responseMessages, cancellationToken); } - private static UsageDetails CreateUsageDetails(Usage usage) - { - return new UsageDetails - { - InputTokenCount = usage.InputTokens, - CachedInputTokenCount = usage.CachedInputTokens, - OutputTokenCount = usage.OutputTokens, - }; - } - private async ValueTask NotifyThreadStartedIfNeededAsync( CodexAgentSession session, string? threadId, diff --git a/src/CodexSdk.MAF/Internal/ThreadEventExtensions.cs b/src/CodexSdk.MAF/Internal/ThreadEventExtensions.cs index d7d9649..9afda80 100644 --- a/src/CodexSdk.MAF/Internal/ThreadEventExtensions.cs +++ b/src/CodexSdk.MAF/Internal/ThreadEventExtensions.cs @@ -35,6 +35,18 @@ internal static class ThreadEventExtensions return update; } + public static UsageDetails ToUsageDetails(this Usage usage) + { + return new UsageDetails + { + InputTokenCount = usage.InputTokens, + CachedInputTokenCount = usage.CachedInputTokens, + OutputTokenCount = usage.OutputTokens, + ReasoningTokenCount = usage.ReasoningOutputTokens, + TotalTokenCount = (long)usage.InputTokens + usage.OutputTokens, + }; + } + private static AgentResponseUpdate CreateLifecycleUpdate(string eventType) { return new AgentResponseUpdate @@ -61,13 +73,7 @@ private static AgentResponseUpdate CreateUsageUpdate(Usage usage) }, Contents = [ - new UsageContent(new UsageDetails - { - InputTokenCount = usage.InputTokens, - CachedInputTokenCount = usage.CachedInputTokens, - OutputTokenCount = usage.OutputTokens, - ReasoningTokenCount = usage.ReasoningOutputTokens, - }) + new UsageContent(usage.ToUsageDetails()) ], }; } diff --git a/tests/CodexSdk.MAF.Tests/ThreadEventExtensionsTests.cs b/tests/CodexSdk.MAF.Tests/ThreadEventExtensionsTests.cs index 2d4266e..c566b80 100644 --- a/tests/CodexSdk.MAF.Tests/ThreadEventExtensionsTests.cs +++ b/tests/CodexSdk.MAF.Tests/ThreadEventExtensionsTests.cs @@ -92,6 +92,34 @@ public void ToAgentResponseUpdate_WhenTurnStarts_ReturnsVisibleSystemContent() Assert.False(update.ShouldSaveAsResponseMessage()); } + [Fact] + public void ToAgentResponseUpdate_WhenTurnCompletes_MapsCompleteUsage() + { + // Arrange + var threadEvent = new TurnCompletedEvent + { + Usage = new Usage + { + InputTokens = 100, + CachedInputTokens = 60, + OutputTokens = 30, + ReasoningOutputTokens = 20, + }, + }; + + // Act + var update = threadEvent.ToAgentResponseUpdate(); + + // Assert + Assert.NotNull(update); + var usage = Assert.IsType(Assert.Single(update.Contents)).Details; + Assert.Equal(100L, usage.InputTokenCount); + Assert.Equal(60L, usage.CachedInputTokenCount); + Assert.Equal(30L, usage.OutputTokenCount); + Assert.Equal(20L, usage.ReasoningTokenCount); + Assert.Equal(130L, usage.TotalTokenCount); + } + [Fact] public void ToAgentResponseUpdate_WhenCommandExecutionCompletedEvent_ReturnsFunctionResultContent() {