diff --git a/CHANGELOG.md b/CHANGELOG.md index 89fefb63..bc76659a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,11 @@ ## Unreleased +### Improvements + +- Cosmos DB data-plane commands now consistently expose their aggregate observed request charge in structured output and connection-scoped `info` telemetry, including metadata/configuration operations, scripts, change feed reads, paginated operations, handled probes, and charged failures. Azure Resource Manager control-plane operations remain uncharged. +- Added `$sessionRequestCharge` and `$sessionChargedOperationCount` as read-only shell variables. Set `$sessionRequestChargeWarningThreshold` to a positive RU threshold to print one warning when the current connection reaches it; `info` reports it as `session.requestChargeWarningThreshold`. + ### Fixes - Local emulator outages are now detected across Cosmos DB commands. Requests fail promptly with an error and return the shell to its disconnected state instead of leaving an unresponsive session labeled as connected. @@ -90,7 +95,8 @@ A focused cycle on top of 1.1.115-preview. New `ttl` and `conflict` commands man ### Improvements - **Structured (JSON) tool results for MCP.** MCP tool results now carry the machine-readable JSON payload (`result`/`outputText`/`error` plus `currentLocation`) as first-class `structuredContent` in addition to the existing JSON text block, so agents can consume structured results directly. The two representations are kept byte-for-byte equivalent, and text-only clients are unaffected. ([#154](https://github.com/Azure/CosmosDBShell/issues/154)) - +- **Request charge in MCP structured results.** Instrumented data-plane commands (`query`, including `--explain`; `print`; container-scoped `ls`; `can-i` probes; `batch run`; `mkitem`; `replace`; `patch`; `rm`; `import`; and `export`) now report the Cosmos DB request charge (in RUs) consumed by the operation as a uniform `requestCharge` field on the MCP tool result, so agents can track observed RU cost consistently across calls. Budget enforcement remains tracked separately in #162. ([#162](https://github.com/Azure/CosmosDBShell/issues/162)) +- **Connection-scoped request-charge totals.** The shell accumulates request charges observed from instrumented commands and reports the total in `info` as `session.requestCharge`, together with the number of positively charged command operations as `session.chargedOperationCount`. A successful `connect` starts new totals; database and container navigation do not reset them. This is usage telemetry, not budget enforcement or billing data. ([#162](https://github.com/Azure/CosmosDBShell/issues/162)) - **Destructive MCP commands now prompt for confirmation instead of being blocked.** When an MCP client invokes `delete`, `rm`, `rmcon`, or `rmdb`, the server sends an elicitation prompt describing the exact command line and only runs it if the user approves; declining, cancelling, or a client that cannot confirm results in nothing being executed. This removes the need for any write opt-in flag. ([#158](https://github.com/Azure/CosmosDBShell/issues/158)) ### Fixes 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/CommandTests/CosmosCommandTests.cs b/CosmosDBShell.Tests/CommandTests/CosmosCommandTests.cs index db152da8..23948011 100644 --- a/CosmosDBShell.Tests/CommandTests/CosmosCommandTests.cs +++ b/CosmosDBShell.Tests/CommandTests/CosmosCommandTests.cs @@ -11,6 +11,67 @@ namespace CosmosShell.Tests.CommandTests; public class CosmosCommandTests { + [Fact] + public async Task ExecuteCosmosCommandAsync_RecordsChargeAndClearsStaleCharge() + { + using var shell = ShellInterpreter.CreateInstance(); + var state = new CommandState { RequestCharge = 9 }; + + state = await shell.ExecuteCosmosCommandAsync(new TestCosmosCommand(2.5), state, string.Empty, CancellationToken.None); + state = await shell.ExecuteCosmosCommandAsync(new TestCosmosCommand(null), state, string.Empty, CancellationToken.None); + + Assert.Null(state.RequestCharge); + Assert.Equal(2.5, shell.SessionRequestCharge); + Assert.Equal(1, shell.SessionChargedOperationCount); + } + + [Fact] + public async Task ExecuteCosmosCommandAsync_CombinesExplicitAndScopedChargesOnce() + { + using var shell = ShellInterpreter.CreateInstance(); + + var state = await shell.ExecuteCosmosCommandAsync( + new TestCosmosCommand(2.5, scopedRequestCharge: 1.25), + new CommandState(), + string.Empty, + CancellationToken.None); + + Assert.Equal(3.75, state.RequestCharge); + Assert.Equal(3.75, shell.SessionRequestCharge); + Assert.Equal(1, shell.SessionChargedOperationCount); + } + + [Fact] + public async Task ExecuteCosmosCommandAsync_ThrownCommandPreservesScopedCharge() + { + using var shell = ShellInterpreter.CreateInstance(); + + var exception = await Assert.ThrowsAsync(() => 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() { @@ -42,8 +103,37 @@ public void CreatePartitionKey_WithHierarchicalIntegerComponents_PreservesIntege private sealed class TestCosmosCommand : CosmosCommand { + private readonly double? requestCharge; + private readonly double scopedRequestCharge; + private readonly bool throwAfterRecording; + private readonly bool cancelAfterRecording; + + public TestCosmosCommand( + double? requestCharge = null, + double scopedRequestCharge = 0, + bool throwAfterRecording = false, + bool cancelAfterRecording = false) + { + this.requestCharge = requestCharge; + this.scopedRequestCharge = scopedRequestCharge; + this.throwAfterRecording = throwAfterRecording; + this.cancelAfterRecording = cancelAfterRecording; + } + public override Task 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 f13243d2..3e8ab943 100644 --- a/CosmosDBShell.Tests/CommandTests/InfoCommandTests.cs +++ b/CosmosDBShell.Tests/CommandTests/InfoCommandTests.cs @@ -4,8 +4,10 @@ namespace CosmosShell.Tests.CommandTests; +using System.Text.Json; using Azure.Data.Cosmos.Shell.Commands; using Azure.Data.Cosmos.Shell.Core; +using Azure.Data.Cosmos.Shell.Parser; using Azure.Data.Cosmos.Shell.States; using Azure.Data.Cosmos.Shell.Util; using Microsoft.Azure.Cosmos; @@ -135,6 +137,101 @@ public void FormatSize_Null_ReturnsNotAvailable() Assert.Equal(MessageService.GetString("command-stats-na"), InfoCommand.FormatSize(null)); } + [Fact] + public void SessionRequestCharge_AccumulatesObservedCharges() + { + using var shell = ShellInterpreter.CreateInstance(); + + shell.RecordRequestCharge(new CommandState { RequestCharge = 1.25 }); + shell.RecordRequestCharge(new CommandState()); + shell.RecordRequestCharge(new CommandState { RequestCharge = 0 }); + shell.RecordRequestCharge(new ErrorCommandState(new InvalidOperationException()) { RequestCharge = 2.5 }); + + Assert.Equal(3.75, shell.SessionRequestCharge); + Assert.Equal(2, shell.SessionChargedOperationCount); + } + + [Fact] + public void Connect_ResetsSessionRequestCharge() + { + using var shell = ShellInterpreter.CreateInstance(); + shell.RecordRequestCharge(new CommandState { RequestCharge = 4.5 }); + + shell.Connect(CreateTestClient(), credentialTypeOverride: "AccountKey"); + + Assert.Equal(0, shell.SessionRequestCharge); + Assert.Equal(0, shell.SessionChargedOperationCount); + } + + [Fact] + public void Disconnect_DoesNotResetSessionRequestCharge() + { + using var shell = ShellInterpreter.CreateInstance(); + shell.Connect(CreateTestClient(), credentialTypeOverride: "AccountKey"); + shell.RecordRequestCharge(new CommandState { RequestCharge = 4.5 }); + + shell.Disconnect(); + + Assert.Equal(4.5, shell.SessionRequestCharge); + Assert.Equal(1, shell.SessionChargedOperationCount); + } + + [Fact] + public void AddSessionUsage_AddsCurrentChargeToStructuredResult() + { + using var shell = ShellInterpreter.CreateInstance(); + shell.RecordRequestCharge(new CommandState { RequestCharge = 3.75 }); + var result = new Dictionary(); + + InfoCommand.AddSessionUsage(shell, result, renderOutput: false); + + var json = JsonSerializer.SerializeToElement(result); + Assert.Equal(3.75, json.GetProperty("session").GetProperty("requestCharge").GetDouble()); + Assert.Equal(1, json.GetProperty("session").GetProperty("chargedOperationCount").GetInt64()); + } + + [Fact] + public void AddSessionUsage_IncludesCurrentCommandCharge() + { + using var shell = ShellInterpreter.CreateInstance(); + shell.RecordRequestCharge(new CommandState { RequestCharge = 3.75 }); + using var scope = RequestChargeContext.Begin(); + RequestChargeContext.Record(1.25); + var result = new Dictionary(); + + InfoCommand.AddSessionUsage(shell, result, renderOutput: false); + + var json = JsonSerializer.SerializeToElement(result); + Assert.Equal(5, json.GetProperty("session").GetProperty("requestCharge").GetDouble()); + Assert.Equal(2, json.GetProperty("session").GetProperty("chargedOperationCount").GetInt64()); + } + + [Fact] + public void AddSessionUsage_IncludesConfiguredMaximum() + { + using var shell = ShellInterpreter.CreateInstance(); + shell.SetVariable("sessionRequestChargeWarningThreshold", new ShellDecimal(25.5)); + var result = new Dictionary(); + + InfoCommand.AddSessionUsage(shell, result, renderOutput: false); + + var session = JsonSerializer.SerializeToElement(result).GetProperty("session"); + Assert.Equal(25.5, session.GetProperty("requestChargeWarningThreshold").GetDouble()); + } + + [Fact] + public void RecordRequestCharge_IgnoresPriorConnectionGeneration() + { + using var shell = ShellInterpreter.CreateInstance(); + var generation = shell.SessionRequestChargeGeneration; + shell.Connect(CreateTestClient(), credentialTypeOverride: "AccountKey"); + + shell.RecordRequestCharge(new CommandState { RequestCharge = 8 }, generation); + + Assert.Equal(0, shell.SessionRequestCharge); + Assert.Equal(0, shell.SessionChargedOperationCount); + } + [Theory] [InlineData("table")] [InlineData("tbl")] diff --git a/CosmosDBShell.Tests/CommandTests/ListCommandTests.cs b/CosmosDBShell.Tests/CommandTests/ListCommandTests.cs index 478ca6e4..579bd1e9 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/SessionRequestChargeTests.cs b/CosmosDBShell.Tests/CommandTests/SessionRequestChargeTests.cs new file mode 100644 index 00000000..bcd025e8 --- /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("sessionRequestChargeWarningThreshold")); + + 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 SessionRequestChargeWarningThreshold_WarnsOnlyOnceWhenReached() + { + using var shell = ShellInterpreter.CreateInstance(); + shell.SetVariable("sessionRequestChargeWarningThreshold", 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 SessionRequestChargeWarningThreshold_ZeroDisablesWarning() + { + using var shell = ShellInterpreter.CreateInstance(); + shell.SetVariable("sessionRequestChargeWarningThreshold", 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("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("sessionRequestChargeWarningThreshold")).Value); + Assert.Equal(1, CountOccurrences(output, "has reached the configured warning threshold")); + } + + [Fact] + public void SessionRequestChargeWarningThreshold_RejectsInvalidValues() + { + using var shell = ShellInterpreter.CreateInstance(); + + Assert.Throws(() => shell.SetVariable("sessionRequestChargeWarningThreshold", new ShellDecimal(-1))); + Assert.Throws(() => shell.SetVariable("sessionRequestChargeWarningThreshold", 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/CommandTests/SprocCommandExecutionTests.cs b/CosmosDBShell.Tests/CommandTests/SprocCommandExecutionTests.cs index 31f4f78f..a7afe48e 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/Integration/BatchOperationTests.cs b/CosmosDBShell.Tests/Integration/BatchOperationTests.cs index 758ba7cd..833d8962 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 is > 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 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 0ced42e7..2ac8ef18 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 is > 0); using var document = JsonDocument.Parse(output); var root = document.RootElement; diff --git a/CosmosDBShell.Tests/Lsp/CosmosShellCompletionHandlerTests.cs b/CosmosDBShell.Tests/Lsp/CosmosShellCompletionHandlerTests.cs index ffc7b2ab..ffab4518 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("$sessionRequestChargeWarningThreshold", labels); + } + [Fact] public async Task VariableCompletion_IgnoresWhenNotVariableContext() { diff --git a/CosmosDBShell.Tests/McpResponseFactoryTests.cs b/CosmosDBShell.Tests/McpResponseFactoryTests.cs index ba6d4d2c..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() { @@ -165,4 +176,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/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/ExportCommand.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ExportCommand.cs index 6cbf24ed..010bf10c 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 88f28009..7178b9ea 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 > 0 ? charge : null, }; } diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/InfoCommand.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/InfoCommand.cs index 8a8e5599..e55b80b4 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,9 +301,53 @@ private static async Task WriteAccountDatabaseBreakdownAsync(ConnectedState stat AnsiConsole.Write(databaseTable); } + 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); + var session = new Dictionary + { + ["requestCharge"] = requestCharge, + ["chargedOperationCount"] = chargedOperationCount, + }; + if (sessionUsage.RequestChargeWarningThreshold > 0) + { + session["requestChargeWarningThreshold"] = sessionUsage.RequestChargeWarningThreshold; + } + + mcpTable["session"] = session; + + 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(requestCharge.ToString("0.##", CultureInfo.InvariantCulture))); + table.AddRow( + MessageService.GetString("command-stats-session-charged-operations"), + Theme.FormatTableValue(chargedOperationCount.ToString(CultureInfo.InvariantCulture))); + if (sessionUsage.RequestChargeWarningThreshold > 0) + { + table.AddRow( + MessageService.GetString("command-stats-session-request-charge-warning-threshold"), + Theme.FormatTableValue(sessionUsage.RequestChargeWarningThreshold.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); + RequestChargeContext.Record(response.RequestCharge); return ParseResourceUsage(response.Headers[ResourceUsageHeader]); } @@ -316,15 +360,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) @@ -490,7 +538,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); @@ -547,7 +597,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; } @@ -655,7 +707,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 +991,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 +1080,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 +1124,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.Commands/ListCommand.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ListCommand.cs index 8bc8e76b..9ac995f7 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ListCommand.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ListCommand.cs @@ -228,24 +228,37 @@ 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) + { + RequestChargeContext.Record(returnState.RequestCharge ?? 0); + throw new CommandException("ls", ex); + } - 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; + } } } @@ -280,6 +293,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 8fb30111..1ed15ff9 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++; @@ -175,6 +175,7 @@ private static async Task WriteItemAsync(Container container, Comm } catch (CosmosException ce) { + RequestChargeContext.Record(ce.RequestCharge); failCount++; ShellInterpreter.WriteLine( MessageService.GetArgsString( diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/PatchCommand.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/PatchCommand.cs index cc684d5e..0a9e2d31 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 b9887990..194cc972 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/PrintCommand.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/PrintCommand.cs @@ -54,6 +54,7 @@ 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) { @@ -61,8 +62,8 @@ private async Task PrintItemAsync(Container container, Cancellatio 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) { @@ -83,6 +84,7 @@ private async Task PrintItemAsync(Container container, Cancellatio } catch (CosmosException ex) { + RequestChargeContext.Record(ex.RequestCharge); throw new CommandException("print", MessageService.GetArgsString("command-print-error-reading_item", "message", CommandException.GetDisplayMessage(ex)), ex); } diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/QueryCommand.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/QueryCommand.cs index 80e3d891..786caa6b 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/QueryCommand.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/QueryCommand.cs @@ -631,13 +631,14 @@ 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) { + RequestChargeContext.Record(requestCharge); await this.ThrowIfRequestFailedAsync(response, shell); } var cumulative = response?.Diagnostics.GetQueryMetrics()?.CumulativeMetrics; - double requestCharge = response?.Diagnostics.GetQueryMetrics()?.TotalRequestCharge ?? 0; var (planAvailable, utilized, potential) = ParseIndexPlan(response?.IndexMetrics); var evaluation = EvaluatePlan( @@ -651,6 +652,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) @@ -677,6 +679,7 @@ private async Task ExecuteQueryAsync(Container container, ShellInt { var returnState = CreateCommandState(this.OutputFormat); var aggregatedDocuments = new List(); + double totalRequestCharge = 0; var options = new QueryRequestOptions { @@ -711,7 +714,12 @@ 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) + { + RequestChargeContext.Record(totalRequestCharge + pageRequestCharge); + await this.ThrowIfRequestFailedAsync(response, shell); + } if (response.Content == null) { @@ -727,11 +735,13 @@ 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. + 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 +761,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 +898,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 fc89552f..3f33c428 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, }; } @@ -108,6 +109,7 @@ private static async Task ReplaceArrayAsync(Container container, } catch (CommandException ex) { + RequestChargeContext.Record(RequestChargeContext.GetCosmosExceptionCharge(ex)); failCount++; ShellInterpreter.WriteLine(ex.Message); } @@ -128,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 df0c9a23..d054c7f1 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,9 +241,15 @@ 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()); + using var queryDocument = JsonDocument.Parse(await streamReader.ReadToEndAsync()); foreach (var element in queryDocument.RootElement.GetProperty("Documents").EnumerateArray()) { @@ -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 > 0 ? totalCharge : null; 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 e9801ca9..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) @@ -446,6 +457,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) @@ -461,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", @@ -489,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) @@ -507,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/CommandState.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/CommandState.cs index 5a206dfe..2dbe33d9 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.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 c4f38f33..b46d1927 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 SessionRequestChargeWarningThresholdVariable = "sessionRequestChargeWarningThreshold"; + internal static readonly ShellInterpreter Instance = new(); private const int MAXHISTORYITEMS = 60; @@ -48,6 +54,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 +72,16 @@ public partial class ShellInterpreter : IDisposable private List history; + private double sessionRequestCharge; + + private long sessionChargedOperationCount; + + private double sessionRequestChargeWarningThreshold; + + private bool sessionRequestChargeWarningIssued; + + private long sessionRequestChargeGeneration; + internal ShellInterpreter(string? configPath = null) { this.State = new DisconnectedState(); @@ -135,6 +153,13 @@ internal static char CSVSeparator } } + internal static IReadOnlyList SessionVariableNames { get; } = + [ + SessionRequestChargeVariable, + SessionChargedOperationCountVariable, + SessionRequestChargeWarningThresholdVariable, + ]; + internal Dictionary Functions { get; } = []; /// @@ -150,6 +175,57 @@ 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; + } + } + } + + /// + /// 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, double RequestChargeWarningThreshold) SessionUsage + { + get + { + lock (this.sessionRequestChargeLock) + { + return (this.sessionRequestCharge, this.sessionChargedOperationCount, this.sessionRequestChargeWarningThreshold); + } + } + } + + internal long SessionRequestChargeGeneration + { + get + { + lock (this.sessionRequestChargeLock) + { + return this.sessionRequestChargeGeneration; + } + } + } + internal string HistoryFile { get; private set; } internal string WelcomeMarkerFile => this.welcomeMarkerFile; @@ -490,7 +566,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; } @@ -658,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, SessionRequestChargeWarningThresholdVariable, StringComparison.OrdinalIgnoreCase)) + { + return new ShellDecimal(this.sessionRequestChargeWarningThreshold); + } + } + var scope = this.GetScope(name); if (scope?.TryGetValue(name, out var value) == true) { @@ -888,6 +985,80 @@ internal async Task RunCommandAsync(CommandState currentState, str return currentState; } + 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; + 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) + { + bool printWarning = false; + double requestChargeTotal = 0; + double requestChargeWarningThreshold = 0; + if (commandState.RequestCharge is { } requestCharge) + { + lock (this.sessionRequestChargeLock) + { + if (generation == this.sessionRequestChargeGeneration) + { + this.sessionRequestCharge += requestCharge; + if (requestCharge > 0) + { + this.sessionChargedOperationCount++; + } + + if (this.sessionRequestChargeWarningThreshold > 0 + && !this.sessionRequestChargeWarningIssued + && this.sessionRequestCharge >= this.sessionRequestChargeWarningThreshold) + { + this.sessionRequestChargeWarningIssued = true; + printWarning = true; + requestChargeTotal = this.sessionRequestCharge; + requestChargeWarningThreshold = this.sessionRequestChargeWarningThreshold; + } + } + } + } + + if (printWarning) + { + this.PrintSessionRequestChargeWarning(requestChargeTotal, requestChargeWarningThreshold); + } + } + 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 +1663,14 @@ internal void Connect(CosmosClient client, ArmCosmosContext? armContext = null, { this.State?.Dispose(); this.State = new ConnectedState(client, armContext); + lock (this.sessionRequestChargeLock) + { + this.sessionRequestCharge = 0; + this.sessionChargedOperationCount = 0; + this.sessionRequestChargeWarningIssued = false; + this.sessionRequestChargeGeneration++; + } + this.activeCredential = credential; this.ActiveCredentialType = credentialTypeOverride ?? credential?.GetType().Name; this.CurrentBatch = null; @@ -1752,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, SessionRequestChargeWarningThresholdVariable, StringComparison.OrdinalIgnoreCase)) + { + this.SetSessionRequestChargeWarningThreshold(value); + return; + } + // Ensure we have at least one variable container (global scope) if (this.VariableContainers.Count == 0) { @@ -1794,6 +1985,70 @@ internal void SetVariable(string variableName, ShellObject value) currentScope.Set(variableName, shellValue); } + private void SetSessionRequestChargeWarningThreshold(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-request-charge-warning-threshold-invalid")); + } + + bool printWarning; + double requestChargeTotal; + lock (this.sessionRequestChargeLock) + { + if (Math.Abs(maximum - this.sessionRequestChargeWarningThreshold) > 1e-9) + { + this.sessionRequestChargeWarningIssued = false; + } + + this.sessionRequestChargeWarningThreshold = 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 warningThreshold) + { + var message = MessageService.GetArgsString( + "warning-session-request-charge-threshold-reached", + "requestCharge", + requestCharge.ToString("0.##", CultureInfo.InvariantCulture), + "warningThreshold", + warningThreshold.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/Azure.Data.Cosmos.Shell.Mcp/McpResponseFactory.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Mcp/McpResponseFactory.cs index 7a197539..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) @@ -83,6 +86,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/CosmosDBShell/Azure.Data.Cosmos.Shell.Mcp/ToolOperations.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Mcp/ToolOperations.cs index 6f39542f..f49f490f 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); } @@ -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/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 43364df4..fae24f8a 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Parser/Statement/CommandStatement.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Parser/Statement/CommandStatement.cs @@ -186,7 +186,7 @@ public override async Task RunAsync(ShellInterpreter shell, Comman } 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); } if (File.Exists(this.Name)) diff --git a/CosmosDBShell/lang/en.ftl b/CosmosDBShell/lang/en.ftl index 726c82f0..53b2521f 100644 --- a/CosmosDBShell/lang/en.ftl +++ b/CosmosDBShell/lang/en.ftl @@ -931,6 +931,13 @@ 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 (RUs) +command-stats-session-charged-operations = Charged operations +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-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 bb9e9090..a8662d3b 100644 --- a/docs/commands.md +++ b/docs/commands.md @@ -1330,7 +1330,11 @@ 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 +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. On serverless accounts, throughput/offer settings are not available, so the scale section reports that throughput settings are not available for serverless @@ -1354,7 +1358,13 @@ 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`, 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. When +`$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 10b89a25..f40ec6df 100644 --- a/docs/mcp.md +++ b/docs/mcp.md @@ -100,8 +100,22 @@ 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` | 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. +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 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 +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 `$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 1b8450a2..c92c38e4 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. | +| `$sessionRequestChargeWarningThreshold` | Configurable warning threshold in RUs. A positive value enables the warning; `0` disables it. | + +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 +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