diff --git a/CHANGELOG.md b/CHANGELOG.md index 89fefb6..32d8d1c 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)) - **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/ListCommandTests.cs b/CosmosDBShell.Tests/CommandTests/ListCommandTests.cs index 478ca6e..579bd1e 100644 --- a/CosmosDBShell.Tests/CommandTests/ListCommandTests.cs +++ b/CosmosDBShell.Tests/CommandTests/ListCommandTests.cs @@ -223,4 +223,15 @@ public async Task ReadQueryResponseAsync_ValidContent_ReturnsJsonDocument() var item = Assert.Single(document.RootElement.GetProperty("Documents").EnumerateArray()); Assert.Equal("1", item.GetProperty("id").GetString()); } + + [Fact] + public void AccumulateRequestCharge_AddsEveryPageCharge() + { + var state = new CommandState(); + + ListCommand.AccumulateRequestCharge(state, 1.25); + ListCommand.AccumulateRequestCharge(state, 2.5); + + Assert.Equal(3.75, state.RequestCharge); + } } diff --git a/CosmosDBShell.Tests/CommandTests/SprocCommandExecutionTests.cs b/CosmosDBShell.Tests/CommandTests/SprocCommandExecutionTests.cs index 31f4f78..a7afe48 100644 --- a/CosmosDBShell.Tests/CommandTests/SprocCommandExecutionTests.cs +++ b/CosmosDBShell.Tests/CommandTests/SprocCommandExecutionTests.cs @@ -206,6 +206,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 ba6d4d2..4957cee 100644 --- a/CosmosDBShell.Tests/McpResponseFactoryTests.cs +++ b/CosmosDBShell.Tests/McpResponseFactoryTests.cs @@ -165,4 +165,52 @@ public void CreateError_PopulatesStructuredContentMatchingTextBlock() Assert.Equal("boom", structured.GetProperty("error").GetString()); Assert.Equal("/TestDatabase", structured.GetProperty("currentLocation").GetString()); } + + [Fact] + public void CreateSuccess_IncludesRequestChargeWhenSet() + { + var commandState = new CommandState + { + Result = new ShellJson(JsonSerializer.SerializeToElement(new { result = "success" })), + RequestCharge = 4.25, + }; + + var result = McpResponseFactory.CreateSuccess(commandState, new ConnectedState(null!)); + + Assert.NotNull(result.StructuredContent); + var structured = result.StructuredContent!.Value; + Assert.True(structured.TryGetProperty("requestCharge", out var requestCharge)); + Assert.Equal(4.25, requestCharge.GetDouble()); + } + + [Fact] + public void CreateSuccess_StructuredError_IncludesRequestChargeWhenSet() + { + var commandState = new StructuredErrorCommandState( + new CommandException("batch", "Batch failed."), + new ShellJson(JsonSerializer.SerializeToElement(new { success = false }))) + { + RequestCharge = 3.5, + }; + + var result = McpResponseFactory.CreateSuccess(commandState, new ConnectedState(null!)); + + Assert.True(result.IsError); + Assert.NotNull(result.StructuredContent); + Assert.Equal(3.5, result.StructuredContent!.Value.GetProperty("requestCharge").GetDouble()); + } + + [Fact] + public void CreateSuccess_OmitsRequestChargeWhenNotSet() + { + var commandState = new CommandState + { + Result = new ShellJson(JsonSerializer.SerializeToElement(new { result = "success" })), + }; + + var result = McpResponseFactory.CreateSuccess(commandState, new ConnectedState(null!)); + + Assert.NotNull(result.StructuredContent); + Assert.False(result.StructuredContent!.Value.TryGetProperty("requestCharge", out _)); + } } \ 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 6cbf24e..010bf10 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ExportCommand.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ExportCommand.cs @@ -111,6 +111,7 @@ public override async Task ExecuteAsync(ShellInterpreter shell, Co exported = count, requestCharge = charge, })), + RequestCharge = charge, }; } diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ImportCommand.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ImportCommand.cs index 88f2800..eb4ab44 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ImportCommand.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ImportCommand.cs @@ -548,6 +548,7 @@ public override async Task ExecuteAsync(ShellInterpreter shell, Co requestCharge = charge, dryRun, })), + RequestCharge = charge, }; } diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ListCommand.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ListCommand.cs index 8bc8e76..81b3ca7 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ListCommand.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ListCommand.cs @@ -229,6 +229,7 @@ private async Task ListContainerItemsAsync(ConnectedState state, S { using var response = await feedIterator.ReadNextAsync(token); using var queryDocument = await ReadQueryResponseAsync(response, token); + AccumulateRequestCharge(returnState, response.Headers.RequestCharge); foreach (var element in queryDocument.RootElement.GetProperty("Documents").EnumerateArray()) { @@ -280,6 +281,11 @@ private async Task ListContainerItemsAsync(ConnectedState state, S return returnState; } + internal static void AccumulateRequestCharge(CommandState commandState, double requestCharge) + { + commandState.RequestCharge = (commandState.RequestCharge ?? 0) + requestCharge; + } + internal static async Task ReadQueryResponseAsync(ResponseMessage response, CancellationToken token) { if (!response.IsSuccessStatusCode) diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/MakeItemCommand.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/MakeItemCommand.cs index 8fb3011..788e1ce 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/MakeItemCommand.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/MakeItemCommand.cs @@ -59,7 +59,7 @@ public async override Task ExecuteAsync(ShellInterpreter shell, Co "mkitem", token); - var summary = await WriteItemAsync(container, commandState, jsonOpt, this.Force == true, token); + var summary = await WriteItemAsync(container, jsonOpt, this.Force == true, token); var returnState = new CommandState(); returnState.Result = new ShellJson(JsonSerializer.SerializeToElement(new @@ -70,6 +70,7 @@ public async override Task ExecuteAsync(ShellInterpreter shell, Co failed = summary.Failed, requestCharge = summary.RequestCharge, })); + returnState.RequestCharge = summary.RequestCharge; return returnState; } @@ -131,7 +132,7 @@ private static object ParseJsonElement(JsonElement element) } } - private static async Task WriteItemAsync(Container container, CommandState commandState, string? jsonOpt, bool force, CancellationToken token) + private static async Task WriteItemAsync(Container container, string? jsonOpt, bool force, CancellationToken token) { if (!string.IsNullOrEmpty(jsonOpt)) { @@ -154,7 +155,6 @@ private static async Task WriteItemAsync(Container container, Comm ? await container.UpsertItemAsync(element, cancellationToken: token) : await container.CreateItemAsync(element, cancellationToken: token); charge += result.RequestCharge; - if (result.StatusCode == System.Net.HttpStatusCode.Created) { createdCount++; diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/PatchCommand.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/PatchCommand.cs index cc684d5..0a9e2d3 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/PatchCommand.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/PatchCommand.cs @@ -133,6 +133,7 @@ public override async Task ExecuteAsync(ShellInterpreter shell, Co patched = true, requestCharge = response.RequestCharge, })), + 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 b988799..9ed343d 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 80e3d89..1dba912 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/QueryCommand.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/QueryCommand.cs @@ -677,6 +677,7 @@ private async Task ExecuteQueryAsync(Container container, ShellInt { var returnState = CreateCommandState(this.OutputFormat); var aggregatedDocuments = new List(); + double totalRequestCharge = 0; var options = new QueryRequestOptions { @@ -727,11 +728,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) - { - 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("F2", CultureInfo.InvariantCulture) } })); var pageDocuments = queryDocument.RootElement.GetProperty("Documents"); var pageExceedsLimit = PageExceedsLimit(aggregatedDocuments.Count, pageDocuments, effectiveMaxItemCount); @@ -751,7 +755,7 @@ private async Task ExecuteQueryAsync(Container container, ShellInt { { "type", "item" }, { "values", aggregatedDocuments }, - { "requestCharge", queryMetrics?.TotalRequestCharge ?? 0 }, + { "requestCharge", totalRequestCharge }, { "queryMetrics", metricProperty }, { "indexMetrics", parsedIndexMetrics ?? new Dictionary() }, }); @@ -888,6 +892,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 fc89552..d3a95ff 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ReplaceCommand.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ReplaceCommand.cs @@ -63,6 +63,7 @@ public override async Task ExecuteAsync(ShellInterpreter shell, Co failed = summary.Failed, requestCharge = summary.RequestCharge, })), + RequestCharge = summary.RequestCharge, }; } diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/RmCommand.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/RmCommand.cs index df0c9a2..87278b9 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/RmCommand.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/RmCommand.cs @@ -118,25 +118,26 @@ private async Task RemoveItemsFromContainerAsync(ConnectedState state, var matchKeyPropertyNames = string.IsNullOrEmpty(this.Key) ? partitionKeyPropertyNames : [this.Key]; var totalCount = 0; + double totalCharge = 0; bool dryRun = this.DryRun == true; // In dry-run mode, count what would be deleted without issuing any delete. - async Task TryDeleteAsync(string id, PartitionKey partitionKey) + async Task<(bool Counted, double RequestCharge)> TryDeleteAsync(string id, PartitionKey partitionKey) { if (dryRun) { - return true; + return (true, 0); } try { - await container.DeleteItemAsync(id, partitionKey, cancellationToken: token); - return true; + var deleteResponse = await container.DeleteItemAsync(id, partitionKey, cancellationToken: token); + return (true, deleteResponse.RequestCharge); } catch (CosmosException ex) when (ex.StatusCode == System.Net.HttpStatusCode.NotFound) { // Item was already deleted, skip - return false; + return (false, ex.RequestCharge); } } @@ -183,7 +184,9 @@ async Task TryDeleteAsync(string id, PartitionKey partitionKey) if (id != null && shouldDelete) { - if (await TryDeleteAsync(id, CreatePartitionKey(pkElements))) + var deleteResult = await TryDeleteAsync(id, CreatePartitionKey(pkElements)); + totalCharge += deleteResult.RequestCharge; + if (deleteResult.Counted) { totalCount++; } @@ -214,7 +217,9 @@ async Task TryDeleteAsync(string id, PartitionKey partitionKey) var id = idElement.GetString(); if (id != null) { - if (await TryDeleteAsync(id, CreatePartitionKey(pkElements))) + var deleteResult = await TryDeleteAsync(id, CreatePartitionKey(pkElements)); + totalCharge += deleteResult.RequestCharge; + if (deleteResult.Counted) { totalCount++; } @@ -236,7 +241,13 @@ async Task TryDeleteAsync(string id, PartitionKey partitionKey) break; } - var response = await feedIterator.ReadNextAsync(token); + using var response = await feedIterator.ReadNextAsync(token); + + // The scan pages that locate matching items consume RUs regardless of whether + // any item is ultimately deleted (including in --dry-run), so include each + // page's request charge from the response headers. + totalCharge += response.Headers.RequestCharge; + using var streamReader = new StreamReader(response.Content); var queryDocument = JsonDocument.Parse(await streamReader.ReadToEndAsync()); @@ -272,7 +283,9 @@ async Task TryDeleteAsync(string id, PartitionKey partitionKey) if (shouldDelete) { - if (await TryDeleteAsync(id, CreatePartitionKey(pkElements))) + var deleteResult = await TryDeleteAsync(id, CreatePartitionKey(pkElements)); + totalCharge += deleteResult.RequestCharge; + if (deleteResult.Counted) { totalCount++; } @@ -296,6 +309,7 @@ async Task TryDeleteAsync(string id, PartitionKey partitionKey) commandState.Result = new ShellJson(JsonSerializer.SerializeToElement(new { type = "item", count = totalCount, dryRun })); commandState.RenderUser = () => AnsiConsole.MarkupLine(renderMessage); + commandState.RequestCharge = totalCharge; return new ExitCode(0); } diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/SprocCommand.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/SprocCommand.cs index e9801ca..4889266 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 5a206df..2dbe33d 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/CommandState.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/CommandState.cs @@ -69,6 +69,12 @@ public OutputFormat OutputFormat /// internal Func? RenderTabular { 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 7a19753..ef02702 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); diff --git a/docs/mcp.md b/docs/mcp.md index 10b89a2..3221ed5 100644 --- a/docs/mcp.md +++ b/docs/mcp.md @@ -100,8 +100,9 @@ Both representations are always byte-for-byte equivalent. | ----- | ------------ | ----------- | | `result` | Commands that produce output | The command result as JSON (objects, arrays, or a scalar). Text-only results are represented as a JSON string. Failed transactional batches include their per-operation summary here alongside `error`. | | `outputText` | CSV output commands with non-empty text | The CSV rendering of the result. Omitted when the CSV output is empty or whitespace. | +| `requestCharge` | Instrumented data-plane command results | The Cosmos DB request charge (in RUs) consumed by the command, as a number. | | `error` | Failed commands | The error message. | | `currentLocation` | Always | The shell's current navigation path (for example `/MyDatabase/MyContainer`), or `null` when disconnected. | -Successful results set `result` (and optionally `outputText`); failed results set `error`, may also include a structured `result`, and mark the tool result as an error. `currentLocation` is always included so a client can track navigation state across calls. +Successful results set `result` (and optionally `outputText`); failed results set `error`, may also include a structured `result`, and mark the tool result as an error. `currentLocation` is always included so a client can track navigation state across calls. Instrumented data-plane commands (`query`, `print`, `ls` for container items, `mkitem`, `replace`, `patch`, `rm`, `import`, and `export`) additionally set `requestCharge` when available so a client can track RU cost across successful and structured-error results.