Add euid cookie value to logs and traces (#2237)

* Add euid cookie value to logs and traces

* Handle CodeQL suggestions

* Revert ElasticsearchGateway.cs

* Potential fix for pull request finding 'Missing Dispose call on local IDisposable'

Co-authored-by: Copilot Autofix powered by AI <223894421+github-code-quality[bot]@users.noreply.github.com>

---------

Co-authored-by: Copilot Autofix powered by AI <223894421+github-code-quality[bot]@users.noreply.github.com>

Author: reakaleek

Directory.Packages.props

@@ -93,6 +93,8 @@
     <PackageVersion Include="FsUnit.xUnit" Version="7.0.1" />
     <PackageVersion Include="GitHubActionsTestLogger" Version="2.4.1" />
     <PackageVersion Include="JetBrains.Annotations" Version="2024.3.0" />
+    <PackageVersion Include="Microsoft.AspNetCore.Mvc.Testing" Version="10.0.0" />
+    <PackageVersion Include="Microsoft.AspNetCore.TestHost" Version="10.0.0" />
     <PackageVersion Include="Microsoft.NET.Test.Sdk" Version="18.0.1" />
     <PackageVersion Include="OpenTelemetry.Exporter.InMemory" Version="1.13.0" />
     <PackageVersion Include="System.IO.Abstractions.TestingHelpers" Version="22.0.16" />

src/api/Elastic.Documentation.Api.Core/AskAi/AskAiUsecase.cs

@@ -13,7 +13,7 @@ public class AskAiUsecase(
 	IStreamTransformer streamTransformer,
 	ILogger<AskAiUsecase> logger)
 {
-	private static readonly ActivitySource AskAiActivitySource = new("Elastic.Documentation.Api.AskAi");
+	private static readonly ActivitySource AskAiActivitySource = new(TelemetryConstants.AskAiSourceName);
 
 	public async Task<Stream> AskAi(AskAiRequest askAiRequest, Cancel ctx)
 	{

src/api/Elastic.Documentation.Api.Core/TelemetryConstants.cs

@@ -0,0 +1,28 @@
+// Licensed to Elasticsearch B.V under one or more agreements.
+// Elasticsearch B.V licenses this file to you under the Apache 2.0 License.
+// See the LICENSE file in the project root for more information
+
+namespace Elastic.Documentation.Api.Core;
+
+/// <summary>
+/// Constants for OpenTelemetry instrumentation in the Docs API.
+/// </summary>
+public static class TelemetryConstants
+{
+	/// <summary>
+	/// ActivitySource name for AskAi operations.
+	/// Used in AskAiUsecase to create spans.
+	/// </summary>
+	public const string AskAiSourceName = "Elastic.Documentation.Api.AskAi";
+
+	/// <summary>
+	/// ActivitySource name for StreamTransformer operations.
+	/// Used in stream transformer implementations to create spans.
+	/// </summary>
+	public const string StreamTransformerSourceName = "Elastic.Documentation.Api.StreamTransformer";
+
+	/// <summary>
+	/// Tag/baggage name used to annotate spans with the user's EUID value.
+	/// </summary>
+	public const string UserEuidAttributeName = "user.euid";
+}

src/api/Elastic.Documentation.Api.Infrastructure/Adapters/AskAi/StreamTransformerBase.cs

@@ -20,7 +20,7 @@ public abstract class StreamTransformerBase(ILogger logger) : IStreamTransformer
 	protected ILogger Logger { get; } = logger;
 
 	// ActivitySource for tracing streaming operations
-	private static readonly ActivitySource StreamTransformerActivitySource = new("Elastic.Documentation.Api.StreamTransformer");
+	private static readonly ActivitySource StreamTransformerActivitySource = new(TelemetryConstants.StreamTransformerSourceName);
 
 	/// <summary>
 	/// Get the agent ID for this transformer

src/api/Elastic.Documentation.Api.Infrastructure/EuidSpanProcessor.cs

@@ -0,0 +1,32 @@
+// Licensed to Elasticsearch B.V under one or more agreements.
+// Elasticsearch B.V licenses this file to you under the Apache 2.0 License.
+// See the LICENSE file in the project root for more information
+
+using System.Diagnostics;
+using Elastic.Documentation.Api.Core;
+using OpenTelemetry;
+
+namespace Elastic.Documentation.Api.Infrastructure;
+
+/// <summary>
+/// OpenTelemetry span processor that automatically adds user.euid tag to all spans
+/// when it exists in the activity baggage.
+/// This ensures the euid is present on all spans (root and children) without manual propagation.
+/// </summary>
+public class EuidSpanProcessor : BaseProcessor<Activity>
+{
+	public override void OnStart(Activity activity)
+	{
+		// Check if euid exists in baggage (set by ASP.NET Core request enrichment)
+		var euid = activity.GetBaggageItem(TelemetryConstants.UserEuidAttributeName);
+		if (!string.IsNullOrEmpty(euid))
+		{
+			// Add as a tag to this span if not already present
+			var hasEuidTag = activity.TagObjects.Any(t => t.Key == TelemetryConstants.UserEuidAttributeName);
+			if (!hasEuidTag)
+			{
+				_ = activity.SetTag(TelemetryConstants.UserEuidAttributeName, euid);
+			}
+		}
+	}
+}

src/api/Elastic.Documentation.Api.Infrastructure/Middleware/EuidLoggingMiddleware.cs

@@ -0,0 +1,45 @@
+// Licensed to Elasticsearch B.V under one or more agreements.
+// Elasticsearch B.V licenses this file to you under the Apache 2.0 License.
+// See the LICENSE file in the project root for more information
+
+using Elastic.Documentation.Api.Core;
+using Microsoft.AspNetCore.Builder;
+using Microsoft.AspNetCore.Http;
+using Microsoft.Extensions.Logging;
+
+namespace Elastic.Documentation.Api.Infrastructure.Middleware;
+
+/// <summary>
+/// Middleware that adds the euid cookie value to the logging scope for all subsequent log entries in the request.
+/// </summary>
+public class EuidLoggingMiddleware(RequestDelegate next, ILogger<EuidLoggingMiddleware> logger)
+{
+	public async Task InvokeAsync(HttpContext context)
+	{
+		// Try to get the euid cookie
+		if (context.Request.Cookies.TryGetValue("euid", out var euid) && !string.IsNullOrEmpty(euid))
+		{
+			// Add euid to logging scope so it appears in all log entries for this request
+			using (logger.BeginScope(new Dictionary<string, object> { [TelemetryConstants.UserEuidAttributeName] = euid }))
+			{
+				await next(context);
+			}
+		}
+		else
+		{
+			await next(context);
+		}
+	}
+}
+
+/// <summary>
+/// Extension methods for registering the EuidLoggingMiddleware.
+/// </summary>
+public static class EuidLoggingMiddlewareExtensions
+{
+	/// <summary>
+	/// Adds the EuidLoggingMiddleware to the application pipeline.
+	/// This middleware enriches logs with the euid cookie value.
+	/// </summary>
+	public static IApplicationBuilder UseEuidLogging(this IApplicationBuilder app) => app.UseMiddleware<EuidLoggingMiddleware>();
+}

src/api/Elastic.Documentation.Api.Infrastructure/OpenTelemetryExtensions.cs

@@ -2,7 +2,10 @@
 // Elasticsearch B.V licenses this file to you under the Apache 2.0 License.
 // See the LICENSE file in the project root for more information
 
+using System.Diagnostics;
+using Elastic.Documentation.Api.Core;
 using Elastic.OpenTelemetry;
+using Microsoft.AspNetCore.Http;
 using Microsoft.Extensions.Hosting;
 using OpenTelemetry;
 using OpenTelemetry.Metrics;
@@ -12,6 +15,35 @@ namespace Elastic.Documentation.Api.Infrastructure;
 
 public static class OpenTelemetryExtensions
 {
+	/// <summary>
+	/// Configures tracing for the Docs API with sources, instrumentation, and enrichment.
+	/// This is the shared configuration used in both production and tests.
+	/// </summary>
+	public static TracerProviderBuilder AddDocsApiTracing(this TracerProviderBuilder builder)
+	{
+		_ = builder
+			.AddSource(TelemetryConstants.AskAiSourceName)
+			.AddSource(TelemetryConstants.StreamTransformerSourceName)
+			.AddAspNetCoreInstrumentation(aspNetCoreOptions =>
+			{
+				// Enrich spans with custom attributes from HTTP context
+				aspNetCoreOptions.EnrichWithHttpRequest = (activity, httpRequest) =>
+				{
+					// Add euid cookie value to span attributes and baggage
+					if (httpRequest.Cookies.TryGetValue("euid", out var euid) && !string.IsNullOrEmpty(euid))
+					{
+						_ = activity.SetTag(TelemetryConstants.UserEuidAttributeName, euid);
+						// Add to baggage so it propagates to all child spans
+						_ = activity.AddBaggage(TelemetryConstants.UserEuidAttributeName, euid);
+					}
+				};
+			})
+			.AddProcessor<EuidSpanProcessor>() // Automatically add euid to all child spans
+			.AddHttpClientInstrumentation();
+
+		return builder;
+	}
+
 	/// <summary>
 	/// Configures Elastic OpenTelemetry (EDOT) for the Docs API.
 	/// Only enables if OTEL_EXPORTER_OTLP_ENDPOINT environment variable is set.
@@ -31,14 +63,7 @@ public static TBuilder AddDocsApiOpenTelemetry<TBuilder>(
 		{
 			_ = edotBuilder
 				.WithLogging()
-				.WithTracing(tracing =>
-				{
-					_ = tracing
-						.AddSource("Elastic.Documentation.Api.AskAi")
-						.AddSource("Elastic.Documentation.Api.StreamTransformer")
-						.AddAspNetCoreInstrumentation()
-						.AddHttpClientInstrumentation();
-				})
+				.WithTracing(tracing => tracing.AddDocsApiTracing())
 				.WithMetrics(metrics =>
 				{
 					_ = metrics

src/api/Elastic.Documentation.Api.Lambda/Program.cs

@@ -8,6 +8,7 @@
 using Elastic.Documentation.Api.Core.AskAi;
 using Elastic.Documentation.Api.Core.Search;
 using Elastic.Documentation.Api.Infrastructure;
+using Elastic.Documentation.Api.Infrastructure.Middleware;
 
 try
 {
@@ -37,6 +38,10 @@
 
 	builder.Services.AddElasticDocsApiUsecases(environment);
 	var app = builder.Build();
+
+	// Add middleware to enrich logs with euid cookie
+	_ = app.UseEuidLogging();
+
 	var v1 = app.MapGroup("/docs/_api/v1");
 	v1.MapElasticDocsApiEndpoints();
 	Console.WriteLine("API endpoints mapped");
@@ -58,3 +63,8 @@
 [JsonSerializable(typeof(SearchRequest))]
 [JsonSerializable(typeof(SearchResponse))]
 internal sealed partial class LambdaJsonSerializerContext : JsonSerializerContext;
+
+// Make the Program class accessible for integration testing
+#pragma warning disable ASP0027
+public partial class Program { }
+#pragma warning restore ASP0027

tests-integration/Elastic.Documentation.Api.IntegrationTests/Elastic.Documentation.Api.IntegrationTests.csproj

@@ -0,0 +1,19 @@
+<Project Sdk="Microsoft.NET.Sdk">
+
+  <PropertyGroup>
+    <TargetFramework>net10.0</TargetFramework>
+  </PropertyGroup>
+
+  <ItemGroup>
+    <ProjectReference Include="$(SolutionRoot)\src\api\Elastic.Documentation.Api.Lambda\Elastic.Documentation.Api.Lambda.csproj"/>
+    <ProjectReference Include="$(SolutionRoot)\src\api\Elastic.Documentation.Api.Infrastructure\Elastic.Documentation.Api.Infrastructure.csproj"/>
+    <ProjectReference Include="$(SolutionRoot)\src\api\Elastic.Documentation.Api.Core\Elastic.Documentation.Api.Core.csproj"/>
+  </ItemGroup>
+
+  <ItemGroup>
+    <PackageReference Include="FakeItEasy" />
+    <PackageReference Include="Microsoft.AspNetCore.Mvc.Testing" />
+    <PackageReference Include="OpenTelemetry.Exporter.InMemory" />
+  </ItemGroup>
+
+</Project>

tests-integration/Elastic.Documentation.Api.IntegrationTests/EuidEnrichmentIntegrationTests.cs

@@ -0,0 +1,85 @@
+// Licensed to Elasticsearch B.V under one or more agreements.
+// Elasticsearch B.V licenses this file to you under the Apache 2.0 License.
+// See the LICENSE file in the project root for more information
+
+using System.Diagnostics;
+using System.Text;
+using Elastic.Documentation.Api.Core;
+using Elastic.Documentation.Api.IntegrationTests.Fixtures;
+using FluentAssertions;
+
+namespace Elastic.Documentation.Api.IntegrationTests;
+
+/// <summary>
+/// Integration tests for euid cookie enrichment in OpenTelemetry traces and logging.
+/// Uses WebApplicationFactory to test the real API configuration with mocked services.
+/// </summary>
+public class EuidEnrichmentIntegrationTests(ApiWebApplicationFactory factory) : IClassFixture<ApiWebApplicationFactory>
+{
+	private readonly ApiWebApplicationFactory _factory = factory;
+
+	/// <summary>
+	/// Test that verifies euid cookie is added to both HTTP span and custom AskAi span,
+	/// and appears in log entries - using the real API configuration.
+	/// </summary>
+	[Fact]
+	public async Task AskAiEndpointPropagatatesEuidToAllSpansAndLogs()
+	{
+		// Arrange
+		const string expectedEuid = "integration-test-euid-12345";
+
+		// Create client
+		using var client = _factory.CreateClient();
+
+		// Act - Make request to /ask-ai/stream with euid cookie
+		using var request = new HttpRequestMessage(HttpMethod.Post, "/docs/_api/v1/ask-ai/stream");
+		request.Headers.Add("Cookie", $"euid={expectedEuid}");
+		request.Content = new StringContent(
+								"""{"message":"test question","conversationId":null}""",
+			Encoding.UTF8,
+			"application/json"
+		);
+
+		using var response = await client.SendAsync(request, TestContext.Current.CancellationToken);
+
+		// Assert - Response is successful
+		response.IsSuccessStatusCode.Should().BeTrue();
+
+		// Assert - Verify spans were captured
+		var activities = _factory.ExportedActivities;
+		activities.Should().NotBeEmpty("OpenTelemetry should have captured activities");
+
+		// Verify HTTP span has euid
+		var httpSpan = activities.FirstOrDefault(a =>
+			a.DisplayName.Contains("POST") && a.DisplayName.Contains("ask-ai"));
+		httpSpan.Should().NotBeNull("Should have captured HTTP request span");
+		var httpEuidTag = httpSpan!.TagObjects.FirstOrDefault(t => t.Key == TelemetryConstants.UserEuidAttributeName);
+		httpEuidTag.Should().NotBeNull("HTTP span should have user.euid tag");
+		httpEuidTag.Value.Should().Be(expectedEuid, "HTTP span euid should match cookie value");
+
+		// Verify custom AskAi span has euid (proves baggage + processor work)
+		var askAiSpan = activities.FirstOrDefault(a => a.Source.Name == TelemetryConstants.AskAiSourceName);
+		askAiSpan.Should().NotBeNull("Should have captured custom AskAi span from AskAiUsecase");
+		var askAiEuidTag = askAiSpan!.TagObjects.FirstOrDefault(t => t.Key == TelemetryConstants.UserEuidAttributeName);
+		askAiEuidTag.Should().NotBeNull("AskAi span should have user.euid tag from baggage");
+		askAiEuidTag.Value.Should().Be(expectedEuid, "AskAi span euid should match cookie value");
+
+		// Assert - Verify logs have euid in scope
+		var logEntries = _factory.LogEntries;
+		logEntries.Should().NotBeEmpty("Should have captured log entries");
+
+		// Find a log entry from AskAiUsecase
+		var askAiLog = logEntries.FirstOrDefault(e =>
+			e.CategoryName.Contains("AskAiUsecase") &&
+			e.Message.Contains("Starting AskAI"));
+		askAiLog.Should().NotBeNull("Should have logged from AskAiUsecase");
+
+		// Verify euid is in the logging scope
+		var hasEuidInScope = askAiLog!.Scopes
+			.Any(scope => scope is IDictionary<string, object> dict &&
+						  dict.TryGetValue(TelemetryConstants.UserEuidAttributeName, out var value) &&
+						  value?.ToString() == expectedEuid);
+
+		hasEuidInScope.Should().BeTrue("Log entry should have user.euid in scope from middleware");
+	}
+}