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 01/16] 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 02/16] 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 03/16] 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 04/16] 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 05/16] 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++; } From 9a2ea3767e86223e783b63c0dc5c8b52a82c438b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mike=20Kr=C3=BCger?= Date: Mon, 31 Aug 2026 12:23:59 +0200 Subject: [PATCH 06/16] Complete RU observability coverage --- CHANGELOG.md | 2 +- .../CommandTests/BatchCommandTests.cs | 17 +++++++ .../CommandTests/CanICommandTests.cs | 22 +++++++++ .../Integration/BatchOperationTests.cs | 18 ++++++- .../Integration/QueryCommandTests.cs | 17 ++++++- .../BatchExecutor.cs | 15 ++++-- .../CanICommand.cs | 49 ++++++++++--------- .../QueryCommand.cs | 3 +- docs/mcp.md | 6 ++- 9 files changed, 116 insertions(+), 33 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 32d8d1c6..1804f85d 100644 --- a/CHANGELOG.md +++ b/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)) +- **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)) - **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 diff --git a/CosmosDBShell.Tests/CommandTests/BatchCommandTests.cs b/CosmosDBShell.Tests/CommandTests/BatchCommandTests.cs index 236e8832..3850dce3 100644 --- a/CosmosDBShell.Tests/CommandTests/BatchCommandTests.cs +++ b/CosmosDBShell.Tests/CommandTests/BatchCommandTests.cs @@ -262,6 +262,23 @@ await Assert.ThrowsAsync( () => 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(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")] diff --git a/CosmosDBShell.Tests/CommandTests/CanICommandTests.cs b/CosmosDBShell.Tests/CommandTests/CanICommandTests.cs index 04159be3..7a8b2da4 100644 --- a/CosmosDBShell.Tests/CommandTests/CanICommandTests.cs +++ b/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(state.Result).Value; + Assert.Equal("probe", json.GetProperty("method").GetString()); } [Fact] diff --git a/CosmosDBShell.Tests/Integration/BatchOperationTests.cs b/CosmosDBShell.Tests/Integration/BatchOperationTests.cs index 758ba7cd..7d14e2bc 100644 --- a/CosmosDBShell.Tests/Integration/BatchOperationTests.cs +++ b/CosmosDBShell.Tests/Integration/BatchOperationTests.cs @@ -33,7 +33,22 @@ public async Task BatchRun_MultipleCreates_CommitsAtomically() "{\"op\":\"create\",\"item\":{\"id\":\"a\",\"pk\":\"" + pk + "\"}}," + "{\"op\":\"create\",\"item\":{\"id\":\"b\",\"pk\":\"" + pk + "\"}}]"; - var output = await ExecuteWithOutputAsync($"batch run '{json}' --partition-key {pk}"); + var outputFile = CreateTempFile(); + Shell.StdOutRedirect = outputFile; + CommandState state; + string output; + try + { + state = await ExecuteAsync($"batch run '{json}' --partition-key {pk}"); + output = await File.ReadAllTextAsync(outputFile, TestContext.Current.CancellationToken); + } + finally + { + Shell.StdOutRedirect = null; + } + + Assert.False(state.IsError, FormatError(state)); + Assert.True(state.RequestCharge > 0); var root = JsonDocument.Parse(output).RootElement; Assert.True(root.GetProperty("success").GetBoolean()); @@ -81,6 +96,7 @@ public async Task BatchRun_FailingOperation_RollsBackEntireBatch() var batchState = await ExecuteAsync($"batch run '{json}' --partition-key {pk}"); Assert.True(batchState.IsError); Assert.Equal(ShellExitCode.GeneralFailure, batchState.ExitCode); + Assert.True(batchState.RequestCharge > 0); var output = await File.ReadAllTextAsync(outputFile, TestContext.Current.CancellationToken); var root = JsonDocument.Parse(output).RootElement; diff --git a/CosmosDBShell.Tests/Integration/QueryCommandTests.cs b/CosmosDBShell.Tests/Integration/QueryCommandTests.cs index 0ced42e7..b9004423 100644 --- a/CosmosDBShell.Tests/Integration/QueryCommandTests.cs +++ b/CosmosDBShell.Tests/Integration/QueryCommandTests.cs @@ -135,7 +135,22 @@ public async Task Query_Explain_ReturnsStructuredResultWithoutDocuments() { var query = $"SELECT * FROM c WHERE c.id = '{this.GetSeedItemId(1)}'"; - var output = await ExecuteWithOutputAsync($"query \"{query}\" --explain"); + var outputFile = CreateTempFile(); + Shell.StdOutRedirect = outputFile; + CommandState state; + string output; + try + { + state = await ExecuteAsync($"query \"{query}\" --explain"); + output = await File.ReadAllTextAsync(outputFile, TestContext.Current.CancellationToken); + } + finally + { + Shell.StdOutRedirect = null; + } + + Assert.False(state.IsError, FormatError(state)); + Assert.True(state.RequestCharge > 0); using var document = JsonDocument.Parse(output); var root = document.RootElement; diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/BatchExecutor.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/BatchExecutor.cs index 67e109d1..68040c96 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/BatchExecutor.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/BatchExecutor.cs @@ -74,7 +74,7 @@ public static async Task ExecuteAsync( operations.Count, "charge", response.RequestCharge.ToString("F2", CultureInfo.InvariantCulture)); - return CreateResultState(summary, successMessage); + return CreateResultState(summary, successMessage, response.RequestCharge); } var errorMessage = MessageService.GetArgsString( @@ -88,14 +88,21 @@ public static async Task ExecuteAsync( commandName, errorMessage, new RequestFailedException((int)response.StatusCode, errorMessage)), - new ShellJson(summary)); + new ShellJson(summary)) + { + RequestCharge = response.RequestCharge, + }; errorState.RenderUser = () => ShellInterpreter.WriteLine(errorMessage); return errorState; } - private static CommandState CreateResultState(JsonElement summary, string message) + internal static CommandState CreateResultState(JsonElement summary, string message, double requestCharge) { - var state = new CommandState { Result = new ShellJson(summary) }; + var state = new CommandState + { + Result = new ShellJson(summary), + RequestCharge = requestCharge, + }; state.RenderUser = () => ShellInterpreter.WriteLine(message); return state; } diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/CanICommand.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/CanICommand.cs index 8df38ded..d303084f 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/CanICommand.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/CanICommand.cs @@ -75,7 +75,7 @@ public async override Task ExecuteAsync(ShellInterpreter shell, Co // mutating or control-plane operation, so it is always reported as indeterminate. if (action == "manage") { - return this.Build(commandState, action, databaseName, containerName, "indeterminate", "none", null, MessageService.GetString("command-can-i-manage-note")); + return this.Build(commandState, action, databaseName, containerName, "indeterminate", "none", null, MessageService.GetString("command-can-i-manage-note"), null); } if (string.IsNullOrEmpty(databaseName) || string.IsNullOrEmpty(containerName)) @@ -86,24 +86,18 @@ public async override Task ExecuteAsync(ShellInterpreter shell, Co // Account-key and emulator connections use a master key, which grants full access. if (shell.ActiveCredential is null) { - return this.Build(commandState, action, databaseName, containerName, "allow", "key", null, MessageService.GetString("command-can-i-key-note")); + return this.Build(commandState, action, databaseName, containerName, "allow", "key", null, MessageService.GetString("command-can-i-key-note"), null); } var container = connectedState.Client.GetContainer(databaseName, containerName); - HttpStatusCode statusCode; - switch (action) + var probeResult = action switch { - case "read": - statusCode = await ProbeReadAsync(container, token); - break; - case "query": - statusCode = await ProbeQueryAsync(container, token); - break; - default: - statusCode = await ProbeWriteAsync(container, token); - break; - } + "read" => await ProbeReadAsync(container, token), + "query" => await ProbeQueryAsync(container, token), + _ => await ProbeWriteAsync(container, token), + }; + var statusCode = probeResult.StatusCode; var (decision, statusNote) = MapDecision(statusCode); string? note = statusNote; @@ -120,20 +114,20 @@ public async override Task ExecuteAsync(ShellInterpreter shell, Co note = MessageService.GetString("command-can-i-write-heuristic-note"); } - return this.Build(commandState, action, databaseName, containerName, decision, "probe", (int)statusCode, note); + return this.Build(commandState, action, databaseName, containerName, decision, "probe", (int)statusCode, note, probeResult.RequestCharge); } - private static async Task ProbeReadAsync(Container container, CancellationToken token) + private static async Task ProbeReadAsync(Container container, CancellationToken token) { using var response = await container.ReadItemStreamAsync( Guid.NewGuid().ToString(), new PartitionKey(Guid.NewGuid().ToString()), requestOptions: null, cancellationToken: token); - return response.StatusCode; + return CreateProbeResult(response); } - private static async Task ProbeQueryAsync(Container container, CancellationToken token) + private static async Task ProbeQueryAsync(Container container, CancellationToken token) { // A minimal TOP 1 query with a single-item page proves query authorization without // forcing a full scan (as an aggregate like COUNT would) on large containers. @@ -142,10 +136,10 @@ private static async Task ProbeQueryAsync(Container container, C new QueryDefinition("SELECT TOP 1 c.id FROM c"), requestOptions: requestOptions); using var response = await iterator.ReadNextAsync(token); - return response.StatusCode; + return CreateProbeResult(response); } - private static async Task ProbeWriteAsync(Container container, CancellationToken token) + private static async Task ProbeWriteAsync(Container container, CancellationToken token) { // Deleting a random, almost-certainly-nonexistent id is non-mutating: an authorized // caller gets 404 NotFound, an unauthorized caller gets 403 Forbidden. The bogus @@ -158,7 +152,12 @@ private static async Task ProbeWriteAsync(Container container, C new PartitionKey(Guid.NewGuid().ToString()), requestOptions: requestOptions, cancellationToken: token); - return response.StatusCode; + return CreateProbeResult(response); + } + + internal static ProbeResult CreateProbeResult(ResponseMessage response) + { + return new ProbeResult(response.StatusCode, response.Headers.RequestCharge); } private static (string Decision, string? Note) MapDecision(HttpStatusCode statusCode) @@ -185,7 +184,7 @@ private static string BuildScope(string databaseName, string? containerName) return string.IsNullOrEmpty(containerName) ? $"/{databaseName}" : $"/{databaseName}/{containerName}"; } - private CommandState Build( + internal CommandState Build( CommandState commandState, string action, string? databaseName, @@ -193,7 +192,8 @@ private CommandState Build( string decision, string method, int? statusCode, - string? note) + string? note, + double? requestCharge) { var result = new Dictionary { @@ -208,6 +208,7 @@ private CommandState Build( commandState.RenderUser = () => this.RenderTable(action, databaseName, containerName, decision, method, statusCode, note); commandState.Result = new ShellJson(JsonSerializer.SerializeToElement(result)); + commandState.RequestCharge = requestCharge; return commandState; } @@ -239,4 +240,6 @@ private void RenderTable(string action, string? databaseName, string? containerN AnsiConsole.MarkupLine(Theme.FormatMuted(note)); } } + + internal readonly record struct ProbeResult(HttpStatusCode StatusCode, double RequestCharge); } diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/QueryCommand.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/QueryCommand.cs index 1dba9129..e085eccf 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/QueryCommand.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/QueryCommand.cs @@ -637,7 +637,7 @@ private async Task ExecuteExplainAsync(Container container, ShellI } var cumulative = response?.Diagnostics.GetQueryMetrics()?.CumulativeMetrics; - double requestCharge = response?.Diagnostics.GetQueryMetrics()?.TotalRequestCharge ?? 0; + double requestCharge = response?.Headers.RequestCharge ?? 0; var (planAvailable, utilized, potential) = ParseIndexPlan(response?.IndexMetrics); var evaluation = EvaluatePlan( @@ -651,6 +651,7 @@ private async Task ExecuteExplainAsync(Container container, ShellI returnState.Result = BuildExplainJson(this.Query, evaluation, requestCharge, messages); returnState.RenderUser = () => RenderExplain(evaluation, requestCharge, messages); + returnState.RequestCharge = requestCharge; return returnState; } catch (OperationCanceledException) when (token.IsCancellationRequested) diff --git a/docs/mcp.md b/docs/mcp.md index 3221ed50..aa35a776 100644 --- a/docs/mcp.md +++ b/docs/mcp.md @@ -100,9 +100,11 @@ 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. | +| `requestCharge` | Instrumented data-plane command results | The Cosmos DB request charge (in RUs) consumed by the command, as a number. This is omitted for commands that do not issue a billable request. | | `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. 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. +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`, including `--explain`; `print`; `ls` for container items; `can-i` data-plane probes; `batch run`; `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. + +This field reports observed cost; it does not enforce an RU budget. Budget guardrails are tracked separately in [#162](https://github.com/Azure/CosmosDBShell/issues/162). From b6d39cd60db6688b999d283fa8968d20118206d4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mike=20Kr=C3=BCger?= Date: Mon, 31 Aug 2026 13:01:35 +0200 Subject: [PATCH 07/16] Track connection-scoped request charges --- CHANGELOG.md | 1 + .../CommandTests/InfoCommandTests.cs | 61 +++++++++++++++++++ .../InfoCommand.cs | 37 +++++++++-- .../ShellInterpreter.cs | 54 ++++++++++++++++ .../Statement/CommandStatement.cs | 9 ++- CosmosDBShell/lang/en.ftl | 2 + docs/commands.md | 8 ++- docs/mcp.md | 6 ++ 8 files changed, 169 insertions(+), 9 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 1804f85d..03067788 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -91,6 +91,7 @@ A focused cycle on top of 1.1.115-preview. New `ttl` and `conflict` commands man - **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`. A successful `connect` starts a new total; database and container navigation do not reset it. 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 diff --git a/CosmosDBShell.Tests/CommandTests/InfoCommandTests.cs b/CosmosDBShell.Tests/CommandTests/InfoCommandTests.cs index f13243d2..529d489a 100644 --- a/CosmosDBShell.Tests/CommandTests/InfoCommandTests.cs +++ b/CosmosDBShell.Tests/CommandTests/InfoCommandTests.cs @@ -4,6 +4,7 @@ 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.States; @@ -135,6 +136,66 @@ 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 ErrorCommandState(new InvalidOperationException()) { RequestCharge = 2.5 }); + + Assert.Equal(3.75, shell.SessionRequestCharge); + } + + [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); + } + + [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); + } + + [Fact] + public void AddSessionUsage_AddsCurrentChargeToStructuredResult() + { + using var shell = ShellInterpreter.CreateInstance(); + shell.RecordRequestCharge(new CommandState { RequestCharge = 3.75 }); + var result = new Dictionary(); + + InfoCommand.AddSessionUsage(shell, result, renderOutput: false); + + var json = JsonSerializer.SerializeToElement(result); + Assert.Equal(3.75, json.GetProperty("session").GetProperty("requestCharge").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); + } + [Theory] [InlineData("table")] [InlineData("tbl")] diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/InfoCommand.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/InfoCommand.cs index 8a8e5599..c065679f 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/InfoCommand.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/InfoCommand.cs @@ -82,17 +82,17 @@ public async override Task ExecuteAsync(ShellInterpreter shell, Co // If both database and container are resolved, show container settings if (!string.IsNullOrEmpty(databaseName) && !string.IsNullOrEmpty(containerName)) { - return await this.ShowContainerSettingsAsync(connectedState, databaseName, containerName, commandState, renderOutput, token); + return await this.ShowContainerSettingsAsync(shell, connectedState, databaseName, containerName, commandState, renderOutput, token); } // If only a database is resolved, show database settings if (!string.IsNullOrEmpty(databaseName)) { - return await this.ShowDatabaseSettingsAsync(connectedState, databaseName, commandState, renderOutput, token); + return await this.ShowDatabaseSettingsAsync(shell, connectedState, databaseName, commandState, renderOutput, token); } // Otherwise show account overview - return await this.PrintOverviewAsync(connectedState, commandState, renderOutput, token); + return await this.PrintOverviewAsync(shell, connectedState, commandState, renderOutput, token); } catch (Exception e) when (e is not OperationCanceledException) { @@ -301,6 +301,28 @@ private static async Task WriteAccountDatabaseBreakdownAsync(ConnectedState stat AnsiConsole.Write(databaseTable); } + internal static void AddSessionUsage(ShellInterpreter shell, Dictionary mcpTable, bool renderOutput) + { + mcpTable["session"] = new Dictionary + { + ["requestCharge"] = shell.SessionRequestCharge, + }; + + if (!renderOutput) + { + return; + } + + AnsiConsole.MarkupLine(Theme.FormatSectionHeader(MessageService.GetString("command-stats-session-heading"))); + var table = new Table(); + table.AddColumns(string.Empty, string.Empty); + table.HideHeaders(); + table.AddRow( + MessageService.GetString("command-stats-session-request-charge"), + Theme.FormatTableValue(shell.SessionRequestCharge.ToString("0.##", CultureInfo.InvariantCulture))); + AnsiConsole.Write(table); + } + private static async Task ReadContainerUsageAsync(Container container, CancellationToken token) { var response = await container.ReadContainerAsync(new ContainerRequestOptions { PopulateQuotaInfo = true }, token); @@ -655,7 +677,7 @@ internal static string FormatSize(long? kilobytes) return string.Create(CultureInfo.InvariantCulture, $"{value:0.##} {units[unit]}"); } - private async Task ShowContainerSettingsAsync(ConnectedState state, string databaseName, string containerName, CommandState commandState, bool renderOutput, CancellationToken token) + private async Task ShowContainerSettingsAsync(ShellInterpreter shell, ConnectedState state, string databaseName, string containerName, CommandState commandState, bool renderOutput, CancellationToken token) { await ValidateContainerExistsAsync(state, databaseName, containerName, "info", token); var view = await CosmosResourceFacade.GetContainerSettingsAsync(state, databaseName, containerName, token); @@ -939,12 +961,13 @@ private async Task ShowContainerSettingsAsync(ConnectedState state mcpTable["topPartitionKeys"] = await WriteTopPartitionKeysAsync(container, view.PartitionKeyPaths, renderOutput, token); } + AddSessionUsage(shell, mcpTable, renderOutput); commandState.Result = new ShellJson(JsonSerializer.SerializeToElement(mcpTable)); commandState.RenderUser = renderOutput ? () => { } : null; return commandState; } - private async Task ShowDatabaseSettingsAsync(ConnectedState state, string databaseName, CommandState commandState, bool renderOutput, CancellationToken token) + private async Task ShowDatabaseSettingsAsync(ShellInterpreter shell, ConnectedState state, string databaseName, CommandState commandState, bool renderOutput, CancellationToken token) { await ValidateDatabaseExistsAsync(state, databaseName, "info", token); @@ -1027,12 +1050,13 @@ private async Task ShowDatabaseSettingsAsync(ConnectedState state, mcpTable["containers"] = perContainer; } + AddSessionUsage(shell, mcpTable, renderOutput); commandState.Result = new ShellJson(JsonSerializer.SerializeToElement(mcpTable)); commandState.RenderUser = renderOutput ? () => { } : null; return commandState; } - private async Task PrintOverviewAsync(ConnectedState state, CommandState commandState, bool renderOutput, CancellationToken token) + private async Task PrintOverviewAsync(ShellInterpreter shell, ConnectedState state, CommandState commandState, bool renderOutput, CancellationToken token) { var client = state.Client; var acc = await client.ReadAccountAsync(); @@ -1070,6 +1094,7 @@ private async Task PrintOverviewAsync(ConnectedState state, Comman await WriteAccountDatabaseBreakdownAsync(state, databaseNames, mcpTable, renderOutput, token); } + AddSessionUsage(shell, mcpTable, renderOutput); commandState.Result = new ShellJson(JsonSerializer.SerializeToElement(mcpTable)); commandState.RenderUser = renderOutput ? () => { } : null; return commandState; diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/ShellInterpreter.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/ShellInterpreter.cs index c4f38f33..3538bb13 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/ShellInterpreter.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/ShellInterpreter.cs @@ -48,6 +48,8 @@ public partial class ShellInterpreter : IDisposable private readonly HashSet diagnosticSecrets = new(StringComparer.Ordinal); + private readonly object sessionRequestChargeLock = new(); + private TokenCredential? activeCredential; private LineEditor? lineEditor; @@ -64,6 +66,10 @@ public partial class ShellInterpreter : IDisposable private List history; + private double sessionRequestCharge; + + private long sessionRequestChargeGeneration; + internal ShellInterpreter(string? configPath = null) { this.State = new DisconnectedState(); @@ -150,6 +156,31 @@ internal static char CSVSeparator /// internal string? ActiveCredentialType { get; private set; } + /// + /// Gets the request charge observed from instrumented commands since the most recent connection. + /// + internal double SessionRequestCharge + { + get + { + lock (this.sessionRequestChargeLock) + { + return this.sessionRequestCharge; + } + } + } + + internal long SessionRequestChargeGeneration + { + get + { + lock (this.sessionRequestChargeLock) + { + return this.sessionRequestChargeGeneration; + } + } + } + internal string HistoryFile { get; private set; } internal string WelcomeMarkerFile => this.welcomeMarkerFile; @@ -888,6 +919,23 @@ internal async Task RunCommandAsync(CommandState currentState, str return currentState; } + internal void RecordRequestCharge(CommandState commandState) + => this.RecordRequestCharge(commandState, this.SessionRequestChargeGeneration); + + internal void RecordRequestCharge(CommandState commandState, long generation) + { + if (commandState.RequestCharge is { } requestCharge) + { + lock (this.sessionRequestChargeLock) + { + if (generation == this.sessionRequestChargeGeneration) + { + this.sessionRequestCharge += requestCharge; + } + } + } + } + internal async Task ConnectAsync(string connectionString, string? loginHint = null, ConnectionMode? mode = null, string? tenantId = null, string? authorityHost = null, string? managedIdentityClientId = null, CredentialMethod credentialMethod = CredentialMethod.Default, string? subscriptionId = null, string? resourceGroupName = null, CancellationToken token = default) { token.ThrowIfCancellationRequested(); @@ -1492,6 +1540,12 @@ internal void Connect(CosmosClient client, ArmCosmosContext? armContext = null, { this.State?.Dispose(); this.State = new ConnectedState(client, armContext); + lock (this.sessionRequestChargeLock) + { + this.sessionRequestCharge = 0; + this.sessionRequestChargeGeneration++; + } + this.activeCredential = credential; this.ActiveCredentialType = credentialTypeOverride ?? credential?.GetType().Name; this.CurrentBatch = null; diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Parser/Statement/CommandStatement.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Parser/Statement/CommandStatement.cs index 43364df4..ddf97a7d 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Parser/Statement/CommandStatement.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Parser/Statement/CommandStatement.cs @@ -185,8 +185,15 @@ public override async Task RunAsync(ShellInterpreter shell, Comman return HelpCommand.PrintCommandHelp(this.Name, shell.App, false); } + // CommandState is intentionally reused by pipelines. Clear the previous command's + // charge before dispatch so it cannot be counted again when the next command does + // not issue an instrumented Cosmos DB request. + commandState.RequestCharge = null; + var requestChargeGeneration = shell.SessionRequestChargeGeneration; var cmd = await this.CreateCommandAsync(factory, shell, commandState, token); - return await cmd.ExecuteAsync(shell, commandState, string.Empty, token); + var result = await cmd.ExecuteAsync(shell, commandState, string.Empty, token); + shell.RecordRequestCharge(result, requestChargeGeneration); + return result; } if (File.Exists(this.Name)) diff --git a/CosmosDBShell/lang/en.ftl b/CosmosDBShell/lang/en.ftl index 726c82f0..aa25c85a 100644 --- a/CosmosDBShell/lang/en.ftl +++ b/CosmosDBShell/lang/en.ftl @@ -931,6 +931,8 @@ command-stats-account-databases-col-containers = Containers command-stats-account-databases-col-count = Documents command-stats-account-databases-col-size = Size command-stats-account-detailed-cost-note = Aggregating account totals reads every container's quota and consumes request units. +command-stats-session-heading = Session Usage +command-stats-session-request-charge = Observed request charge (RU) command-version-description = Displays the version of Cosmos DB Shell. command-version = Cosmos Shell version: { $version } diff --git a/docs/commands.md b/docs/commands.md index bb9e9090..5988532f 100644 --- a/docs/commands.md +++ b/docs/commands.md @@ -1330,7 +1330,10 @@ data/total storage size. Use `index show` for the full indexing policy JSON. Whe only a database is in scope it reports the container count, aggregate document count, total storage, and shared throughput. When neither is in scope (the account root) it reports account metadata: read/write regions and the database -count. +count. Every scope also includes the cumulative request charge observed from +instrumented commands during the current connection. The total resets to zero +after each successful `connect`; changing database or container scope does not +reset it. It is session telemetry, not a budget or billing total. On serverless accounts, throughput/offer settings are not available, so the scale section reports that throughput settings are not available for serverless @@ -1354,7 +1357,8 @@ with redirection, the report is written to the file as a plain-text grid instead the rich console layout. The `--partitions` and `--detailed` options issue queries against the data and therefore consume request units; at the account root, `--detailed` aggregates every container's -storage and document count across all databases. This command is read-only. +storage and document count across all databases. In JSON output the connection +total is available as `session.requestCharge`. This command is read-only. ### help diff --git a/docs/mcp.md b/docs/mcp.md index aa35a776..0cbff146 100644 --- a/docs/mcp.md +++ b/docs/mcp.md @@ -108,3 +108,9 @@ Successful results set `result` (and optionally `outputText`); failed results se This field reports observed cost; it does not enforce an RU budget. Budget guardrails are tracked separately in [#162](https://github.com/Azure/CosmosDBShell/issues/162). +The `info` command result also includes `session.requestCharge`, the cumulative +charge observed from instrumented commands during the current connection. A +successful `connect` starts a new total; navigation between databases and +containers does not reset it. This session value is telemetry rather than a +budget or billing total. + From 1af804157daec9dae52ca92875671f78cd9b19f4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mike=20Kr=C3=BCger?= Date: Mon, 31 Aug 2026 13:08:42 +0200 Subject: [PATCH 08/16] Address request charge review feedback --- .../Integration/BatchOperationTests.cs | 4 ++-- .../Integration/QueryCommandTests.cs | 2 +- .../ImportCommand.cs | 2 +- .../PrintCommand.cs | 19 ++++++++++++++----- .../RmCommand.cs | 2 +- 5 files changed, 19 insertions(+), 10 deletions(-) diff --git a/CosmosDBShell.Tests/Integration/BatchOperationTests.cs b/CosmosDBShell.Tests/Integration/BatchOperationTests.cs index 7d14e2bc..833d8962 100644 --- a/CosmosDBShell.Tests/Integration/BatchOperationTests.cs +++ b/CosmosDBShell.Tests/Integration/BatchOperationTests.cs @@ -48,7 +48,7 @@ public async Task BatchRun_MultipleCreates_CommitsAtomically() } Assert.False(state.IsError, FormatError(state)); - Assert.True(state.RequestCharge > 0); + Assert.True(state.RequestCharge is > 0); var root = JsonDocument.Parse(output).RootElement; Assert.True(root.GetProperty("success").GetBoolean()); @@ -96,7 +96,7 @@ public async Task BatchRun_FailingOperation_RollsBackEntireBatch() var batchState = await ExecuteAsync($"batch run '{json}' --partition-key {pk}"); Assert.True(batchState.IsError); Assert.Equal(ShellExitCode.GeneralFailure, batchState.ExitCode); - Assert.True(batchState.RequestCharge > 0); + Assert.True(batchState.RequestCharge is > 0); var output = await File.ReadAllTextAsync(outputFile, TestContext.Current.CancellationToken); var root = JsonDocument.Parse(output).RootElement; diff --git a/CosmosDBShell.Tests/Integration/QueryCommandTests.cs b/CosmosDBShell.Tests/Integration/QueryCommandTests.cs index b9004423..2ac8ef18 100644 --- a/CosmosDBShell.Tests/Integration/QueryCommandTests.cs +++ b/CosmosDBShell.Tests/Integration/QueryCommandTests.cs @@ -150,7 +150,7 @@ public async Task Query_Explain_ReturnsStructuredResultWithoutDocuments() } Assert.False(state.IsError, FormatError(state)); - Assert.True(state.RequestCharge > 0); + Assert.True(state.RequestCharge is > 0); using var document = JsonDocument.Parse(output); var root = document.RootElement; diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ImportCommand.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ImportCommand.cs index eb4ab44c..7178b9ea 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ImportCommand.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ImportCommand.cs @@ -548,7 +548,7 @@ public override async Task ExecuteAsync(ShellInterpreter shell, Co requestCharge = charge, dryRun, })), - RequestCharge = charge, + RequestCharge = charge > 0 ? charge : null, }; } diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/PrintCommand.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/PrintCommand.cs index 9ed343d8..98f89e83 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/PrintCommand.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/PrintCommand.cs @@ -67,24 +67,33 @@ private async Task PrintItemAsync(Container container, Cancellatio } else if (response.StatusCode == System.Net.HttpStatusCode.NotFound) { - throw new CommandException("print", MessageService.GetString("command-print-error-item_not_found", new Dictionary + return new ErrorCommandState(new CommandException("print", MessageService.GetString("command-print-error-item_not_found", new Dictionary { { "id", this.Id ?? "(null)" }, { "partitionKey", this.PartitionKey ?? "(null)" }, - })); + }))) + { + RequestCharge = response.Headers.RequestCharge, + }; } else { - throw new CommandException("print", MessageService.GetString("command-print-error-request_failed", new Dictionary + return new ErrorCommandState(new CommandException("print", MessageService.GetString("command-print-error-request_failed", new Dictionary { { "id", this.Id ?? "(null)" }, { "status", (int)response.StatusCode }, - })); + }))) + { + RequestCharge = response.Headers.RequestCharge, + }; } } catch (CosmosException ex) { - throw new CommandException("print", MessageService.GetArgsString("command-print-error-reading_item", "message", CommandException.GetDisplayMessage(ex)), ex); + return new ErrorCommandState(new CommandException("print", MessageService.GetArgsString("command-print-error-reading_item", "message", CommandException.GetDisplayMessage(ex)), ex)) + { + RequestCharge = ex.RequestCharge > 0 ? ex.RequestCharge : null, + }; } return commandState; diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/RmCommand.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/RmCommand.cs index 87278b99..6cdf9656 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/RmCommand.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/RmCommand.cs @@ -309,7 +309,7 @@ private async Task RemoveItemsFromContainerAsync(ConnectedState state, commandState.Result = new ShellJson(JsonSerializer.SerializeToElement(new { type = "item", count = totalCount, dryRun })); commandState.RenderUser = () => AnsiConsole.MarkupLine(renderMessage); - commandState.RequestCharge = totalCharge; + commandState.RequestCharge = totalCharge > 0 ? totalCharge : null; return new ExitCode(0); } From c04214f7705a1615755038641c906b98cc8c063b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mike=20Kr=C3=BCger?= Date: Mon, 31 Aug 2026 13:24:47 +0200 Subject: [PATCH 09/16] Centralize session request charge accounting --- .../CommandTests/CosmosCommandTests.cs | 21 ++++++++++++++ .../QueryCommand.cs | 28 +++++++++++++++---- .../ShellInterpreter.cs | 15 ++++++++++ .../ToolOperations.cs | 2 +- .../Expression/CommandExpression.cs | 2 +- .../Statement/CommandStatement.cs | 9 +----- CosmosDBShell/lang/en.ftl | 2 +- 7 files changed, 63 insertions(+), 16 deletions(-) diff --git a/CosmosDBShell.Tests/CommandTests/CosmosCommandTests.cs b/CosmosDBShell.Tests/CommandTests/CosmosCommandTests.cs index db152da8..11f6f4f4 100644 --- a/CosmosDBShell.Tests/CommandTests/CosmosCommandTests.cs +++ b/CosmosDBShell.Tests/CommandTests/CosmosCommandTests.cs @@ -11,6 +11,19 @@ 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); + } + [Fact] public void CreatePartitionKey_WithHierarchicalIntegerComponents_PreservesIntegerTypes() { @@ -42,8 +55,16 @@ public void CreatePartitionKey_WithHierarchicalIntegerComponents_PreservesIntege private sealed class TestCosmosCommand : CosmosCommand { + private readonly double? requestCharge; + + public TestCosmosCommand(double? requestCharge = null) + { + this.requestCharge = requestCharge; + } + public override Task ExecuteAsync(ShellInterpreter shell, CommandState commandState, string commandText, CancellationToken token) { + commandState.RequestCharge = this.requestCharge; return Task.FromResult(commandState); } diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/QueryCommand.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/QueryCommand.cs index e085eccf..d7633874 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/QueryCommand.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/QueryCommand.cs @@ -631,13 +631,20 @@ private async Task ExecuteExplainAsync(Container container, ShellI using ResponseMessage? response = feedIterator.HasMoreResults ? await feedIterator.ReadNextAsync(token) : null; - if (response is not null) + double requestCharge = response?.Headers.RequestCharge ?? 0; + if (response is not null && !response.IsSuccessStatusCode) { - await this.ThrowIfRequestFailedAsync(response, shell); + try + { + await this.ThrowIfRequestFailedAsync(response, shell); + } + catch (Exception ex) when (ex is not OperationCanceledException) + { + return new ErrorCommandState(ex) { RequestCharge = requestCharge > 0 ? requestCharge : null }; + } } var cumulative = response?.Diagnostics.GetQueryMetrics()?.CumulativeMetrics; - double requestCharge = response?.Headers.RequestCharge ?? 0; var (planAvailable, utilized, potential) = ParseIndexPlan(response?.IndexMetrics); var evaluation = EvaluatePlan( @@ -713,7 +720,19 @@ private async Task ExecuteQueryAsync(Container container, ShellInt using var response = await feedIterator.ReadNextAsync(token); - await this.ThrowIfRequestFailedAsync(response, shell); + var pageRequestCharge = response.Headers.RequestCharge; + if (!response.IsSuccessStatusCode) + { + try + { + await this.ThrowIfRequestFailedAsync(response, shell); + } + catch (Exception ex) when (ex is not OperationCanceledException) + { + var failedCharge = totalRequestCharge + pageRequestCharge; + return new ErrorCommandState(ex) { RequestCharge = failedCharge > 0 ? failedCharge : null }; + } + } if (response.Content == null) { @@ -734,7 +753,6 @@ private async Task ExecuteQueryAsync(Container container, ShellInt // 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("F2", CultureInfo.InvariantCulture) } })); diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/ShellInterpreter.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/ShellInterpreter.cs index 3538bb13..14fc061d 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/ShellInterpreter.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/ShellInterpreter.cs @@ -922,6 +922,21 @@ internal async Task RunCommandAsync(CommandState currentState, str internal void RecordRequestCharge(CommandState commandState) => this.RecordRequestCharge(commandState, this.SessionRequestChargeGeneration); + internal async Task ExecuteCosmosCommandAsync( + CosmosCommand command, + CommandState commandState, + string commandText, + CancellationToken token) + { + // CommandState is intentionally reused by pipelines and expressions. Clear the + // previous command's charge so an uninstrumented command cannot count it again. + commandState.RequestCharge = null; + var generation = this.SessionRequestChargeGeneration; + var result = await command.ExecuteAsync(this, commandState, commandText, token); + this.RecordRequestCharge(result, generation); + return result; + } + internal void RecordRequestCharge(CommandState commandState, long generation) { if (commandState.RequestCharge is { } requestCharge) diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Mcp/ToolOperations.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Mcp/ToolOperations.cs index 6f39542f..92cad62d 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Mcp/ToolOperations.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Mcp/ToolOperations.cs @@ -491,7 +491,7 @@ private async ValueTask OnCallToolsAsync( try { ShellInterpreter.Instance.PrintCommand(sb.ToString()); - var response = await cmd.ExecuteAsync(ShellInterpreter.Instance, new CommandState(), command.CommandName, cancellationToken); + var response = await ShellInterpreter.Instance.ExecuteCosmosCommandAsync(cmd, new CommandState(), command.CommandName, cancellationToken); ShellInterpreter.Instance.CancelPrompt(); return McpResponseFactory.CreateSuccess(response, ShellInterpreter.Instance.State); } diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Parser/Expression/CommandExpression.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Parser/Expression/CommandExpression.cs index 50d118f0..1d9c3e73 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Parser/Expression/CommandExpression.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Parser/Expression/CommandExpression.cs @@ -137,7 +137,7 @@ internal async Task ExecuteCommandAsync(ShellInterpreter shell, Co if (shell.App.Commands.TryGetValue(this.Name, out var factory)) { var cmd = await this.CreateCommandAsync(factory, shell, commandState, token); - return await cmd.ExecuteAsync(shell, commandState, string.Empty, token); + return await shell.ExecuteCosmosCommandAsync(cmd, commandState, string.Empty, token); } // Check for script files diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Parser/Statement/CommandStatement.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Parser/Statement/CommandStatement.cs index ddf97a7d..fae24f8a 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Parser/Statement/CommandStatement.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Parser/Statement/CommandStatement.cs @@ -185,15 +185,8 @@ public override async Task RunAsync(ShellInterpreter shell, Comman return HelpCommand.PrintCommandHelp(this.Name, shell.App, false); } - // CommandState is intentionally reused by pipelines. Clear the previous command's - // charge before dispatch so it cannot be counted again when the next command does - // not issue an instrumented Cosmos DB request. - commandState.RequestCharge = null; - var requestChargeGeneration = shell.SessionRequestChargeGeneration; var cmd = await this.CreateCommandAsync(factory, shell, commandState, token); - var result = await cmd.ExecuteAsync(shell, commandState, string.Empty, token); - shell.RecordRequestCharge(result, requestChargeGeneration); - return result; + return await shell.ExecuteCosmosCommandAsync(cmd, commandState, string.Empty, token); } if (File.Exists(this.Name)) diff --git a/CosmosDBShell/lang/en.ftl b/CosmosDBShell/lang/en.ftl index aa25c85a..b747cdb5 100644 --- a/CosmosDBShell/lang/en.ftl +++ b/CosmosDBShell/lang/en.ftl @@ -932,7 +932,7 @@ command-stats-account-databases-col-count = Documents command-stats-account-databases-col-size = Size command-stats-account-detailed-cost-note = Aggregating account totals reads every container's quota and consumes request units. command-stats-session-heading = Session Usage -command-stats-session-request-charge = Observed request charge (RU) +command-stats-session-request-charge = Observed request charge (RUs) command-version-description = Displays the version of Cosmos DB Shell. command-version = Cosmos Shell version: { $version } From e2fbd175623e79bf8282283f27112a59f92cfabb Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mike=20Kr=C3=BCger?= Date: Mon, 31 Aug 2026 13:47:45 +0200 Subject: [PATCH 10/16] Preserve request charge for failed listings --- .../ListCommand.cs | 37 ++++++++++++------- 1 file changed, 24 insertions(+), 13 deletions(-) diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ListCommand.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ListCommand.cs index 81b3ca7b..975c24e0 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ListCommand.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ListCommand.cs @@ -228,25 +228,36 @@ private async Task ListContainerItemsAsync(ConnectedState state, S while (feedIterator.HasMoreResults) { using var response = await feedIterator.ReadNextAsync(token); - using var queryDocument = await ReadQueryResponseAsync(response, token); AccumulateRequestCharge(returnState, response.Headers.RequestCharge); + JsonDocument queryDocument; + try + { + queryDocument = await ReadQueryResponseAsync(response, token); + } + catch (Exception ex) when (ex is not OperationCanceledException) + { + return new ErrorCommandState(ex) { RequestCharge = returnState.RequestCharge }; + } - foreach (var element in queryDocument.RootElement.GetProperty("Documents").EnumerateArray()) + using (queryDocument) { - // Check if pattern matches - bool shouldList = this.matcher == null; + foreach (var element in queryDocument.RootElement.GetProperty("Documents").EnumerateArray()) + { + // Check if pattern matches + bool shouldList = this.matcher == null; - shouldList = shouldList || MatchesAnyPath(element, matchKeyPropertyNames, this.matcher!); + shouldList = shouldList || MatchesAnyPath(element, matchKeyPropertyNames, this.matcher!); - if (shouldList) - { - list.Add(element.Clone()); - } + if (shouldList) + { + list.Add(element.Clone()); + } - if (ResultLimit.IsLimitReached(list.Count, effectiveMaxItemCount)) - { - limitReached = ShouldReportLimitReached(list.Count, effectiveMaxItemCount, usesServerSideTop, feedIterator.HasMoreResults); - break; + if (ResultLimit.IsLimitReached(list.Count, effectiveMaxItemCount)) + { + limitReached = ShouldReportLimitReached(list.Count, effectiveMaxItemCount, usesServerSideTop, feedIterator.HasMoreResults); + break; + } } } From d69c70f3b5a64c1ac9a0c8b7783b8651695d5f6a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mike=20Kr=C3=BCger?= Date: Mon, 31 Aug 2026 13:51:21 +0200 Subject: [PATCH 11/16] Count charged session operations --- CHANGELOG.md | 2 +- .../CommandTests/CosmosCommandTests.cs | 1 + .../CommandTests/InfoCommandTests.cs | 6 ++++ .../InfoCommand.cs | 9 +++-- .../ShellInterpreter.cs | 33 +++++++++++++++++++ CosmosDBShell/lang/en.ftl | 1 + docs/commands.md | 4 ++- docs/mcp.md | 4 ++- 8 files changed, 55 insertions(+), 5 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 03067788..92e39eb3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -91,7 +91,7 @@ A focused cycle on top of 1.1.115-preview. New `ttl` and `conflict` commands man - **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`. A successful `connect` starts a new total; database and container navigation do not reset it. This is usage telemetry, not budget enforcement or billing data. ([#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 diff --git a/CosmosDBShell.Tests/CommandTests/CosmosCommandTests.cs b/CosmosDBShell.Tests/CommandTests/CosmosCommandTests.cs index 11f6f4f4..8f651325 100644 --- a/CosmosDBShell.Tests/CommandTests/CosmosCommandTests.cs +++ b/CosmosDBShell.Tests/CommandTests/CosmosCommandTests.cs @@ -22,6 +22,7 @@ public async Task ExecuteCosmosCommandAsync_RecordsChargeAndClearsStaleCharge() Assert.Null(state.RequestCharge); Assert.Equal(2.5, shell.SessionRequestCharge); + Assert.Equal(1, shell.SessionChargedOperationCount); } [Fact] diff --git a/CosmosDBShell.Tests/CommandTests/InfoCommandTests.cs b/CosmosDBShell.Tests/CommandTests/InfoCommandTests.cs index 529d489a..2d0a64b1 100644 --- a/CosmosDBShell.Tests/CommandTests/InfoCommandTests.cs +++ b/CosmosDBShell.Tests/CommandTests/InfoCommandTests.cs @@ -143,9 +143,11 @@ public void SessionRequestCharge_AccumulatesObservedCharges() 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] @@ -157,6 +159,7 @@ public void Connect_ResetsSessionRequestCharge() shell.Connect(CreateTestClient(), credentialTypeOverride: "AccountKey"); Assert.Equal(0, shell.SessionRequestCharge); + Assert.Equal(0, shell.SessionChargedOperationCount); } [Fact] @@ -169,6 +172,7 @@ public void Disconnect_DoesNotResetSessionRequestCharge() shell.Disconnect(); Assert.Equal(4.5, shell.SessionRequestCharge); + Assert.Equal(1, shell.SessionChargedOperationCount); } [Fact] @@ -182,6 +186,7 @@ public void AddSessionUsage_AddsCurrentChargeToStructuredResult() 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] @@ -194,6 +199,7 @@ public void RecordRequestCharge_IgnoresPriorConnectionGeneration() shell.RecordRequestCharge(new CommandState { RequestCharge = 8 }, generation); Assert.Equal(0, shell.SessionRequestCharge); + Assert.Equal(0, shell.SessionChargedOperationCount); } [Theory] diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/InfoCommand.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/InfoCommand.cs index c065679f..f441fada 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/InfoCommand.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/InfoCommand.cs @@ -303,9 +303,11 @@ private static async Task WriteAccountDatabaseBreakdownAsync(ConnectedState stat internal static void AddSessionUsage(ShellInterpreter shell, Dictionary mcpTable, bool renderOutput) { + var sessionUsage = shell.SessionUsage; mcpTable["session"] = new Dictionary { - ["requestCharge"] = shell.SessionRequestCharge, + ["requestCharge"] = sessionUsage.RequestCharge, + ["chargedOperationCount"] = sessionUsage.ChargedOperationCount, }; if (!renderOutput) @@ -319,7 +321,10 @@ internal static void AddSessionUsage(ShellInterpreter shell, Dictionary + /// Gets the number of command operations that reported a positive request charge + /// since the most recent connection. + /// + internal long SessionChargedOperationCount + { + get + { + lock (this.sessionRequestChargeLock) + { + return this.sessionChargedOperationCount; + } + } + } + + internal (double RequestCharge, long ChargedOperationCount) SessionUsage + { + get + { + lock (this.sessionRequestChargeLock) + { + return (this.sessionRequestCharge, this.sessionChargedOperationCount); + } + } + } + internal long SessionRequestChargeGeneration { get @@ -946,6 +974,10 @@ internal void RecordRequestCharge(CommandState commandState, long generation) if (generation == this.sessionRequestChargeGeneration) { this.sessionRequestCharge += requestCharge; + if (requestCharge > 0) + { + this.sessionChargedOperationCount++; + } } } } @@ -1558,6 +1590,7 @@ internal void Connect(CosmosClient client, ArmCosmosContext? armContext = null, lock (this.sessionRequestChargeLock) { this.sessionRequestCharge = 0; + this.sessionChargedOperationCount = 0; this.sessionRequestChargeGeneration++; } diff --git a/CosmosDBShell/lang/en.ftl b/CosmosDBShell/lang/en.ftl index b747cdb5..a62b07f0 100644 --- a/CosmosDBShell/lang/en.ftl +++ b/CosmosDBShell/lang/en.ftl @@ -933,6 +933,7 @@ command-stats-account-databases-col-size = Size command-stats-account-detailed-cost-note = Aggregating account totals reads every container's quota and consumes request units. command-stats-session-heading = Session Usage command-stats-session-request-charge = Observed request charge (RUs) +command-stats-session-charged-operations = Charged operations command-version-description = Displays the version of Cosmos DB Shell. command-version = Cosmos Shell version: { $version } diff --git a/docs/commands.md b/docs/commands.md index 5988532f..b468274f 100644 --- a/docs/commands.md +++ b/docs/commands.md @@ -1358,7 +1358,9 @@ the rich console layout. The `--partitions` and `--detailed` options issue queries against the data and therefore consume request units; at the account root, `--detailed` aggregates every container's storage and document count across all databases. In JSON output the connection -total is available as `session.requestCharge`. This command is read-only. +total is available as `session.requestCharge`, and `session.chargedOperationCount` +counts instrumented command operations that reported a positive charge. This +command is read-only. ### help diff --git a/docs/mcp.md b/docs/mcp.md index 0cbff146..7f8e69ac 100644 --- a/docs/mcp.md +++ b/docs/mcp.md @@ -112,5 +112,7 @@ The `info` command result also includes `session.requestCharge`, the cumulative charge observed from instrumented commands during the current connection. A successful `connect` starts a new total; navigation between databases and containers does not reset it. This session value is telemetry rather than a -budget or billing total. +budget or billing total. `session.chargedOperationCount` counts instrumented +command operations that reported a positive request charge; it counts command +operations rather than individual query pages or transactional batch items. From 9ff84a0443533d86aee856152e7a1404f150dd5a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mike=20Kr=C3=BCger?= Date: Mon, 31 Aug 2026 14:19:40 +0200 Subject: [PATCH 12/16] Ensure all data-plane operations report request charge --- CHANGELOG.md | 4 + .../CommandTests/CosmosCommandTests.cs | 70 +++++++++++++- .../CommandTests/InfoCommandTests.cs | 16 ++++ .../InfoCommand.cs | 24 +++-- .../MakeItemCommand.cs | 1 + .../ReplaceCommand.cs | 1 + .../SprocCommand.cs | 18 +++- .../TriggerCommand.cs | 18 +++- .../UdfCommand.cs | 18 +++- .../WatchCommand.cs | 1 + .../DataPlaneCosmosResourceOperations.cs | 40 +++++++- .../RequestChargeContext.cs | 93 +++++++++++++++++++ .../ShellInterpreter.cs | 31 ++++++- .../McpResponseFactory.cs | 19 ++-- .../ToolOperations.cs | 5 +- docs/commands.md | 7 +- docs/mcp.md | 9 +- 17 files changed, 339 insertions(+), 36 deletions(-) create mode 100644 CosmosDBShell/Azure.Data.Cosmos.Shell.Core/RequestChargeContext.cs diff --git a/CHANGELOG.md b/CHANGELOG.md index 92e39eb3..0a7619d1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,10 @@ ## 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. + ### 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. diff --git a/CosmosDBShell.Tests/CommandTests/CosmosCommandTests.cs b/CosmosDBShell.Tests/CommandTests/CosmosCommandTests.cs index 8f651325..23948011 100644 --- a/CosmosDBShell.Tests/CommandTests/CosmosCommandTests.cs +++ b/CosmosDBShell.Tests/CommandTests/CosmosCommandTests.cs @@ -25,6 +25,53 @@ public async Task ExecuteCosmosCommandAsync_RecordsChargeAndClearsStaleCharge() 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(() => 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(() => 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() { @@ -57,14 +104,35 @@ 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) + 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 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); } diff --git a/CosmosDBShell.Tests/CommandTests/InfoCommandTests.cs b/CosmosDBShell.Tests/CommandTests/InfoCommandTests.cs index 2d0a64b1..7981673d 100644 --- a/CosmosDBShell.Tests/CommandTests/InfoCommandTests.cs +++ b/CosmosDBShell.Tests/CommandTests/InfoCommandTests.cs @@ -189,6 +189,22 @@ public void AddSessionUsage_AddsCurrentChargeToStructuredResult() 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(); + + 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 RecordRequestCharge_IgnoresPriorConnectionGeneration() { diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/InfoCommand.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/InfoCommand.cs index f441fada..32378e95 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/InfoCommand.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/InfoCommand.cs @@ -304,10 +304,13 @@ private static async Task WriteAccountDatabaseBreakdownAsync(ConnectedState stat internal static void AddSessionUsage(ShellInterpreter shell, Dictionary mcpTable, bool renderOutput) { var sessionUsage = shell.SessionUsage; + var currentRequestCharge = RequestChargeContext.CurrentRequestCharge; + var requestCharge = sessionUsage.RequestCharge + currentRequestCharge; + var chargedOperationCount = sessionUsage.ChargedOperationCount + (currentRequestCharge > 0 ? 1 : 0); mcpTable["session"] = new Dictionary { - ["requestCharge"] = sessionUsage.RequestCharge, - ["chargedOperationCount"] = sessionUsage.ChargedOperationCount, + ["requestCharge"] = requestCharge, + ["chargedOperationCount"] = chargedOperationCount, }; if (!renderOutput) @@ -321,16 +324,17 @@ internal static void AddSessionUsage(ShellInterpreter shell, Dictionary ReadContainerUsageAsync(Container container, CancellationToken token) { var response = await container.ReadContainerAsync(new ContainerRequestOptions { PopulateQuotaInfo = true }, token); + RequestChargeContext.Record(response.RequestCharge); return ParseResourceUsage(response.Headers[ResourceUsageHeader]); } @@ -343,15 +347,19 @@ private static async Task WriteDatabaseThroughputAsync(Database database, Dictio try { var throughput = await database.ReadThroughputAsync(new RequestOptions(), token); + RequestChargeContext.Record(throughput.RequestCharge); min = throughput.MinThroughput; max = throughput.Resource?.AutoscaleMaxThroughput ?? throughput.Resource?.Throughput ?? min; } catch (CosmosException ex) when (ex.StatusCode == System.Net.HttpStatusCode.NotFound) { + RequestChargeContext.Record(ex.RequestCharge); + // Database has no shared throughput; containers provide their own. } catch (CosmosException ex) when (ex.StatusCode == System.Net.HttpStatusCode.BadRequest && ThroughputErrors.IsServerlessThroughputError(ex.Message)) { + RequestChargeContext.Record(ex.RequestCharge); serverless = true; } catch (Exception ex) when (ex is not OperationCanceledException) @@ -517,7 +525,9 @@ private static async Task WriteDatabaseThroughputAsync(Database database, Dictio using var iterator = container.GetItemQueryIterator(new QueryDefinition(queryText)); while (iterator.HasMoreResults) { - foreach (var element in await iterator.ReadNextAsync(token)) + var response = await iterator.ReadNextAsync(token); + RequestChargeContext.Record(response.RequestCharge); + foreach (var element in response) { long count = element.TryGetProperty("count", out var countProperty) && countProperty.TryGetInt64(out var parsed) ? parsed : 0; var keyParts = new List(partitionKeyPaths.Count); @@ -574,7 +584,9 @@ private static async Task CountFeedRangeAsync(Container container, FeedRan new QueryDefinition("SELECT VALUE COUNT(1) FROM c")); while (iterator.HasMoreResults) { - foreach (var value in await iterator.ReadNextAsync(token)) + var response = await iterator.ReadNextAsync(token); + RequestChargeContext.Record(response.RequestCharge); + foreach (var value in response) { count += value; } diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/MakeItemCommand.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/MakeItemCommand.cs index 788e1cef..1ed15ff9 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/MakeItemCommand.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/MakeItemCommand.cs @@ -175,6 +175,7 @@ private static async Task WriteItemAsync(Container container, stri } catch (CosmosException ce) { + RequestChargeContext.Record(ce.RequestCharge); failCount++; ShellInterpreter.WriteLine( MessageService.GetArgsString( diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ReplaceCommand.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ReplaceCommand.cs index d3a95fff..4eeebc0c 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ReplaceCommand.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ReplaceCommand.cs @@ -109,6 +109,7 @@ private static async Task ReplaceArrayAsync(Container container, } catch (CommandException ex) { + RequestChargeContext.Record(RequestChargeContext.GetCosmosExceptionCharge(ex)); failCount++; ShellInterpreter.WriteLine(ex.Message); } diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/SprocCommand.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/SprocCommand.cs index 48892668..e0fefdff 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/SprocCommand.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/SprocCommand.cs @@ -212,7 +212,9 @@ internal async Task ListAsync(Container container, ShellInterprete using var iterator = container.Scripts.GetStoredProcedureQueryIterator(); while (iterator.HasMoreResults) { - foreach (var properties in await iterator.ReadNextAsync(token)) + var response = await iterator.ReadNextAsync(token); + RequestChargeContext.Record(response.RequestCharge); + foreach (var properties in response) { items.Add(new { @@ -277,6 +279,7 @@ internal async Task ShowAsync(Container container, CommandState co try { var response = await container.Scripts.ReadStoredProcedureAsync(name, cancellationToken: token); + RequestChargeContext.Record(response.RequestCharge); commandState.Result = new ShellText(response.Resource.Body ?? string.Empty) { Highlighter = JavaScriptOutputHighlighter.BuildMarkup }; return commandState; } @@ -293,11 +296,13 @@ internal async Task ExistsAsync(Container container, CommandState bool exists; try { - await container.Scripts.ReadStoredProcedureAsync(name, cancellationToken: token); + var response = await container.Scripts.ReadStoredProcedureAsync(name, cancellationToken: token); + RequestChargeContext.Record(response.RequestCharge); exists = true; } catch (CosmosException ex) when (ex.StatusCode == HttpStatusCode.NotFound) { + RequestChargeContext.Record(ex.RequestCharge); exists = false; } @@ -337,10 +342,12 @@ private async Task CreateAsync(Container container, ShellInterpret try { var read = await container.Scripts.ReadStoredProcedureAsync(name, cancellationToken: token); + RequestChargeContext.Record(read.RequestCharge); existingBody = read.Resource.Body ?? string.Empty; } catch (CosmosException ex) when (ex.StatusCode == HttpStatusCode.NotFound) { + RequestChargeContext.Record(ex.RequestCharge); existingBody = null; } @@ -384,11 +391,14 @@ internal async Task WriteCreateAsync(Container container, CommandS try { response = await container.Scripts.ReplaceStoredProcedureAsync(properties, cancellationToken: token); + RequestChargeContext.Record(response.RequestCharge); replaced = true; } catch (CosmosException ex) when (ex.StatusCode == HttpStatusCode.NotFound) { + RequestChargeContext.Record(ex.RequestCharge); response = await container.Scripts.CreateStoredProcedureAsync(properties, cancellationToken: token); + RequestChargeContext.Record(response.RequestCharge); replaced = false; } } @@ -397,6 +407,7 @@ internal async Task WriteCreateAsync(Container container, CommandS try { response = await container.Scripts.CreateStoredProcedureAsync(properties, cancellationToken: token); + RequestChargeContext.Record(response.RequestCharge); replaced = false; } catch (CosmosException ex) when (ex.StatusCode == HttpStatusCode.Conflict) @@ -462,6 +473,7 @@ internal async Task DeleteAsync(Container container, CommandState try { var response = await container.Scripts.DeleteStoredProcedureAsync(name, cancellationToken: token); + RequestChargeContext.Record(response.RequestCharge); commandState.Result = new ShellJson(JsonSerializer.SerializeToElement(new { type = "sproc", id = name, deleted = true })); commandState.RenderUser = () => ShellInterpreter.WriteLine(MessageService.GetArgsString( "command-sproc-deleted", @@ -490,6 +502,7 @@ private async Task EditAsync(Container container, ShellInterpreter try { var read = await container.Scripts.ReadStoredProcedureAsync(name, cancellationToken: token); + RequestChargeContext.Record(read.RequestCharge); existingBody = read.Resource.Body ?? string.Empty; } catch (CosmosException ex) when (ex.StatusCode == HttpStatusCode.NotFound) @@ -508,6 +521,7 @@ private async Task EditAsync(Container container, ShellInterpreter var properties = new StoredProcedureProperties { Id = name, Body = newBody }; var response = await container.Scripts.ReplaceStoredProcedureAsync(properties, cancellationToken: token); + RequestChargeContext.Record(response.RequestCharge); commandState.Result = new ShellJson(JsonSerializer.SerializeToElement(new { type = "sproc", id = name, changed = true })); commandState.RenderUser = () => ShellInterpreter.WriteLine(MessageService.GetArgsString( diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/TriggerCommand.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/TriggerCommand.cs index add55d18..ae78b696 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/TriggerCommand.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/TriggerCommand.cs @@ -182,7 +182,9 @@ internal async Task ListAsync(Container container, ShellInterprete using var iterator = container.Scripts.GetTriggerQueryIterator(); while (iterator.HasMoreResults) { - foreach (var properties in await iterator.ReadNextAsync(token)) + var response = await iterator.ReadNextAsync(token); + RequestChargeContext.Record(response.RequestCharge); + foreach (var properties in response) { items.Add(new { @@ -252,6 +254,7 @@ internal async Task ShowAsync(Container container, CommandState co try { var response = await container.Scripts.ReadTriggerAsync(name, cancellationToken: token); + RequestChargeContext.Record(response.RequestCharge); commandState.Result = new ShellText(response.Resource.Body ?? string.Empty) { Highlighter = JavaScriptOutputHighlighter.BuildMarkup }; return commandState; } @@ -268,11 +271,13 @@ internal async Task ExistsAsync(Container container, CommandState bool exists; try { - await container.Scripts.ReadTriggerAsync(name, cancellationToken: token); + var response = await container.Scripts.ReadTriggerAsync(name, cancellationToken: token); + RequestChargeContext.Record(response.RequestCharge); exists = true; } catch (CosmosException ex) when (ex.StatusCode == HttpStatusCode.NotFound) { + RequestChargeContext.Record(ex.RequestCharge); exists = false; } @@ -320,10 +325,12 @@ private async Task CreateAsync(Container container, ShellInterpret try { var read = await container.Scripts.ReadTriggerAsync(name, cancellationToken: token); + RequestChargeContext.Record(read.RequestCharge); existingBody = read.Resource.Body ?? string.Empty; } catch (CosmosException ex) when (ex.StatusCode == HttpStatusCode.NotFound) { + RequestChargeContext.Record(ex.RequestCharge); existingBody = null; } @@ -373,11 +380,14 @@ internal async Task WriteCreateAsync(Container container, CommandS try { response = await container.Scripts.ReplaceTriggerAsync(properties, cancellationToken: token); + RequestChargeContext.Record(response.RequestCharge); replaced = true; } catch (CosmosException ex) when (ex.StatusCode == HttpStatusCode.NotFound) { + RequestChargeContext.Record(ex.RequestCharge); response = await container.Scripts.CreateTriggerAsync(properties, cancellationToken: token); + RequestChargeContext.Record(response.RequestCharge); replaced = false; } } @@ -386,6 +396,7 @@ internal async Task WriteCreateAsync(Container container, CommandS try { response = await container.Scripts.CreateTriggerAsync(properties, cancellationToken: token); + RequestChargeContext.Record(response.RequestCharge); replaced = false; } catch (CosmosException ex) when (ex.StatusCode == HttpStatusCode.Conflict) @@ -421,6 +432,7 @@ internal async Task DeleteAsync(Container container, CommandState try { var response = await container.Scripts.DeleteTriggerAsync(name, cancellationToken: token); + RequestChargeContext.Record(response.RequestCharge); commandState.Result = new ShellJson(JsonSerializer.SerializeToElement(new { type = "trigger", id = name, deleted = true })); commandState.RenderUser = () => ShellInterpreter.WriteLine(MessageService.GetArgsString( "command-trigger-deleted", @@ -449,6 +461,7 @@ private async Task EditAsync(Container container, ShellInterpreter try { var read = await container.Scripts.ReadTriggerAsync(name, cancellationToken: token); + RequestChargeContext.Record(read.RequestCharge); existing = read.Resource; } catch (CosmosException ex) when (ex.StatusCode == HttpStatusCode.NotFound) @@ -474,6 +487,7 @@ private async Task EditAsync(Container container, ShellInterpreter TriggerOperation = existing.TriggerOperation, }; var response = await container.Scripts.ReplaceTriggerAsync(properties, cancellationToken: token); + RequestChargeContext.Record(response.RequestCharge); commandState.Result = new ShellJson(JsonSerializer.SerializeToElement(new { type = "trigger", id = name, changed = true })); commandState.RenderUser = () => ShellInterpreter.WriteLine(MessageService.GetArgsString( diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/UdfCommand.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/UdfCommand.cs index 6ab0c887..e0067d51 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/UdfCommand.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/UdfCommand.cs @@ -139,7 +139,9 @@ internal async Task ListAsync(Container container, ShellInterprete using var iterator = container.Scripts.GetUserDefinedFunctionQueryIterator(); while (iterator.HasMoreResults) { - foreach (var properties in await iterator.ReadNextAsync(token)) + var response = await iterator.ReadNextAsync(token); + RequestChargeContext.Record(response.RequestCharge); + foreach (var properties in response) { items.Add(new { @@ -199,6 +201,7 @@ internal async Task ShowAsync(Container container, CommandState co try { var response = await container.Scripts.ReadUserDefinedFunctionAsync(name, cancellationToken: token); + RequestChargeContext.Record(response.RequestCharge); commandState.Result = new ShellText(response.Resource.Body ?? string.Empty) { Highlighter = JavaScriptOutputHighlighter.BuildMarkup }; return commandState; } @@ -215,11 +218,13 @@ internal async Task ExistsAsync(Container container, CommandState bool exists; try { - await container.Scripts.ReadUserDefinedFunctionAsync(name, cancellationToken: token); + var response = await container.Scripts.ReadUserDefinedFunctionAsync(name, cancellationToken: token); + RequestChargeContext.Record(response.RequestCharge); exists = true; } catch (CosmosException ex) when (ex.StatusCode == HttpStatusCode.NotFound) { + RequestChargeContext.Record(ex.RequestCharge); exists = false; } @@ -259,10 +264,12 @@ private async Task CreateAsync(Container container, ShellInterpret try { var read = await container.Scripts.ReadUserDefinedFunctionAsync(name, cancellationToken: token); + RequestChargeContext.Record(read.RequestCharge); existingBody = read.Resource.Body ?? string.Empty; } catch (CosmosException ex) when (ex.StatusCode == HttpStatusCode.NotFound) { + RequestChargeContext.Record(ex.RequestCharge); existingBody = null; } @@ -306,11 +313,14 @@ internal async Task WriteCreateAsync(Container container, CommandS try { response = await container.Scripts.ReplaceUserDefinedFunctionAsync(properties, cancellationToken: token); + RequestChargeContext.Record(response.RequestCharge); replaced = true; } catch (CosmosException ex) when (ex.StatusCode == HttpStatusCode.NotFound) { + RequestChargeContext.Record(ex.RequestCharge); response = await container.Scripts.CreateUserDefinedFunctionAsync(properties, cancellationToken: token); + RequestChargeContext.Record(response.RequestCharge); replaced = false; } } @@ -319,6 +329,7 @@ internal async Task WriteCreateAsync(Container container, CommandS try { response = await container.Scripts.CreateUserDefinedFunctionAsync(properties, cancellationToken: token); + RequestChargeContext.Record(response.RequestCharge); replaced = false; } catch (CosmosException ex) when (ex.StatusCode == HttpStatusCode.Conflict) @@ -347,6 +358,7 @@ internal async Task DeleteAsync(Container container, CommandState try { var response = await container.Scripts.DeleteUserDefinedFunctionAsync(name, cancellationToken: token); + RequestChargeContext.Record(response.RequestCharge); commandState.Result = new ShellJson(JsonSerializer.SerializeToElement(new { type = "udf", id = name, deleted = true })); commandState.RenderUser = () => ShellInterpreter.WriteLine(MessageService.GetArgsString( "command-udf-deleted", @@ -375,6 +387,7 @@ private async Task EditAsync(Container container, ShellInterpreter try { var read = await container.Scripts.ReadUserDefinedFunctionAsync(name, cancellationToken: token); + RequestChargeContext.Record(read.RequestCharge); existingBody = read.Resource.Body ?? string.Empty; } catch (CosmosException ex) when (ex.StatusCode == HttpStatusCode.NotFound) @@ -393,6 +406,7 @@ private async Task EditAsync(Container container, ShellInterpreter var properties = new UserDefinedFunctionProperties { Id = name, Body = newBody }; var response = await container.Scripts.ReplaceUserDefinedFunctionAsync(properties, cancellationToken: token); + RequestChargeContext.Record(response.RequestCharge); commandState.Result = new ShellJson(JsonSerializer.SerializeToElement(new { type = "udf", id = name, changed = true })); commandState.RenderUser = () => ShellInterpreter.WriteLine(MessageService.GetArgsString( diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/WatchCommand.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/WatchCommand.cs index f186adb6..c97a7a22 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/WatchCommand.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/WatchCommand.cs @@ -173,6 +173,7 @@ private async Task WatchAsync(ShellInterpreter shell, Container co while (!token.IsCancellationRequested) { using var response = await iterator.ReadNextAsync(token); + RequestChargeContext.Record(response.Headers.RequestCharge); if (response.StatusCode == HttpStatusCode.NotModified) { diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/DataPlaneCosmosResourceOperations.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/DataPlaneCosmosResourceOperations.cs index 8805e651..115b5881 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/DataPlaneCosmosResourceOperations.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/DataPlaneCosmosResourceOperations.cs @@ -27,6 +27,7 @@ public async IAsyncEnumerable GetDatabaseNamesAsync([EnumeratorCancellat while (iterator.HasMoreResults) { var page = await iterator.ReadNextAsync(token); + RequestChargeContext.Record(page.RequestCharge); foreach (var database in page) { yield return database.Id; @@ -41,6 +42,7 @@ public async IAsyncEnumerable GetContainerNamesAsync(string databaseName while (iterator.HasMoreResults) { var page = await iterator.ReadNextAsync(token); + RequestChargeContext.Record(page.RequestCharge); foreach (var container in page) { yield return container.Id; @@ -53,10 +55,12 @@ public async Task DatabaseExistsAsync(string databaseName, CancellationTok try { var response = await client.GetDatabase(databaseName).ReadAsync(cancellationToken: token); + RequestChargeContext.Record(response.RequestCharge); return response.StatusCode == HttpStatusCode.OK; } catch (CosmosException ex) when (ex.StatusCode == HttpStatusCode.NotFound) { + RequestChargeContext.Record(ex.RequestCharge); return false; } } @@ -66,10 +70,12 @@ public async Task ContainerExistsAsync(string databaseName, string contain try { var response = await client.GetDatabase(databaseName).GetContainer(containerName).ReadContainerAsync(cancellationToken: token); + RequestChargeContext.Record(response.RequestCharge); return response.StatusCode == HttpStatusCode.OK; } catch (CosmosException ex) when (ex.StatusCode == HttpStatusCode.NotFound) { + RequestChargeContext.Record(ex.RequestCharge); return false; } } @@ -78,6 +84,7 @@ public async Task CreateDatabaseAsync(string databaseName, string? scale { var throughput = CreateThroughputProperties(scale, maxRu); var response = await client.CreateDatabaseIfNotExistsAsync(databaseName, throughput, cancellationToken: token); + RequestChargeContext.Record(response.RequestCharge); return response.Database.Id; } @@ -114,22 +121,26 @@ public async Task CreateContainerAsync( var throughput = CreateThroughputProperties(scale, maxRu); var database = client.GetDatabase(databaseName); var response = await database.CreateContainerIfNotExistsAsync(props, throughput, cancellationToken: token); + RequestChargeContext.Record(response.RequestCharge); return response.Container.Id; } - public Task DeleteDatabaseAsync(string databaseName, CancellationToken token) + public async Task DeleteDatabaseAsync(string databaseName, CancellationToken token) { - return client.GetDatabase(databaseName).DeleteAsync(cancellationToken: token); + var response = await client.GetDatabase(databaseName).DeleteAsync(cancellationToken: token); + RequestChargeContext.Record(response.RequestCharge); } - public Task DeleteContainerAsync(string databaseName, string containerName, CancellationToken token) + public async Task DeleteContainerAsync(string databaseName, string containerName, CancellationToken token) { - return client.GetDatabase(databaseName).GetContainer(containerName).DeleteContainerAsync(cancellationToken: token); + var response = await client.GetDatabase(databaseName).GetContainer(containerName).DeleteContainerAsync(cancellationToken: token); + RequestChargeContext.Record(response.RequestCharge); } public async Task> GetPartitionKeyPathsAsync(string databaseName, string containerName, CancellationToken token) { var response = await client.GetDatabase(databaseName).GetContainer(containerName).ReadContainerAsync(cancellationToken: token); + RequestChargeContext.Record(response.RequestCharge); var properties = response.Resource; if (properties == null) { @@ -147,6 +158,7 @@ public async Task> GetPartitionKeyPathsAsync(string databa public async Task GetContainerSettingsAsync(string databaseName, string containerName, CancellationToken token) { var dpResponse = await client.GetDatabase(databaseName).GetContainer(containerName).ReadContainerAsync(cancellationToken: token); + RequestChargeContext.Record(dpResponse.RequestCharge); var properties = GetContainerPropertiesOrThrow(dpResponse); int? dpMin = null; int? dpMax = null; @@ -155,19 +167,23 @@ public async Task GetContainerSettingsAsync(string databa try { var throughputResponse = await client.GetDatabase(databaseName).GetContainer(containerName).ReadThroughputAsync(new RequestOptions(), token); + RequestChargeContext.Record(throughputResponse.RequestCharge); dpMin = throughputResponse.MinThroughput; dpMax = throughputResponse.Resource?.AutoscaleMaxThroughput ?? throughputResponse.Resource?.Throughput ?? dpMin; } catch (CosmosException ex) when (ex.StatusCode == HttpStatusCode.NotFound) { + RequestChargeContext.Record(ex.RequestCharge); dpAvailability = ThroughputAvailability.NotConfigured; } catch (CosmosException ex) when (ex.StatusCode == HttpStatusCode.BadRequest && ThroughputErrors.IsServerlessThroughputError(ex.Message)) { + RequestChargeContext.Record(ex.RequestCharge); dpAvailability = ThroughputAvailability.Serverless; } catch (Exception ex) { + RequestChargeContext.Record(RequestChargeContext.GetCosmosExceptionCharge(ex)); dpAvailability = ThroughputAvailability.Unavailable; dpError = ex.Message; } @@ -217,6 +233,7 @@ public async Task GetContainerSettingsAsync(string databa public async Task GetIndexingPolicyJsonAsync(string databaseName, string containerName, CancellationToken token) { var response = await client.GetDatabase(databaseName).GetContainer(containerName).ReadContainerAsync(cancellationToken: token); + RequestChargeContext.Record(response.RequestCharge); var properties = GetContainerPropertiesOrThrow(response); var policy = properties.IndexingPolicy ?? throw new IndexPolicyMissingException(); @@ -227,16 +244,19 @@ public async Task ReplaceIndexingPolicyAsync(string databaseName, string { var container = client.GetDatabase(databaseName).GetContainer(containerName); var current = await container.ReadContainerAsync(cancellationToken: token); + RequestChargeContext.Record(current.RequestCharge); var props = GetContainerPropertiesOrThrow(current); var policy = ParseIndexingPolicy(indexPolicyJson); props.IndexingPolicy = policy; var replaced = await container.ReplaceContainerAsync(props, cancellationToken: token); + RequestChargeContext.Record(replaced.RequestCharge); return JsonSerializer.Serialize(replaced.Resource?.IndexingPolicy ?? policy, IndexingPolicyJsonOptions); } public async Task GetTimeToLiveAsync(string databaseName, string containerName, CancellationToken token) { var response = await client.GetDatabase(databaseName).GetContainer(containerName).ReadContainerAsync(cancellationToken: token); + RequestChargeContext.Record(response.RequestCharge); var props = GetContainerPropertiesOrThrow(response); return new ContainerTtlView(props.DefaultTimeToLive); } @@ -245,9 +265,11 @@ public async Task ReplaceTimeToLiveAsync(string databaseName, { var container = client.GetDatabase(databaseName).GetContainer(containerName); var current = await container.ReadContainerAsync(cancellationToken: token); + RequestChargeContext.Record(current.RequestCharge); var props = GetContainerPropertiesOrThrow(current); props.DefaultTimeToLive = defaultTimeToLive; var replaced = await container.ReplaceContainerAsync(props, cancellationToken: token); + RequestChargeContext.Record(replaced.RequestCharge); var updated = GetContainerPropertiesOrThrow(replaced); return new ContainerTtlView(updated.DefaultTimeToLive); } @@ -255,6 +277,7 @@ public async Task ReplaceTimeToLiveAsync(string databaseName, public async Task GetAnalyticalTimeToLiveAsync(string databaseName, string containerName, CancellationToken token) { var response = await client.GetDatabase(databaseName).GetContainer(containerName).ReadContainerAsync(cancellationToken: token); + RequestChargeContext.Record(response.RequestCharge); var props = GetContainerPropertiesOrThrow(response); return new ContainerAnalyticalTtlView(props.AnalyticalStoreTimeToLiveInSeconds); } @@ -263,9 +286,11 @@ public async Task ReplaceAnalyticalTimeToLiveAsync(s { var container = client.GetDatabase(databaseName).GetContainer(containerName); var current = await container.ReadContainerAsync(cancellationToken: token); + RequestChargeContext.Record(current.RequestCharge); var props = GetContainerPropertiesOrThrow(current); props.AnalyticalStoreTimeToLiveInSeconds = analyticalTimeToLive; var replaced = await container.ReplaceContainerAsync(props, cancellationToken: token); + RequestChargeContext.Record(replaced.RequestCharge); var updated = GetContainerPropertiesOrThrow(replaced); return new ContainerAnalyticalTtlView(updated.AnalyticalStoreTimeToLiveInSeconds); } @@ -273,6 +298,7 @@ public async Task ReplaceAnalyticalTimeToLiveAsync(s public async Task GetConflictResolutionPolicyAsync(string databaseName, string containerName, CancellationToken token) { var response = await client.GetDatabase(databaseName).GetContainer(containerName).ReadContainerAsync(cancellationToken: token); + RequestChargeContext.Record(response.RequestCharge); var props = GetContainerPropertiesOrThrow(response); return ToConflictResolutionView(props.ConflictResolutionPolicy); } @@ -281,9 +307,11 @@ public async Task ReplaceConflictResolutionPolicyAsync(s { var container = client.GetDatabase(databaseName).GetContainer(containerName); var current = await container.ReadContainerAsync(cancellationToken: token); + RequestChargeContext.Record(current.RequestCharge); var props = GetContainerPropertiesOrThrow(current); props.ConflictResolutionPolicy = BuildConflictResolutionPolicy(update); var replaced = await container.ReplaceContainerAsync(props, cancellationToken: token); + RequestChargeContext.Record(replaced.RequestCharge); var updated = GetContainerPropertiesOrThrow(replaced); return ToConflictResolutionView(updated.ConflictResolutionPolicy); } @@ -298,10 +326,12 @@ public async Task GetThroughputAsync(string databaseName, string var throughputResponse = !isContainer ? await client.GetDatabase(databaseName).ReadThroughputAsync(new RequestOptions(), token) : await client.GetDatabase(databaseName).GetContainer(containerName).ReadThroughputAsync(new RequestOptions(), token); + RequestChargeContext.Record(throughputResponse.RequestCharge); return BuildThroughputView(scope, resourceName, throughputResponse); } catch (CosmosException ex) when (ex.StatusCode is HttpStatusCode.NotFound or HttpStatusCode.BadRequest) { + RequestChargeContext.Record(ex.RequestCharge); return new ThroughputView(scope, resourceName, false, null, null, null, ThroughputAvailability.NotConfigured, null); } } @@ -318,6 +348,7 @@ public async Task ReplaceThroughputAsync(string databaseName, st currentResponse = !isContainer ? await client.GetDatabase(databaseName).ReadThroughputAsync(new RequestOptions(), token) : await client.GetDatabase(databaseName).GetContainer(containerName).ReadThroughputAsync(new RequestOptions(), token); + RequestChargeContext.Record(currentResponse.RequestCharge); } catch (CosmosException ex) when (ex.StatusCode is HttpStatusCode.NotFound or HttpStatusCode.BadRequest) { @@ -341,6 +372,7 @@ public async Task ReplaceThroughputAsync(string databaseName, st var throughputResponse = !isContainer ? await client.GetDatabase(databaseName).ReplaceThroughputAsync(properties, cancellationToken: token) : await client.GetDatabase(databaseName).GetContainer(containerName).ReplaceThroughputAsync(properties, cancellationToken: token); + RequestChargeContext.Record(throughputResponse.RequestCharge); return BuildThroughputView(scope, resourceName, throughputResponse); } catch (CosmosException ex) when (ex.StatusCode is HttpStatusCode.NotFound or HttpStatusCode.BadRequest) diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/RequestChargeContext.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/RequestChargeContext.cs new file mode 100644 index 00000000..8724d157 --- /dev/null +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/RequestChargeContext.cs @@ -0,0 +1,93 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace Azure.Data.Cosmos.Shell.Core; + +using System.Threading; +using Microsoft.Azure.Cosmos; + +/// +/// Collects request charges from shared data-plane helpers during one command execution. +/// +internal static class RequestChargeContext +{ + private const string ExceptionChargeKey = "CosmosDBShell.RequestCharge"; + + private static readonly AsyncLocal CurrentScope = new(); + + internal static double CurrentRequestCharge => CurrentScope.Value?.RequestCharge ?? 0; + + internal static Scope Begin() + { + var scope = new Scope(CurrentScope.Value); + CurrentScope.Value = scope; + return scope; + } + + internal static void Record(double requestCharge) + { + if (requestCharge > 0 && CurrentScope.Value is { } scope) + { + scope.RequestCharge += requestCharge; + } + } + + internal static double GetCosmosExceptionCharge(Exception exception) + { + for (Exception? current = exception; current is not null; current = current.InnerException) + { + if (current is CosmosException cosmosException && cosmosException.RequestCharge > 0) + { + return cosmosException.RequestCharge; + } + } + + return 0; + } + + internal static void SetExceptionCharge(Exception exception, double requestCharge) + { + if (requestCharge > 0) + { + exception.Data[ExceptionChargeKey] = requestCharge; + } + } + + internal static double? GetExceptionCharge(Exception exception) + { + for (Exception? current = exception; current is not null; current = current.InnerException) + { + if (current.Data[ExceptionChargeKey] is double requestCharge) + { + return requestCharge; + } + } + + return null; + } + + internal sealed class Scope : IDisposable + { + private readonly Scope? parent; + private bool disposed; + + internal Scope(Scope? parent) + { + this.parent = parent; + } + + internal double RequestCharge { get; set; } + + public void Dispose() + { + if (this.disposed) + { + return; + } + + this.disposed = true; + CurrentScope.Value = this.parent; + } + } +} diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/ShellInterpreter.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/ShellInterpreter.cs index 9cb65b6a..de7120e8 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/ShellInterpreter.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/ShellInterpreter.cs @@ -549,7 +549,10 @@ public async Task ExecuteCommandAsync(string command, Cancellation this.ReportExecutionError(e, command); this.DisconnectLocalEmulatorAfterConnectivityFailure(e); var inner = e is PositionalException pe ? (pe.InnerException ?? pe) : e; - result = new ErrorCommandState(inner); + result = new ErrorCommandState(inner) + { + RequestCharge = RequestChargeContext.GetExceptionCharge(e), + }; return result; } @@ -960,9 +963,29 @@ internal async Task ExecuteCosmosCommandAsync( // previous command's charge so an uninstrumented command cannot count it again. commandState.RequestCharge = null; var generation = this.SessionRequestChargeGeneration; - var result = await command.ExecuteAsync(this, commandState, commandText, token); - this.RecordRequestCharge(result, generation); - return result; + using var requestChargeScope = RequestChargeContext.Begin(); + try + { + var result = await command.ExecuteAsync(this, commandState, commandText, token); + if (requestChargeScope.RequestCharge > 0) + { + result.RequestCharge = (result.RequestCharge ?? 0) + requestChargeScope.RequestCharge; + } + + this.RecordRequestCharge(result, generation); + return result; + } + catch (Exception ex) + { + var requestCharge = requestChargeScope.RequestCharge + RequestChargeContext.GetCosmosExceptionCharge(ex); + if (requestCharge > 0) + { + RequestChargeContext.SetExceptionCharge(ex, requestCharge); + this.RecordRequestCharge(new CommandState { RequestCharge = requestCharge }, generation); + } + + throw; + } } internal void RecordRequestCharge(CommandState commandState, long generation) diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Mcp/McpResponseFactory.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Mcp/McpResponseFactory.cs index ef02702e..d659cb2b 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Mcp/McpResponseFactory.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Mcp/McpResponseFactory.cs @@ -27,15 +27,18 @@ public static CallToolResult CreateSuccess(CommandState commandState, State shel return CreateResponse(CreateSuccessPayload(commandState), shellState, commandState.IsError); } - public static CallToolResult CreateError(string message, State shellState) + public static CallToolResult CreateError(string message, State shellState, double? requestCharge = null) { - return CreateResponse( - new JsonObject - { - ["error"] = message, - }, - shellState, - isError: true); + var payload = new JsonObject + { + ["error"] = message, + }; + if (requestCharge is not null) + { + payload["requestCharge"] = requestCharge.Value; + } + + return CreateResponse(payload, shellState, isError: true); } internal static string? GetCurrentLocation(State shellState) diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Mcp/ToolOperations.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Mcp/ToolOperations.cs index 92cad62d..f49f490f 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Mcp/ToolOperations.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Mcp/ToolOperations.cs @@ -499,7 +499,10 @@ private async ValueTask OnCallToolsAsync( { this.logger?.LogError(ex, $"An exception occurred running '{command.CommandName}'. "); - return McpResponseFactory.CreateError($"Error executing command '{command.CommandName}': {ex.Message}", ShellInterpreter.Instance.State); + return McpResponseFactory.CreateError( + $"Error executing command '{command.CommandName}': {ex.Message}", + ShellInterpreter.Instance.State, + RequestChargeContext.GetExceptionCharge(ex)); } finally { diff --git a/docs/commands.md b/docs/commands.md index b468274f..cf66a06c 100644 --- a/docs/commands.md +++ b/docs/commands.md @@ -1331,7 +1331,8 @@ only a database is in scope it reports the container count, aggregate document count, total storage, and shared throughput. When neither is in scope (the account root) it reports account metadata: read/write regions and the database count. Every scope also includes the cumulative request charge observed from -instrumented commands during the current connection. The total resets to zero +Cosmos DB data-plane requests during the current connection, including the +requests made by the current `info` command. The total resets to zero after each successful `connect`; changing database or container scope does not reset it. It is session telemetry, not a budget or billing total. @@ -1359,7 +1360,9 @@ the rich console layout. The `--partitions` and request units; at the account root, `--detailed` aggregates every container's storage and document count across all databases. In JSON output the connection total is available as `session.requestCharge`, and `session.chargedOperationCount` -counts instrumented command operations that reported a positive charge. This +counts command operations that reported a positive charge. Paginated and +multi-request commands contribute their aggregate observed charge. Azure Resource +Manager control-plane operations do not consume Cosmos DB request units. This command is read-only. ### help diff --git a/docs/mcp.md b/docs/mcp.md index 7f8e69ac..86180483 100644 --- a/docs/mcp.md +++ b/docs/mcp.md @@ -100,19 +100,20 @@ 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. This is omitted for commands that do not issue a billable request. | +| `requestCharge` | Charged data-plane command results | The Cosmos DB request charge (in RUs) consumed by the command, as a number. This is omitted for commands that do not issue a billable request. | | `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. Instrumented data-plane commands (`query`, including `--explain`; `print`; `ls` for container items; `can-i` data-plane probes; `batch run`; `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. +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. Commands report `requestCharge` whenever their Cosmos DB data-plane requests expose one, including paginated reads, metadata and configuration operations, scripts, change feed reads, handled probes, and charged failures. Multi-request commands aggregate the observed charges. Azure Resource Manager control-plane operations do not consume or report Cosmos DB request units. This field reports observed cost; it does not enforce an RU budget. Budget guardrails are tracked separately in [#162](https://github.com/Azure/CosmosDBShell/issues/162). The `info` command result also includes `session.requestCharge`, the cumulative -charge observed from instrumented commands during the current connection. A +charge observed from data-plane commands during the current connection, +including the current `info` request cost. A successful `connect` starts a new total; navigation between databases and containers does not reset it. This session value is telemetry rather than a -budget or billing total. `session.chargedOperationCount` counts instrumented +budget or billing total. `session.chargedOperationCount` counts command operations that reported a positive request charge; it counts command operations rather than individual query pages or transactional batch items. From 4a9aa31d440142c94f1ee4c42b69054f231b04cd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mike=20Kr=C3=BCger?= Date: Mon, 31 Aug 2026 14:24:28 +0200 Subject: [PATCH 13/16] Test request charge on generic MCP errors --- CosmosDBShell.Tests/McpResponseFactoryTests.cs | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/CosmosDBShell.Tests/McpResponseFactoryTests.cs b/CosmosDBShell.Tests/McpResponseFactoryTests.cs index 4957cee9..53057422 100644 --- a/CosmosDBShell.Tests/McpResponseFactoryTests.cs +++ b/CosmosDBShell.Tests/McpResponseFactoryTests.cs @@ -114,6 +114,17 @@ public void CreateError_WrapsMessageWithCurrentLocation() Assert.Equal("boom", document.RootElement.GetProperty("error").GetString()); } + [Fact] + public void CreateError_WithRequestCharge_IncludesChargeInEquivalentPayloads() + { + var result = McpResponseFactory.CreateError("boom", new DatabaseState("TestDatabase", null!), 2.5); + var text = Assert.IsType(Assert.Single(result.Content)).Text; + + using var document = JsonDocument.Parse(text); + Assert.Equal(2.5, document.RootElement.GetProperty("requestCharge").GetDouble()); + Assert.Equal(document.RootElement.GetRawText(), result.StructuredContent?.GetRawText()); + } + [Fact] public void CreateSuccess_WhenDisconnected_UsesNullCurrentLocation() { From bf65ce4eb86427103f481427c432cc9fa2742558 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mike=20Kr=C3=BCger?= Date: Mon, 31 Aug 2026 14:45:15 +0200 Subject: [PATCH 14/16] Add session request charge shell variables --- CHANGELOG.md | 1 + .../CommandTests/InfoCommandTests.cs | 14 ++ .../CommandTests/SessionRequestChargeTests.cs | 130 +++++++++++++++++ .../Lsp/CosmosShellCompletionHandlerTests.cs | 11 ++ .../InfoCommand.cs | 15 +- .../ShellInterpreter.cs | 134 +++++++++++++++++- .../CosmosShellCompletionHandler.cs | 15 ++ CosmosDBShell/lang/en.ftl | 4 + docs/commands.md | 5 +- docs/mcp.md | 4 +- docs/programming.md | 17 +++ 11 files changed, 344 insertions(+), 6 deletions(-) create mode 100644 CosmosDBShell.Tests/CommandTests/SessionRequestChargeTests.cs diff --git a/CHANGELOG.md b/CHANGELOG.md index 0a7619d1..9c857a62 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,7 @@ ### 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 `$sessionMaxRequestCharge` to a positive RU threshold to print one warning when the current connection reaches it; `info` reports the configured maximum as `session.maxRequestCharge`. ### Fixes diff --git a/CosmosDBShell.Tests/CommandTests/InfoCommandTests.cs b/CosmosDBShell.Tests/CommandTests/InfoCommandTests.cs index 7981673d..c276e6e1 100644 --- a/CosmosDBShell.Tests/CommandTests/InfoCommandTests.cs +++ b/CosmosDBShell.Tests/CommandTests/InfoCommandTests.cs @@ -7,6 +7,7 @@ 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; @@ -205,6 +206,19 @@ public void AddSessionUsage_IncludesCurrentCommandCharge() Assert.Equal(2, json.GetProperty("session").GetProperty("chargedOperationCount").GetInt64()); } + [Fact] + public void AddSessionUsage_IncludesConfiguredMaximum() + { + using var shell = ShellInterpreter.CreateInstance(); + shell.SetVariable("sessionMaxRequestCharge", new ShellDecimal(25.5)); + var result = new Dictionary(); + + InfoCommand.AddSessionUsage(shell, result, renderOutput: false); + + var session = JsonSerializer.SerializeToElement(result).GetProperty("session"); + Assert.Equal(25.5, session.GetProperty("maxRequestCharge").GetDouble()); + } + [Fact] public void RecordRequestCharge_IgnoresPriorConnectionGeneration() { diff --git a/CosmosDBShell.Tests/CommandTests/SessionRequestChargeTests.cs b/CosmosDBShell.Tests/CommandTests/SessionRequestChargeTests.cs new file mode 100644 index 00000000..c8943f55 --- /dev/null +++ b/CosmosDBShell.Tests/CommandTests/SessionRequestChargeTests.cs @@ -0,0 +1,130 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace CosmosShell.Tests.CommandTests; + +using Azure.Data.Cosmos.Shell.Core; +using Azure.Data.Cosmos.Shell.Parser; +using Microsoft.Azure.Cosmos; +using Spectre.Console; + +[Collection(ConsoleOutputTestCollection.Name)] +public class SessionRequestChargeTests +{ + [Fact] + public void SessionVariables_ReflectCurrentUsage() + { + using var shell = ShellInterpreter.CreateInstance(); + shell.RecordRequestCharge(new CommandState { RequestCharge = 2.5 }); + shell.RecordRequestCharge(new CommandState { RequestCharge = 1.25 }); + + var charge = Assert.IsType(shell.GetVariable("sessionRequestCharge")); + var operationCount = Assert.IsType(shell.GetVariable("sessionChargedOperationCount")); + var maximum = Assert.IsType(shell.GetVariable("sessionMaxRequestCharge")); + + Assert.Equal(3.75, charge.Value); + Assert.Equal(2, operationCount.Value); + Assert.Equal(0, maximum.Value); + } + + [Fact] + public void SessionUsageVariables_AreReadOnly() + { + using var shell = ShellInterpreter.CreateInstance(); + + Assert.Throws(() => shell.SetVariable("sessionRequestCharge", new ShellNumber(1))); + Assert.Throws(() => shell.SetVariable("sessionChargedOperationCount", new ShellNumber(1))); + } + + [Fact] + public void SessionMaxRequestCharge_WarnsOnlyOnceWhenReached() + { + using var shell = ShellInterpreter.CreateInstance(); + shell.SetVariable("sessionMaxRequestCharge", new ShellDecimal(3)); + + var output = CaptureConsole(() => + { + shell.RecordRequestCharge(new CommandState { RequestCharge = 2 }); + shell.RecordRequestCharge(new CommandState { RequestCharge = 1 }); + shell.RecordRequestCharge(new CommandState { RequestCharge = 1 }); + }); + + Assert.Equal(1, CountOccurrences(output, "has reached the configured warning threshold")); + } + + [Fact] + public void SessionMaxRequestCharge_ZeroDisablesWarning() + { + using var shell = ShellInterpreter.CreateInstance(); + shell.SetVariable("sessionMaxRequestCharge", new ShellNumber(0)); + + var output = CaptureConsole(() => shell.RecordRequestCharge(new CommandState { RequestCharge = 10 })); + + Assert.Empty(output); + } + + [Fact] + public void Connect_PreservesMaximumAndRearmsWarning() + { + using var shell = ShellInterpreter.CreateInstance(); + shell.SetVariable("sessionMaxRequestCharge", new ShellNumber(2)); + _ = CaptureConsole(() => shell.RecordRequestCharge(new CommandState { RequestCharge = 2 })); + + shell.Connect(CreateTestClient(), credentialTypeOverride: "AccountKey"); + var output = CaptureConsole(() => shell.RecordRequestCharge(new CommandState { RequestCharge = 2 })); + + Assert.Equal(2, Assert.IsType(shell.GetVariable("sessionMaxRequestCharge")).Value); + Assert.Equal(1, CountOccurrences(output, "has reached the configured warning threshold")); + } + + [Fact] + public void SessionMaxRequestCharge_RejectsInvalidValues() + { + using var shell = ShellInterpreter.CreateInstance(); + + Assert.Throws(() => shell.SetVariable("sessionMaxRequestCharge", new ShellDecimal(-1))); + Assert.Throws(() => shell.SetVariable("sessionMaxRequestCharge", new ShellText("ten"))); + } + + private static CosmosClient CreateTestClient() => new( + "https://localhost:8081", + Convert.ToBase64String(new byte[64]), + new CosmosClientOptions()); + + private static int CountOccurrences(string value, string substring) + { + int count = 0; + int index = 0; + while ((index = value.IndexOf(substring, index, StringComparison.Ordinal)) >= 0) + { + count++; + index += substring.Length; + } + + return count; + } + + private static string CaptureConsole(Action action) + { + var saved = AnsiConsole.Console; + using var writer = new StringWriter(); + try + { + AnsiConsole.Console = AnsiConsole.Create(new AnsiConsoleSettings + { + Ansi = AnsiSupport.No, + ColorSystem = ColorSystemSupport.NoColors, + Out = new AnsiConsoleOutput(writer), + }); + + action(); + } + finally + { + AnsiConsole.Console = saved; + } + + return writer.ToString(); + } +} diff --git a/CosmosDBShell.Tests/Lsp/CosmosShellCompletionHandlerTests.cs b/CosmosDBShell.Tests/Lsp/CosmosShellCompletionHandlerTests.cs index ffc7b2ab..6775dba6 100644 --- a/CosmosDBShell.Tests/Lsp/CosmosShellCompletionHandlerTests.cs +++ b/CosmosDBShell.Tests/Lsp/CosmosShellCompletionHandlerTests.cs @@ -87,6 +87,17 @@ public async Task VariableCompletion_SuggestsVariables() Assert.Contains("$foobar", labels); } + [Fact] + public async Task VariableCompletion_SuggestsSessionVariables() + { + var completions = await GetCompletionsAsync("echo $session", 0, 13); + var labels = Labels(completions); + + Assert.Contains("$sessionRequestCharge", labels); + Assert.Contains("$sessionChargedOperationCount", labels); + Assert.Contains("$sessionMaxRequestCharge", labels); + } + [Fact] public async Task VariableCompletion_IgnoresWhenNotVariableContext() { diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/InfoCommand.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/InfoCommand.cs index 32378e95..db3ff111 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/InfoCommand.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/InfoCommand.cs @@ -307,11 +307,17 @@ internal static void AddSessionUsage(ShellInterpreter shell, Dictionary 0 ? 1 : 0); - mcpTable["session"] = new Dictionary + var session = new Dictionary { ["requestCharge"] = requestCharge, ["chargedOperationCount"] = chargedOperationCount, }; + if (sessionUsage.MaxRequestCharge > 0) + { + session["maxRequestCharge"] = sessionUsage.MaxRequestCharge; + } + + mcpTable["session"] = session; if (!renderOutput) { @@ -328,6 +334,13 @@ internal static void AddSessionUsage(ShellInterpreter shell, Dictionary 0) + { + table.AddRow( + MessageService.GetString("command-stats-session-max-request-charge"), + Theme.FormatTableValue(sessionUsage.MaxRequestCharge.ToString("0.##", CultureInfo.InvariantCulture))); + } + AnsiConsole.Write(table); } diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/ShellInterpreter.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/ShellInterpreter.cs index de7120e8..772e1fdd 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/ShellInterpreter.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/ShellInterpreter.cs @@ -25,6 +25,12 @@ namespace Azure.Data.Cosmos.Shell.Core; /// public partial class ShellInterpreter : IDisposable { + private const string SessionRequestChargeVariable = "sessionRequestCharge"; + + private const string SessionChargedOperationCountVariable = "sessionChargedOperationCount"; + + private const string SessionMaxRequestChargeVariable = "sessionMaxRequestCharge"; + internal static readonly ShellInterpreter Instance = new(); private const int MAXHISTORYITEMS = 60; @@ -70,6 +76,10 @@ public partial class ShellInterpreter : IDisposable private long sessionChargedOperationCount; + private double sessionMaxRequestCharge; + + private bool sessionRequestChargeWarningIssued; + private long sessionRequestChargeGeneration; internal ShellInterpreter(string? configPath = null) @@ -143,6 +153,13 @@ internal static char CSVSeparator } } + internal static IReadOnlyList SessionVariableNames { get; } = + [ + SessionRequestChargeVariable, + SessionChargedOperationCountVariable, + SessionMaxRequestChargeVariable, + ]; + internal Dictionary Functions { get; } = []; /// @@ -187,13 +204,13 @@ internal long SessionChargedOperationCount } } - internal (double RequestCharge, long ChargedOperationCount) SessionUsage + internal (double RequestCharge, long ChargedOperationCount, double MaxRequestCharge) SessionUsage { get { lock (this.sessionRequestChargeLock) { - return (this.sessionRequestCharge, this.sessionChargedOperationCount); + return (this.sessionRequestCharge, this.sessionChargedOperationCount, this.sessionMaxRequestCharge); } } } @@ -720,6 +737,24 @@ internal static void ReportError(string message, params object[] par) internal ShellObject GetVariable(string name) { + lock (this.sessionRequestChargeLock) + { + if (string.Equals(name, SessionRequestChargeVariable, StringComparison.OrdinalIgnoreCase)) + { + return new ShellDecimal(this.sessionRequestCharge); + } + + if (string.Equals(name, SessionChargedOperationCountVariable, StringComparison.OrdinalIgnoreCase)) + { + return new ShellDecimal(this.sessionChargedOperationCount); + } + + if (string.Equals(name, SessionMaxRequestChargeVariable, StringComparison.OrdinalIgnoreCase)) + { + return new ShellDecimal(this.sessionMaxRequestCharge); + } + } + var scope = this.GetScope(name); if (scope?.TryGetValue(name, out var value) == true) { @@ -990,6 +1025,9 @@ internal async Task ExecuteCosmosCommandAsync( internal void RecordRequestCharge(CommandState commandState, long generation) { + bool printWarning = false; + double requestChargeTotal = 0; + double requestChargeMaximum = 0; if (commandState.RequestCharge is { } requestCharge) { lock (this.sessionRequestChargeLock) @@ -1001,9 +1039,24 @@ internal void RecordRequestCharge(CommandState commandState, long generation) { this.sessionChargedOperationCount++; } + + if (this.sessionMaxRequestCharge > 0 + && !this.sessionRequestChargeWarningIssued + && this.sessionRequestCharge >= this.sessionMaxRequestCharge) + { + this.sessionRequestChargeWarningIssued = true; + printWarning = true; + requestChargeTotal = this.sessionRequestCharge; + requestChargeMaximum = this.sessionMaxRequestCharge; + } } } } + + if (printWarning) + { + this.PrintSessionRequestChargeWarning(requestChargeTotal, requestChargeMaximum); + } } internal async Task ConnectAsync(string connectionString, string? loginHint = null, ConnectionMode? mode = null, string? tenantId = null, string? authorityHost = null, string? managedIdentityClientId = null, CredentialMethod credentialMethod = CredentialMethod.Default, string? subscriptionId = null, string? resourceGroupName = null, CancellationToken token = default) @@ -1614,6 +1667,7 @@ internal void Connect(CosmosClient client, ArmCosmosContext? armContext = null, { this.sessionRequestCharge = 0; this.sessionChargedOperationCount = 0; + this.sessionRequestChargeWarningIssued = false; this.sessionRequestChargeGeneration++; } @@ -1877,6 +1931,18 @@ internal void DeclareFunction(DefStatement defStatement) internal void SetVariable(string variableName, ShellObject value) { + if (string.Equals(variableName, SessionRequestChargeVariable, StringComparison.OrdinalIgnoreCase) + || string.Equals(variableName, SessionChargedOperationCountVariable, StringComparison.OrdinalIgnoreCase)) + { + throw new ShellException(MessageService.GetArgsString("error-session-variable-read-only", "name", variableName)); + } + + if (string.Equals(variableName, SessionMaxRequestChargeVariable, StringComparison.OrdinalIgnoreCase)) + { + this.SetSessionMaxRequestCharge(value); + return; + } + // Ensure we have at least one variable container (global scope) if (this.VariableContainers.Count == 0) { @@ -1919,6 +1985,70 @@ internal void SetVariable(string variableName, ShellObject value) currentScope.Set(variableName, shellValue); } + private void SetSessionMaxRequestCharge(ShellObject value) + { + double maximum = value switch + { + ShellNumber number => number.Value, + ShellDecimal decimalValue => decimalValue.Value, + _ => double.NaN, + }; + + if (!double.IsFinite(maximum) || maximum < 0) + { + throw new ShellException(MessageService.GetString("error-session-max-request-charge-invalid")); + } + + bool printWarning; + double requestChargeTotal; + lock (this.sessionRequestChargeLock) + { + if (maximum != this.sessionMaxRequestCharge) + { + this.sessionRequestChargeWarningIssued = false; + } + + this.sessionMaxRequestCharge = maximum; + printWarning = maximum > 0 + && !this.sessionRequestChargeWarningIssued + && this.sessionRequestCharge >= maximum; + if (printWarning) + { + this.sessionRequestChargeWarningIssued = true; + } + + requestChargeTotal = this.sessionRequestCharge; + } + + if (printWarning) + { + this.PrintSessionRequestChargeWarning(requestChargeTotal, maximum); + } + } + + private void PrintSessionRequestChargeWarning(double requestCharge, double maximum) + { + var message = MessageService.GetArgsString( + "warning-session-max-request-charge-reached", + "requestCharge", + requestCharge.ToString("0.##", CultureInfo.InvariantCulture), + "maximum", + maximum.ToString("0.##", CultureInfo.InvariantCulture)); + + if (this.IsMachineMode) + { + Console.Error.WriteLine(JsonSerializer.Serialize(new Dictionary + { + ["status"] = "warning", + ["warning"] = message, + })); + } + else + { + AnsiConsole.MarkupLine(Theme.FormatWarning(message)); + } + } + /// /// Releases the unmanaged resources used by the and optionally releases the managed resources. /// diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Lsp/CosmosShellCompletionHandler.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Lsp/CosmosShellCompletionHandler.cs index bb4067e1..379d6442 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Lsp/CosmosShellCompletionHandler.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Lsp/CosmosShellCompletionHandler.cs @@ -142,6 +142,21 @@ private static void AddVariableCompletions(List items, string pa } var seen = new HashSet(StringComparer.OrdinalIgnoreCase); + foreach (var name in ShellInterpreter.SessionVariableNames) + { + string variableName = "$" + name; + if (seen.Add(name) && variableName.StartsWith(partial, StringComparison.OrdinalIgnoreCase)) + { + items.Add(new CompletionItem + { + Label = variableName, + Kind = CompletionItemKind.Variable, + InsertText = variableName, + SortText = "0_" + name, + }); + } + } + foreach (var container in ShellInterpreter.Instance.VariableContainers.Reverse()) { foreach (var name in container.Variables.Keys) diff --git a/CosmosDBShell/lang/en.ftl b/CosmosDBShell/lang/en.ftl index a62b07f0..98d96d19 100644 --- a/CosmosDBShell/lang/en.ftl +++ b/CosmosDBShell/lang/en.ftl @@ -934,6 +934,10 @@ command-stats-account-detailed-cost-note = Aggregating account totals reads ever command-stats-session-heading = Session Usage command-stats-session-request-charge = Observed request charge (RUs) command-stats-session-charged-operations = Charged operations +command-stats-session-max-request-charge = Warning threshold (RUs) +warning-session-max-request-charge-reached = Session request charge { $requestCharge } RUs has reached the configured warning threshold of { $maximum } RUs. +error-session-variable-read-only = Variable '${ $name }' is read-only. +error-session-max-request-charge-invalid = Variable '$sessionMaxRequestCharge' must be a non-negative number. Set it to 0 to disable the warning. command-version-description = Displays the version of Cosmos DB Shell. command-version = Cosmos Shell version: { $version } diff --git a/docs/commands.md b/docs/commands.md index cf66a06c..61e5b2df 100644 --- a/docs/commands.md +++ b/docs/commands.md @@ -1362,8 +1362,9 @@ storage and document count across all databases. In JSON output the connection total is available as `session.requestCharge`, and `session.chargedOperationCount` counts command operations that reported a positive charge. Paginated and multi-request commands contribute their aggregate observed charge. Azure Resource -Manager control-plane operations do not consume Cosmos DB request units. This -command is read-only. +Manager control-plane operations do not consume Cosmos DB request units. When +`$sessionMaxRequestCharge` is positive, `session.maxRequestCharge` reports the +configured warning threshold. This command is read-only. ### help diff --git a/docs/mcp.md b/docs/mcp.md index 86180483..3d7b5ce5 100644 --- a/docs/mcp.md +++ b/docs/mcp.md @@ -115,5 +115,7 @@ successful `connect` starts a new total; navigation between databases and containers does not reset it. This session value is telemetry rather than a budget or billing total. `session.chargedOperationCount` counts command operations that reported a positive request charge; it counts command -operations rather than individual query pages or transactional batch items. +operations rather than individual query pages or transactional batch items. If +the shell variable `$sessionMaxRequestCharge` is set to a positive value, the +session object also includes `session.maxRequestCharge`. diff --git a/docs/programming.md b/docs/programming.md index 1b8450a2..723d7316 100644 --- a/docs/programming.md +++ b/docs/programming.md @@ -72,6 +72,23 @@ echo $"Hello $name" # interpolate echo "Hello $name" # print $name literally ``` +### Session request charge variables + +The shell provides three built-in session variables: + +| Variable | Description | +| --- | --- | +| `$sessionRequestCharge` | Read-only cumulative request charge observed during the current connection. | +| `$sessionChargedOperationCount` | Read-only number of command operations that reported a positive charge. | +| `$sessionMaxRequestCharge` | Configurable warning threshold in RUs. A positive value enables the warning; `0` disables it. | + +For example, `$sessionMaxRequestCharge = 100` prints one warning when the +current connection reaches or exceeds 100 observed RUs. The warning is emitted +only once for that threshold. A successful `connect` resets the accumulated +charge and operation count and rearms the warning, while preserving the +configured maximum. Assigning a different positive maximum also rearms the +warning for the new threshold. + For script positional parameters, see [Writing and Running Scripts](#writing-and-running-scripts). ## Writing and Running Scripts From 41ea35015decce23eb149bead0f8d3ff41754d92 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mike=20Kr=C3=BCger?= Date: Mon, 31 Aug 2026 15:06:47 +0200 Subject: [PATCH 15/16] Address request charge review findings --- .../ListCommand.cs | 3 ++- .../PrintCommand.cs | 26 +++++++------------ .../QueryCommand.cs | 21 +++------------ .../ReplaceCommand.cs | 1 + .../RmCommand.cs | 2 +- .../ShellInterpreter.cs | 2 +- 6 files changed, 18 insertions(+), 37 deletions(-) diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ListCommand.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ListCommand.cs index 975c24e0..9ac995f7 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ListCommand.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ListCommand.cs @@ -236,7 +236,8 @@ private async Task ListContainerItemsAsync(ConnectedState state, S } catch (Exception ex) when (ex is not OperationCanceledException) { - return new ErrorCommandState(ex) { RequestCharge = returnState.RequestCharge }; + RequestChargeContext.Record(returnState.RequestCharge ?? 0); + throw new CommandException("ls", ex); } using (queryDocument) diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/PrintCommand.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/PrintCommand.cs index 98f89e83..194cc972 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/PrintCommand.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/PrintCommand.cs @@ -54,46 +54,38 @@ private async Task PrintItemAsync(Container container, Cancellatio try { using var response = await container.ReadItemStreamAsync(this.Id, new PartitionKey(this.PartitionKey), cancellationToken: token); + RequestChargeContext.Record(response.Headers.RequestCharge); if (response.IsSuccessStatusCode) { - commandState.RequestCharge = response.Headers.RequestCharge; using var reader = new StreamReader(response.Content); var content = await reader.ReadToEndAsync(); // Parse the content as JSON for structured output - var jsonDocument = System.Text.Json.JsonDocument.Parse(content); - commandState.Result = new ShellJson(jsonDocument.RootElement); + using var jsonDocument = System.Text.Json.JsonDocument.Parse(content); + commandState.Result = new ShellJson(jsonDocument.RootElement.Clone()); } else if (response.StatusCode == System.Net.HttpStatusCode.NotFound) { - return new ErrorCommandState(new CommandException("print", MessageService.GetString("command-print-error-item_not_found", new Dictionary + throw new CommandException("print", MessageService.GetString("command-print-error-item_not_found", new Dictionary { { "id", this.Id ?? "(null)" }, { "partitionKey", this.PartitionKey ?? "(null)" }, - }))) - { - RequestCharge = response.Headers.RequestCharge, - }; + })); } else { - return new ErrorCommandState(new CommandException("print", MessageService.GetString("command-print-error-request_failed", new Dictionary + throw new CommandException("print", MessageService.GetString("command-print-error-request_failed", new Dictionary { { "id", this.Id ?? "(null)" }, { "status", (int)response.StatusCode }, - }))) - { - RequestCharge = response.Headers.RequestCharge, - }; + })); } } catch (CosmosException ex) { - return new ErrorCommandState(new CommandException("print", MessageService.GetArgsString("command-print-error-reading_item", "message", CommandException.GetDisplayMessage(ex)), ex)) - { - RequestCharge = ex.RequestCharge > 0 ? ex.RequestCharge : null, - }; + RequestChargeContext.Record(ex.RequestCharge); + throw new CommandException("print", MessageService.GetArgsString("command-print-error-reading_item", "message", CommandException.GetDisplayMessage(ex)), ex); } return commandState; diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/QueryCommand.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/QueryCommand.cs index d7633874..786caa6b 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/QueryCommand.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/QueryCommand.cs @@ -634,14 +634,8 @@ private async Task ExecuteExplainAsync(Container container, ShellI double requestCharge = response?.Headers.RequestCharge ?? 0; if (response is not null && !response.IsSuccessStatusCode) { - try - { - await this.ThrowIfRequestFailedAsync(response, shell); - } - catch (Exception ex) when (ex is not OperationCanceledException) - { - return new ErrorCommandState(ex) { RequestCharge = requestCharge > 0 ? requestCharge : null }; - } + RequestChargeContext.Record(requestCharge); + await this.ThrowIfRequestFailedAsync(response, shell); } var cumulative = response?.Diagnostics.GetQueryMetrics()?.CumulativeMetrics; @@ -723,15 +717,8 @@ private async Task ExecuteQueryAsync(Container container, ShellInt var pageRequestCharge = response.Headers.RequestCharge; if (!response.IsSuccessStatusCode) { - try - { - await this.ThrowIfRequestFailedAsync(response, shell); - } - catch (Exception ex) when (ex is not OperationCanceledException) - { - var failedCharge = totalRequestCharge + pageRequestCharge; - return new ErrorCommandState(ex) { RequestCharge = failedCharge > 0 ? failedCharge : null }; - } + RequestChargeContext.Record(totalRequestCharge + pageRequestCharge); + await this.ThrowIfRequestFailedAsync(response, shell); } if (response.Content == null) diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ReplaceCommand.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ReplaceCommand.cs index 4eeebc0c..3f33c428 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ReplaceCommand.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ReplaceCommand.cs @@ -130,6 +130,7 @@ private static async Task ReplaceArrayAsync(Container container, if (failCount > 0) { + RequestChargeContext.Record(charge); throw new CommandException( "replace", MessageService.GetArgsString( diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/RmCommand.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/RmCommand.cs index 6cdf9656..d054c7f1 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/RmCommand.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/RmCommand.cs @@ -249,7 +249,7 @@ private async Task RemoveItemsFromContainerAsync(ConnectedState state, totalCharge += response.Headers.RequestCharge; using var streamReader = new StreamReader(response.Content); - var queryDocument = JsonDocument.Parse(await streamReader.ReadToEndAsync()); + using var queryDocument = JsonDocument.Parse(await streamReader.ReadToEndAsync()); foreach (var element in queryDocument.RootElement.GetProperty("Documents").EnumerateArray()) { diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/ShellInterpreter.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/ShellInterpreter.cs index 772e1fdd..84156f07 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/ShellInterpreter.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/ShellInterpreter.cs @@ -2003,7 +2003,7 @@ private void SetSessionMaxRequestCharge(ShellObject value) double requestChargeTotal; lock (this.sessionRequestChargeLock) { - if (maximum != this.sessionMaxRequestCharge) + if (Math.Abs(maximum - this.sessionMaxRequestCharge) > 1e-9) { this.sessionRequestChargeWarningIssued = false; } From ae3cdb185d9a4536de895207ead60c2094fe651f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mike=20Kr=C3=BCger?= Date: Tue, 1 Sep 2026 12:00:40 +0200 Subject: [PATCH 16/16] Clarify session request charge warning threshold Rename the unreleased `$sessionMaxRequestCharge` setting to `$sessionRequestChargeWarningThreshold` so it cannot be mistaken for a hard cap. Rename the corresponding `info` structured field to `session.requestChargeWarningThreshold` and update completion, localization, tests, changelog, and documentation. --- CHANGELOG.md | 2 +- .../CommandTests/InfoCommandTests.cs | 4 +- .../CommandTests/SessionRequestChargeTests.cs | 20 ++++----- .../Lsp/CosmosShellCompletionHandlerTests.cs | 2 +- .../InfoCommand.cs | 10 ++--- .../ShellInterpreter.cs | 44 +++++++++---------- CosmosDBShell/lang/en.ftl | 6 +-- docs/commands.md | 2 +- docs/mcp.md | 4 +- docs/programming.md | 4 +- 10 files changed, 49 insertions(+), 49 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 9c857a62..bc76659a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,7 +5,7 @@ ### 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 `$sessionMaxRequestCharge` to a positive RU threshold to print one warning when the current connection reaches it; `info` reports the configured maximum as `session.maxRequestCharge`. +- 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 diff --git a/CosmosDBShell.Tests/CommandTests/InfoCommandTests.cs b/CosmosDBShell.Tests/CommandTests/InfoCommandTests.cs index c276e6e1..3e8ab943 100644 --- a/CosmosDBShell.Tests/CommandTests/InfoCommandTests.cs +++ b/CosmosDBShell.Tests/CommandTests/InfoCommandTests.cs @@ -210,13 +210,13 @@ public void AddSessionUsage_IncludesCurrentCommandCharge() public void AddSessionUsage_IncludesConfiguredMaximum() { using var shell = ShellInterpreter.CreateInstance(); - shell.SetVariable("sessionMaxRequestCharge", new ShellDecimal(25.5)); + shell.SetVariable("sessionRequestChargeWarningThreshold", new ShellDecimal(25.5)); var result = new Dictionary(); InfoCommand.AddSessionUsage(shell, result, renderOutput: false); var session = JsonSerializer.SerializeToElement(result).GetProperty("session"); - Assert.Equal(25.5, session.GetProperty("maxRequestCharge").GetDouble()); + Assert.Equal(25.5, session.GetProperty("requestChargeWarningThreshold").GetDouble()); } [Fact] diff --git a/CosmosDBShell.Tests/CommandTests/SessionRequestChargeTests.cs b/CosmosDBShell.Tests/CommandTests/SessionRequestChargeTests.cs index c8943f55..bcd025e8 100644 --- a/CosmosDBShell.Tests/CommandTests/SessionRequestChargeTests.cs +++ b/CosmosDBShell.Tests/CommandTests/SessionRequestChargeTests.cs @@ -21,7 +21,7 @@ public void SessionVariables_ReflectCurrentUsage() var charge = Assert.IsType(shell.GetVariable("sessionRequestCharge")); var operationCount = Assert.IsType(shell.GetVariable("sessionChargedOperationCount")); - var maximum = Assert.IsType(shell.GetVariable("sessionMaxRequestCharge")); + var maximum = Assert.IsType(shell.GetVariable("sessionRequestChargeWarningThreshold")); Assert.Equal(3.75, charge.Value); Assert.Equal(2, operationCount.Value); @@ -38,10 +38,10 @@ public void SessionUsageVariables_AreReadOnly() } [Fact] - public void SessionMaxRequestCharge_WarnsOnlyOnceWhenReached() + public void SessionRequestChargeWarningThreshold_WarnsOnlyOnceWhenReached() { using var shell = ShellInterpreter.CreateInstance(); - shell.SetVariable("sessionMaxRequestCharge", new ShellDecimal(3)); + shell.SetVariable("sessionRequestChargeWarningThreshold", new ShellDecimal(3)); var output = CaptureConsole(() => { @@ -54,10 +54,10 @@ public void SessionMaxRequestCharge_WarnsOnlyOnceWhenReached() } [Fact] - public void SessionMaxRequestCharge_ZeroDisablesWarning() + public void SessionRequestChargeWarningThreshold_ZeroDisablesWarning() { using var shell = ShellInterpreter.CreateInstance(); - shell.SetVariable("sessionMaxRequestCharge", new ShellNumber(0)); + shell.SetVariable("sessionRequestChargeWarningThreshold", new ShellNumber(0)); var output = CaptureConsole(() => shell.RecordRequestCharge(new CommandState { RequestCharge = 10 })); @@ -68,23 +68,23 @@ public void SessionMaxRequestCharge_ZeroDisablesWarning() public void Connect_PreservesMaximumAndRearmsWarning() { using var shell = ShellInterpreter.CreateInstance(); - shell.SetVariable("sessionMaxRequestCharge", new ShellNumber(2)); + shell.SetVariable("sessionRequestChargeWarningThreshold", new ShellNumber(2)); _ = CaptureConsole(() => shell.RecordRequestCharge(new CommandState { RequestCharge = 2 })); shell.Connect(CreateTestClient(), credentialTypeOverride: "AccountKey"); var output = CaptureConsole(() => shell.RecordRequestCharge(new CommandState { RequestCharge = 2 })); - Assert.Equal(2, Assert.IsType(shell.GetVariable("sessionMaxRequestCharge")).Value); + Assert.Equal(2, Assert.IsType(shell.GetVariable("sessionRequestChargeWarningThreshold")).Value); Assert.Equal(1, CountOccurrences(output, "has reached the configured warning threshold")); } [Fact] - public void SessionMaxRequestCharge_RejectsInvalidValues() + public void SessionRequestChargeWarningThreshold_RejectsInvalidValues() { using var shell = ShellInterpreter.CreateInstance(); - Assert.Throws(() => shell.SetVariable("sessionMaxRequestCharge", new ShellDecimal(-1))); - Assert.Throws(() => shell.SetVariable("sessionMaxRequestCharge", new ShellText("ten"))); + Assert.Throws(() => shell.SetVariable("sessionRequestChargeWarningThreshold", new ShellDecimal(-1))); + Assert.Throws(() => shell.SetVariable("sessionRequestChargeWarningThreshold", new ShellText("ten"))); } private static CosmosClient CreateTestClient() => new( diff --git a/CosmosDBShell.Tests/Lsp/CosmosShellCompletionHandlerTests.cs b/CosmosDBShell.Tests/Lsp/CosmosShellCompletionHandlerTests.cs index 6775dba6..ffab4518 100644 --- a/CosmosDBShell.Tests/Lsp/CosmosShellCompletionHandlerTests.cs +++ b/CosmosDBShell.Tests/Lsp/CosmosShellCompletionHandlerTests.cs @@ -95,7 +95,7 @@ public async Task VariableCompletion_SuggestsSessionVariables() Assert.Contains("$sessionRequestCharge", labels); Assert.Contains("$sessionChargedOperationCount", labels); - Assert.Contains("$sessionMaxRequestCharge", labels); + Assert.Contains("$sessionRequestChargeWarningThreshold", labels); } [Fact] diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/InfoCommand.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/InfoCommand.cs index db3ff111..e55b80b4 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/InfoCommand.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/InfoCommand.cs @@ -312,9 +312,9 @@ internal static void AddSessionUsage(ShellInterpreter shell, Dictionary 0) + if (sessionUsage.RequestChargeWarningThreshold > 0) { - session["maxRequestCharge"] = sessionUsage.MaxRequestCharge; + session["requestChargeWarningThreshold"] = sessionUsage.RequestChargeWarningThreshold; } mcpTable["session"] = session; @@ -334,11 +334,11 @@ internal static void AddSessionUsage(ShellInterpreter shell, Dictionary 0) + if (sessionUsage.RequestChargeWarningThreshold > 0) { table.AddRow( - MessageService.GetString("command-stats-session-max-request-charge"), - Theme.FormatTableValue(sessionUsage.MaxRequestCharge.ToString("0.##", CultureInfo.InvariantCulture))); + MessageService.GetString("command-stats-session-request-charge-warning-threshold"), + Theme.FormatTableValue(sessionUsage.RequestChargeWarningThreshold.ToString("0.##", CultureInfo.InvariantCulture))); } AnsiConsole.Write(table); diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/ShellInterpreter.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/ShellInterpreter.cs index 84156f07..b46d1927 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/ShellInterpreter.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/ShellInterpreter.cs @@ -29,7 +29,7 @@ public partial class ShellInterpreter : IDisposable private const string SessionChargedOperationCountVariable = "sessionChargedOperationCount"; - private const string SessionMaxRequestChargeVariable = "sessionMaxRequestCharge"; + private const string SessionRequestChargeWarningThresholdVariable = "sessionRequestChargeWarningThreshold"; internal static readonly ShellInterpreter Instance = new(); @@ -76,7 +76,7 @@ public partial class ShellInterpreter : IDisposable private long sessionChargedOperationCount; - private double sessionMaxRequestCharge; + private double sessionRequestChargeWarningThreshold; private bool sessionRequestChargeWarningIssued; @@ -157,7 +157,7 @@ internal static char CSVSeparator [ SessionRequestChargeVariable, SessionChargedOperationCountVariable, - SessionMaxRequestChargeVariable, + SessionRequestChargeWarningThresholdVariable, ]; internal Dictionary Functions { get; } = []; @@ -204,13 +204,13 @@ internal long SessionChargedOperationCount } } - internal (double RequestCharge, long ChargedOperationCount, double MaxRequestCharge) SessionUsage + internal (double RequestCharge, long ChargedOperationCount, double RequestChargeWarningThreshold) SessionUsage { get { lock (this.sessionRequestChargeLock) { - return (this.sessionRequestCharge, this.sessionChargedOperationCount, this.sessionMaxRequestCharge); + return (this.sessionRequestCharge, this.sessionChargedOperationCount, this.sessionRequestChargeWarningThreshold); } } } @@ -749,9 +749,9 @@ internal ShellObject GetVariable(string name) return new ShellDecimal(this.sessionChargedOperationCount); } - if (string.Equals(name, SessionMaxRequestChargeVariable, StringComparison.OrdinalIgnoreCase)) + if (string.Equals(name, SessionRequestChargeWarningThresholdVariable, StringComparison.OrdinalIgnoreCase)) { - return new ShellDecimal(this.sessionMaxRequestCharge); + return new ShellDecimal(this.sessionRequestChargeWarningThreshold); } } @@ -1027,7 +1027,7 @@ internal void RecordRequestCharge(CommandState commandState, long generation) { bool printWarning = false; double requestChargeTotal = 0; - double requestChargeMaximum = 0; + double requestChargeWarningThreshold = 0; if (commandState.RequestCharge is { } requestCharge) { lock (this.sessionRequestChargeLock) @@ -1040,14 +1040,14 @@ internal void RecordRequestCharge(CommandState commandState, long generation) this.sessionChargedOperationCount++; } - if (this.sessionMaxRequestCharge > 0 + if (this.sessionRequestChargeWarningThreshold > 0 && !this.sessionRequestChargeWarningIssued - && this.sessionRequestCharge >= this.sessionMaxRequestCharge) + && this.sessionRequestCharge >= this.sessionRequestChargeWarningThreshold) { this.sessionRequestChargeWarningIssued = true; printWarning = true; requestChargeTotal = this.sessionRequestCharge; - requestChargeMaximum = this.sessionMaxRequestCharge; + requestChargeWarningThreshold = this.sessionRequestChargeWarningThreshold; } } } @@ -1055,7 +1055,7 @@ internal void RecordRequestCharge(CommandState commandState, long generation) if (printWarning) { - this.PrintSessionRequestChargeWarning(requestChargeTotal, requestChargeMaximum); + this.PrintSessionRequestChargeWarning(requestChargeTotal, requestChargeWarningThreshold); } } @@ -1937,9 +1937,9 @@ internal void SetVariable(string variableName, ShellObject value) throw new ShellException(MessageService.GetArgsString("error-session-variable-read-only", "name", variableName)); } - if (string.Equals(variableName, SessionMaxRequestChargeVariable, StringComparison.OrdinalIgnoreCase)) + if (string.Equals(variableName, SessionRequestChargeWarningThresholdVariable, StringComparison.OrdinalIgnoreCase)) { - this.SetSessionMaxRequestCharge(value); + this.SetSessionRequestChargeWarningThreshold(value); return; } @@ -1985,7 +1985,7 @@ internal void SetVariable(string variableName, ShellObject value) currentScope.Set(variableName, shellValue); } - private void SetSessionMaxRequestCharge(ShellObject value) + private void SetSessionRequestChargeWarningThreshold(ShellObject value) { double maximum = value switch { @@ -1996,19 +1996,19 @@ private void SetSessionMaxRequestCharge(ShellObject value) if (!double.IsFinite(maximum) || maximum < 0) { - throw new ShellException(MessageService.GetString("error-session-max-request-charge-invalid")); + throw new ShellException(MessageService.GetString("error-session-request-charge-warning-threshold-invalid")); } bool printWarning; double requestChargeTotal; lock (this.sessionRequestChargeLock) { - if (Math.Abs(maximum - this.sessionMaxRequestCharge) > 1e-9) + if (Math.Abs(maximum - this.sessionRequestChargeWarningThreshold) > 1e-9) { this.sessionRequestChargeWarningIssued = false; } - this.sessionMaxRequestCharge = maximum; + this.sessionRequestChargeWarningThreshold = maximum; printWarning = maximum > 0 && !this.sessionRequestChargeWarningIssued && this.sessionRequestCharge >= maximum; @@ -2026,14 +2026,14 @@ private void SetSessionMaxRequestCharge(ShellObject value) } } - private void PrintSessionRequestChargeWarning(double requestCharge, double maximum) + private void PrintSessionRequestChargeWarning(double requestCharge, double warningThreshold) { var message = MessageService.GetArgsString( - "warning-session-max-request-charge-reached", + "warning-session-request-charge-threshold-reached", "requestCharge", requestCharge.ToString("0.##", CultureInfo.InvariantCulture), - "maximum", - maximum.ToString("0.##", CultureInfo.InvariantCulture)); + "warningThreshold", + warningThreshold.ToString("0.##", CultureInfo.InvariantCulture)); if (this.IsMachineMode) { diff --git a/CosmosDBShell/lang/en.ftl b/CosmosDBShell/lang/en.ftl index 98d96d19..53b2521f 100644 --- a/CosmosDBShell/lang/en.ftl +++ b/CosmosDBShell/lang/en.ftl @@ -934,10 +934,10 @@ command-stats-account-detailed-cost-note = Aggregating account totals reads ever command-stats-session-heading = Session Usage command-stats-session-request-charge = Observed request charge (RUs) command-stats-session-charged-operations = Charged operations -command-stats-session-max-request-charge = Warning threshold (RUs) -warning-session-max-request-charge-reached = Session request charge { $requestCharge } RUs has reached the configured warning threshold of { $maximum } RUs. +command-stats-session-request-charge-warning-threshold = Warning threshold (RUs) +warning-session-request-charge-threshold-reached = Session request charge { $requestCharge } RUs has reached the configured warning threshold of { $warningThreshold } RUs. error-session-variable-read-only = Variable '${ $name }' is read-only. -error-session-max-request-charge-invalid = Variable '$sessionMaxRequestCharge' must be a non-negative number. Set it to 0 to disable the warning. +error-session-request-charge-warning-threshold-invalid = Variable '$sessionRequestChargeWarningThreshold' must be a non-negative number. Set it to 0 to disable the warning. command-version-description = Displays the version of Cosmos DB Shell. command-version = Cosmos Shell version: { $version } diff --git a/docs/commands.md b/docs/commands.md index 61e5b2df..a8662d3b 100644 --- a/docs/commands.md +++ b/docs/commands.md @@ -1363,7 +1363,7 @@ total is available as `session.requestCharge`, and `session.chargedOperationCoun counts command operations that reported a positive charge. Paginated and multi-request commands contribute their aggregate observed charge. Azure Resource Manager control-plane operations do not consume Cosmos DB request units. When -`$sessionMaxRequestCharge` is positive, `session.maxRequestCharge` reports the +`$sessionRequestChargeWarningThreshold` is positive, `session.requestChargeWarningThreshold` reports the configured warning threshold. This command is read-only. ### help diff --git a/docs/mcp.md b/docs/mcp.md index 3d7b5ce5..f40ec6df 100644 --- a/docs/mcp.md +++ b/docs/mcp.md @@ -116,6 +116,6 @@ containers does not reset it. This session value is telemetry rather than a budget or billing total. `session.chargedOperationCount` counts command operations that reported a positive request charge; it counts command operations rather than individual query pages or transactional batch items. If -the shell variable `$sessionMaxRequestCharge` is set to a positive value, the -session object also includes `session.maxRequestCharge`. +the shell variable `$sessionRequestChargeWarningThreshold` is set to a positive value, the +session object also includes `session.requestChargeWarningThreshold`. diff --git a/docs/programming.md b/docs/programming.md index 723d7316..c92c38e4 100644 --- a/docs/programming.md +++ b/docs/programming.md @@ -80,9 +80,9 @@ The shell provides three built-in session variables: | --- | --- | | `$sessionRequestCharge` | Read-only cumulative request charge observed during the current connection. | | `$sessionChargedOperationCount` | Read-only number of command operations that reported a positive charge. | -| `$sessionMaxRequestCharge` | Configurable warning threshold in RUs. A positive value enables the warning; `0` disables it. | +| `$sessionRequestChargeWarningThreshold` | Configurable warning threshold in RUs. A positive value enables the warning; `0` disables it. | -For example, `$sessionMaxRequestCharge = 100` prints one warning when the +For example, `$sessionRequestChargeWarningThreshold = 100` prints one warning when the current connection reaches or exceeds 100 observed RUs. The warning is emitted only once for that threshold. A successful `connect` resets the accumulated charge and operation count and rearms the warning, while preserving the