Surface Cosmos request charge (RUs) in MCP structured results

[Mike Krüger (mkrueger)]

Fixes #162

Summary

Surfaces the aggregate request charge observed by Cosmos DB data-plane commands as a uniform top-level numeric requestCharge field in MCP and structured output. Commands that issue multiple requests aggregate their observed charges; ARM control-plane operations do not report Cosmos DB RUs.

Coverage

Coverage includes queries and --explain, point reads and item writes, transactional batches, import/export, ls and resource discovery, cd validation, info, indexing/TTL/conflict/throughput metadata, stored procedures, UDFs, triggers, change feed watching, permission probes, pagination, cancellation, and charged failures.

Generic shell and MCP failures preserve charges already observed before the error. Existing command-specific result shapes remain compatible.

Session usage and guardrail

• $sessionRequestCharge exposes the cumulative observed request charge for the current connection.
• $sessionChargedOperationCount exposes the number of command operations that reported a positive charge.
• $sessionMaxRequestCharge configures an optional RU warning threshold. The shell warns once when usage reaches or exceeds it; 0 disables the warning.
• A successful connect resets usage and rearms the warning while preserving the configured maximum. Navigation and disconnect do not reset usage.
• info reports session.requestCharge, session.chargedOperationCount, and, when enabled, session.maxRequestCharge.

This satisfies #162 using its allowed warning semantics. It intentionally does not implement hard-abort behavior or a per-call --max-ru option.

Implementation

• CommandState.RequestCharge carries the common command-level value.
• A per-command RequestChargeContext collects charges from shared helpers, pagination, handled failures, and cancellation paths.
• A connection-generation guard prevents commands from an earlier connection contributing to a newer session.
• MCP structured content and its JSON text representation remain equivalent.

Validation

• Full Release suite: 2,354 total, 2,256 passed, 98 emulator-dependent skipped, 0 failed.
• Focused review-fix suite: 161 passed, 16 emulator-dependent skipped, 0 failed.
• Release build: 0 warnings, 0 errors.

[github-code-quality]

Code Coverage Overview

Languages: C#

C# / code-coverage/dotnet

The overall line coverage in commit 41ea350 in the dev/mkrueger/request... branch remains at 61%, unchanged from commit 4bc063e in the main branch.

Show a line coverage summary of the most impacted files.

File	main 4bc063e	dev/mkrueger/request... 41ea350	+/-
D:\a\CosmosDBSh...ds\RmCommand.cs	19%	17%	-2%
D:\a\CosmosDBSh...olOperations.cs	90%	89%	-1%
D:\a\CosmosDBSh...QueryCommand.cs	41%	41%	0%
D:\a\CosmosDBSh...s\UdfCommand.cs	43%	44%	+1%
D:\a\CosmosDBSh...iggerCommand.cs	47%	48%	+1%
D:\a\CosmosDBSh...SprocCommand.cs	51%	52%	+1%
D:\a\CosmosDBSh...\InfoCommand.cs	12%	14%	+2%
D:\a\CosmosDBSh...lInterpreter.cs	63%	66%	+3%
D:\a\CosmosDBSh...atchExecutor.cs	8%	17%	+9%
D:\a\CosmosDBSh...hargeContext.cs	0%	93%	+93%

---

_{Updated August 31, 2026 13:13 UTC}

[Copilot]

Pull request overview

This PR adds a consistent way for data-plane commands to report Cosmos DB request charge (RUs) by storing it on CommandState and emitting it as a top-level requestCharge field in MCP tool-result payloads (both structuredContent and the legacy JSON text block).

Changes:

• Add CommandState.RequestCharge and include it as requestCharge in MCP success payloads when set.
• Retrofit multiple data-plane commands (query, print, mkitem, replace, patch, rm, import, export, sproc exec) to capture RU charge.
• Document the MCP payload shape and add/extend tests validating structuredContent parity and requestCharge behavior.

Reviewed changes

Copilot reviewed 15 out of 15 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
docs/mcp.md	Documents tool-result payload shape and the new requestCharge metadata field.
CosmosDBShell/Azure.Data.Cosmos.Shell.Mcp/McpResponseFactory.cs	Emits StructuredContent and adds requestCharge to the MCP success payload when available.
CosmosDBShell/Azure.Data.Cosmos.Shell.Core/CommandState.cs	Introduces the nullable RequestCharge field for commands to populate.
CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/QueryCommand.cs	Attempts to aggregate query RU charges and store them on CommandState.
CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/PrintCommand.cs	Captures point-read RU charge from response headers.
CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/MakeItemCommand.cs	Accumulates RU charges across single/bulk writes and propagates to returned state.
CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ReplaceCommand.cs	Returns total replace charge (single/bulk) and sets it on the returned state.
CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/PatchCommand.cs	Sets request charge from patch response.
CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/RmCommand.cs	Tracks delete RU charges and sets them on the command state.
CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ImportCommand.cs	Sets aggregated RU charge on success.
CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ExportCommand.cs	Sets aggregated RU charge on success.
CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/SprocCommand.cs	Captures RU charge from stored procedure execution.
CosmosDBShell.Tests/McpResponseFactoryTests.cs	Adds tests for StructuredContent population/parity and requestCharge presence/omission.
CosmosDBShell.Tests/CommandTests/SprocCommandExecutionTests.cs	Extends sproc exec test to assert RequestCharge is captured.
CHANGELOG.md	Adds entries describing structured MCP results and RU surfacing.

[Copilot]

Pull request overview

Copilot reviewed 15 out of 15 changed files in this pull request and generated no new comments.

[Copilot]

Pull request overview

Copilot reviewed 15 out of 15 changed files in this pull request and generated 3 comments.

Suppressed comments (5)

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

• SprocCommand is marked Restricted = true without Confirmable, and ToolOperations rejects the sproc tool before invoking ExecAsync. Therefore this assignment is never observed by McpResponseFactory, while the PR documentation claims sproc exec reports requestCharge in MCP results. Either make this command callable/confirmable over MCP or remove it from this MCP-specific contract and the associated documentation.

            commandState.RequestCharge = response.RequestCharge;

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

• Container-scoped ls is also an MCP-exposed data-plane read: ListContainerItemsAsync runs GetItemQueryStreamIterator and consumes query pages, but it never populates CommandState.RequestCharge. As a result, the new factory still omits requestCharge for ls, so the documented uniform data-plane contract is incomplete. Accumulate each page's request charge and assign it to the returned state (with a focused test).

        if (commandState.RequestCharge.HasValue)
        {
            payload["requestCharge"] = commandState.RequestCharge.Value;
        }

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

• Error responses bypass this block: ToolOperations catches command exceptions and calls CreateError, whose payload contains only error and currentLocation. A CosmosException can carry a non-zero RequestCharge even when a read or write fails, so those MCP results still omit the cost despite the contract describing requestCharge for data-plane operations. Propagate the charge through the error path, or explicitly document that the field is success-only.

        if (commandState.RequestCharge.HasValue)
        {
            payload["requestCharge"] = commandState.RequestCharge.Value;
        }

docs/mcp.md:105

• The documented list is incomplete for the stated uniform data-plane contract: ls item listing, watch, and info --partitions/--detailed also issue item/change-feed queries, but their returned states never populate CommandState.RequestCharge (see ListCommand.cs:228-230, WatchCommand.cs:173-175, and InfoCommand.cs:490-550). Clients therefore still receive no RU field for these MCP calls. Instrument those paths as well, or explicitly narrow this contract to the commands listed here.

Successful results set `result` (and optionally `outputText`); failed results set `error` and mark the tool result as an error. `currentLocation` is always included so a client can track navigation state across calls. Data-plane commands (`query`, `print`, `mkitem`, `replace`, `patch`, `rm`, `import`, `export`, and `sproc exec`) additionally set `requestCharge` so a client can track RU cost across calls.

docs/mcp.md:105

• sproc exec cannot be invoked through MCP: SprocCommand is marked restricted without a confirmable path, and ToolOperations rejects it before execution. This sentence therefore overstates the client-facing MCP contract by saying clients can track requestCharge for sproc exec; remove it from this MCP list or change the command's MCP availability consistently.

Successful results set `result` (and optionally `outputText`); failed results set `error` and mark the tool result as an error. `currentLocation` is always included so a client can track navigation state across calls. Data-plane commands (`query`, `print`, `mkitem`, `replace`, `patch`, `rm`, `import`, `export`, and `sproc exec`) additionally set `requestCharge` so a client can track RU cost across calls.

[Copilot]

Pull request overview

Copilot reviewed 17 out of 17 changed files in this pull request and generated 4 comments.

[Copilot]

Pull request overview

Copilot reviewed 17 out of 17 changed files in this pull request and generated no new comments.

Suppressed comments (1)

Previously missed (1) — in code that hasn't changed since the last review.

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

• requestCharge is only emitted for non-error command states because the commandState.IsError early-return happens before the RequestCharge block. If a command captures RU cost but returns an error CommandState (for example, partial failures that still have a request charge), the MCP payload will omit requestCharge, undermining the goal of uniform RU reporting.

        if (commandState.RequestCharge.HasValue)
        {
            payload["requestCharge"] = commandState.RequestCharge.Value;
        }

[Copilot]

Pull request overview

Copilot reviewed 29 out of 29 changed files in this pull request and generated 1 comment.

Suppressed comments (1)

Previously missed (1) — in code that hasn't changed since the last review.

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

• JsonDocument.Parse(...) allocates pooled buffers and should be disposed. In the scan loop this document is currently never disposed, which can increase memory usage for large rm scans.

                using var streamReader = new StreamReader(response.Content);
                var queryDocument = JsonDocument.Parse(await streamReader.ReadToEndAsync());

[Copilot]

Pull request overview

Copilot reviewed 32 out of 32 changed files in this pull request and generated 1 comment.

[Copilot]

Pull request overview

Copilot reviewed 32 out of 32 changed files in this pull request and generated no new comments.

Suppressed comments (1)

Previously missed (1) — in code that hasn't changed since the last review.

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

• In ExecuteExplainAsync, the HasMoreResults ? ReadNextAsync(...) : null pattern can yield response == null, which would return a successful explain result with requestCharge implicitly 0 even though no request was made (and the rest of the method assumes a response exists to compute plan/index metrics). Since EXPLAIN must execute at least one page, always issue the first request and remove the nullable response path.

            using ResponseMessage? response = feedIterator.HasMoreResults
                ? await feedIterator.ReadNextAsync(token)
                : null;
            double requestCharge = response?.Headers.RequestCharge ?? 0;
            if (response is not null && !response.IsSuccessStatusCode)

[Copilot]

Pull request overview

Copilot reviewed 32 out of 32 changed files in this pull request and generated no new comments.

Suppressed comments (4)

Previously missed (4) — in code that hasn't changed since the last review.

CosmosDBShell/Azure.Data.Cosmos.Shell.Mcp/ToolOperations.cs:495

• The MCP exception path returns McpResponseFactory.CreateError(...), which drops CommandState.RequestCharge and any structured error details for commands that throw (e.g., patch/replace/mkitem failures). This means MCP responses can omit the RU charge even though Cosmos returned it in response headers/exceptions, and session totals won’t account for it either.

        try
        {
            ShellInterpreter.Instance.PrintCommand(sb.ToString());
            var response = await ShellInterpreter.Instance.ExecuteCosmosCommandAsync(cmd, new CommandState(), command.CommandName, cancellationToken);
            ShellInterpreter.Instance.CancelPrompt();

CosmosDBShell/Azure.Data.Cosmos.Shell.Core/ShellInterpreter.cs:966

• ExecuteCosmosCommandAsync only records request charge on successful completion. If a command throws after issuing a billable request (common for data-plane errors like 404/412/etc.), session.requestCharge and session.chargedOperationCount will underreport because RecordRequestCharge is never reached.

        // CommandState is intentionally reused by pipelines and expressions. Clear the
        // previous command's charge so an uninstrumented command cannot count it again.
        commandState.RequestCharge = null;
        var generation = this.SessionRequestChargeGeneration;
        var result = await command.ExecuteAsync(this, commandState, commandText, token);

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

• ReplaceCommand now surfaces summary.RequestCharge on success, but array-mode failures throw CommandException after accumulating charge (ReplaceArrayAsync), so MCP/tool callers will lose the observed RU cost and session totals won’t include it for partial/failed replaces. This contradicts the new observability contract for instrumented data-plane commands.

                failed = summary.Failed,
                requestCharge = summary.RequestCharge,
            })),
            RequestCharge = summary.RequestCharge,
        };

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

• In array mode, WriteItemAsync accumulates RU charge as it writes items, but throws CommandException when any item fails (lines 260-270 in this file). Because the exception doesn’t carry the accumulated charge, MCP/tool callers will lose requestCharge and session totals won’t include the cost for partial/failed mkitem runs.

            replaced = summary.Replaced,
            failed = summary.Failed,
            requestCharge = summary.RequestCharge,
        }));
        returnState.RequestCharge = summary.RequestCharge;

[Copilot]

Pull request overview

Copilot reviewed 37 out of 37 changed files in this pull request and generated 5 comments.

Suppressed comments (3)

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

• Returning an ErrorCommandState here will suppress query failures in the interactive shell because ErrorCommandState carries no Result/RenderUser and PrintState won’t render the exception message. Prefer letting ThrowIfRequestFailedAsync throw so ShellInterpreter can report the error, but record the accumulated RU charge into RequestChargeContext first so MCP/session accounting still captures it.

                    catch (Exception ex) when (ex is not OperationCanceledException)
                    {
                        var failedCharge = totalRequestCharge + pageRequestCharge;
                        return new ErrorCommandState(ex) { RequestCharge = failedCharge > 0 ? failedCharge : null };
                    }

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

• These branches return ErrorCommandState, which won’t print anything in interactive mode (no Result/RenderUser), effectively hiding print errors. Record the response-header RU charge into RequestChargeContext and throw the CommandException so ShellInterpreter renders the error while MCP can still include requestCharge from the exception context.

            else if (response.StatusCode == System.Net.HttpStatusCode.NotFound)
            {
                return new ErrorCommandState(new CommandException("print", MessageService.GetString("command-print-error-item_not_found", new Dictionary<string, object>
                {
                    { "id", this.Id ?? "(null)" },

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

• Catching CosmosException and returning ErrorCommandState here will also suppress error output in interactive mode. Throw the CommandException instead so ShellInterpreter reports it; ExecuteCosmosCommandAsync can still extract RU charge from the CosmosException for MCP/session accounting.

        catch (CosmosException ex)
        {
            return new ErrorCommandState(new CommandException("print", MessageService.GetArgsString("command-print-error-reading_item", "message", CommandException.GetDisplayMessage(ex)), ex))
            {
                RequestCharge = ex.RequestCharge > 0 ? ex.RequestCharge : null,
            };

[Copilot]

Pull request overview

Copilot reviewed 37 out of 37 changed files in this pull request and generated no new comments.

Suppressed comments (1)

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

• In the array replace path, successful item replacements add to the local charge total but (on partial failure) that total is lost when the method throws. Only failure charges recorded via RequestChargeContext.Record(...) survive into the charged error, so requestCharge underreports the RU cost when some items succeed before the array ultimately fails.

            catch (CommandException ex)
            {
                RequestChargeContext.Record(RequestChargeContext.GetCosmosExceptionCharge(ex));
                failCount++;
                ShellInterpreter.WriteLine(ex.Message);

[Copilot]

Pull request overview

Copilot reviewed 41 out of 41 changed files in this pull request and generated no new comments.

Suppressed comments (4)

Previously missed (2) — in code that hasn't changed since the last review.

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

• returnState.RequestCharge is always set, which can force MCP to emit a top-level requestCharge: 0 when the input is an empty array (no Cosmos requests issued). Per the documented contract, the top-level field should be omitted when no billable request was observed.

This issue also appears on line 175 of the same file.

            failed = summary.Failed,
            requestCharge = summary.RequestCharge,
        }));
        returnState.RequestCharge = summary.RequestCharge;
        return returnState;

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

• RequestCharge = summary.RequestCharge can cause MCP to emit a top-level requestCharge: 0 when the input array is empty (no Cosmos requests issued). The top-level field should be omitted when no billable request was observed.

                replaced = summary.Replaced,
                failed = summary.Failed,
                requestCharge = summary.RequestCharge,
            })),
            RequestCharge = summary.RequestCharge,

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

• In replace array mode, the command throws after iterating when any item fails. Because successful item charges are only kept in the local charge variable (not recorded in RequestChargeContext), the structured error/session requestCharge will miss the RU cost of successful items on partial failures.

                successCount++;
            }
            catch (CommandException ex)
            {
                RequestChargeContext.Record(RequestChargeContext.GetCosmosExceptionCharge(ex));

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

• In the bulk mkitem path, array failures are surfaced by throwing after the loop. Since only CosmosException charges are recorded into RequestChargeContext, the successful per-item request charges accumulated in charge are not included in the eventual structured error/session requestCharge, underreporting RUs for partial failures.

                        }
                        catch (CosmosException ce)
                        {
                            RequestChargeContext.Record(ce.RequestCharge);
                            failCount++;

[Copilot]

Pull request overview

Copilot reviewed 41 out of 41 changed files in this pull request and generated no new comments.