Add Engine Telemetry - Phase 1

.gitignore

@@ -38,5 +38,11 @@ dab-config*.json
 .env
 /docs/design/McpToolRegistryHotReload.md
 
+# Internal telemetry design drafts (keep local copies out of publication)
+/docs/design/engine-telemetry-functional-spec.md
+/docs/design/engine-telemetry-implementation-handoff.md
+/docs/design/engine-telemetry-requirements-review.md
+/docs/design/engine-telemetry-technical-design.md
+
 # Verify test files
 *.received.*

src/Azure.DataApiBuilder.Mcp/Core/McpEndpointRouteBuilderExtensions.cs

@@ -4,8 +4,12 @@
 using System.Diagnostics.CodeAnalysis;
 using Azure.DataApiBuilder.Config.ObjectModel;
 using Azure.DataApiBuilder.Core.Configurations;
+using Azure.DataApiBuilder.Core.Telemetry.Product;
+using Azure.DataApiBuilder.Mcp.Utils;
 using Microsoft.AspNetCore.Builder;
+using Microsoft.AspNetCore.Http;
 using Microsoft.AspNetCore.Routing;
+using Microsoft.Extensions.DependencyInjection;
 
 namespace Azure.DataApiBuilder.Mcp.Core
 {
@@ -37,7 +41,35 @@ public static IEndpointRouteBuilder MapDabMcp(
             string mcpPath = mcpOptions.Path ?? McpRuntimeOptions.DEFAULT_PATH;
 
             // Map the MCP endpoint
-            endpoints.MapMcp(mcpPath);
+            endpoints.MapMcp(mcpPath).Add(builder =>
+            {
+                RequestDelegate? next = builder.RequestDelegate;
+                if (next is null)
+                {
+                    return;
+                }
+
+                builder.RequestDelegate = async context =>
+                {
+                    if (context.RequestServices.GetService<EngineTelemetrySession>()?.IsEnabled != true)
+                    {
+                        await next(context);
+                        return;
+                    }
+
+                    Stream original = context.Response.Body;
+                    using McpProductResponseStream observer = new(original);
+                    context.Response.Body = observer;
+                    try
+                    {
+                        await next(context);
+                    }
+                    finally
+                    {
+                        context.Response.Body = original;
+                    }
+                };
+            });
 
             return endpoints;
         }

src/Azure.DataApiBuilder.Mcp/Core/McpServerConfiguration.cs

@@ -76,13 +76,14 @@ internal static IServiceCollection ConfigureMcpServer(this IServiceCollection se
                     }
 
                     return await McpTelemetryHelper.ExecuteWithTelemetryAsync(
-                        tool, toolName, arguments, request.Services, ct);
+                        tool, toolName, arguments, request.Services, ct, sdkResponseItems: request.Items);
                 }
                 finally
                 {
                     arguments?.Dispose();
                 }
             })
+            .WithMessageFilters(filters => filters.AddOutgoingFilter(McpProductResponseCompletion.Filter))
             .WithHttpTransport();
 
             // Configure underlying MCP server options

src/Azure.DataApiBuilder.Mcp/Core/McpStdioServer.cs

@@ -483,7 +483,6 @@ private async Task HandleCallToolAsync(JsonElement? id, JsonElement root, Cancel
                 // helpers that read IHttpContextAccessor will see the role. We also ensure the
                 // Simulator authentication handler can authenticate the user by flowing the
                 // Authorization header commonly used in tests/simulator scenarios.
-                CallToolResult callResult;
                 IConfiguration? configuration = _serviceProvider.GetService<IConfiguration>();
                 string? stdioRole = configuration?.GetValue<string>("MCP:Role");
                 if (!string.IsNullOrWhiteSpace(stdioRole))
@@ -512,8 +511,10 @@ private async Task HandleCallToolAsync(JsonElement? id, JsonElement root, Cancel
                     try
                     {
                         // Execute the tool with the scoped service provider so any scoped services resolve correctly.
-                        callResult = await McpTelemetryHelper.ExecuteWithTelemetryAsync(
-                            tool, toolName!, argsDoc, scopedProvider, ct);
+                        // Product completion must follow the successful stdout write, not just tool execution.
+                        await McpTelemetryHelper.ExecuteWithTelemetryAsync(
+                            tool, toolName!, argsDoc, scopedProvider, ct,
+                            writeStdioResponse: result => HandleCallToolAsync(id ?? default, result));
                     }
                     finally
                     {
@@ -526,11 +527,10 @@ private async Task HandleCallToolAsync(JsonElement? id, JsonElement root, Cancel
                 }
                 else
                 {
-                    callResult = await McpTelemetryHelper.ExecuteWithTelemetryAsync(
-                        tool, toolName!, argsDoc, _serviceProvider, ct);
+                    await McpTelemetryHelper.ExecuteWithTelemetryAsync(
+                        tool, toolName!, argsDoc, _serviceProvider, ct,
+                        writeStdioResponse: result => HandleCallToolAsync(id ?? default, result));
                 }
-
-                await HandleCallToolAsync(id ?? default, callResult);
             }
             finally
             {

src/Azure.DataApiBuilder.Mcp/Core/McpToolRegistryRefreshService.cs

@@ -3,6 +3,7 @@
 
 using Azure.DataApiBuilder.Config;
 using Azure.DataApiBuilder.Config.ObjectModel;
+using Azure.DataApiBuilder.Config.Telemetry;
 using Azure.DataApiBuilder.Core.Configurations;
 using Azure.DataApiBuilder.Core.Services.MetadataProviders;
 using Azure.DataApiBuilder.Mcp.Model;
@@ -138,6 +139,9 @@ private void OnConfigChanged(object? sender, HotReloadEventArgs args)
             }
             catch (Exception ex)
             {
+                // Keep serving the previous registry, while preserving the failed attempt's
+                // telemetry epoch rather than accepting a partially refreshed generation.
+                TelemetryFailureContext.Current?.RecordFailure(TelemetryFailureStage.Serving);
                 _logger.LogError(
                     ex,
                     "Failed to refresh the MCP tool registry after a runtime configuration change. " +

src/Azure.DataApiBuilder.Mcp/Utils/McpProductResponseCompletion.cs

@@ -0,0 +1,126 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+using Azure.DataApiBuilder.Core.Telemetry.Product;
+using Microsoft.AspNetCore.Http;
+using ModelContextProtocol.Protocol;
+using ModelContextProtocol.Server;
+
+namespace Azure.DataApiBuilder.Mcp.Utils;
+
+/// <summary>
+/// Correlates one tool response through the SDK's nonserialized Items, without retaining an
+/// ID or payload. A byte-opaque stream observer confirms that this response actually flushed.
+/// </summary>
+internal sealed class McpProductResponseCompletion
+{
+    private const string ITEM_KEY = "DAB.ProductTelemetry.ResponseCompletion";
+    private static readonly AsyncLocal<McpProductResponseCompletion?> _sending = new();
+    private readonly EngineTelemetryRequestScope _request;
+    private CancellationTokenRegistration _cancellation;
+    private int _completed;
+    private int _wrote;
+    private int _flushed;
+
+    internal McpProductResponseCompletion(EngineTelemetryRequestScope request, CancellationToken cancellationToken)
+    {
+        _request = request;
+        _cancellation = cancellationToken.UnsafeRegister(static value =>
+            ((McpProductResponseCompletion)value!).Complete(EngineTelemetryOutcome.Canceled), this);
+        if (Volatile.Read(ref _completed) != 0)
+        {
+            _cancellation.Unregister();
+        }
+    }
+
+    internal static McpProductResponseCompletion Attach(IDictionary<string, object?> items, EngineTelemetryRequestScope request,
+        CancellationToken cancellationToken = default)
+    {
+        McpProductResponseCompletion completion = new(request, cancellationToken);
+        items[ITEM_KEY] = completion;
+        return completion;
+    }
+
+    internal void Abandon(IDictionary<string, object?> items)
+    {
+        _cancellation.Unregister();
+        items.Remove(ITEM_KEY);
+    }
+
+    internal static void Wrote() => _sending.Value?.RecordWrite();
+
+    internal static void Flushed() => _sending.Value?.RecordFlush();
+
+    internal static void WriteFailed(bool canceled) => _sending.Value?.Complete(canceled
+        ? EngineTelemetryOutcome.Canceled : EngineTelemetryOutcome.Failure);
+
+    internal static McpMessageHandler Filter(McpMessageHandler next) => async (context, cancellationToken) =>
+    {
+        // The SDK copies the request context onto ordinary responses. SDK-generated error
+        // messages may not preserve Items; the tool wrapper already records thrown failures.
+        if (context.JsonRpcMessage is not JsonRpcResponse ||
+            context.JsonRpcMessage.Context?.Items is not { } items ||
+            !items.TryGetValue(ITEM_KEY, out object? value) || value is not McpProductResponseCompletion completion)
+        {
+            await next(context, cancellationToken).ConfigureAwait(false);
+            return;
+        }
+
+        McpProductResponseCompletion? previous = _sending.Value;
+        _sending.Value = completion;
+        try
+        {
+            await next(context, cancellationToken).ConfigureAwait(false);
+            if (!completion._request.IsCompleted)
+            {
+                // Disposed transports can silently drop a send; an SDK return alone is not
+                // delivery. Preserve an unknown outcome if no write+flush was observed.
+                completion.Complete(Volatile.Read(ref completion._flushed) != 0
+                    ? completion._request.Outcome : EngineTelemetryOutcome.Unknown);
+            }
+        }
+        catch (OperationCanceledException)
+        {
+            completion.Complete(EngineTelemetryOutcome.Canceled);
+            throw;
+        }
+        catch (Exception)
+        {
+            completion.Complete(EngineTelemetryOutcome.Failure);
+            throw;
+        }
+        finally
+        {
+            _sending.Value = previous;
+            items.Remove(ITEM_KEY);
+        }
+    };
+
+    private void RecordWrite() => Volatile.Write(ref _wrote, 1);
+
+    private void RecordFlush()
+    {
+        if (Volatile.Read(ref _wrote) != 0)
+        {
+            Volatile.Write(ref _flushed, 1);
+        }
+    }
+
+    private void Complete(EngineTelemetryOutcome outcome)
+    {
+        if (Interlocked.Exchange(ref _completed, 1) != 0)
+        {
+            return;
+        }
+
+        _cancellation.Unregister();
+        try
+        {
+            _request.Complete(outcome, Volatile.Read(ref _flushed) != 0 ? StatusCodes.Status200OK : null);
+        }
+        catch (Exception)
+        {
+            // Optional product recording must never change SDK response delivery.
+        }
+    }
+}

src/Azure.DataApiBuilder.Mcp/Utils/McpProductResponseStream.cs

@@ -0,0 +1,87 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+namespace Azure.DataApiBuilder.Mcp.Utils;
+
+/// <summary>
+/// Delegates every byte unchanged. Only successful write/flush boundaries and closed failure
+/// categories are observed; no buffer, message text, request ID or exception is retained.
+/// </summary>
+internal sealed class McpProductResponseStream(Stream inner) : Stream
+{
+    public override bool CanRead => inner.CanRead;
+    public override bool CanSeek => inner.CanSeek;
+    public override bool CanWrite => inner.CanWrite;
+    public override long Length => inner.Length;
+    public override long Position { get => inner.Position; set => inner.Position = value; }
+
+    public override void Flush()
+    {
+        try
+        {
+            inner.Flush();
+            McpProductResponseCompletion.Flushed();
+        }
+        catch (Exception exception)
+        {
+            McpProductResponseCompletion.WriteFailed(exception is OperationCanceledException);
+            throw;
+        }
+    }
+
+    public override async Task FlushAsync(CancellationToken cancellationToken)
+    {
+        try
+        {
+            await inner.FlushAsync(cancellationToken).ConfigureAwait(false);
+            McpProductResponseCompletion.Flushed();
+        }
+        catch (Exception exception)
+        {
+            McpProductResponseCompletion.WriteFailed(exception is OperationCanceledException || cancellationToken.IsCancellationRequested);
+            throw;
+        }
+    }
+
+    public override void Write(byte[] buffer, int offset, int count)
+    {
+        try
+        {
+            inner.Write(buffer, offset, count);
+            if (count > 0)
+            {
+                McpProductResponseCompletion.Wrote();
+            }
+        }
+        catch (Exception exception)
+        {
+            McpProductResponseCompletion.WriteFailed(exception is OperationCanceledException);
+            throw;
+        }
+    }
+
+    public override Task WriteAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken)
+        => WriteAsync(buffer.AsMemory(offset, count), cancellationToken).AsTask();
+
+    public override async ValueTask WriteAsync(ReadOnlyMemory<byte> buffer, CancellationToken cancellationToken = default)
+    {
+        try
+        {
+            await inner.WriteAsync(buffer, cancellationToken).ConfigureAwait(false);
+            if (!buffer.IsEmpty)
+            {
+                McpProductResponseCompletion.Wrote();
+            }
+        }
+        catch (Exception exception)
+        {
+            McpProductResponseCompletion.WriteFailed(exception is OperationCanceledException || cancellationToken.IsCancellationRequested);
+            throw;
+        }
+    }
+
+    public override int Read(byte[] buffer, int offset, int count) => inner.Read(buffer, offset, count);
+    public override long Seek(long offset, SeekOrigin origin) => inner.Seek(offset, origin);
+    public override void SetLength(long value) => inner.SetLength(value);
+    // The ASP.NET response owns the inner stream; this observer must never dispose it.
+}