diff --git a/CHANGELOG.md b/CHANGELOG.md index d25a26a..22d8cdf 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,10 +6,27 @@ - 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`. +- Destructive MCP confirmations now identify their target. The elicitation prompt adds the connected account endpoint and the current database/container location, and notes that explicit `--db`/`--con` arguments override that location. ([#207](https://github.com/Azure/CosmosDBShell/pull/207)) +- Import and export no longer hold entire files in memory. CSV imports are parsed incrementally, and CSV exports spool documents to a private temporary file to determine the complete column set, so transfers no longer scale with document count. Allow temporary disk space for the CSV export spool in addition to the destination file. ([#207](https://github.com/Azure/CosmosDBShell/pull/207)) + +### Breaking changes + +- Malformed CSV files are now rejected instead of being silently misread. An unterminated or misplaced quote previously caused the remainder of the file to be absorbed into a single field, so the import reported success while writing corrupted items. Such files now abort with `Invalid CSV record at line `. Imports that previously appeared to succeed may now fail and require the source file to be corrected. ([#207](https://github.com/Azure/CosmosDBShell/pull/207)) ### 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. +- A failed or cancelled export no longer destroys its destination file. Exports are written to a temporary file in the destination directory and moved into place only after they complete, so an existing file survives query failures, write failures, and cancellation. An abrupt process termination can leave an unfinished `.cosmos-export-*.tmp` file behind. ([#207](https://github.com/Azure/CosmosDBShell/pull/207)) +- `export --max` no longer requests a further query page once the limit is reached, so the reported request charge no longer includes a page whose items were discarded. Query iterators are now disposed. ([#207](https://github.com/Azure/CosmosDBShell/pull/207)) +- Shell and MCP command execution is serialized, including nested shell calls, so concurrent requests can no longer interleave and corrupt the shared connection and navigation state. Waiting for a destructive confirmation does not hold the execution lock. ([#207](https://github.com/Azure/CosmosDBShell/pull/207)) +- A destructive MCP command is refused when the connection or navigation context changes while its confirmation is pending, including navigating away and back. It previously ran against the changed context. ([#207](https://github.com/Azure/CosmosDBShell/pull/207)) +- Echoing an MCP command line no longer fails the command it announces on hosts without an ANSI terminal, which previously reported `Terminal does not support ANSI` instead of running it. ([#207](https://github.com/Azure/CosmosDBShell/pull/207)) +- MCP command lines now list positional arguments in the order the command binds them. A destructive confirmation and the recorded history entry previously followed the client's argument order, so `rmdb` could display its `force` flag in place of the database name. A call that supplies a positional argument while omitting an earlier one is now rejected, because the shell cannot express that call and the recorded command would bind differently on replay. ([#207](https://github.com/Azure/CosmosDBShell/pull/207)) +- MCP invocations are now saved to the history file as they run and are bounded by the history size limit. They were previously saved only when a later interactive command was entered. On Linux and macOS, the history file is now restricted to its owner, including an existing file that was previously readable by other users. ([#207](https://github.com/Azure/CosmosDBShell/pull/207)) + +### Build & pipeline + +- Added a dependency on CsvHelper 33.1.0 for CSV parsing. ([#207](https://github.com/Azure/CosmosDBShell/pull/207)) ## 1.1.209-preview — 2026-08-26 diff --git a/CosmosDBShell.Tests/CommandTests/ExportCommandTests.cs b/CosmosDBShell.Tests/CommandTests/ExportCommandTests.cs index d56747b..a30d647 100644 --- a/CosmosDBShell.Tests/CommandTests/ExportCommandTests.cs +++ b/CosmosDBShell.Tests/CommandTests/ExportCommandTests.cs @@ -10,6 +10,8 @@ namespace CosmosShell.Tests.CommandTests; using System.Threading; using System.Threading.Tasks; using Azure.Data.Cosmos.Shell.Commands; +using Microsoft.Azure.Cosmos; +using NSubstitute; public class ExportCommandTests { @@ -213,6 +215,180 @@ public async Task WriteCsvAsync_WithNoItems_ProducesEmptyOutput() Assert.Equal(string.Empty, writer.ToString()); } + [Theory] + [InlineData(0)] + [InlineData(1)] + [InlineData(2)] + public async Task WriteFileAsync_FailurePreservesExistingFileAndRemovesTemporaryFile(int format) + { + var directory = Path.Join(Path.GetTempPath(), Guid.NewGuid().ToString("N")); + Directory.CreateDirectory(directory); + var path = Path.Join(directory, "export.json"); + try + { + await File.WriteAllTextAsync(path, "previous export", TestContext.Current.CancellationToken); + await Assert.ThrowsAsync(() => ExportCommand.WriteFileAsync( + FailingItemsAsync(), (ExportFormat)format, path, true, CancellationToken.None)); + Assert.Equal("previous export", await File.ReadAllTextAsync(path, TestContext.Current.CancellationToken)); + Assert.Single(Directory.GetFiles(directory)); + } + finally + { + Directory.Delete(directory, recursive: true); + } + } + + [Theory] + [InlineData(true)] + [InlineData(false)] + public async Task WriteFileAsync_RespectsOverwriteFlag(bool overwrite) + { + var directory = Path.Join(Path.GetTempPath(), Guid.NewGuid().ToString("N")); + Directory.CreateDirectory(directory); + var path = Path.Join(directory, "export.json"); + try + { + await File.WriteAllTextAsync(path, "previous export", TestContext.Current.CancellationToken); + Task ExportAsync() => ExportCommand.WriteFileAsync( + ToAsyncEnumerableAsync(JsonSerializer.SerializeToElement(new { id = "1" })), + ExportFormat.JsonLines, path, overwrite, CancellationToken.None); + if (overwrite) + { + Assert.Equal(1, await ExportAsync()); + Assert.Equal("{\"id\":\"1\"}\n", await File.ReadAllTextAsync(path, TestContext.Current.CancellationToken)); + } + else + { + await Assert.ThrowsAsync(ExportAsync); + Assert.Equal("previous export", await File.ReadAllTextAsync(path, TestContext.Current.CancellationToken)); + } + + Assert.Single(Directory.GetFiles(directory)); + } + finally + { + Directory.Delete(directory, recursive: true); + } + } + + [Fact] + public async Task WriteFileAsync_CancellationPreservesExistingFile() + { + var directory = Path.Join(Path.GetTempPath(), Guid.NewGuid().ToString("N")); + Directory.CreateDirectory(directory); + var path = Path.Join(directory, "export.json"); + using var cancellation = new CancellationTokenSource(); + try + { + await File.WriteAllTextAsync(path, "previous export", TestContext.Current.CancellationToken); + await cancellation.CancelAsync(); + await Assert.ThrowsAnyAsync(() => ExportCommand.WriteFileAsync( + ToAsyncEnumerableAsync(JsonSerializer.SerializeToElement(new { id = "1" })), + ExportFormat.Array, path, true, cancellation.Token)); + Assert.Equal("previous export", await File.ReadAllTextAsync(path, TestContext.Current.CancellationToken)); + Assert.Single(Directory.GetFiles(directory)); + } + finally + { + Directory.Delete(directory, recursive: true); + } + } + + [Fact] + public async Task EnumerateAsync_LimitAtPageBoundaryDoesNotFetchAnotherPage() + { + using var iterator = Substitute.For>(); + var response = Substitute.For>(); + response.GetEnumerator().Returns(_ => ((IEnumerable)new[] { JsonSerializer.SerializeToElement(new { id = "1" }) }).GetEnumerator()); + response.RequestCharge.Returns(3); + iterator.HasMoreResults.Returns(true); + iterator.ReadNextAsync(Arg.Any()).Returns(Task.FromResult(response)); + var charge = 0.0; + var count = 0; + await foreach (var item in ExportCommand.EnumerateAsync(iterator, 1, value => charge += value, CancellationToken.None)) + { + count++; + } + + Assert.Equal(1, count); + Assert.Equal(3, charge); + await iterator.Received(1).ReadNextAsync(Arg.Any()); + } + + [Fact] + public async Task WriteCsvAsync_DoesNotRetainSourceDocuments() + { + using var writer = new StringWriter(); + Assert.Equal(200, await ExportCommand.WriteCsvAsync(TransientItemsAsync(), writer, ',', TestContext.Current.CancellationToken)); + Assert.Contains("199", writer.ToString()); + } + + [Fact] + public void DeleteTemporaryFile_DoesNotThrowWhenDirectoryIsMissing() + { + var directory = Directory.CreateTempSubdirectory("cosmos-export-test-"); + var path = Path.Join(directory.FullName, "export.tmp"); + directory.Delete(); + + Assert.IsAssignableFrom(new DirectoryNotFoundException()); + ExportCommand.DeleteTemporaryFile(path); + } + + [Fact] + public void DeleteTemporaryFile_DoesNotThrowWhenPathIsDirectory() + { + var directory = Directory.CreateTempSubdirectory("cosmos-export-test-"); + try + { + ExportCommand.DeleteTemporaryFile(directory.FullName); + Assert.True(directory.Exists); + } + finally + { + directory.Delete(); + } + } + + [Fact] + public async Task WriteArrayAsync_FlushesIncrementallyWithoutFlushingEachItem() + { + using var stream = new MemoryStream(); + var item = JsonSerializer.SerializeToElement(new { value = new string('x', 1024) }); + async IAsyncEnumerable ItemsAsync() + { + yield return item; + Assert.Equal(0, stream.Length); + for (var index = 0; index < 128; index++) + { + yield return item; + } + + Assert.True(stream.Length >= 64 * 1024); + await Task.Yield(); + } + + Assert.Equal(129, await ExportCommand.WriteArrayAsync(ItemsAsync(), stream, TestContext.Current.CancellationToken)); + using var result = JsonDocument.Parse(stream.ToArray()); + Assert.Equal(129, result.RootElement.GetArrayLength()); + } + + private static async IAsyncEnumerable TransientItemsAsync() + { + for (var index = 0; index < 200; index++) + { + using var document = JsonDocument.Parse($"{{\"id\":{index}}}"); + yield return document.RootElement; + await Task.Yield(); + } + } + + private static async IAsyncEnumerable FailingItemsAsync() + { + yield return JsonSerializer.SerializeToElement(new { id = "1" }); + await Task.Yield(); + throw new IOException("simulated read failure"); + } + private static async IAsyncEnumerable ToAsyncEnumerableAsync(params JsonElement[] items) { foreach (var item in items) diff --git a/CosmosDBShell.Tests/CommandTests/ImportCommandTests.cs b/CosmosDBShell.Tests/CommandTests/ImportCommandTests.cs index 3f6f0dd..6fbccca 100644 --- a/CosmosDBShell.Tests/CommandTests/ImportCommandTests.cs +++ b/CosmosDBShell.Tests/CommandTests/ImportCommandTests.cs @@ -347,6 +347,117 @@ public void BuildCsvObject_MapsColumnsToStringProperties() Assert.Equal("Alice", element.GetProperty("name").GetString()); } + [Fact] + public void ReadCsvRecords_ReadsValidRowsBeforeReportingMalformedRecord() + { + using var reader = new StringReader("id,name\n1,Alice\n2,\"unterminated"); + using var records = ImportCommand.ReadCsvRecords(reader, ',', TestContext.Current.CancellationToken).GetEnumerator(); + Assert.True(records.MoveNext()); + Assert.True(records.MoveNext()); + Assert.Equal(new[] { "1", "Alice" }, records.Current.Fields); + var error = Assert.Throws(() => records.MoveNext()); + Assert.Contains("3", error.Message); + } + + [Fact] + public void ReadCsvRecords_CancellationStopsBetweenRecords() + { + using var reader = new StringReader("id,name\n1,Alice\n2,Bob"); + using var cancellation = new CancellationTokenSource(); + using var records = ImportCommand.ReadCsvRecords(reader, ',', cancellation.Token).GetEnumerator(); + Assert.True(records.MoveNext()); + cancellation.Cancel(); + Assert.ThrowsAny(() => records.MoveNext()); + } + + [Fact] + public async Task ReadCsvRecordsAsync_MatchesSynchronousRecordsAndStartLines() + { + const string content = "id,name\n1,\"multi\nline\"\n2,Bob\n"; + using var reader = new StringReader(content); + var actual = new List<(int StartLine, List Fields)>(); + await foreach (var record in ImportCommand.ReadCsvRecordsAsync(reader, ',', TestContext.Current.CancellationToken)) + { + actual.Add(record); + } + + var expected = ImportCommand.ParseCsvWithLines(content, ','); + Assert.Equal(expected.Select(r => r.StartLine), actual.Select(r => r.StartLine)); + Assert.Equal(expected.Select(r => r.Fields), actual.Select(r => r.Fields)); + } + + [Fact] + public async Task ReadCsvRecordsAsync_ReportsMalformedRecordWithPhysicalLine() + { + using var reader = new StringReader("id,name\n1,Alice\n2,\"unterminated"); + var error = await Assert.ThrowsAsync(async () => + { + await foreach (var _ in ImportCommand.ReadCsvRecordsAsync(reader, ',', TestContext.Current.CancellationToken)) + { + // Enumerate only to drive parsing until the malformed record throws. + } + }); + Assert.Contains("3", error.Message); + } + + [Fact] + public async Task ReadCsvRecordsAsync_CancellationStopsBetweenRecords() + { + using var reader = new StringReader("id,name\n1,Alice\n2,Bob"); + using var cancellation = new CancellationTokenSource(); + await using var records = ImportCommand.ReadCsvRecordsAsync(reader, ',', cancellation.Token).GetAsyncEnumerator(TestContext.Current.CancellationToken); + Assert.True(await records.MoveNextAsync()); + await cancellation.CancelAsync(); + await Assert.ThrowsAnyAsync(async () => await records.MoveNextAsync()); + } + + [Fact] + public async Task ReadCsvRecordsAsync_CancellationInterruptsBlockedRecordRead() + { + using var reader = new BlockingAfterPrefixReader("id,name\n1,\"multi"); + using var cancellation = new CancellationTokenSource(); + await using var records = ImportCommand.ReadCsvRecordsAsync(reader, ',', cancellation.Token).GetAsyncEnumerator(TestContext.Current.CancellationToken); + Assert.True(await records.MoveNextAsync()); + var move = records.MoveNextAsync().AsTask(); + await reader.Blocked.Task.WaitAsync(TimeSpan.FromSeconds(10), TestContext.Current.CancellationToken); + Assert.False(move.IsCompleted); + + await cancellation.CancelAsync(); + + await Assert.ThrowsAnyAsync(() => move.WaitAsync(TimeSpan.FromSeconds(10), TestContext.Current.CancellationToken)); + } + + private sealed class BlockingAfterPrefixReader(string prefix) : TextReader + { + private bool prefixReturned; + + public TaskCompletionSource Blocked { get; } = new(TaskCreationOptions.RunContinuationsAsynchronously); + + public override async ValueTask ReadAsync(Memory buffer, CancellationToken cancellationToken = default) + { + if (!this.prefixReturned) + { + this.prefixReturned = true; + prefix.AsSpan().CopyTo(buffer.Span); + return prefix.Length; + } + + this.Blocked.TrySetResult(); + await Task.Delay(Timeout.Infinite, cancellationToken); + return 0; + } + + public override Task ReadAsync(char[] buffer, int index, int count) + => this.ReadAsync(buffer.AsMemory(index, count)).AsTask(); + } + + [Fact] + public void ParseCsvWithLines_SkipsBlankLinesWithoutLosingPhysicalLineNumbers() + { + var records = ImportCommand.ParseCsvWithLines("id,name\n\n\n1,Alice\n", ','); + Assert.Equal(4, records[1].StartLine); + } + [Fact] public void BuildCsvObject_SingleSegmentPartitionKey_StaysTopLevel() { diff --git a/CosmosDBShell.Tests/Runtime/SerializedExecutionTests.cs b/CosmosDBShell.Tests/Runtime/SerializedExecutionTests.cs new file mode 100644 index 0000000..1daeac0 --- /dev/null +++ b/CosmosDBShell.Tests/Runtime/SerializedExecutionTests.cs @@ -0,0 +1,210 @@ +namespace CosmosShell.Tests.Runtime; + +using Azure.Data.Cosmos.Shell.Core; + +public class SerializedExecutionTests +{ + [Fact] + public async Task PrintCommand_ConcurrentWithHistorySnapshots_DoesNotCorruptHistory() + { + using var shell = ShellInterpreter.CreateInstance(); + using var cancellation = new CancellationTokenSource(TimeSpan.FromSeconds(10)); + + var writer = Task.Run( + () => + { + for (var i = 0; i < 500; i++) + { + shell.PrintCommand($"echo {i}"); + } + }, + cancellation.Token); + + var reader = Task.Run( + () => + { + while (!writer.IsCompleted) + { + foreach (var entry in shell.History) + { + Assert.NotNull(entry); + } + } + }, + cancellation.Token); + + await Task.WhenAll(writer, reader); + Assert.Equal("echo 499", shell.History[^1]); + } + + [Fact] + public void PrintCommand_PersistsBoundedHistory() + { + var configPath = Path.Join(Path.GetTempPath(), $"cosmosshell-history-{Guid.NewGuid():N}"); + try + { + using (var shell = new ShellInterpreter(configPath)) + { + for (var i = 0; i < 70; i++) + { + shell.PrintCommand($"echo {i}"); + } + } + + var persisted = File.ReadAllLines(Path.Join(configPath, "cmd_history")); + Assert.Equal(60, persisted.Length); + Assert.Equal("echo 69", persisted[^1]); + + using var restarted = new ShellInterpreter(configPath); + Assert.Equal("echo 69", restarted.History[^1]); + } + finally + { + if (Directory.Exists(configPath)) + { + Directory.Delete(configPath, recursive: true); + } + } + } + + [Fact] + public void PrintCommand_RestrictsHistoryFileToOwnerOnUnix() + { + if (OperatingSystem.IsWindows()) + { + Assert.Skip("Unix file modes do not apply on Windows."); + return; + } + + var configPath = Path.Join(Path.GetTempPath(), $"cosmosshell-history-{Guid.NewGuid():N}"); + try + { + Directory.CreateDirectory(configPath); + var historyFile = Path.Join(configPath, "cmd_history"); + File.WriteAllText(historyFile, string.Empty); + File.SetUnixFileMode(historyFile, UnixFileMode.UserRead | UnixFileMode.UserWrite | UnixFileMode.GroupRead | UnixFileMode.OtherRead); + + using (var shell = new ShellInterpreter(configPath)) + { + shell.PrintCommand("echo 1"); + } + + Assert.Equal(UnixFileMode.UserRead | UnixFileMode.UserWrite, File.GetUnixFileMode(historyFile)); + + File.Delete(historyFile); + using (var shell = new ShellInterpreter(configPath)) + { + shell.PrintCommand("echo 2"); + } + + Assert.Equal(UnixFileMode.UserRead | UnixFileMode.UserWrite, File.GetUnixFileMode(historyFile)); + } + finally + { + if (Directory.Exists(configPath)) + { + Directory.Delete(configPath, recursive: true); + } + } + } + + [Fact] + public async Task Dispose_ReleasesExecutionGateAndIsIdempotent() + { + var shell = ShellInterpreter.CreateInstance(); + Assert.Equal(42, await shell.RunSerializedAsync(() => Task.FromResult(42), TestContext.Current.CancellationToken)); + shell.Dispose(); + shell.Dispose(); + var executed = false; + await Assert.ThrowsAsync(() => shell.RunSerializedAsync(() => + { + executed = true; + return Task.FromResult(1); + }, CancellationToken.None)); + Assert.False(executed); + } + + [Fact] + public async Task RunSerializedAsync_WaitsForOtherExecutionButAllowsNestedCalls() + { + using var shell = ShellInterpreter.CreateInstance(); + var entered = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var release = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var first = shell.RunSerializedAsync(async () => + { + entered.SetResult(); + await release.Task; + return await shell.RunSerializedAsync(() => Task.FromResult(42), CancellationToken.None); + }, CancellationToken.None); + await entered.Task; + var secondEntered = false; + var second = shell.RunSerializedAsync(() => + { + secondEntered = true; + return Task.FromResult(7); + }, CancellationToken.None); + try + { + Assert.False(secondEntered); + Assert.False(second.IsCompleted); + } + finally + { + release.SetResult(); + } + + Assert.Equal(42, await first.WaitAsync(TimeSpan.FromSeconds(5), TestContext.Current.CancellationToken)); + Assert.Equal(7, await second.WaitAsync(TimeSpan.FromSeconds(5), TestContext.Current.CancellationToken)); + } + + [Fact] + public async Task RunSerializedAsync_ReleasesGateAfterFailure() + { + using var shell = ShellInterpreter.CreateInstance(); + await Assert.ThrowsAsync(() => shell.RunSerializedAsync( + () => throw new InvalidOperationException(), CancellationToken.None)); + Assert.Equal(42, await shell.RunSerializedAsync(() => Task.FromResult(42), CancellationToken.None)); + } + + [Fact] + public async Task RunSerializedAsync_CancelledWaiterDoesNotExecuteOrReleaseAnotherOwnersGate() + { + using var shell = ShellInterpreter.CreateInstance(); + using var cancellation = new CancellationTokenSource(); + var release = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var first = shell.RunSerializedAsync(async () => + { + await release.Task; + return 1; + }, TestContext.Current.CancellationToken); + var executed = false; + var waiter = shell.RunSerializedAsync(() => + { + executed = true; + return Task.FromResult(2); + }, cancellation.Token); + try + { + await cancellation.CancelAsync(); + try + { + await waiter; + Assert.Fail("The queued operation should have been cancelled."); + } + catch (OperationCanceledException) + { + Assert.True(cancellation.IsCancellationRequested); + } + + Assert.False(executed); + Assert.False(first.IsCompleted); + } + finally + { + release.SetResult(); + await first; + } + + Assert.Equal(3, await shell.RunSerializedAsync(() => Task.FromResult(3), TestContext.Current.CancellationToken)); + } +} \ No newline at end of file diff --git a/CosmosDBShell.Tests/TestEnvironmentInitializer.cs b/CosmosDBShell.Tests/TestEnvironmentInitializer.cs new file mode 100644 index 0000000..9071d5f --- /dev/null +++ b/CosmosDBShell.Tests/TestEnvironmentInitializer.cs @@ -0,0 +1,19 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace CosmosShell.Tests; + +using System.Runtime.CompilerServices; + +internal static class TestEnvironmentInitializer +{ + // Shells created by tests persist command history; always override, since a pre-set value may be the developer's real config. + [ModuleInitializer] + internal static void IsolateShellConfigDirectory() + { + Environment.SetEnvironmentVariable( + "COSMOSDB_SHELL_CONFIG_DIR", + Path.Join(Path.GetTempPath(), $"cosmosshell-tests-{Guid.NewGuid():N}")); + } +} diff --git a/CosmosDBShell.Tests/ToolOperationsCallToolTests.cs b/CosmosDBShell.Tests/ToolOperationsCallToolTests.cs index 0cef79e..406ab92 100644 --- a/CosmosDBShell.Tests/ToolOperationsCallToolTests.cs +++ b/CosmosDBShell.Tests/ToolOperationsCallToolTests.cs @@ -12,7 +12,9 @@ namespace CosmosShell.Tests; using System.Threading.Tasks; using Azure.Data.Cosmos.Shell.Core; +using Azure.Data.Cosmos.Shell.Commands; using Azure.Data.Cosmos.Shell.Mcp; +using Azure.Data.Cosmos.Shell.States; using Microsoft.Extensions.Logging.Abstractions; @@ -146,6 +148,94 @@ public async Task ConfirmDestructive_UserAccepts_ReturnsNull() Assert.Null(result); } + [Fact] + public async Task ExecuteTool_ContextChangesDuringConfirmation_RefusesExecution() + { + var shell = ShellInterpreter.Instance; + var originalState = shell.State; + var command = new TrackingCommand(); + try + { + var result = await CreateToolOperations().ExecuteToolAsync( + shell.App.Commands["rm"], command, "rm test-*", + (request, _) => + { + Assert.Contains("Account:", request.Message); + Assert.Contains("Current location:", request.Message); + shell.State = new DisconnectedState(); + shell.State = originalState; + return new ValueTask(new ElicitResult { Action = "accept" }); + }, CancellationToken.None); + + Assert.True(result.IsError); + Assert.False(command.Executed); + Assert.Contains("context changed", Assert.IsType(Assert.Single(result.Content)).Text); + } + finally + { + shell.State = originalState; + } + } + + private sealed class TrackingCommand : CosmosCommand + { + public bool Executed { get; private set; } + + public override Task ExecuteAsync(ShellInterpreter shell, CommandState commandState, string commandText, CancellationToken token) + { + this.Executed = true; + return Task.FromResult(commandState); + } + } + + [Fact] + public async Task ExecuteTool_UnchangedContextAfterConfirmation_ExecutesCommand() + { + var command = new TrackingCommand(); + var result = await CreateToolOperations().ExecuteToolAsync( + ShellInterpreter.Instance.App.Commands["rm"], command, "rm test-*", + (_, _) => new ValueTask(new ElicitResult { Action = "accept" }), + TestContext.Current.CancellationToken); + Assert.False(result.IsError == true); + Assert.True(command.Executed); + } + + [Fact] + public async Task ExecuteTool_WithoutAnsiTerminal_EchoesPlainlyAndStillExecutes() + { + var command = new TrackingCommand(); + using var plain = new StringWriter(); + + var savedConsole = AnsiConsole.Console; + var savedOut = Console.Out; + try + { + AnsiConsole.Console = AnsiConsole.Create(new AnsiConsoleSettings + { + Ansi = AnsiSupport.No, + ColorSystem = ColorSystemSupport.NoColors, + Out = new AnsiConsoleOutput(plain), + }); + Console.SetOut(plain); + + var result = await CreateToolOperations().ExecuteToolAsync( + ShellInterpreter.Instance.App.Commands["rm"], command, "rm test-*", + (_, _) => new ValueTask(new ElicitResult { Action = "accept" }), + TestContext.Current.CancellationToken); + + Assert.False(result.IsError == true); + Assert.True(command.Executed); + } + finally + { + Console.SetOut(savedOut); + AnsiConsole.Console = savedConsole; + } + + Assert.Contains("rm test-*", plain.ToString(), StringComparison.Ordinal); + Assert.Equal("rm test-*", ShellInterpreter.Instance.History.ToArray()[^1]); + } + [Theory] [InlineData("decline")] [InlineData("cancel")] @@ -193,6 +283,69 @@ public async Task ConfirmDestructive_ElicitationThrows_ReturnsErrorAndDoesNotExe } } + [Fact] + public async Task CallTool_PositionalArgumentsOutOfOrder_DisplaysDeclarationOrder() + { + var tool = CreateToolOperations(); + var arguments = new Dictionary + { + ["force"] = Json("true"), + ["name"] = Json("\"OldDb\""), + }; + + var result = await tool.CallToolHandler(CallContext("rmdb", arguments), CancellationToken.None); + + var (isError, root, document) = ReadResult(result); + using (document) + { + Assert.True(isError); + Assert.Contains("rmdb \"OldDb\" \"True\"", root.GetProperty("error").GetString(), StringComparison.Ordinal); + } + } + + [Fact] + public async Task CallTool_PositionalGap_ReturnsErrorWithoutExecuting() + { + var tool = CreateToolOperations(); + var arguments = new Dictionary + { + ["path"] = Json("\"dark.json\""), + }; + + var result = await tool.CallToolHandler(CallContext("theme", arguments), CancellationToken.None); + + var (isError, root, document) = ReadResult(result); + using (document) + { + Assert.True(isError); + var error = root.GetProperty("error").GetString(); + Assert.Contains("'path'", error, StringComparison.Ordinal); + Assert.Contains("'action'", error, StringComparison.Ordinal); + } + } + + [Fact] + public void FindPositionalGap_DetectsOmittedAndNullPredecessors() + { + var parameters = ShellInterpreter.Instance.App.Commands["theme"].Parameters; + + Assert.NotNull(ToolOperations.FindPositionalGap(parameters, new Dictionary { [parameters[2]] = "dark.json" })); + Assert.NotNull(ToolOperations.FindPositionalGap(parameters, new Dictionary { [parameters[0]] = null, [parameters[1]] = "dark" })); + Assert.Null(ToolOperations.FindPositionalGap(parameters, new Dictionary { [parameters[0]] = "show", [parameters[1]] = "dark" })); + } + + [Fact] + public void FormatPositionalsForHistory_RendersContiguousValuesAndExpandsArrays() + { + var parameters = ShellInterpreter.Instance.App.Commands["theme"].Parameters; + var values = new Dictionary { [parameters[0]] = "show", [parameters[1]] = "dark" }; + Assert.Equal(" \"show\" \"dark\"", ToolOperations.FormatPositionalsForHistory(parameters, values)); + + var echoParameters = ShellInterpreter.Instance.App.Commands["echo"].Parameters; + var echoValues = new Dictionary { [echoParameters[0]] = new[] { "hello", "world" } }; + Assert.Equal(" \"hello\" \"world\"", ToolOperations.FormatPositionalsForHistory(echoParameters, echoValues)); + } + [Fact] public async Task CallTool_UnknownArgument_ReturnsErrorListingKnownArguments() { @@ -353,6 +506,7 @@ public async Task CallTool_InvalidContinuationType_ReturnsError(string continuat public async Task CallTool_EchoCommand_ReturnsSuccessResult() { var tool = CreateToolOperations(); + using var output = new StringWriter(); var arguments = new Dictionary { ["messages"] = Json("[\"hello\", \"world\"]"), @@ -365,11 +519,16 @@ public async Task CallTool_EchoCommand_ReturnsSuccessResult() { Ansi = AnsiSupport.Yes, ColorSystem = ColorSystemSupport.NoColors, - Out = new AnsiConsoleOutput(new StringWriter()), + Out = new AnsiConsoleOutput(output), }); var result = await tool.CallToolHandler(CallContext("echo", arguments), CancellationToken.None); + Assert.Contains("echo", output.ToString(), StringComparison.Ordinal); + var recorded = ShellInterpreter.Instance.History.ToArray(); + Assert.Equal("echo \"hello\" \"world\"", recorded[^1]); + Assert.Single(recorded, entry => entry == recorded[^1]); + var (isError, root, document) = ReadResult(result); using (document) { diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ExportCommand.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ExportCommand.cs index 0fade21..57152c4 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ExportCommand.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ExportCommand.cs @@ -91,7 +91,7 @@ public override async Task ExecuteAsync(ShellInterpreter shell, Co var max = ResultLimit.ResolveMaxItemCount(this.Max, defaultMaxItemCount: null); var format = this.Format ?? ExportFormat.JsonLines; - var (count, charge) = await ExecuteExportAsync(container, query, max, format, filePath, token); + var (count, charge) = await ExecuteExportAsync(container, query, max, format, filePath, this.Force == true, token); ShellInterpreter.WriteLine(MessageService.GetArgsString( "command-export-success", @@ -171,6 +171,11 @@ internal static async Task WriteArrayAsync(IAsyncEnumerable it await foreach (var item in items.WithCancellation(token)) { item.WriteTo(writer); + if (writer.BytesPending >= 64 * 1024) + { + await writer.FlushAsync(token); + } + count++; } @@ -183,7 +188,7 @@ internal static async Task WriteArrayAsync(IAsyncEnumerable it /// Writes a sequence of items to as CSV. The header row is the /// union of all top-level property names (in first-seen order); each subsequent row /// contains the corresponding values. Nested objects and arrays are written as compact - /// JSON. Items are buffered to compute the column set. + /// JSON. Items are spooled to disk to compute the column set. /// /// The items to write. /// The destination writer. @@ -192,16 +197,25 @@ internal static async Task WriteArrayAsync(IAsyncEnumerable it /// The number of data rows written. internal static async Task WriteCsvAsync(IAsyncEnumerable items, TextWriter writer, char separator, CancellationToken token) { - var buffered = new List(); - await foreach (var item in items.WithCancellation(token)) + var spoolOptions = new FileStreamOptions + { + Mode = FileMode.CreateNew, + Access = FileAccess.ReadWrite, + Share = FileShare.None, + Options = FileOptions.Asynchronous | FileOptions.DeleteOnClose, + }; + if (!OperatingSystem.IsWindows()) { - buffered.Add(item); + spoolOptions.UnixCreateMode = UnixFileMode.UserRead | UnixFileMode.UserWrite; } + await using var spool = new FileStream(Path.Join(Path.GetTempPath(), $"cosmos-csv-{Guid.NewGuid():N}.tmp"), spoolOptions); + using var spoolWriter = new StreamWriter(spool, new UTF8Encoding(false), leaveOpen: true); var headers = new List(); var headerSet = new HashSet(StringComparer.Ordinal); - foreach (var item in buffered) + await foreach (var item in items.WithCancellation(token)) { + await spoolWriter.WriteLineAsync(SerializeJsonLine(item).AsMemory(), token); if (item.ValueKind != JsonValueKind.Object) { continue; @@ -216,6 +230,9 @@ internal static async Task WriteCsvAsync(IAsyncEnumerable item } } + await spoolWriter.FlushAsync(token); + spool.Position = 0; + using var spoolReader = new StreamReader(spool, leaveOpen: true); var sb = new StringBuilder(); if (headers.Count > 0) { @@ -233,8 +250,10 @@ internal static async Task WriteCsvAsync(IAsyncEnumerable item } var count = 0; - foreach (var item in buffered) + while (await spoolReader.ReadLineAsync(token) is { } line) { + using var document = JsonDocument.Parse(line); + var item = document.RootElement; sb.Clear(); for (var i = 0; i < headers.Count; i++) { @@ -268,14 +287,9 @@ internal static async Task WriteCsvAsync(IAsyncEnumerable item int? max, ExportFormat format, string filePath, + bool overwrite, CancellationToken token) { - var directory = Path.GetDirectoryName(Path.GetFullPath(filePath)); - if (!string.IsNullOrEmpty(directory) && !Directory.Exists(directory)) - { - Directory.CreateDirectory(directory); - } - var options = new QueryRequestOptions(); if (max is int explicitMax && explicitMax > 0) { @@ -285,30 +299,10 @@ internal static async Task WriteCsvAsync(IAsyncEnumerable item try { var totalCharge = 0.0; - var iterator = container.GetItemQueryIterator(query, requestOptions: options); - - if (format == ExportFormat.Array) - { - await using var stream = new FileStream(filePath, FileMode.Create, FileAccess.Write, FileShare.None); - var count = await WriteArrayAsync(EnumerateAsync(iterator, max, charge => totalCharge += charge, token), stream, token); - return (count, totalCharge); - } - else if (format == ExportFormat.Csv) - { - await using var stream = new FileStream(filePath, FileMode.Create, FileAccess.Write, FileShare.None); - using var writer = new StreamWriter(stream, new UTF8Encoding(encoderShouldEmitUTF8Identifier: false)); - writer.NewLine = "\n"; - var count = await WriteCsvAsync(EnumerateAsync(iterator, max, charge => totalCharge += charge, token), writer, ShellInterpreter.CSVSeparator, token); - return (count, totalCharge); - } - else - { - await using var stream = new FileStream(filePath, FileMode.Create, FileAccess.Write, FileShare.None); - using var writer = new StreamWriter(stream, new UTF8Encoding(encoderShouldEmitUTF8Identifier: false)); - writer.NewLine = "\n"; - var count = await WriteJsonLinesAsync(EnumerateAsync(iterator, max, charge => totalCharge += charge, token), writer, token); - return (count, totalCharge); - } + using var iterator = container.GetItemQueryIterator(query, requestOptions: options); + var count = await WriteFileAsync( + EnumerateAsync(iterator, max, charge => totalCharge += charge, token), format, filePath, overwrite, token); + return (count, totalCharge); } catch (CosmosException ce) { @@ -324,14 +318,78 @@ internal static async Task WriteCsvAsync(IAsyncEnumerable item } } - private static async IAsyncEnumerable EnumerateAsync( + internal static async Task WriteFileAsync( + IAsyncEnumerable items, + ExportFormat format, + string filePath, + bool overwrite, + CancellationToken token) + { + var destination = Path.GetFullPath(filePath); + var directory = Path.GetDirectoryName(destination)!; + Directory.CreateDirectory(directory); + var temporary = Path.Join(directory, $".cosmos-export-{Guid.NewGuid():N}.tmp"); + try + { + int count; + var fileOptions = new FileStreamOptions + { + Mode = FileMode.CreateNew, + Access = FileAccess.Write, + Share = FileShare.None, + Options = FileOptions.Asynchronous, + }; + if (!OperatingSystem.IsWindows()) + { + fileOptions.UnixCreateMode = UnixFileMode.UserRead | UnixFileMode.UserWrite; + } + + await using (var stream = new FileStream(temporary, fileOptions)) + { + if (format == ExportFormat.Array) + { + count = await WriteArrayAsync(items, stream, token); + } + else + { + using var writer = new StreamWriter(stream, new UTF8Encoding(encoderShouldEmitUTF8Identifier: false)); + writer.NewLine = "\n"; + count = format == ExportFormat.Csv + ? await WriteCsvAsync(items, writer, ShellInterpreter.CSVSeparator, token) + : await WriteJsonLinesAsync(items, writer, token); + } + } + + token.ThrowIfCancellationRequested(); + System.IO.File.Move(temporary, destination, overwrite); + return count; + } + finally + { + DeleteTemporaryFile(temporary); + } + } + + internal static void DeleteTemporaryFile(string path) + { + try + { + System.IO.File.Delete(path); + } + catch (Exception exception) when (exception is IOException or UnauthorizedAccessException) + { + System.Diagnostics.Trace.TraceWarning("Export temporary-file cleanup failed ({0}).", exception.GetType().Name); + } + } + + internal static async IAsyncEnumerable EnumerateAsync( FeedIterator iterator, int? max, Action recordCharge, [System.Runtime.CompilerServices.EnumeratorCancellation] CancellationToken token) { var emitted = 0; - while (iterator.HasMoreResults) + while ((!max.HasValue || emitted < max.Value) && iterator.HasMoreResults) { FeedResponse response; try diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ImportCommand.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ImportCommand.cs index bb82728..466b7e5 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ImportCommand.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ImportCommand.cs @@ -13,6 +13,8 @@ namespace Azure.Data.Cosmos.Shell.Commands; using Azure.Data.Cosmos.Shell.Mcp; using Azure.Data.Cosmos.Shell.Parser; using Azure.Data.Cosmos.Shell.Util; +using CsvHelper; +using CsvHelper.Configuration; using global::Azure.Data.Cosmos.Shell.Core; using global::Azure.Data.Cosmos.Shell.States; @@ -256,100 +258,80 @@ internal static List> ParseCsv(string content, char separator) /// A list of (1-based start line, field values) records. internal static List<(int StartLine, List Fields)> ParseCsvWithLines(string content, char separator) { - var records = new List<(int StartLine, List Fields)>(); - var record = new List(); - var field = new StringBuilder(); - var inQuotes = false; - var hasContent = false; - var physicalLine = 1; - var recordStartLine = 0; - - for (var i = 0; i < content.Length; i++) - { - var c = content[i]; - if (inQuotes) - { - if (c == '"') - { - if (i + 1 < content.Length && content[i + 1] == '"') - { - field.Append('"'); - i++; - } - else - { - inQuotes = false; - } - } - else - { - if (c == '\n') - { - physicalLine++; - } + using var reader = new StringReader(content); + return ReadCsvRecords(reader, separator, CancellationToken.None).ToList(); + } - field.Append(c); - } - } - else if (c == '"') + internal static IEnumerable<(int StartLine, List Fields)> ReadCsvRecords(TextReader reader, char separator, CancellationToken token) + { + using var parser = CreateCsvParser(reader, separator, token); + while (true) + { + token.ThrowIfCancellationRequested(); + var startLine = parser.RawRow + 1; + string[]? fields; + try { - if (recordStartLine == 0) + if (!parser.Read()) { - recordStartLine = physicalLine; + yield break; } - inQuotes = true; - hasContent = true; + fields = parser.Record; } - else if (c == separator) + catch (CsvHelperException ex) { - if (recordStartLine == 0) - { - recordStartLine = physicalLine; - } - - record.Add(field.ToString()); - field.Clear(); - hasContent = true; + throw new CommandException("import", MessageService.GetArgsString("command-import-error-invalid_csv", "line", startLine), ex); } - else if (c == '\r') + + if (fields is not null && parser.RawRecord.TrimEnd('\r', '\n').Length > 0) { - // Ignored; line breaks are handled on '\n'. + yield return (startLine, fields.ToList()); } - else if (c == '\n') + } + } + + internal static async IAsyncEnumerable<(int StartLine, List Fields)> ReadCsvRecordsAsync( + TextReader reader, + char separator, + [EnumeratorCancellation] CancellationToken token) + { + using var parser = CreateCsvParser(reader, separator, token); + while (true) + { + token.ThrowIfCancellationRequested(); + var startLine = parser.RawRow + 1; + string[]? fields; + try { - record.Add(field.ToString()); - field.Clear(); - if (hasContent || record.Count > 1) + if (!await parser.ReadAsync()) { - records.Add((recordStartLine == 0 ? physicalLine : recordStartLine, record)); + yield break; } - record = new List(); - hasContent = false; - recordStartLine = 0; - physicalLine++; + fields = parser.Record; } - else + catch (CsvHelperException ex) { - if (recordStartLine == 0) - { - recordStartLine = physicalLine; - } + throw new CommandException("import", MessageService.GetArgsString("command-import-error-invalid_csv", "line", startLine), ex); + } - field.Append(c); - hasContent = true; + if (fields is not null && parser.RawRecord.TrimEnd('\r', '\n').Length > 0) + { + yield return (startLine, fields.ToList()); } } + } - if (hasContent || field.Length > 0 || record.Count > 0) + // CsvParser.ReadAsync accepts no token, so cancellation is enforced at the reader to + // interrupt a large or unterminated record while it is still being read. + private static CsvParser CreateCsvParser(TextReader reader, char separator, CancellationToken token) + => new(new CancellationAwareTextReader(reader, token), new CsvConfiguration(CultureInfo.InvariantCulture) { - record.Add(field.ToString()); - records.Add((recordStartLine == 0 ? physicalLine : recordStartLine, record)); - } - - return records; - } + Delimiter = separator.ToString(), + IgnoreBlankLines = false, + ExceptionMessagesContainRawData = false, + }); /// /// Builds a JSON object from a CSV header row and a value row. Every column becomes a @@ -433,17 +415,17 @@ internal static JsonElement BuildCsvObject(IReadOnlyList headers, IReadO string[]? partitionKeySegments, [EnumeratorCancellation] CancellationToken token) { - var content = await System.IO.File.ReadAllTextAsync(filePath, token); - var records = ParseCsvWithLines(content, ShellInterpreter.CSVSeparator); - if (records.Count == 0) + await using var stream = new FileStream(filePath, FileMode.Open, FileAccess.Read, FileShare.Read); + using var reader = new StreamReader(stream); + List? headers = null; + await foreach (var (startLine, fields) in ReadCsvRecordsAsync(reader, ShellInterpreter.CSVSeparator, token)) { - yield break; - } + if (headers is null) + { + headers = fields; + continue; + } - var headers = records[0].Fields; - for (var r = 1; r < records.Count; r++) - { - var (startLine, fields) = records[r]; yield return (startLine, BuildCsvObject(headers, fields, partitionKeySegments)); } } @@ -676,4 +658,56 @@ private static async Task ResolveFormatAsync(string filePath, Impo return (success, failed, charge); } + + private sealed class CancellationAwareTextReader(TextReader inner, CancellationToken token) : TextReader + { + public override int Peek() + { + token.ThrowIfCancellationRequested(); + return inner.Peek(); + } + + public override int Read() + { + token.ThrowIfCancellationRequested(); + return inner.Read(); + } + + public override int Read(char[] buffer, int index, int count) + { + token.ThrowIfCancellationRequested(); + return inner.Read(buffer, index, count); + } + + public override int Read(Span buffer) + { + token.ThrowIfCancellationRequested(); + return inner.Read(buffer); + } + + public override Task ReadAsync(char[] buffer, int index, int count) + => this.ReadAsync(buffer.AsMemory(index, count), CancellationToken.None).AsTask(); + + public override async ValueTask ReadAsync(Memory buffer, CancellationToken cancellationToken = default) + { + token.ThrowIfCancellationRequested(); + if (!cancellationToken.CanBeCanceled) + { + return await inner.ReadAsync(buffer, token).ConfigureAwait(false); + } + + using var linked = CancellationTokenSource.CreateLinkedTokenSource(token, cancellationToken); + return await inner.ReadAsync(buffer, linked.Token).ConfigureAwait(false); + } + + protected override void Dispose(bool disposing) + { + if (disposing) + { + inner.Dispose(); + } + + base.Dispose(disposing); + } + } } diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/ShellInterpreter.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/ShellInterpreter.cs index 711e52d..f164092 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/ShellInterpreter.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/ShellInterpreter.cs @@ -31,6 +31,9 @@ public partial class ShellInterpreter : IDisposable private const string SessionRequestChargeWarningThresholdVariable = "sessionRequestChargeWarningThreshold"; + // Declared before Instance: static initializers run in order and the constructor reads history. + private static readonly object HistoryFileLock = new(); + internal static readonly ShellInterpreter Instance = new(); private const int MAXHISTORYITEMS = 60; @@ -56,6 +59,14 @@ public partial class ShellInterpreter : IDisposable private readonly object sessionRequestChargeLock = new(); + private readonly object historyLock = new(); + + private readonly SemaphoreSlim executionGate = new(1, 1); + + private readonly AsyncLocal ownsExecutionGate = new(); + + private long stateVersion; + private TokenCredential? activeCredential; private LineEditor? lineEditor; @@ -103,11 +114,16 @@ internal ShellInterpreter(string? configPath = null) if (File.Exists(this.HistoryFile)) { - foreach (var line in File.ReadAllLines(this.HistoryFile)) + string[] lines; + lock (HistoryFileLock) + { + lines = File.ReadAllLines(this.HistoryFile); + } + + foreach (var line in lines) { var decoded = DecodeHistoryLine(line); - this.history.Remove(decoded); - this.history.Add(decoded); + this.RecordHistoryEntry(decoded); } } @@ -233,7 +249,17 @@ internal long SessionRequestChargeGeneration internal Func IsInteractiveSession { get; set; } = static () => !Console.IsInputRedirected && !Console.IsOutputRedirected; - internal IReadOnlyList History => this.history; + internal IReadOnlyList History + { + get + { + // Snapshot: interactive readers must not enumerate the list while an MCP echo mutates it. + lock (this.historyLock) + { + return this.history.ToArray(); + } + } + } internal string? LastBuffer { get; set; } @@ -253,7 +279,17 @@ internal long SessionRequestChargeGeneration internal bool AppendErrRedirection { get; set; } - internal State State { get; set; } + internal long StateVersion => Interlocked.Read(ref this.stateVersion); + + internal State State + { + get; + set + { + field = value; + Interlocked.Increment(ref this.stateVersion); + } + } internal Program.CosmosShellOptions? Options { get; set; } @@ -507,6 +543,39 @@ public void CancelPrompt() /// A cancellation token to observe while waiting for the task to complete. /// A representing the result of the command execution. public async Task ExecuteCommandAsync(string command, CancellationToken token) + { + try + { + return await this.RunSerializedAsync(() => this.ExecuteCommandCoreAsync(command, token), token); + } + catch (OperationCanceledException) when (token.IsCancellationRequested) + { + this.Diagnostics?.LogCancelled(0, command); + return new CommandState(); + } + } + + internal async Task RunSerializedAsync(Func> operation, CancellationToken token) + { + if (this.ownsExecutionGate.Value) + { + return await operation(); + } + + await this.executionGate.WaitAsync(token); + try + { + this.ownsExecutionGate.Value = true; + return await operation(); + } + finally + { + this.ownsExecutionGate.Value = false; + this.executionGate.Release(); + } + } + + private async Task ExecuteCommandCoreAsync(string command, CancellationToken token) { using var activity = TracingBootstrap.StartCommandActivity("cosmosdbshell.command"); var isLocalEmulatorOperation = this.State is ConnectedState connectedState @@ -890,8 +959,7 @@ internal async Task RunAsync() if (!string.IsNullOrWhiteSpace(command)) { - this.history.Remove(command); - this.history.Add(command); + this.RecordHistoryEntry(command); this.SaveHistory(); CancellationToken token = UserCancellationTokenSource.Token; await this.ExecuteCommandAsync(command, token); @@ -988,7 +1056,14 @@ internal async Task RunCommandAsync(CommandState currentState, str internal void RecordRequestCharge(CommandState commandState) => this.RecordRequestCharge(commandState, this.SessionRequestChargeGeneration); - internal async Task ExecuteCosmosCommandAsync( + internal Task ExecuteCosmosCommandAsync( + CosmosCommand command, + CommandState commandState, + string commandText, + CancellationToken token) + => this.RunSerializedAsync(() => this.ExecuteCosmosCommandCoreAsync(command, commandState, commandText, token), token); + + private async Task ExecuteCosmosCommandCoreAsync( CosmosCommand command, CommandState commandState, string commandText, @@ -1758,13 +1833,31 @@ internal void PrintCommand(string cmdString) // Print the shell prompt similar to how it appears when typing command // AnsiConsole.Markup(new CosmosShellPrompt(this).GetPromptString()); // AnsiConsole.Write(" "); - var txt = ((IHighlighter)Instance).BuildHighlightedText(cmdString); - AnsiConsole.Write(txt); - AnsiConsole.WriteLine(); // Ensure the next output starts on a new line + this.RecordHistoryEntry(cmdString); - this.history.Remove(cmdString); - this.history.Add(cmdString); - this.Editor?.History.Add(cmdString); + try + { + this.SaveHistory(); + } + catch (Exception ex) when (ex is IOException or UnauthorizedAccessException) + { + // History is best-effort; an unwritable history file must not fail the command. + System.Diagnostics.Debug.WriteLine(ex); + } + + // Echoing and the line editor both need an ANSI terminal, which an MCP host may not + // provide. Neither may fail the command being announced. + try + { + var txt = ((IHighlighter)Instance).BuildHighlightedText(cmdString); + AnsiConsole.Write(txt); + AnsiConsole.WriteLine(); // Ensure the next output starts on a new line + this.Editor?.History.Add(cmdString); + } + catch (NotSupportedException) + { + Console.Out.WriteLine(cmdString); + } } internal CommandState PrintState(CommandState state, bool markAsRendered = false) @@ -2066,6 +2159,7 @@ protected virtual void Dispose(bool disposing) this.editorCancelTokenSource?.Dispose(); this.State?.Dispose(); this.Diagnostics?.Dispose(); + this.executionGate.Dispose(); } this.disposedValue = true; @@ -2132,7 +2226,7 @@ private LineEditor CreateLineEditor() lineEditor.KeyBindings.Add(ConsoleKey.S, ConsoleModifiers.Control, () => new ReverseSearchHistoryCommand(this, startsForward: true)); lineEditor.KeyBindings.Add(ConsoleKey.Tab, () => new CosmosCompleteCommand(this, AutoComplete.Next)); lineEditor.KeyBindings.Add(ConsoleKey.Tab, ConsoleModifiers.Control, () => new CosmosCompleteCommand(this, AutoComplete.Previous)); - foreach (var line in this.history) + foreach (var line in this.History) { lineEditor.History.Add(line); } @@ -2175,12 +2269,44 @@ private void Console_CancelKeyPress(object? sender, ConsoleCancelEventArgs e) private void SaveHistory() { - if (this.history.Count > MAXHISTORYITEMS) + lock (this.historyLock) { - this.history = [.. this.history.Skip(this.history.Count - MAXHISTORYITEMS)]; + if (this.history.Count > MAXHISTORYITEMS) + { + this.history = [.. this.history.Skip(this.history.Count - MAXHISTORYITEMS)]; + } + + // Written under the locks so concurrent interactive and MCP saves, and shells + // sharing the history file, cannot interleave. + lock (HistoryFileLock) + { + var options = new FileStreamOptions { Mode = FileMode.Create, Access = FileAccess.Write, Share = FileShare.Read }; + if (!OperatingSystem.IsWindows()) + { + // History can contain connection secrets; UnixCreateMode covers only new files. + options.UnixCreateMode = UnixFileMode.UserRead | UnixFileMode.UserWrite; + if (File.Exists(this.HistoryFile)) + { + File.SetUnixFileMode(this.HistoryFile, UnixFileMode.UserRead | UnixFileMode.UserWrite); + } + } + + using var writer = new StreamWriter(this.HistoryFile, new System.Text.UTF8Encoding(encoderShouldEmitUTF8Identifier: false), options); + foreach (var line in this.history) + { + writer.WriteLine(EncodeHistoryLine(line)); + } + } } + } - File.WriteAllLines(this.HistoryFile, this.history.Select(EncodeHistoryLine)); + private void RecordHistoryEntry(string entry) + { + lock (this.historyLock) + { + this.history.Remove(entry); + this.history.Add(entry); + } } [System.Diagnostics.CodeAnalysis.SuppressMessage("StyleCop.CSharp.OrderingRules", "SA1204", Justification = "History helpers are grouped with SaveHistory for cohesion.")] diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Mcp/ToolOperations.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Mcp/ToolOperations.cs index 3d46801..3336f74 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Mcp/ToolOperations.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Mcp/ToolOperations.cs @@ -10,6 +10,7 @@ namespace Azure.Data.Cosmos.Shell.Mcp; using System.Text.Json.Nodes; using Azure.Data.Cosmos.Shell.Commands; using Azure.Data.Cosmos.Shell.Core; +using Azure.Data.Cosmos.Shell.States; using Azure.Data.Cosmos.Shell.Util; using Microsoft.AspNetCore.Http.Metadata; using Microsoft.Extensions.Logging; @@ -208,6 +209,59 @@ internal static string FormatOptionForHistory(Option option, object? value) return $" --{option.Name[0]} {ShellLiteral.Quote(value?.ToString())}"; } + // Shell syntax cannot skip a positional, so a later value would bind to the omitted slot on replay. + internal static string? FindPositionalGap(IReadOnlyList parameters, IReadOnlyDictionary values) + { + Parameter? firstOmitted = null; + foreach (var parameter in parameters) + { + if (!IsPositionalSupplied(values, parameter)) + { + firstOmitted ??= parameter; + } + else if (firstOmitted != null) + { + return $"Parameter '{parameter.Name[0]}' requires the preceding positional parameter '{firstOmitted.Name[0]}'. Supply '{firstOmitted.Name[0]}' as well."; + } + } + + return null; + } + + internal static string FormatPositionalsForHistory(IReadOnlyList parameters, IReadOnlyDictionary values) + { + var sb = new StringBuilder(); + foreach (var parameter in parameters) + { + if (!IsPositionalSupplied(values, parameter)) + { + break; + } + + var value = values[parameter]; + if (value is Array array) + { + foreach (var element in array) + { + sb.Append(' ').Append(ShellLiteral.Quote(element?.ToString())); + } + } + else + { + sb.Append(' ').Append(ShellLiteral.Quote(value?.ToString())); + } + } + + return sb.ToString(); + } + + private static bool IsPositionalSupplied(IReadOnlyDictionary values, Parameter parameter) + { + return values.TryGetValue(parameter, out var value) + && value != null + && (value is not Array array || array.Length > 0); + } + internal static void ConfigurePaging(object command) { if (command is IPagedCommand paged) @@ -457,6 +511,8 @@ private async ValueTask OnCallToolsAsync( var cmd = command.CreateCommand(); ConfigurePaging(cmd); var suppliedParameters = new HashSet(StringComparer.OrdinalIgnoreCase); + var positionalValues = new Dictionary(); + var optionText = new StringBuilder(); if (parameters.Params.Arguments != null) { @@ -477,7 +533,7 @@ private async ValueTask OnCallToolsAsync( memberKind: "option", memberDisplay: $"--{option.Name[0]}", commandName: command.CommandName, - appendToHistory: value => sb.Append(FormatOptionForHistory(option, value))); + appendToHistory: value => optionText.Append(FormatOptionForHistory(option, value))); if (bindError != null) { return bindError; @@ -496,7 +552,7 @@ private async ValueTask OnCallToolsAsync( memberKind: "parameter", memberDisplay: parameter.Name[0], commandName: command.CommandName, - appendToHistory: value => sb.Append(' ').Append(ShellLiteral.Quote(value?.ToString()))); + appendToHistory: value => positionalValues[parameter] = value); if (bindError != null) { return bindError; @@ -549,6 +605,14 @@ private async ValueTask OnCallToolsAsync( return McpResponseFactory.CreateError(missingMessage, ShellInterpreter.Instance.State); } + var positionalGap = FindPositionalGap(command.Parameters, positionalValues); + if (positionalGap != null) + { + var gapMessage = $"Invalid positional arguments for command '{command.CommandName}': {positionalGap}"; + this.logger?.LogWarning("{Message}", gapMessage); + return McpResponseFactory.CreateError(gapMessage, ShellInterpreter.Instance.State); + } + var batchSubcommand = (cmd as BatchCommand)?.Subcommand?.Trim(); if (!string.IsNullOrEmpty(batchSubcommand) && !string.Equals(batchSubcommand, "run", StringComparison.OrdinalIgnoreCase)) @@ -558,27 +622,58 @@ private async ValueTask OnCallToolsAsync( return McpResponseFactory.CreateError(errorMessage, ShellInterpreter.Instance.State); } + // MCP argument order is not semantic, so render positionals in the order the shell binds them. + sb.Append(FormatPositionalsForHistory(command.Parameters, positionalValues)); + sb.Append(optionText); + + var server = parameters.Server; + Func>? elicit = + server?.ClientCapabilities?.Elicitation != null ? server.ElicitAsync : null; + return await this.ExecuteToolAsync(command, cmd, sb.ToString(), elicit, cancellationToken); + } + + internal async Task ExecuteToolAsync( + CommandFactory command, + CosmosCommand cmd, + string commandLine, + Func>? elicit, + CancellationToken cancellationToken) + { + var shell = ShellInterpreter.Instance; + long? confirmedVersion = null; if (RequiresConfirmation(command)) { - var server = parameters.Server; - Func>? elicit = - server?.ClientCapabilities?.Elicitation != null ? server.ElicitAsync : null; - - var confirmation = await this.ConfirmDestructiveAsync(elicit, command.CommandName, sb.ToString(), cancellationToken); + var snapshot = await shell.RunSerializedAsync( + () => Task.FromResult((Version: shell.StateVersion, Context: DescribeContext(shell.State))), cancellationToken); + var confirmation = await this.ConfirmDestructiveAsync( + elicit, command.CommandName, commandLine, cancellationToken, snapshot.Context); if (confirmation != null) { return confirmation; } + + confirmedVersion = snapshot.Version; } - this.logger?.LogTrace($"Invoking '{command.CommandName}'."); + this.logger?.LogTrace($"Requested '{command.CommandName}'."); try { - ShellInterpreter.Instance.PrintCommand(sb.ToString()); - var response = await ShellInterpreter.Instance.ExecuteCosmosCommandAsync(cmd, new CommandState(), command.CommandName, cancellationToken); - ShellInterpreter.Instance.CancelPrompt(); - return McpResponseFactory.CreateSuccess(response, ShellInterpreter.Instance.State); + return await shell.RunSerializedAsync( + async () => + { + if (confirmedVersion.HasValue && confirmedVersion.Value != shell.StateVersion) + { + return McpResponseFactory.CreateError( + "The shell context changed while awaiting confirmation. Nothing was executed. Retry the command and confirm its current target.", shell.State); + } + + shell.PrintCommand(commandLine); + var response = await shell.ExecuteCosmosCommandAsync(cmd, new CommandState(), command.CommandName, cancellationToken); + shell.CancelPrompt(); + return McpResponseFactory.CreateSuccess(response, shell.State); + }, + cancellationToken); } catch (Exception ex) { @@ -591,10 +686,16 @@ private async ValueTask OnCallToolsAsync( } finally { - this.logger?.LogTrace($"Finished executing '{command.CommandName}'."); + this.logger?.LogTrace($"Finished handling request for '{command.CommandName}'."); } } + private static string DescribeContext(State state) + { + var endpoint = state is ConnectedState connected ? connected.Client.Endpoint.ToString() : "(disconnected)"; + return $"Account: {endpoint}\nCurrent location: {McpResponseFactory.GetCurrentLocation(state) ?? "(none)"}\nExplicit database/container arguments in the command override this location."; + } + // Gates a destructive command behind an MCP elicitation confirmation. Returns // null when the operation is approved and should proceed; otherwise returns the // CallToolResult to send back (refusal, denial, or a failed confirmation). @@ -604,7 +705,8 @@ private async ValueTask OnCallToolsAsync( Func>? elicit, string commandName, string commandLine, - CancellationToken cancellationToken) + CancellationToken cancellationToken, + string? context = null) { if (elicit == null) { @@ -620,6 +722,7 @@ private async ValueTask OnCallToolsAsync( { Message = $"Confirm destructive operation. The agent wants to run: {commandLine}\n" + + (context is null ? string.Empty : context + "\n") + "This can permanently change or delete data in the connected Azure Cosmos DB account and cannot be undone. Approve this operation?", RequestedSchema = new ElicitRequestParams.RequestSchema(), }; diff --git a/CosmosDBShell/CosmosDBShell.csproj b/CosmosDBShell/CosmosDBShell.csproj index 1264dd3..85095bb 100644 --- a/CosmosDBShell/CosmosDBShell.csproj +++ b/CosmosDBShell/CosmosDBShell.csproj @@ -91,6 +91,7 @@ + diff --git a/CosmosDBShell/lang/en.ftl b/CosmosDBShell/lang/en.ftl index e97f96b..92b8642 100644 --- a/CosmosDBShell/lang/en.ftl +++ b/CosmosDBShell/lang/en.ftl @@ -476,6 +476,7 @@ command-import-dry-run-success = Dry run: { $count } valid { $count -> *[other] items } command-import-error-missing_file = A source file path is required. +command-import-error-invalid_csv = Invalid CSV record at line { $line }. command-import-error-file_not_found = File '{ $file }' was not found. command-import-error-blank_line = Line { $line } is blank. command-import-error-not_object = Line { $line } is not a JSON object. diff --git a/Directory.Packages.props b/Directory.Packages.props index eab837e..c555da6 100644 --- a/Directory.Packages.props +++ b/Directory.Packages.props @@ -10,6 +10,7 @@ + diff --git a/README.md b/README.md index 16071b6..5ad6316 100644 --- a/README.md +++ b/README.md @@ -33,6 +33,10 @@ A terminal-native shell for Azure Cosmos DB — navigate databases like a filesy - MCP server for AI/tool integration - Distributed tracing via OpenTelemetry (`--otel`): emits a sampled W3C `traceparent` on Cosmos requests, with optional OTLP export +Exports replace their destination only after successful completion, preserving an existing file on failure or cancellation. Imports stream records; CSV exports use temporary disk storage to discover columns without retaining all documents in memory. See [import/export](docs/commands.md#export). + +MCP command execution is serialized with the shell, and destructive confirmations are invalidated by connection or navigation changes. MCP invocations are echoed in the shell so their activity stays visible, and they are recorded in history alongside interactive commands. History remains fully replayable, including connection strings; treat its file as sensitive. See [MCP security](docs/mcp.md#security) and [history](docs/navigation.md#history). + ## Quick Start **Requirements:** .NET SDK 10.0+. diff --git a/docs/commands.md b/docs/commands.md index ac3ad34..495261c 100644 --- a/docs/commands.md +++ b/docs/commands.md @@ -652,7 +652,9 @@ Examples: ### export -Stream items from a container to a local file. Default format is JSON Lines (one compact JSON object per line); pass `--format=array` for a single JSON array, or `--format=csv` for CSV. Items are streamed end-to-end for JSON formats; CSV buffers items to compute the column set. The CSV separator follows the `COSMOSDB_SHELL_CSVSEP` environment variable (default `;`). +Stream items from a container to a local file. Default format is JSON Lines (one compact JSON object per line); pass `--format=array` for a single JSON array, or `--format=csv` for CSV. JSON formats stream incrementally. CSV spools documents to a private temporary file to compute the complete column set, keeping only the column names and current record in memory. Allow enough temporary disk space for the JSON spool as well as the destination export. The CSV separator follows the `COSMOSDB_SHELL_CSVSEP` environment variable (default `;`). + +All formats write to a temporary file in the destination directory and move it into place only after successful completion. An existing destination requires `--force` and is preserved if reading, writing, or cancellation interrupts the export. Temporary files are removed on normal completion and handled failures; an abrupt process termination can leave an unfinished destination-directory temporary file. Once `--max` items have been emitted, no further query pages are requested. ```text Usage: export [options] @@ -680,7 +682,7 @@ The summary line reports the number of items written and the total RU charge. ### import -Bulk-load items from a JSON Lines, JSON array, or CSV file into a container. Format is auto-detected: a `.csv` extension selects CSV, otherwise the first non-whitespace character is inspected (`[` ⇒ array, otherwise JSON Lines). It can be forced with `--format`. Default mode is `insert`; pass `--mode=upsert` to replace items that already exist. For CSV, the header row defines property names and every value is imported as a string; the CSV separator follows `COSMOSDB_SHELL_CSVSEP` (default `;`). JSON Lines and JSON array inputs are streamed item-by-item, but CSV import reads and parses the entire file into memory before importing, so very large CSV files can cause a significant memory spike. +Bulk-load items from a JSON Lines, JSON array, or CSV file into a container. Format is auto-detected: a `.csv` extension selects CSV, otherwise the first non-whitespace character is inspected (`[` ⇒ array, otherwise JSON Lines). It can be forced with `--format`. Default mode is `insert`; pass `--mode=upsert` to replace items that already exist. For CSV, the header row defines property names and every value is imported as a string; the CSV separator follows `COSMOSDB_SHELL_CSVSEP` (default `;`). All formats are read incrementally rather than loading the complete file into memory. CSV supports quoted separators, escaped quotes, and multiline fields; malformed records abort the import with their physical start line. Earlier writes are not rolled back, so use `--dry-run` first when the entire file must be validated before any writes. ```text Usage: import [options] diff --git a/docs/mcp.md b/docs/mcp.md index e05b7ea..8f8d1b6 100644 --- a/docs/mcp.md +++ b/docs/mcp.md @@ -81,6 +81,10 @@ Destructive commands (`delete`, `rm`, `rmcon`, `rmdb`) are gated behind an expli This replaces any opt-in write flag: destructive commands are always allowed to be invoked, but always require confirmation. +Confirmation includes the connected account endpoint and current navigation location alongside the command and its explicit target arguments. If the connection or navigation state changes while confirmation is pending, the approved command is refused without executing; retry it to confirm the new context. Even navigating away and back invalidates the pending confirmation. + +Shell and MCP command execution is serialized against the shared interpreter. Confirmation prompts do not hold the execution lock, so the shell remains usable while waiting. Clients still share a connection and navigation context: pass explicit `database` and `container` arguments for independent operations rather than relying on an earlier `cd` call. + The MCP confirmation applies even when a command is invoked with a force / no-prompt argument (for example `rmdb OldDB true`). That argument only skips the *interactive shell* prompt; it does not bypass the MCP elicitation gate. Database and container resource actions are executed through Azure Resource Manager when an ARM context is attached (Entra ID connections). MCP sessions connected with account keys, emulator credentials, or static data-plane tokens fall back to the Cosmos DB data plane for these actions. @@ -89,6 +93,10 @@ For deterministic ARM routing in multi-subscription environments, start the shel ### Data Exposure +MCP tool invocations are echoed as command lines in the shell window, so anyone watching the terminal can see what a connected client is doing. They are also recorded in the shell history. History entries are complete and replayable, including any supplied connection strings. Protect the history file accordingly. On Linux and macOS, the shell restricts the history file to its owner. + +Positional arguments must be supplied without gaps: a call that provides a positional parameter while omitting an earlier one is rejected, because the equivalent shell command line would bind the value to the omitted slot. + Your MCP client may use a remote LLM. Command outputs, query results, and file contents could be transmitted to external services. **Treat all shell output as potentially shared.** ### Best Practices diff --git a/docs/navigation.md b/docs/navigation.md index 6878e16..2af408a 100644 --- a/docs/navigation.md +++ b/docs/navigation.md @@ -236,6 +236,8 @@ There is no separate "enter multi-line mode" command — the shell enters and le Multi-line commands are saved to history as a single entry. When you recall one with `Up` / `Ctrl+P` or reverse-search (`Ctrl+R`), the full multi-line text is restored. History files written by older versions of the shell continue to load unchanged. +Commands are stored in full so they can be executed again, including connection strings containing account keys. Treat the `cmd_history` file in the shell configuration directory as sensitive: protect it with your user account's file permissions and do not share it. Use Entra ID to avoid storing account keys, or `--clear-history` to clear the saved history. MCP tool invocations are echoed as command lines and recorded in the same history. + ## Keyboard Shortcuts Available at the interactive prompt: diff --git a/l10n/CosmosDBShell.json b/l10n/CosmosDBShell.json index 82d03cd..05904bc 100644 --- a/l10n/CosmosDBShell.json +++ b/l10n/CosmosDBShell.json @@ -415,6 +415,7 @@ "command-import-error-blank_line": "Line {0} is blank.", "command-import-error-csv_pk_conflict": "CSV column \u0027{0}\u0027 conflicts with the partition key path \u0027{1}\u0027: the column holds a scalar value but the path requires it to be a nested object. Rename the column or choose a different partition key path.", "command-import-error-file_not_found": "File \u0027{0}\u0027 was not found.", + "command-import-error-invalid_csv": "Invalid CSV record at line {0}.", "command-import-error-invalid_line_json": "Line {0} is not valid JSON: {1}", "command-import-error-item_failed": "Line {0}: {1} - {2}", "command-import-error-item_status": "Line {0}: item returned status {1}.",