Surface Cosmos request charge (RUs) in MCP structured results

CHANGELOG.md

@@ -2,6 +2,11 @@
 
 ## Unreleased
 
+### Improvements
+
+- Cosmos DB data-plane commands now consistently expose their aggregate observed request charge in structured output and connection-scoped `info` telemetry, including metadata/configuration operations, scripts, change feed reads, paginated operations, handled probes, and charged failures. Azure Resource Manager control-plane operations remain uncharged.
+- Added `$sessionRequestCharge` and `$sessionChargedOperationCount` as read-only shell variables. Set `$sessionRequestChargeWarningThreshold` to a positive RU threshold to print one warning when the current connection reaches it; `info` reports it as `session.requestChargeWarningThreshold`.
+
 ### Fixes
 
 - Local emulator outages are now detected across Cosmos DB commands. Requests fail promptly with an error and return the shell to its disconnected state instead of leaving an unresponsive session labeled as connected.
@@ -90,7 +95,8 @@ A focused cycle on top of 1.1.115-preview. New `ttl` and `conflict` commands man
 ### Improvements
 
 - **Structured (JSON) tool results for MCP.** MCP tool results now carry the machine-readable JSON payload (`result`/`outputText`/`error` plus `currentLocation`) as first-class `structuredContent` in addition to the existing JSON text block, so agents can consume structured results directly. The two representations are kept byte-for-byte equivalent, and text-only clients are unaffected. ([#154](https://github.com/Azure/CosmosDBShell/issues/154))
-
+- **Request charge in MCP structured results.** Instrumented data-plane commands (`query`, including `--explain`; `print`; container-scoped `ls`; `can-i` probes; `batch run`; `mkitem`; `replace`; `patch`; `rm`; `import`; and `export`) now report the Cosmos DB request charge (in RUs) consumed by the operation as a uniform `requestCharge` field on the MCP tool result, so agents can track observed RU cost consistently across calls. Budget enforcement remains tracked separately in #162. ([#162](https://github.com/Azure/CosmosDBShell/issues/162))
+- **Connection-scoped request-charge totals.** The shell accumulates request charges observed from instrumented commands and reports the total in `info` as `session.requestCharge`, together with the number of positively charged command operations as `session.chargedOperationCount`. A successful `connect` starts new totals; database and container navigation do not reset them. This is usage telemetry, not budget enforcement or billing data. ([#162](https://github.com/Azure/CosmosDBShell/issues/162))
 - **Destructive MCP commands now prompt for confirmation instead of being blocked.** When an MCP client invokes `delete`, `rm`, `rmcon`, or `rmdb`, the server sends an elicitation prompt describing the exact command line and only runs it if the user approves; declining, cancelling, or a client that cannot confirm results in nothing being executed. This removes the need for any write opt-in flag. ([#158](https://github.com/Azure/CosmosDBShell/issues/158))
 
 ### Fixes

CosmosDBShell.Tests/CommandTests/BatchCommandTests.cs

@@ -262,6 +262,23 @@ await Assert.ThrowsAsync<CommandException>(
             () => BatchExecutor.ExecuteAsync("batch", null!, default, operations, CancellationToken.None));
     }
 
+    [Fact]
+    public void CreateResultState_PreservesBatchResultAndRequestCharge()
+    {
+        var summary = JsonSerializer.SerializeToElement(new
+        {
+            success = true,
+            requestCharge = 4.25,
+        });
+
+        var state = BatchExecutor.CreateResultState(summary, "Batch committed.", 4.25);
+
+        Assert.Equal(4.25, state.RequestCharge);
+        var result = Assert.IsType<ShellJson>(state.Result).Value;
+        Assert.True(result.GetProperty("success").GetBoolean());
+        Assert.Equal(4.25, result.GetProperty("requestCharge").GetDouble());
+    }
+
     [Theory]
     [InlineData(1, "1 operation")]
     [InlineData(2, "2 operations")]

CosmosDBShell.Tests/CommandTests/CanICommandTests.cs

@@ -92,6 +92,28 @@ public async Task CanI_KeyAuth_ReportsAllow()
         Assert.Equal("key", json.GetProperty("method").GetString());
         Assert.Equal("MyDB", json.GetProperty("database").GetString());
         Assert.Equal("Products", json.GetProperty("container").GetString());
+        Assert.Null(state.RequestCharge);
+    }
+
+    [Fact]
+    public void Build_ProbeResultSetsRequestCharge()
+    {
+        var command = new CanICommand();
+
+        var state = command.Build(
+            new CommandState(),
+            "query",
+            "MyDB",
+            "Products",
+            "allow",
+            "probe",
+            200,
+            null,
+            1.75);
+
+        Assert.Equal(1.75, state.RequestCharge);
+        var json = Assert.IsType<ShellJson>(state.Result).Value;
+        Assert.Equal("probe", json.GetProperty("method").GetString());
     }
 
     [Fact]

CosmosDBShell.Tests/CommandTests/CosmosCommandTests.cs

@@ -11,6 +11,67 @@ namespace CosmosShell.Tests.CommandTests;
 
 public class CosmosCommandTests
 {
+    [Fact]
+    public async Task ExecuteCosmosCommandAsync_RecordsChargeAndClearsStaleCharge()
+    {
+        using var shell = ShellInterpreter.CreateInstance();
+        var state = new CommandState { RequestCharge = 9 };
+
+        state = await shell.ExecuteCosmosCommandAsync(new TestCosmosCommand(2.5), state, string.Empty, CancellationToken.None);
+        state = await shell.ExecuteCosmosCommandAsync(new TestCosmosCommand(null), state, string.Empty, CancellationToken.None);
+
+        Assert.Null(state.RequestCharge);
+        Assert.Equal(2.5, shell.SessionRequestCharge);
+        Assert.Equal(1, shell.SessionChargedOperationCount);
+    }
+
+    [Fact]
+    public async Task ExecuteCosmosCommandAsync_CombinesExplicitAndScopedChargesOnce()
+    {
+        using var shell = ShellInterpreter.CreateInstance();
+
+        var state = await shell.ExecuteCosmosCommandAsync(
+            new TestCosmosCommand(2.5, scopedRequestCharge: 1.25),
+            new CommandState(),
+            string.Empty,
+            CancellationToken.None);
+
+        Assert.Equal(3.75, state.RequestCharge);
+        Assert.Equal(3.75, shell.SessionRequestCharge);
+        Assert.Equal(1, shell.SessionChargedOperationCount);
+    }
+
+    [Fact]
+    public async Task ExecuteCosmosCommandAsync_ThrownCommandPreservesScopedCharge()
+    {
+        using var shell = ShellInterpreter.CreateInstance();
+
+        var exception = await Assert.ThrowsAsync<InvalidOperationException>(() => shell.ExecuteCosmosCommandAsync(
+            new TestCosmosCommand(null, scopedRequestCharge: 1.5, throwAfterRecording: true),
+            new CommandState(),
+            string.Empty,
+            CancellationToken.None));
+
+        Assert.Equal(1.5, RequestChargeContext.GetExceptionCharge(exception));
+        Assert.Equal(1.5, shell.SessionRequestCharge);
+        Assert.Equal(1, shell.SessionChargedOperationCount);
+    }
+
+    [Fact]
+    public async Task ExecuteCosmosCommandAsync_CancelledCommandPreservesScopedCharge()
+    {
+        using var shell = ShellInterpreter.CreateInstance();
+
+        await Assert.ThrowsAsync<OperationCanceledException>(() => shell.ExecuteCosmosCommandAsync(
+            new TestCosmosCommand(null, scopedRequestCharge: 1.5, cancelAfterRecording: true),
+            new CommandState(),
+            string.Empty,
+            CancellationToken.None));
+
+        Assert.Equal(1.5, shell.SessionRequestCharge);
+        Assert.Equal(1, shell.SessionChargedOperationCount);
+    }
+
     [Fact]
     public void CreatePartitionKey_WithHierarchicalIntegerComponents_PreservesIntegerTypes()
     {
@@ -42,8 +103,37 @@ public void CreatePartitionKey_WithHierarchicalIntegerComponents_PreservesIntege
 
     private sealed class TestCosmosCommand : CosmosCommand
     {
+        private readonly double? requestCharge;
+        private readonly double scopedRequestCharge;
+        private readonly bool throwAfterRecording;
+        private readonly bool cancelAfterRecording;
+
+        public TestCosmosCommand(
+            double? requestCharge = null,
+            double scopedRequestCharge = 0,
+            bool throwAfterRecording = false,
+            bool cancelAfterRecording = false)
+        {
+            this.requestCharge = requestCharge;
+            this.scopedRequestCharge = scopedRequestCharge;
+            this.throwAfterRecording = throwAfterRecording;
+            this.cancelAfterRecording = cancelAfterRecording;
+        }
+
         public override Task<CommandState> ExecuteAsync(ShellInterpreter shell, CommandState commandState, string commandText, CancellationToken token)
         {
+            RequestChargeContext.Record(this.scopedRequestCharge);
+            if (this.throwAfterRecording)
+            {
+                throw new InvalidOperationException("test");
+            }
+
+            if (this.cancelAfterRecording)
+            {
+                throw new OperationCanceledException();
+            }
+
+            commandState.RequestCharge = this.requestCharge;
             return Task.FromResult(commandState);
         }
 

CosmosDBShell.Tests/CommandTests/InfoCommandTests.cs

@@ -4,8 +4,10 @@
 
 namespace CosmosShell.Tests.CommandTests;
 
+using System.Text.Json;
 using Azure.Data.Cosmos.Shell.Commands;
 using Azure.Data.Cosmos.Shell.Core;
+using Azure.Data.Cosmos.Shell.Parser;
 using Azure.Data.Cosmos.Shell.States;
 using Azure.Data.Cosmos.Shell.Util;
 using Microsoft.Azure.Cosmos;
@@ -135,6 +137,101 @@ public void FormatSize_Null_ReturnsNotAvailable()
         Assert.Equal(MessageService.GetString("command-stats-na"), InfoCommand.FormatSize(null));
     }
 
+    [Fact]
+    public void SessionRequestCharge_AccumulatesObservedCharges()
+    {
+        using var shell = ShellInterpreter.CreateInstance();
+
+        shell.RecordRequestCharge(new CommandState { RequestCharge = 1.25 });
+        shell.RecordRequestCharge(new CommandState());
+        shell.RecordRequestCharge(new CommandState { RequestCharge = 0 });
+        shell.RecordRequestCharge(new ErrorCommandState(new InvalidOperationException()) { RequestCharge = 2.5 });
+
+        Assert.Equal(3.75, shell.SessionRequestCharge);
+        Assert.Equal(2, shell.SessionChargedOperationCount);
+    }
+
+    [Fact]
+    public void Connect_ResetsSessionRequestCharge()
+    {
+        using var shell = ShellInterpreter.CreateInstance();
+        shell.RecordRequestCharge(new CommandState { RequestCharge = 4.5 });
+
+        shell.Connect(CreateTestClient(), credentialTypeOverride: "AccountKey");
+
+        Assert.Equal(0, shell.SessionRequestCharge);
+        Assert.Equal(0, shell.SessionChargedOperationCount);
+    }
+
+    [Fact]
+    public void Disconnect_DoesNotResetSessionRequestCharge()
+    {
+        using var shell = ShellInterpreter.CreateInstance();
+        shell.Connect(CreateTestClient(), credentialTypeOverride: "AccountKey");
+        shell.RecordRequestCharge(new CommandState { RequestCharge = 4.5 });
+
+        shell.Disconnect();
+
+        Assert.Equal(4.5, shell.SessionRequestCharge);
+        Assert.Equal(1, shell.SessionChargedOperationCount);
+    }
+
+    [Fact]
+    public void AddSessionUsage_AddsCurrentChargeToStructuredResult()
+    {
+        using var shell = ShellInterpreter.CreateInstance();
+        shell.RecordRequestCharge(new CommandState { RequestCharge = 3.75 });
+        var result = new Dictionary<string, object?>();
+
+        InfoCommand.AddSessionUsage(shell, result, renderOutput: false);
+
+        var json = JsonSerializer.SerializeToElement(result);
+        Assert.Equal(3.75, json.GetProperty("session").GetProperty("requestCharge").GetDouble());
+        Assert.Equal(1, json.GetProperty("session").GetProperty("chargedOperationCount").GetInt64());
+    }
+
+    [Fact]
+    public void AddSessionUsage_IncludesCurrentCommandCharge()
+    {
+        using var shell = ShellInterpreter.CreateInstance();
+        shell.RecordRequestCharge(new CommandState { RequestCharge = 3.75 });
+        using var scope = RequestChargeContext.Begin();
+        RequestChargeContext.Record(1.25);
+        var result = new Dictionary<string, object?>();
+
+        InfoCommand.AddSessionUsage(shell, result, renderOutput: false);
+
+        var json = JsonSerializer.SerializeToElement(result);
+        Assert.Equal(5, json.GetProperty("session").GetProperty("requestCharge").GetDouble());
+        Assert.Equal(2, json.GetProperty("session").GetProperty("chargedOperationCount").GetInt64());
+    }
+
+    [Fact]
+    public void AddSessionUsage_IncludesConfiguredMaximum()
+    {
+        using var shell = ShellInterpreter.CreateInstance();
+        shell.SetVariable("sessionRequestChargeWarningThreshold", new ShellDecimal(25.5));
+        var result = new Dictionary<string, object?>();
+
+        InfoCommand.AddSessionUsage(shell, result, renderOutput: false);
+
+        var session = JsonSerializer.SerializeToElement(result).GetProperty("session");
+        Assert.Equal(25.5, session.GetProperty("requestChargeWarningThreshold").GetDouble());
+    }
+
+    [Fact]
+    public void RecordRequestCharge_IgnoresPriorConnectionGeneration()
+    {
+        using var shell = ShellInterpreter.CreateInstance();
+        var generation = shell.SessionRequestChargeGeneration;
+        shell.Connect(CreateTestClient(), credentialTypeOverride: "AccountKey");
+
+        shell.RecordRequestCharge(new CommandState { RequestCharge = 8 }, generation);
+
+        Assert.Equal(0, shell.SessionRequestCharge);
+        Assert.Equal(0, shell.SessionChargedOperationCount);
+    }
+
     [Theory]
     [InlineData("table")]
     [InlineData("tbl")]

CosmosDBShell.Tests/CommandTests/ListCommandTests.cs

@@ -223,4 +223,15 @@ public async Task ReadQueryResponseAsync_ValidContent_ReturnsJsonDocument()
         var item = Assert.Single(document.RootElement.GetProperty("Documents").EnumerateArray());
         Assert.Equal("1", item.GetProperty("id").GetString());
     }
+
+    [Fact]
+    public void AccumulateRequestCharge_AddsEveryPageCharge()
+    {
+        var state = new CommandState();
+
+        ListCommand.AccumulateRequestCharge(state, 1.25);
+        ListCommand.AccumulateRequestCharge(state, 2.5);
+
+        Assert.Equal(3.75, state.RequestCharge);
+    }
 }