[BUG] `Cosmos` recordings do not work properly

[scbedd]

Describe the bug

During playback for database query operations, cosmos test recording is not working the way it should. Due to the headers of the response of the queryIterator here AND that the response from the service encodes a binary content blob (which the cosmos client understands how to parse), the test-proxy is breaking the content by storing on disk improperly.

The test works properly during playback if the result content is actually JSON:

I had to manually fix up the content you see above to make playback work.

Expected behavior

Recordings should be written with proper JSON content.

Actual behavior

Instead, due to the way the queryIterator works, the content comes back in CosmosBinary format. However, the content is marked as Content-Type: json in the response headers, so the proxy doesn't understand that it needs to Base64 encode this pure binary content to store it at rest, so it comes out all wonky:

However, I believe during the record/playback round-trip, the stream is broken in a way the queryIterator is not expecting, so we get back a null content stream according to the query iterator, and it throws with the follwing test error detail:

 Azure.Mcp.Tools.Cosmos.LiveTests.CosmosCommandTests.Should_list_and_query_multiple_databases_and_containers
   Source: CosmosCommandTests.cs line 169
   Duration: 6.7 sec

  Message: 
Property 'items' not found. Full element: '{"message":"Error querying items: Value cannot be null. (Parameter \u0027utf8Json\u0027)","stackTrace":"   at Azure.Mcp.Tools.Cosmos.Services.CosmosService.QueryItems(String accountName, String databaseName, String containerName, String query, String subscription, AuthMethod authMethod, String tenant, RetryPolicyOptions retryPolicy, CancellationToken cancellationToken) in C:\\repo\\mcp\\tools\\Azure.Mcp.Tools.Cosmos\\src\\Services\\CosmosService.cs:line 340\r\n   at Azure.Mcp.Tools.Cosmos.Commands.ItemQueryCommand.ExecuteAsync(CommandContext context, ParseResult parseResult, CancellationToken cancellationToken) in C:\\repo\\mcp\\tools\\Azure.Mcp.Tools.Cosmos\\src\\Commands\\ItemQueryCommand.cs:line 71","type":"Exception"}'

  Stack Trace: 
TestExtensions.AssertProperty(JsonElement element, String propertyName) line 38
TestExtensions.AssertProperty(Nullable`1 element, String propertyName) line 33
CosmosCommandTests.Should_list_and_query_multiple_databases_and_containers() line 214
--- End of stack trace from previous location ---

  Standard Output: 
Applying custom matcher to global settings: CompareBodies=, IgnoreQueryOrdering=, IgnoredHeaders=x-ms-activity-id,x-ms-cosmos-correlated-activityid
Attempting to start MCP Client
MCP client initialized successfully
[Playback] Session file: tools/Azure.Mcp.Tools.Cosmos/tests/Azure.Mcp.Tools.Cosmos.LiveTests/SessionRecords/Azure.Mcp.Tools.Cosmos.LiveTests.CosmosCommandTests.Should_list_and_query_multiple_databases_and_containers.json
[Playback] Recording ID: 4fbe7d7b-f4ab-401b-bbf7-c38be94d0cef
[Playback] Variable from recording: database0 = ToDoList
[Playback] Variable from recording: db0/container0 = Items
request: {"command":"cosmos_database_list","parameters":{"subscription":"00000000-0000-0000-0000-000000000000","account":"Sanitized"}}
response: {
  "status": 200,
  "message": "Success",
  "results": {
    "databases": [
      "ToDoList"
    ]
  },
  "duration": 0
}
request: {"command":"cosmos_database_container_list","parameters":{"subscription":"00000000-0000-0000-0000-000000000000","account":"Sanitized","database":"ToDoList"}}
response: {
  "status": 200,
  "message": "Success",
  "results": {
    "containers": [
      "Items"
    ]
  },
  "duration": 0
}
request: {"command":"cosmos_database_container_item_query","parameters":{"subscription":"00000000-0000-0000-0000-000000000000","account":"Sanitized","database":"ToDoList","container":"Items"}}
response: {
  "status": 500,
  "message": "Error querying items: Value cannot be null. (Parameter \u0027utf8Json\u0027). To mitigate this issue, please refer to the troubleshooting guidelines here at https://aka.ms/azmcp/troubleshooting.",
  "results": {
    "message": "Error querying items: Value cannot be null. (Parameter \u0027utf8Json\u0027)",
    "stackTrace": "   at Azure.Mcp.Tools.Cosmos.Services.CosmosService.QueryItems(String accountName, String databaseName, String containerName, String query, String subscription, AuthMethod authMethod, String tenant, RetryPolicyOptions retryPolicy, CancellationToken cancellationToken) in C:\\repo\\mcp\\tools\\Azure.Mcp.Tools.Cosmos\\src\\Services\\CosmosService.cs:line 340\r\n   at Azure.Mcp.Tools.Cosmos.Commands.ItemQueryCommand.ExecuteAsync(CommandContext context, ParseResult parseResult, CancellationToken cancellationToken) in C:\\repo\\mcp\\tools\\Azure.Mcp.Tools.Cosmos\\src\\Commands\\ItemQueryCommand.cs:line 71",
    "type": "Exception"
  },
  "duration": 0
}

Reproduction Steps

Re-record one of these three tests:

• CosmosCommandTests.Should_query_cosmos_database_container_items
• CosmosCommandTests.Should_list_and_query_multiple_databases_and_containers
• CosmosCommandTests.Should_show_single_item_from_cosmos_account

And check the output recording to see the bad output. Then swap the TestMode to playback and re-run. You'll immediately see the error detail above.

Environment

Standard dev env. Nothing special.

[joshfree]

@anuchandy please help drive

[anuchandy]

Hi @scbedd, I spent some time digging into this, and the issue appears to be that the Cosmos library prefers CosmosBinary for query responses by default. While it can encode JSON, the first byte 128 (0x80) in the binary payload (see here) indicates this special serialization type, that Cosmos library knows how to decode.

The reason server returns query responses in this format is that the client when running query internally sets the request header
x-ms-cosmos-supported-serialization-formats: JsonText,CosmosBinary, which expresses a preference for CosmosBinary.

In record mode, if we modify the request (using an http handler) to remove the x-ms-cosmos-supported-serialization-formats header, the Cosmos service falls back to returning plain text JSON instead of the binary format that test proxy can handle.

Would this be an acceptable approach to unblock record–playback?

(below is a query request with this header removed and response in plain json)