From 044df912ba0d5413ec1d1dced34129599cfad37e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mike=20Kr=C3=BCger?= Date: Wed, 8 Jul 2026 09:35:10 +0200 Subject: [PATCH 1/5] Surface Cosmos request charge (RUs) in MCP structured results Add a uniform RequestCharge property to CommandState and emit it as a top-level 'requestCharge' field in the MCP tool result payload. Retrofit data-plane commands (query, print, mkitem, replace, patch, rm, import, export, and sproc exec) to record the request units consumed so agents can track RU cost consistently across calls. The document result shape is unchanged; requestCharge is a sibling metadata field. Addresses part (a) of #162. --- CHANGELOG.md | 1 + .../SprocCommandExecutionTests.cs | 1 + .../McpResponseFactoryTests.cs | 31 +++++++++++++++++++ .../ExportCommand.cs | 1 + .../ImportCommand.cs | 1 + .../MakeItemCommand.cs | 7 +++++ .../PatchCommand.cs | 1 + .../PrintCommand.cs | 1 + .../QueryCommand.cs | 3 ++ .../ReplaceCommand.cs | 14 +++++---- .../RmCommand.cs | 11 +++++-- .../SprocCommand.cs | 1 + .../CommandState.cs | 6 ++++ .../McpResponseFactory.cs | 5 +++ docs/mcp.md | 3 +- 15 files changed, 77 insertions(+), 10 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 8a62b482..847f0323 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,7 @@ ### 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.** Data-plane commands (`query`, `print`, `mkitem`, `replace`, `patch`, `rm`, `import`, `export`, and `sproc exec`) 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)) ## 1.1.115-preview — 2026-07-01 diff --git a/CosmosDBShell.Tests/CommandTests/SprocCommandExecutionTests.cs b/CosmosDBShell.Tests/CommandTests/SprocCommandExecutionTests.cs index 2e6a6c34..e497075a 100644 --- a/CosmosDBShell.Tests/CommandTests/SprocCommandExecutionTests.cs +++ b/CosmosDBShell.Tests/CommandTests/SprocCommandExecutionTests.cs @@ -204,6 +204,7 @@ public async Task ExecAsync_ReturnsResource() var json = Assert.IsType(state.Result!.ConvertShellObject(DataType.Json)); Assert.True(json.GetProperty("ok").GetBoolean()); + Assert.Equal(2.0, state.RequestCharge); } [Fact] diff --git a/CosmosDBShell.Tests/McpResponseFactoryTests.cs b/CosmosDBShell.Tests/McpResponseFactoryTests.cs index 4a9088bd..490a22d9 100644 --- a/CosmosDBShell.Tests/McpResponseFactoryTests.cs +++ b/CosmosDBShell.Tests/McpResponseFactoryTests.cs @@ -144,4 +144,35 @@ 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_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 _)); + } } \ No newline at end of file diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ExportCommand.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ExportCommand.cs index e78daf8d..2fbd7f04 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ExportCommand.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ExportCommand.cs @@ -107,6 +107,7 @@ public override async Task ExecuteAsync(ShellInterpreter shell, Co return new CommandState { Result = new ShellJson(SuccessDocument.RootElement.Clone()), + RequestCharge = charge, }; } diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ImportCommand.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ImportCommand.cs index 4074588a..92695145 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ImportCommand.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ImportCommand.cs @@ -542,6 +542,7 @@ public override async Task ExecuteAsync(ShellInterpreter shell, Co return new CommandState { Result = new ShellJson(SuccessDocument.RootElement.Clone()), + RequestCharge = charge, }; } diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/MakeItemCommand.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/MakeItemCommand.cs index 923f33dd..c535911e 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/MakeItemCommand.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/MakeItemCommand.cs @@ -65,6 +65,7 @@ public async override Task ExecuteAsync(ShellInterpreter shell, Co var returnState = new CommandState(); returnState.Result = new ShellJson(SuccessDocument.RootElement.Clone()); + returnState.RequestCharge = commandState.RequestCharge; return returnState; } @@ -130,6 +131,7 @@ private static async Task WriteItemAsync(Container container, CommandState comma { if (!string.IsNullOrEmpty(jsonOpt)) { + double totalCharge = 0; try { using var doc = JsonDocument.Parse(jsonOpt); @@ -149,6 +151,7 @@ private static async Task WriteItemAsync(Container container, CommandState comma ? await container.UpsertItemAsync(element, cancellationToken: token) : await container.CreateItemAsync(element, cancellationToken: token); charge += result.RequestCharge; + totalCharge += result.RequestCharge; if (result.StatusCode == System.Net.HttpStatusCode.Created) { @@ -273,6 +276,8 @@ private static async Task WriteItemAsync(Container container, CommandState comma ? await container.UpsertItemAsync(root, cancellationToken: token) : await container.CreateItemAsync(root, cancellationToken: token); + totalCharge += result.RequestCharge; + if (result.StatusCode == System.Net.HttpStatusCode.Created) { var key = force ? "command-mkitem-upserted-created" : "command-mkitem-created-success"; @@ -314,6 +319,8 @@ private static async Task WriteItemAsync(Container container, CommandState comma { throw new CommandException("mkitem", MessageService.GetArgsString("json_error_parsing_arg", "message", ex.Message), ex); } + + commandState.RequestCharge = totalCharge; } } } \ No newline at end of file diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/PatchCommand.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/PatchCommand.cs index d9d65786..c5d892b7 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/PatchCommand.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/PatchCommand.cs @@ -130,6 +130,7 @@ public override async Task ExecuteAsync(ShellInterpreter shell, Co return new CommandState { Result = new ShellJson(SuccessDocument.RootElement.Clone()), + RequestCharge = response.RequestCharge, }; } catch (CosmosException ce) when (ce.StatusCode == System.Net.HttpStatusCode.NotFound) diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/PrintCommand.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/PrintCommand.cs index b9887990..9ed343d8 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/PrintCommand.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/PrintCommand.cs @@ -57,6 +57,7 @@ private async Task PrintItemAsync(Container container, Cancellatio if (response.IsSuccessStatusCode) { + commandState.RequestCharge = response.Headers.RequestCharge; using var reader = new StreamReader(response.Content); var content = await reader.ReadToEndAsync(); diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/QueryCommand.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/QueryCommand.cs index 72a6d8e9..ab0a447b 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/QueryCommand.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/QueryCommand.cs @@ -284,6 +284,7 @@ private async Task ExecuteQueryAsync(Container container, ShellInt var returnState = new CommandState(); returnState.SetFormat(this.OutputFormat ?? Environment.GetEnvironmentVariable("COSMOSDB_SHELL_FORMAT")); var aggregatedDocuments = new List(); + double totalRequestCharge = 0; try { @@ -367,6 +368,7 @@ private async Task ExecuteQueryAsync(Container container, ShellInt var queryMetrics = response.Diagnostics.GetQueryMetrics(); if (queryMetrics != null) { + totalRequestCharge += queryMetrics.TotalRequestCharge; AnsiConsole.MarkupLine(MessageService.GetString("command-query-request_charge", new Dictionary { { "charge", queryMetrics.TotalRequestCharge.ToString() } })); } @@ -524,6 +526,7 @@ private async Task ExecuteQueryAsync(Container container, ShellInt AnsiConsole.MarkupLine(MessageService.GetString("command-results-limit_reached", new Dictionary { { "count", effectiveMaxItemCount.Value } })); } + returnState.RequestCharge = totalRequestCharge; return returnState; } catch (OperationCanceledException) when (token.IsCancellationRequested) diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ReplaceCommand.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ReplaceCommand.cs index fc38f8a9..fd0a4fdb 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ReplaceCommand.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ReplaceCommand.cs @@ -54,15 +54,16 @@ public override async Task ExecuteAsync(ShellInterpreter shell, Co var partitionKeyPaths = await CosmosResourceFacade.GetPartitionKeyPathsAsync(connectedState, databaseName!, containerName!, token); - await ReplaceItemsAsync(container, partitionKeyPaths, jsonOpt, this.ETag, token); + var totalCharge = await ReplaceItemsAsync(container, partitionKeyPaths, jsonOpt, this.ETag, token); return new CommandState { Result = new ShellJson(SuccessDocument.RootElement.Clone()), + RequestCharge = totalCharge, }; } - private static async Task ReplaceItemsAsync(Container container, IReadOnlyList partitionKeyPaths, string jsonInput, string? etag, CancellationToken token) + private static async Task ReplaceItemsAsync(Container container, IReadOnlyList partitionKeyPaths, string jsonInput, string? etag, CancellationToken token) { try { @@ -76,11 +77,10 @@ private static async Task ReplaceItemsAsync(Container container, IReadOnlyList partitionKeyPaths, JsonElement arrayRoot, CancellationToken token) + private static async Task ReplaceArrayAsync(Container container, IReadOnlyList partitionKeyPaths, JsonElement arrayRoot, CancellationToken token) { int successCount = 0; int failCount = 0; @@ -133,6 +133,8 @@ private static async Task ReplaceArrayAsync(Container container, IReadOnlyList ReplaceOneAsync(Container container, IReadOnlyList partitionKeyPaths, JsonElement item, string? etag, CancellationToken token, bool printSuccess) diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/RmCommand.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/RmCommand.cs index e082c7d8..9d335090 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/RmCommand.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/RmCommand.cs @@ -114,6 +114,7 @@ private async Task RemoveItemsFromContainerAsync(ConnectedState state, var matchKeyPropertyNames = string.IsNullOrEmpty(this.Key) ? partitionKeyPropertyNames : [this.Key]; var totalCount = 0; + double totalCharge = 0; // Process pipe input if available if (hasPipeInput && commandState.Result is ShellJson jsonResult) @@ -160,7 +161,8 @@ private async Task RemoveItemsFromContainerAsync(ConnectedState state, { try { - await container.DeleteItemAsync(id, CreatePartitionKey(pkElements), cancellationToken: token); + var deleteResponse = await container.DeleteItemAsync(id, CreatePartitionKey(pkElements), cancellationToken: token); + totalCharge += deleteResponse.RequestCharge; totalCount++; } catch (CosmosException ex) when (ex.StatusCode == System.Net.HttpStatusCode.NotFound) @@ -196,7 +198,8 @@ private async Task RemoveItemsFromContainerAsync(ConnectedState state, { try { - await container.DeleteItemAsync(id, CreatePartitionKey(pkElements), cancellationToken: token); + var deleteResponse = await container.DeleteItemAsync(id, CreatePartitionKey(pkElements), cancellationToken: token); + totalCharge += deleteResponse.RequestCharge; totalCount++; } catch (CosmosException ex) when (ex.StatusCode == System.Net.HttpStatusCode.NotFound) @@ -259,7 +262,8 @@ private async Task RemoveItemsFromContainerAsync(ConnectedState state, { try { - await container.DeleteItemAsync(id, CreatePartitionKey(pkElements), cancellationToken: token); + var deleteResponse = await container.DeleteItemAsync(id, CreatePartitionKey(pkElements), cancellationToken: token); + totalCharge += deleteResponse.RequestCharge; totalCount++; } catch (CosmosException ex) when (ex.StatusCode == System.Net.HttpStatusCode.NotFound) @@ -290,6 +294,7 @@ private async Task RemoveItemsFromContainerAsync(ConnectedState state, })); } + commandState.RequestCharge = totalCharge; return new ExitCode(0); } diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/SprocCommand.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/SprocCommand.cs index 04f41bc8..c5f31058 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/SprocCommand.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/SprocCommand.cs @@ -446,6 +446,7 @@ internal async Task 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) diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/CommandState.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/CommandState.cs index 98720118..884a6004 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/CommandState.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/CommandState.cs @@ -28,6 +28,12 @@ public partial class CommandState internal bool IsPrinted { get; set; } + /// + /// 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. + /// + internal double? RequestCharge { get; set; } + internal bool BreakBlock { get; set; } = false; internal bool ContinueBlock { get; set; } = false; diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Mcp/McpResponseFactory.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Mcp/McpResponseFactory.cs index c7c4249e..ab9e52f3 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Mcp/McpResponseFactory.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Mcp/McpResponseFactory.cs @@ -96,6 +96,11 @@ private static JsonObject CreateSuccessPayload(CommandState commandState) payload["result"] = resultNode; } + if (commandState.RequestCharge.HasValue) + { + payload["requestCharge"] = commandState.RequestCharge.Value; + } + if (commandState.OutputFormat == OutputFormat.CSV) { var outputText = commandState.GenerateOutputText(); diff --git a/docs/mcp.md b/docs/mcp.md index d8f4d00f..d05acfa9 100644 --- a/docs/mcp.md +++ b/docs/mcp.md @@ -86,8 +86,9 @@ Both representations are always byte-for-byte equivalent. | ----- | ------------ | ----------- | | `result` | Successful commands that produce output | The command result as JSON (objects, arrays, or a scalar). Text-only results are represented as a JSON string. | | `outputText` | CSV output commands with non-empty text | The CSV rendering of the result. Omitted when the CSV output is empty or whitespace. | +| `requestCharge` | Data-plane commands that consume request units | 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` 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` and mark the tool result as an error. `currentLocation` is always included so a client can track navigation state across calls. Data-plane commands (`query`, `print`, `mkitem`, `replace`, `patch`, `rm`, `import`, `export`, and `sproc exec`) additionally set `requestCharge` so a client can track RU cost across calls. From 9576858a6702b028beae99e8fdaf9a39a5b02916 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mike=20Kr=C3=BCger?= Date: Mon, 13 Jul 2026 15:38:57 +0200 Subject: [PATCH 2/5] Account request charge from response headers in query/rm (always present, include scan pages) --- .../QueryCommand.cs | 16 +++++++++------- .../RmCommand.cs | 6 ++++++ 2 files changed, 15 insertions(+), 7 deletions(-) diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/QueryCommand.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/QueryCommand.cs index ab0a447b..17b224d4 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/QueryCommand.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/QueryCommand.cs @@ -365,12 +365,14 @@ private async Task ExecuteQueryAsync(Container container, ShellInt using var queryDocument = JsonDocument.Parse(responseContent); ShellInterpreter.WriteLine(MessageService.GetString("command-query-fetched", new Dictionary { { "count", queryDocument.RootElement.GetProperty("_count").ToString() } })); - var queryMetrics = response.Diagnostics.GetQueryMetrics(); - if (queryMetrics != null) - { - totalRequestCharge += queryMetrics.TotalRequestCharge; - AnsiConsole.MarkupLine(MessageService.GetString("command-query-request_charge", new Dictionary { { "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 { { "charge", pageRequestCharge.ToString() } })); var pageDocuments = queryDocument.RootElement.GetProperty("Documents"); var pageExceedsLimit = PageExceedsLimit(aggregatedDocuments.Count, pageDocuments, effectiveMaxItemCount); @@ -389,7 +391,7 @@ private async Task ExecuteQueryAsync(Container container, ShellInt new Dictionary() { { "documents", aggregatedDocuments }, - { "requestCharge", queryMetrics?.TotalRequestCharge ?? 0 }, + { "requestCharge", pageRequestCharge }, { "queryMetrics", metricProperty }, { "indexMetrics", parsedIndexMetrics ?? new Dictionary() }, }); diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/RmCommand.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/RmCommand.cs index 55a4f41b..72e61618 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/RmCommand.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/RmCommand.cs @@ -242,6 +242,12 @@ private async Task RemoveItemsFromContainerAsync(ConnectedState state, } 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()); From 7a89a57c454d0e4e9c49f651524ea8401f2dbd3d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mike=20Kr=C3=BCger?= Date: Fri, 28 Aug 2026 16:16:54 +0200 Subject: [PATCH 3/5] Polish request charge reporting --- CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/QueryCommand.cs | 2 +- CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/RmCommand.cs | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/QueryCommand.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/QueryCommand.cs index 04de9f40..1dba9129 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/QueryCommand.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/QueryCommand.cs @@ -735,7 +735,7 @@ private async Task ExecuteQueryAsync(Container container, ShellInt // 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 { { "charge", pageRequestCharge.ToString() } })); + AnsiConsole.MarkupLine(MessageService.GetString("command-query-request_charge", new Dictionary { { "charge", pageRequestCharge.ToString("F2", CultureInfo.InvariantCulture) } })); var pageDocuments = queryDocument.RootElement.GetProperty("Documents"); var pageExceedsLimit = PageExceedsLimit(aggregatedDocuments.Count, pageDocuments, effectiveMaxItemCount); diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/RmCommand.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/RmCommand.cs index 5f6a86e5..0635f9d6 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/RmCommand.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/RmCommand.cs @@ -241,7 +241,7 @@ private async Task RemoveItemsFromContainerAsync(ConnectedState state, 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 From 4caa0fc3f4f3deaccaf313fb182c114eb4698337 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mike=20Kr=C3=BCger?= Date: Fri, 28 Aug 2026 16:29:02 +0200 Subject: [PATCH 4/5] Preserve request charge on MCP errors --- CosmosDBShell.Tests/McpResponseFactoryTests.cs | 17 +++++++++++++++++ .../McpResponseFactory.cs | 10 +++++----- docs/mcp.md | 4 ++-- 3 files changed, 24 insertions(+), 7 deletions(-) diff --git a/CosmosDBShell.Tests/McpResponseFactoryTests.cs b/CosmosDBShell.Tests/McpResponseFactoryTests.cs index 45f502e9..4957cee9 100644 --- a/CosmosDBShell.Tests/McpResponseFactoryTests.cs +++ b/CosmosDBShell.Tests/McpResponseFactoryTests.cs @@ -183,6 +183,23 @@ public void CreateSuccess_IncludesRequestChargeWhenSet() 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() { diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Mcp/McpResponseFactory.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Mcp/McpResponseFactory.cs index ad56be77..ef02702e 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Mcp/McpResponseFactory.cs +++ b/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); @@ -102,11 +107,6 @@ private static JsonObject CreateSuccessPayload(CommandState commandState) payload["result"] = resultNode; } - if (commandState.RequestCharge.HasValue) - { - payload["requestCharge"] = commandState.RequestCharge.Value; - } - if (commandState.OutputFormat == OutputFormat.CSV) { var outputText = commandState.GenerateOutputText(); diff --git a/docs/mcp.md b/docs/mcp.md index d4eabd07..3221ed50 100644 --- a/docs/mcp.md +++ b/docs/mcp.md @@ -100,9 +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` | Successful instrumented data-plane commands | The Cosmos DB request charge (in RUs) consumed by the command, as a number. | +| `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 data-plane commands (`query`, `print`, `ls` for container items, `mkitem`, `replace`, `patch`, `rm`, `import`, and `export`) additionally set `requestCharge` so a client can track RU cost 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. From 515ab16f6fe51588246c91760d8a1db904a86e8e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mike=20Kr=C3=BCger?= Date: Fri, 28 Aug 2026 16:38:57 +0200 Subject: [PATCH 5/5] Clarify dry-run delete counting --- .../Azure.Data.Cosmos.Shell.Commands/RmCommand.cs | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/RmCommand.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/RmCommand.cs index 0635f9d6..87278b99 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/RmCommand.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/RmCommand.cs @@ -122,7 +122,7 @@ private async Task RemoveItemsFromContainerAsync(ConnectedState state, bool dryRun = this.DryRun == true; // In dry-run mode, count what would be deleted without issuing any delete. - async Task<(bool Deleted, double RequestCharge)> TryDeleteAsync(string id, PartitionKey partitionKey) + async Task<(bool Counted, double RequestCharge)> TryDeleteAsync(string id, PartitionKey partitionKey) { if (dryRun) { @@ -186,7 +186,7 @@ private async Task RemoveItemsFromContainerAsync(ConnectedState state, { var deleteResult = await TryDeleteAsync(id, CreatePartitionKey(pkElements)); totalCharge += deleteResult.RequestCharge; - if (deleteResult.Deleted) + if (deleteResult.Counted) { totalCount++; } @@ -219,7 +219,7 @@ private async Task RemoveItemsFromContainerAsync(ConnectedState state, { var deleteResult = await TryDeleteAsync(id, CreatePartitionKey(pkElements)); totalCharge += deleteResult.RequestCharge; - if (deleteResult.Deleted) + if (deleteResult.Counted) { totalCount++; } @@ -285,7 +285,7 @@ private async Task RemoveItemsFromContainerAsync(ConnectedState state, { var deleteResult = await TryDeleteAsync(id, CreatePartitionKey(pkElements)); totalCharge += deleteResult.RequestCharge; - if (deleteResult.Deleted) + if (deleteResult.Counted) { totalCount++; }