Surface Cosmos request charge (RUs) in MCP structured results

CHANGELOG.md

@@ -90,7 +90,7 @@ 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.** Successful data-plane commands (`query`, `print`, container-scoped `ls`, `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 RU cost consistently across calls. ([#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/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);
+    }
 }

CosmosDBShell.Tests/CommandTests/SprocCommandExecutionTests.cs

@@ -206,6 +206,7 @@ public async Task ExecAsync_ReturnsResource()
 
         var json = Assert.IsType<JsonElement>(state.Result!.ConvertShellObject(DataType.Json));
         Assert.True(json.GetProperty("ok").GetBoolean());
+        Assert.Equal(2.0, state.RequestCharge);
     }
 
     [Fact]

CosmosDBShell.Tests/McpResponseFactoryTests.cs

@@ -165,4 +165,52 @@ public void CreateError_PopulatesStructuredContentMatchingTextBlock()
         Assert.Equal("boom", structured.GetProperty("error").GetString());
         Assert.Equal("/TestDatabase", structured.GetProperty("currentLocation").GetString());
     }
+
+    [Fact]
+    public void CreateSuccess_IncludesRequestChargeWhenSet()
+    {
+        var commandState = new CommandState
+        {
+            Result = new ShellJson(JsonSerializer.SerializeToElement(new { result = "success" })),
+            RequestCharge = 4.25,
+        };
+
+        var result = McpResponseFactory.CreateSuccess(commandState, new ConnectedState(null!));
+
+        Assert.NotNull(result.StructuredContent);
+        var structured = result.StructuredContent!.Value;
+        Assert.True(structured.TryGetProperty("requestCharge", out var requestCharge));
+        Assert.Equal(4.25, requestCharge.GetDouble());
+    }
+
+    [Fact]
+    public void CreateSuccess_StructuredError_IncludesRequestChargeWhenSet()
+    {
+        var commandState = new StructuredErrorCommandState(
+            new CommandException("batch", "Batch failed."),
+            new ShellJson(JsonSerializer.SerializeToElement(new { success = false })))
+        {
+            RequestCharge = 3.5,
+        };
+
+        var result = McpResponseFactory.CreateSuccess(commandState, new ConnectedState(null!));
+
+        Assert.True(result.IsError);
+        Assert.NotNull(result.StructuredContent);
+        Assert.Equal(3.5, result.StructuredContent!.Value.GetProperty("requestCharge").GetDouble());
+    }
+
+    [Fact]
+    public void CreateSuccess_OmitsRequestChargeWhenNotSet()
+    {
+        var commandState = new CommandState
+        {
+            Result = new ShellJson(JsonSerializer.SerializeToElement(new { result = "success" })),
+        };
+
+        var result = McpResponseFactory.CreateSuccess(commandState, new ConnectedState(null!));
+
+        Assert.NotNull(result.StructuredContent);
+        Assert.False(result.StructuredContent!.Value.TryGetProperty("requestCharge", out _));
+    }
 }

CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ExportCommand.cs

@@ -111,6 +111,7 @@ public override async Task<CommandState> ExecuteAsync(ShellInterpreter shell, Co
                 exported = count,
                 requestCharge = charge,
             })),
+            RequestCharge = charge,
         };
     }
 

CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ImportCommand.cs

@@ -548,6 +548,7 @@ public override async Task<CommandState> ExecuteAsync(ShellInterpreter shell, Co
                 requestCharge = charge,
                 dryRun,
             })),
+            RequestCharge = charge,
         };
     }
 

CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ListCommand.cs

@@ -229,6 +229,7 @@ private async Task<CommandState> ListContainerItemsAsync(ConnectedState state, S
         {
             using var response = await feedIterator.ReadNextAsync(token);
             using var queryDocument = await ReadQueryResponseAsync(response, token);
+            AccumulateRequestCharge(returnState, response.Headers.RequestCharge);
 
             foreach (var element in queryDocument.RootElement.GetProperty("Documents").EnumerateArray())
             {
@@ -280,6 +281,11 @@ private async Task<CommandState> ListContainerItemsAsync(ConnectedState state, S
         return returnState;
     }
 
+    internal static void AccumulateRequestCharge(CommandState commandState, double requestCharge)
+    {
+        commandState.RequestCharge = (commandState.RequestCharge ?? 0) + requestCharge;
+    }
+
     internal static async Task<JsonDocument> ReadQueryResponseAsync(ResponseMessage response, CancellationToken token)
     {
         if (!response.IsSuccessStatusCode)

CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/MakeItemCommand.cs

@@ -59,7 +59,7 @@ public async override Task<CommandState> ExecuteAsync(ShellInterpreter shell, Co
             "mkitem",
             token);
 
-        var summary = await WriteItemAsync(container, commandState, jsonOpt, this.Force == true, token);
+        var summary = await WriteItemAsync(container, jsonOpt, this.Force == true, token);
 
         var returnState = new CommandState();
         returnState.Result = new ShellJson(JsonSerializer.SerializeToElement(new
@@ -70,6 +70,7 @@ public async override Task<CommandState> ExecuteAsync(ShellInterpreter shell, Co
             failed = summary.Failed,
             requestCharge = summary.RequestCharge,
         }));
+        returnState.RequestCharge = summary.RequestCharge;
         return returnState;
     }
 
@@ -131,7 +132,7 @@ private static object ParseJsonElement(JsonElement element)
         }
     }
 
-    private static async Task<WriteSummary> WriteItemAsync(Container container, CommandState commandState, string? jsonOpt, bool force, CancellationToken token)
+    private static async Task<WriteSummary> WriteItemAsync(Container container, string? jsonOpt, bool force, CancellationToken token)
     {
         if (!string.IsNullOrEmpty(jsonOpt))
         {
@@ -154,7 +155,6 @@ private static async Task<WriteSummary> WriteItemAsync(Container container, Comm
                                 ? await container.UpsertItemAsync(element, cancellationToken: token)
                                 : await container.CreateItemAsync(element, cancellationToken: token);
                             charge += result.RequestCharge;
-
                             if (result.StatusCode == System.Net.HttpStatusCode.Created)
                             {
                                 createdCount++;

CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/PatchCommand.cs

@@ -133,6 +133,7 @@ public override async Task<CommandState> ExecuteAsync(ShellInterpreter shell, Co
                     patched = true,
                     requestCharge = response.RequestCharge,
                 })),
+                RequestCharge = response.RequestCharge,
             };
         }
         catch (CosmosException ce) when (ce.StatusCode == System.Net.HttpStatusCode.NotFound)

CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/PrintCommand.cs

@@ -57,6 +57,7 @@ private async Task<CommandState> PrintItemAsync(Container container, Cancellatio
 
             if (response.IsSuccessStatusCode)
             {
+                commandState.RequestCharge = response.Headers.RequestCharge;
                 using var reader = new StreamReader(response.Content);
                 var content = await reader.ReadToEndAsync();
 

CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/QueryCommand.cs

@@ -677,6 +677,7 @@ private async Task<CommandState> ExecuteQueryAsync(Container container, ShellInt
         {
             var returnState = CreateCommandState(this.OutputFormat);
             var aggregatedDocuments = new List<JsonElement>();
+            double totalRequestCharge = 0;
 
             var options = new QueryRequestOptions
             {
@@ -727,11 +728,14 @@ private async Task<CommandState> ExecuteQueryAsync(Container container, ShellInt
 
                 using var queryDocument = JsonDocument.Parse(responseContent);
                 ShellInterpreter.WriteLine(MessageService.GetString("command-query-fetched", new Dictionary<string, object> { { "count", queryDocument.RootElement.GetProperty("_count").ToString() } }));
-                var queryMetrics = response.Diagnostics.GetQueryMetrics();
-                if (queryMetrics != null)
-                {
-                    AnsiConsole.MarkupLine(MessageService.GetString("command-query-request_charge", new Dictionary<string, object> { { "charge", queryMetrics.TotalRequestCharge.ToString() } }));
-                }
+
+                // Cosmos always returns the RU cost in the response headers, whereas query
+                // metrics (and their TotalRequestCharge) can be null when diagnostics are
+                // unavailable. Accumulate and report from the headers so the charge is always
+                // correct; the detailed metrics payload is built separately from the response.
+                var pageRequestCharge = response.Headers.RequestCharge;
+                totalRequestCharge += pageRequestCharge;
+                AnsiConsole.MarkupLine(MessageService.GetString("command-query-request_charge", new Dictionary<string, object> { { "charge", pageRequestCharge.ToString("F2", CultureInfo.InvariantCulture) } }));
 
                 var pageDocuments = queryDocument.RootElement.GetProperty("Documents");
                 var pageExceedsLimit = PageExceedsLimit(aggregatedDocuments.Count, pageDocuments, effectiveMaxItemCount);
@@ -751,7 +755,7 @@ private async Task<CommandState> ExecuteQueryAsync(Container container, ShellInt
                             {
                                 { "type", "item" },
                                 { "values", aggregatedDocuments },
-                                { "requestCharge", queryMetrics?.TotalRequestCharge ?? 0 },
+                                { "requestCharge", totalRequestCharge },
                                 { "queryMetrics", metricProperty },
                                 { "indexMetrics", parsedIndexMetrics ?? new Dictionary<string, object>() },
                             });
@@ -888,6 +892,7 @@ private async Task<CommandState> ExecuteQueryAsync(Container container, ShellInt
                 AnsiConsole.MarkupLine(MessageService.GetString("command-results-limit_reached", new Dictionary<string, object> { { "count", effectiveMaxItemCount.Value } }));
             }
 
+            returnState.RequestCharge = totalRequestCharge;
             return returnState;
         }
         catch (OperationCanceledException) when (token.IsCancellationRequested)

CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ReplaceCommand.cs

@@ -63,6 +63,7 @@ public override async Task<CommandState> ExecuteAsync(ShellInterpreter shell, Co
                 failed = summary.Failed,
                 requestCharge = summary.RequestCharge,
             })),
+            RequestCharge = summary.RequestCharge,
         };
     }
 

CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/RmCommand.cs

@@ -118,25 +118,26 @@ private async Task<ExitCode> RemoveItemsFromContainerAsync(ConnectedState state,
         var matchKeyPropertyNames = string.IsNullOrEmpty(this.Key) ? partitionKeyPropertyNames : [this.Key];
 
         var totalCount = 0;
+        double totalCharge = 0;
         bool dryRun = this.DryRun == true;
 
         // In dry-run mode, count what would be deleted without issuing any delete.
-        async Task<bool> TryDeleteAsync(string id, PartitionKey partitionKey)
+        async Task<(bool Counted, double RequestCharge)> TryDeleteAsync(string id, PartitionKey partitionKey)
         {
             if (dryRun)
             {
-                return true;
+                return (true, 0);
             }
 
             try
             {
-                await container.DeleteItemAsync<object>(id, partitionKey, cancellationToken: token);
-                return true;
+                var deleteResponse = await container.DeleteItemAsync<object>(id, partitionKey, cancellationToken: token);
+                return (true, deleteResponse.RequestCharge);
             }
             catch (CosmosException ex) when (ex.StatusCode == System.Net.HttpStatusCode.NotFound)
             {
                 // Item was already deleted, skip
-                return false;
+                return (false, ex.RequestCharge);
             }
         }
 
@@ -183,7 +184,9 @@ async Task<bool> TryDeleteAsync(string id, PartitionKey partitionKey)
 
                         if (id != null && shouldDelete)
                         {
-                            if (await TryDeleteAsync(id, CreatePartitionKey(pkElements)))
+                            var deleteResult = await TryDeleteAsync(id, CreatePartitionKey(pkElements));
+                            totalCharge += deleteResult.RequestCharge;
+                            if (deleteResult.Counted)
                             {
                                 totalCount++;
                             }
@@ -214,7 +217,9 @@ async Task<bool> TryDeleteAsync(string id, PartitionKey partitionKey)
                         var id = idElement.GetString();
                         if (id != null)
                         {
-                            if (await TryDeleteAsync(id, CreatePartitionKey(pkElements)))
+                            var deleteResult = await TryDeleteAsync(id, CreatePartitionKey(pkElements));
+                            totalCharge += deleteResult.RequestCharge;
+                            if (deleteResult.Counted)
                             {
                                 totalCount++;
                             }
@@ -236,7 +241,13 @@ async Task<bool> TryDeleteAsync(string id, PartitionKey partitionKey)
                     break;
                 }
 
-                var response = await feedIterator.ReadNextAsync(token);
+                using var response = await feedIterator.ReadNextAsync(token);
+
+                // The scan pages that locate matching items consume RUs regardless of whether
+                // any item is ultimately deleted (including in --dry-run), so include each
+                // page's request charge from the response headers.
+                totalCharge += response.Headers.RequestCharge;
+
                 using var streamReader = new StreamReader(response.Content);
                 var queryDocument = JsonDocument.Parse(await streamReader.ReadToEndAsync());
 
@@ -272,7 +283,9 @@ async Task<bool> TryDeleteAsync(string id, PartitionKey partitionKey)
 
                     if (shouldDelete)
                     {
-                        if (await TryDeleteAsync(id, CreatePartitionKey(pkElements)))
+                        var deleteResult = await TryDeleteAsync(id, CreatePartitionKey(pkElements));
+                        totalCharge += deleteResult.RequestCharge;
+                        if (deleteResult.Counted)
                         {
                             totalCount++;
                         }
@@ -296,6 +309,7 @@ async Task<bool> TryDeleteAsync(string id, PartitionKey partitionKey)
         commandState.Result = new ShellJson(JsonSerializer.SerializeToElement(new { type = "item", count = totalCount, dryRun }));
         commandState.RenderUser = () => AnsiConsole.MarkupLine(renderMessage);
 
+        commandState.RequestCharge = totalCharge;
         return new ExitCode(0);
     }
 

CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/SprocCommand.cs

@@ -446,6 +446,7 @@ internal async Task<CommandState> ExecAsync(Container container, CommandState co
                 response.RequestCharge.ToString("F2")));
 
             commandState.Result = new ShellJson(response.Resource.Clone());
+            commandState.RequestCharge = response.RequestCharge;
             return commandState;
         }
         catch (CosmosException ex) when (ex.StatusCode == HttpStatusCode.NotFound)

CosmosDBShell/Azure.Data.Cosmos.Shell.Core/CommandState.cs

@@ -69,6 +69,12 @@ public OutputFormat OutputFormat
     /// </summary>
     internal Func<TabularData>? RenderTabular { get; set; }
 
+    /// <summary>
+    /// Gets or sets the Cosmos DB request charge (in RUs) consumed by the command, when applicable.
+    /// Data-plane commands set this so consumers such as the MCP structured payload can report cost uniformly.
+    /// </summary>
+    internal double? RequestCharge { get; set; }
+
     internal bool BreakBlock { get; set; } = false;
 
     internal bool ContinueBlock { get; set; } = false;

CosmosDBShell/Azure.Data.Cosmos.Shell.Mcp/McpResponseFactory.cs

@@ -83,6 +83,11 @@ private static JsonObject CreateSuccessPayload(CommandState commandState)
     {
         var payload = new JsonObject();
 
+        if (commandState.RequestCharge.HasValue)
+        {
+            payload["requestCharge"] = commandState.RequestCharge.Value;
+        }
+
         if (commandState.IsError)
         {
             payload["error"] = GetErrorPayloadMessage(commandState);

docs/mcp.md

@@ -100,8 +100,9 @@ Both representations are always byte-for-byte equivalent.
 | ----- | ------------ | ----------- |
 | `result` | Commands that produce output | The command result as JSON (objects, arrays, or a scalar). Text-only results are represented as a JSON string. Failed transactional batches include their per-operation summary here alongside `error`. |
 | `outputText` | CSV output commands with non-empty text | The CSV rendering of the result. Omitted when the CSV output is empty or whitespace. |
+| `requestCharge` | Instrumented data-plane command results | The Cosmos DB request charge (in RUs) consumed by the command, as a number. |
 | `error` | Failed commands | The error message. |
 | `currentLocation` | Always | The shell's current navigation path (for example `/MyDatabase/MyContainer`), or `null` when disconnected. |
 
-Successful results set `result` (and optionally `outputText`); failed results set `error`, may also include a structured `result`, and mark the tool result as an error. `currentLocation` is always included so a client can track navigation state across calls.
+Successful results set `result` (and optionally `outputText`); failed results set `error`, may also include a structured `result`, and mark the tool result as an error. `currentLocation` is always included so a client can track navigation state across calls. Instrumented data-plane commands (`query`, `print`, `ls` for container items, `mkitem`, `replace`, `patch`, `rm`, `import`, and `export`) additionally set `requestCharge` when available so a client can track RU cost across successful and structured-error results.