Add --otel distributed tracing support

CosmosDBShell.Tests/Runtime/TracingBootstrapTests.cs

@@ -0,0 +1,43 @@
+// ------------------------------------------------------------
+// Copyright (c) Microsoft Corporation.  All rights reserved.
+// ------------------------------------------------------------
+
+namespace CosmosShell.Tests;
+
+using System.Diagnostics;
+
+using Azure.Data.Cosmos.Shell.Core;
+
+using Xunit;
+
+public class TracingBootstrapTests
+{
+    [Fact]
+    public void StartCommandActivity_WithoutProvider_ReturnsNull()
+    {
+        using var activity = TracingBootstrap.StartCommandActivity("cosmosdbshell.command");
+
+        Assert.Null(activity);
+    }
+
+    [Fact]
+    public void StartCommandActivity_WhenInitialized_ProducesRecordedActivity()
+    {
+        using var tracing = TracingBootstrap.Initialize(otlpEndpoint: null);
+
+        using var activity = TracingBootstrap.StartCommandActivity("cosmosdbshell.command");
+
+        Assert.NotNull(activity);
+        Assert.True(activity!.Recorded);
+        Assert.True(activity.ActivityTraceFlags.HasFlag(ActivityTraceFlags.Recorded));
+    }
+
+    [Fact]
+    public void Initialize_SetsAzureActivitySourceSwitch()
+    {
+        using var tracing = TracingBootstrap.Initialize(otlpEndpoint: null);
+
+        Assert.True(
+            AppContext.TryGetSwitch("Azure.Experimental.EnableActivitySource", out var enabled) && enabled);
+    }
+}

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

@@ -265,6 +265,7 @@ public void CancelPrompt()
     /// <returns>A <see cref="CommandState"/> representing the result of the command execution.</returns>
     public async Task<CommandState> ExecuteCommandAsync(string command, CancellationToken token)
     {
+        using var activity = TracingBootstrap.StartCommandActivity("cosmosdbshell.command");
         var state = new CommandState();
         state.SetFormat(Environment.GetEnvironmentVariable("COSMOSDB_SHELL_FORMAT"));
 
@@ -1316,7 +1317,10 @@ private static CosmosClientOptions CreateClientOptions(string connectionString,
         {
             ApplicationName = "CosmosDBShell",
             ConnectionMode = requestedMode,
-            CosmosClientTelemetryOptions = new CosmosClientTelemetryOptions(),
+            CosmosClientTelemetryOptions = new CosmosClientTelemetryOptions
+            {
+                DisableDistributedTracing = false,
+            },
             UseSystemTextJsonSerializerWithOptions = new JsonSerializerOptions()
             {
                 DefaultIgnoreCondition = JsonIgnoreCondition.WhenWritingNull,

CosmosDBShell/Azure.Data.Cosmos.Shell.Core/TracingBootstrap.cs

@@ -0,0 +1,79 @@
+// ------------------------------------------------------------
+// Copyright (c) Microsoft Corporation.  All rights reserved.
+// ------------------------------------------------------------
+
+namespace Azure.Data.Cosmos.Shell.Core;
+
+using System.Diagnostics;
+using OpenTelemetry;
+using OpenTelemetry.Resources;
+using OpenTelemetry.Trace;
+
+/// <summary>
+/// Owns the distributed-tracing lifecycle for the shell. When enabled it registers
+/// a <see cref="TracerProvider"/> that records every activity, which causes the
+/// Azure Cosmos DB SDK to emit a sampled W3C <c>traceparent</c>
+/// (the <c>-01</c> flag) on its outgoing requests. An OTLP exporter is added only
+/// when an endpoint is supplied.
+/// </summary>
+public sealed class TracingBootstrap : IDisposable
+{
+    /// <summary>
+    /// Name of the <see cref="ActivitySource"/> used for per-command root activities.
+    /// </summary>
+    public const string ActivitySourceName = "CosmosDBShell";
+
+    private const string CosmosOperationSourceName = "Azure.Cosmos.Operation";
+
+    private static readonly ActivitySource SharedSource = new(ActivitySourceName);
+
+    private readonly TracerProvider provider;
+
+    private TracingBootstrap(TracerProvider provider)
+    {
+        this.provider = provider;
+    }
+
+    /// <summary>
+    /// Enables distributed tracing for the current process. Sets the Azure SDK
+    /// experimental switch required to emit <c>traceparent</c> headers and builds a
+    /// tracer provider that records all activities.
+    /// </summary>
+    /// <param name="otlpEndpoint">Optional OTLP endpoint to export spans to. When null or empty, no exporter is added and tracing only propagates a sampled <c>traceparent</c> on the wire.</param>
+    /// <returns>A <see cref="TracingBootstrap"/> that must be disposed to flush and tear down the provider.</returns>
+    public static TracingBootstrap Initialize(string? otlpEndpoint)
+    {
+        // Required so the Azure.Core HTTP pipeline writes a W3C traceparent header.
+        AppContext.SetSwitch("Azure.Experimental.EnableActivitySource", true);
+
+        var builder = Sdk.CreateTracerProviderBuilder()
+            .SetResourceBuilder(ResourceBuilder.CreateDefault().AddService(ActivitySourceName))
+            .SetSampler(new AlwaysOnSampler())
+            .AddSource(CosmosOperationSourceName)
+            .AddSource(ActivitySourceName);
+
+        if (!string.IsNullOrWhiteSpace(otlpEndpoint))
+        {
+            builder.AddOtlpExporter(options => options.Endpoint = new Uri(otlpEndpoint));
+        }
+
+        return new TracingBootstrap(builder.Build());
+    }
+
+    /// <summary>
+    /// Starts a root activity for a shell command. Returns null when tracing is not
+    /// enabled, so callers incur no overhead in the common case.
+    /// </summary>
+    /// <param name="name">The activity name.</param>
+    /// <returns>The started activity, or null when no tracer is listening.</returns>
+    public static Activity? StartCommandActivity(string name)
+    {
+        return SharedSource.StartActivity(name, ActivityKind.Client);
+    }
+
+    /// <inheritdoc />
+    public void Dispose()
+    {
+        this.provider.Dispose();
+    }
+}

CosmosDBShell/CosmosDBShell.csproj

@@ -95,6 +95,8 @@
     <PackageReference Include="ModelContextProtocol" />
     <PackageReference Include="ModelContextProtocol.AspNetCore" />
     <PackageReference Include="OmniSharp.Extensions.LanguageServer" />
+    <PackageReference Include="OpenTelemetry" />
+    <PackageReference Include="OpenTelemetry.Exporter.OpenTelemetryProtocol" />
     <PackageReference Include="RadLine" />
     <PackageReference Include="Azure.Core" />
     <PackageReference Include="Newtonsoft.Json" />

CosmosDBShell/Program.cs

@@ -39,6 +39,7 @@ public static async Task Main(string[] args)
         }
 
         IHost? host = null;
+        TracingBootstrap? tracing = null;
         try
         {
             // --help / --version handled manually so we can render our own
@@ -108,6 +109,31 @@ public static async Task Main(string[] args)
                 o.McpPort = mcpValue ?? DefaultMcpPort;
             }
 
+            // --otel supports an optional value: when present without an endpoint,
+            // tracing is still enabled (emitting a sampled traceparent) and the
+            // OTLP endpoint, if any, falls back to the standard environment variable.
+            var otelResult = parseResult.FindResultFor(optionMap.Otel);
+            if (otelResult is not null)
+            {
+                o.EnableTracing = true;
+                var otelValue = parseResult.GetValueForOption(optionMap.Otel);
+                o.OtlpEndpoint = string.IsNullOrWhiteSpace(otelValue)
+                    ? Environment.GetEnvironmentVariable("OTEL_EXPORTER_OTLP_ENDPOINT")
+                    : otelValue;
+
+                // Validate the endpoint up front so a malformed --otel value (or
+                // OTEL_EXPORTER_OTLP_ENDPOINT) yields a clean error instead of an
+                // unhandled exception when the exporter is created.
+                if (!string.IsNullOrWhiteSpace(o.OtlpEndpoint)
+                    && !Uri.TryCreate(o.OtlpEndpoint, UriKind.Absolute, out _))
+                {
+                    Environment.ExitCode = 1;
+                    ShellInterpreter.WriteLine(MessageService.GetArgsString(
+                        "otel-error-invalid-endpoint", "endpoint", o.OtlpEndpoint));
+                    return;
+                }
+            }
+
             if (o.StartLspServer)
             {
                 // Already handled above, but keep for completeness
@@ -149,6 +175,13 @@ public static async Task Main(string[] args)
 
             ShellInterpreter.Instance.Options = o;
 
+            // Enable distributed tracing before any CosmosClient is created so the
+            // Azure SDK pipeline emits a sampled W3C traceparent on its requests.
+            if (o.EnableTracing)
+            {
+                tracing = TracingBootstrap.Initialize(o.OtlpEndpoint);
+            }
+
             if (o.ConnectionString != null)
             {
                 using var connectTokenSource = ShellInterpreter.UserCancellationTokenSource;
@@ -286,6 +319,7 @@ await ShellInterpreter.Instance.ConnectAsync(
         {
             ShellInterpreter.Instance.Dispose();
             host?.Dispose();
+            tracing?.Dispose();
         }
     }
 
@@ -451,6 +485,11 @@ private static (RootCommand Command, OptionMap Map) BuildRootCommand()
         var verbose = new Option<bool>("--verbose", MessageService.GetString("help-Verbose"));
         var theme = new Option<string?>("--theme", MessageService.GetString("help-Theme"));
 
+        var otel = new Option<string?>("--otel", MessageService.GetString("help-Otel"))
+        {
+            Arity = ArgumentArity.ZeroOrOne,
+        };
+
         var root = new RootCommand("Cosmos DB Shell")
         {
             colorSystem,
@@ -471,6 +510,7 @@ private static (RootCommand Command, OptionMap Map) BuildRootCommand()
             lspStdio,
             verbose,
             theme,
+            otel,
         };
 
         var map = new OptionMap(
@@ -491,7 +531,8 @@ private static (RootCommand Command, OptionMap Map) BuildRootCommand()
             startLspServer,
             lspStdio,
             verbose,
-            theme);
+            theme,
+            otel);
 
         return (root, map);
     }
@@ -517,6 +558,7 @@ private static string BuildHelpText()
             [map.ConnectResourceGroup] = "<name>",
             [map.McpPort] = "[<port>]",
             [map.Theme] = "<name>",
+            [map.Otel] = "[<endpoint>]",
         };
 
         var rows = new List<(string Label, string? Description)>();
@@ -620,7 +662,8 @@ private sealed record OptionMap(
         Option<bool> StartLspServer,
         Option<bool> LspStdio,
         Option<bool> Verbose,
-        Option<string?> Theme);
+        Option<string?> Theme,
+        Option<string?> Otel);
 
     /// <summary>
     /// Maps the most common <c>System.CommandLine</c> parse error messages
@@ -698,5 +741,9 @@ public class CosmosShellOptions
         public bool Verbose { get; set; }
 
         public string? Theme { get; set; }
+
+        public bool EnableTracing { get; set; }
+
+        public string? OtlpEndpoint { get; set; }
     }
 }

CosmosDBShell/lang/en.ftl

@@ -622,7 +622,9 @@ help-EnableLspServer = Enable Language Server Protocol (LSP) server for editor i
 help-McpPort = Enable MCP HTTP server. Optionally specify a port with --mcp <port>; default is 6128.
 help-Verbose = Print full exception details instead of only the message.
 help-Theme = Color theme profile to apply at startup. Falls back to the COSMOSDB_SHELL_THEME environment variable.
+help-Otel = Enable distributed tracing so requests carry a sampled W3C traceparent. Optionally specify an OTLP endpoint with --otel <endpoint>; falls back to the OTEL_EXPORTER_OTLP_ENDPOINT environment variable.
 mcp-error-invalid-port = Error: --mcp port must be greater than 0.
+otel-error-invalid-endpoint = Error: --otel endpoint '{ $endpoint }' is not a valid absolute URI.
 
 warning-unknown-theme = Unknown theme '{ $name }'. Available themes: { $themes }. Falling back to default.
 

Directory.Packages.props

@@ -19,6 +19,8 @@
     <PackageVersion Include="Nerdbank.GitVersioning" Version="3.9.50" />
     <PackageVersion Include="Newtonsoft.Json" Version="13.0.5-beta1" />
     <PackageVersion Include="Microsoft.Extensions.Hosting" Version="11.0.0-preview.2.26159.112" />
+    <PackageVersion Include="OpenTelemetry" Version="1.15.3" />
+    <PackageVersion Include="OpenTelemetry.Exporter.OpenTelemetryProtocol" Version="1.15.3" />
     <PackageVersion Include="ModelContextProtocol" Version="1.1.0" />
     <PackageVersion Include="ModelContextProtocol.AspNetCore" Version="1.1.0" />
     <PackageVersion Include="OmniSharp.Extensions.LanguageServer" Version="0.19.9" />

README.md

@@ -15,6 +15,7 @@ Lightweight CLI for Azure Cosmos DB.
 - Pipelines and scripting with variables, loops, functions
 - Multi-line input at the prompt — automatic continuation for unclosed blocks/strings, plus explicit `\` line continuation ([docs](docs/navigation.md#multi-line-input))
 - MCP server for AI/tool integration
+- Distributed tracing via OpenTelemetry (`--otel`): emits a sampled W3C `traceparent` on Cosmos requests, with optional OTLP export
 
 ## Quick Start
 
@@ -130,6 +131,7 @@ Packaging runs produce preview versions in the form `1.0.<run>-preview.<branch>`
 | `--connect-subscription <id>` | Azure subscription ID for ARM database and container operations |
 | `--connect-resource-group <name>` | Azure resource group name for ARM database and container operations |
 | `--mcp [port]` | Enable MCP server on the given port, or `6128` by default |
+| `--otel [endpoint]` | Enable distributed tracing (sampled W3C `traceparent`); optional OTLP `endpoint`, else `OTEL_EXPORTER_OTLP_ENDPOINT` |
 | `--verbose` | Print full exception details |
 | `--color-system <n>` | Colors: 0=off, 1=standard, 2=truecolor (alias: `--cs`) |
 | `--theme <name>` | Color theme profile to apply at startup (`default`, `light`, `dark`, `monochrome`). Falls back to `COSMOSDB_SHELL_THEME`. |

docs/navigation.md

@@ -269,6 +269,7 @@ Start the shell with options to customize behavior:
 | `--connect-subscription <id>` | Azure subscription ID for ARM database and container operations at startup |
 | `--connect-resource-group <name>` | Azure resource group name for ARM database and container operations at startup |
 | `--mcp [port]` | Enable MCP (Model Context Protocol) server on the given port, or `6128` by default |
+| `--otel [endpoint]` | Enable distributed tracing so requests carry a sampled W3C `traceparent`. Optionally export spans to an OTLP `endpoint`; falls back to the `OTEL_EXPORTER_OTLP_ENDPOINT` environment variable |
 | `--color-system <n>` | Color scheme: 0=off, 1=standard, 2=truecolor (alias: `--cs`) |
 | `--clear-history` | Clear command history on start |
 | `--help` | Show usage information |
@@ -282,6 +283,7 @@ Start the shell with options to customize behavior:
 | `COSMOSDB_SHELL_ACCOUNT_KEY` | Account key for authentication |
 | `COSMOSDB_SHELL_FORMAT` | Default output format |
 | `COSMOSDB_SHELL_CSVSEP` | CSV column separator |
+| `OTEL_EXPORTER_OTLP_ENDPOINT` | Default OTLP endpoint used by `--otel` when no endpoint is supplied |
 
 **Examples:**
 
@@ -300,4 +302,10 @@ cosmosdbshell --mcp
 
 # Start with MCP server enabled on a custom port
 cosmosdbshell --mcp 5050
+
+# Enable distributed tracing (emits a sampled traceparent on Cosmos requests)
+cosmosdbshell --otel
+
+# Enable distributed tracing and export spans to an OTLP collector
+cosmosdbshell --otel http://localhost:4317
 ```