Propagate trace context for entity operations

src/Grpc/orchestrator_service.proto

@@ -182,6 +182,7 @@ message EntityOperationSignaledEvent {
     google.protobuf.Timestamp scheduledTime = 3;
     google.protobuf.StringValue input = 4;
     google.protobuf.StringValue targetInstanceId = 5; // used only within histories, null in messages
+    TraceContext parentTraceContext = 6;
 }
 
 message EntityOperationCalledEvent {
@@ -192,6 +193,7 @@ message EntityOperationCalledEvent {
     google.protobuf.StringValue parentInstanceId = 5;  // used only within messages, null in histories
     google.protobuf.StringValue parentExecutionId = 6;  // used only within messages, null in histories
     google.protobuf.StringValue targetInstanceId = 7; // used only within histories, null in messages
+    TraceContext parentTraceContext = 8;
 }
 
 message EntityLockRequestedEvent {
@@ -318,6 +320,8 @@ message SendEntityMessageAction {
         EntityLockRequestedEvent entityLockRequested = 3;
         EntityUnlockSentEvent entityUnlockSent = 4;
     }
+
+    TraceContext parentTraceContext = 5;
 }
 
 message OrchestratorAction {

src/Shared/Grpc/ProtoUtils.cs

@@ -415,6 +415,7 @@ internal static P.OrchestratorResponse ConstructOrchestratorResponse(
                             out string requestId);
 
                         entityConversionState.EntityRequestIds.Add(requestId);
+                        sendAction.ParentTraceContext = CreateTraceContext();
 
                         switch (sendAction.EntityMessageTypeCase)
                         {
@@ -636,6 +637,9 @@ internal static void ToEntityBatchRequest(
                     Id = Guid.Parse(op.EntityOperationSignaled.RequestId),
                     Operation = op.EntityOperationSignaled.Operation,
                     Input = op.EntityOperationSignaled.Input,
+                    TraceContext = op.EntityOperationSignaled.ParentTraceContext is { } signalTc
+                        ? new DistributedTraceContext(signalTc.TraceParent, signalTc.TraceState)
+                        : null,
                 });
                 operationInfos.Add(new P.OperationInfo
                 {
@@ -650,6 +654,9 @@ internal static void ToEntityBatchRequest(
                     Id = Guid.Parse(op.EntityOperationCalled.RequestId),
                     Operation = op.EntityOperationCalled.Operation,
                     Input = op.EntityOperationCalled.Input,
+                    TraceContext = op.EntityOperationCalled.ParentTraceContext is { } calledTc
+                        ? new DistributedTraceContext(calledTc.TraceParent, calledTc.TraceState)
+                        : null,
                 });
                 operationInfos.Add(new P.OperationInfo
                 {

src/Shared/Grpc/Tracing/Schema.cs

@@ -60,5 +60,10 @@ public static class Task
         /// The time at which the timer is scheduled to fire.
         /// </summary>
         public const string FireAt = "durabletask.fire_at";
+
+        /// <summary>
+        /// The name of the entity operation being executed.
+        /// </summary>
+        public const string EntityOperationName = "durabletask.entity.operation";
     }
 }

src/Shared/Grpc/Tracing/TraceActivityConstants.cs

@@ -24,6 +24,11 @@ static class TraceActivityConstants
     /// </summary>
     public const string Event = "event";
 
+    /// <summary>
+    /// The name of the activity that represents entity operation execution.
+    /// </summary>
+    public const string EntityOperation = "entity_operation";
+
     /// <summary>
     /// The name of the activity that represents timer operations.
     /// </summary>

src/Shared/Grpc/Tracing/TraceHelper.cs

@@ -156,6 +156,58 @@ static class TraceHelper
         return newActivity;
     }
 
+    /// <summary>
+    /// Starts a new trace activity for executing an entity operation.
+    /// </summary>
+    /// <param name="entityName">The name of the entity.</param>
+    /// <param name="operationName">The name of the operation being executed.</param>
+    /// <param name="instanceId">The instance ID of the entity.</param>
+    /// <param name="traceParent">The W3C traceparent header value from the parent orchestration.</param>
+    /// <param name="traceState">The W3C tracestate header value from the parent orchestration.</param>
+    /// <returns>
+    /// Returns a newly started <see cref="Activity"/> with entity operation metadata, or null if tracing is not enabled.
+    /// </returns>
+    public static Activity? StartTraceActivityForEntityOperation(
+        string entityName,
+        string? operationName,
+        string instanceId,
+        string? traceParent,
+        string? traceState)
+    {
+        if (traceParent is null || !ActivityContext.TryParse(
+                traceParent,
+                traceState,
+                out ActivityContext activityContext))
+        {
+            return null;
+        }
+
+        string spanName = string.IsNullOrEmpty(operationName)
+            ? $"{TraceActivityConstants.EntityOperation}:{entityName}"
+            : $"{TraceActivityConstants.EntityOperation}:{entityName}:{operationName}";
+
+        Activity? newActivity = ActivityTraceSource.StartActivity(
+            spanName,
+            kind: ActivityKind.Server,
+            parentContext: activityContext);
+
+        if (newActivity == null)
+        {
+            return null;
+        }
+
+        newActivity.SetTag(Schema.Task.Type, TraceActivityConstants.EntityOperation);
+        newActivity.SetTag(Schema.Task.Name, entityName);
+        newActivity.SetTag(Schema.Task.InstanceId, instanceId);
+
+        if (!string.IsNullOrEmpty(operationName))
+        {
+            newActivity.SetTag(Schema.Task.EntityOperationName, operationName);
+        }
+
+        return newActivity;
+    }
+
     /// <summary>
     /// Emits a new trace activity for a (task) activity that successfully completes.
     /// </summary>
@@ -202,6 +254,51 @@ public static void EmitTraceActivityForTaskFailed(
         activity?.Dispose();
     }
 
+    /// <summary>
+    /// Emits a new trace activity for an entity operation that successfully completes.
+    /// </summary>
+    /// <param name="instanceId">The ID of the associated orchestration.</param>
+    /// <param name="historyEvent">The associated <see cref="P.HistoryEvent" />.</param>
+    /// <param name="calledEvent">The associated <see cref="P.EntityOperationCalledEvent"/>.</param>
+    public static void EmitTraceActivityForEntityOperationCompleted(
+        string? instanceId,
+        P.HistoryEvent? historyEvent,
+        P.EntityOperationCalledEvent? calledEvent)
+    {
+        Activity? activity = StartTraceActivityForSchedulingEntityOperation(instanceId, historyEvent, calledEvent);
+
+        activity?.Dispose();
+    }
+
+    /// <summary>
+    /// Emits a new trace activity for an entity operation that fails.
+    /// </summary>
+    /// <param name="instanceId">The ID of the associated orchestration.</param>
+    /// <param name="historyEvent">The associated <see cref="P.HistoryEvent" />.</param>
+    /// <param name="calledEvent">The associated <see cref="P.EntityOperationCalledEvent"/>.</param>
+    /// <param name="failedEvent">The associated <see cref="P.EntityOperationFailedEvent"/>.</param>
+    public static void EmitTraceActivityForEntityOperationFailed(
+        string? instanceId,
+        P.HistoryEvent? historyEvent,
+        P.EntityOperationCalledEvent? calledEvent,
+        P.EntityOperationFailedEvent? failedEvent)
+    {
+        Activity? activity = StartTraceActivityForSchedulingEntityOperation(instanceId, historyEvent, calledEvent);
+
+        if (activity is null)
+        {
+            return;
+        }
+
+        if (failedEvent != null)
+        {
+            string statusDescription = failedEvent.FailureDetails?.ErrorMessage ?? "Unspecified entity operation failure";
+            activity.SetStatus(ActivityStatusCode.Error, statusDescription);
+        }
+
+        activity.Dispose();
+    }
+
     /// <summary>
     /// Emits a new trace activity for sub-orchestration execution when the sub-orchestration
     /// completes successfully.
@@ -409,6 +506,74 @@ static string CreateSpanName(string spanDescription, string? taskName, string? t
         return newActivity;
     }
 
+    /// <summary>
+    /// Starts a new trace activity for scheduling an entity operation. Represents the time between
+    /// enqueuing the entity operation message and it completing.
+    /// </summary>
+    /// <param name="instanceId">The ID of the associated orchestration.</param>
+    /// <param name="historyEvent">The associated <see cref="P.HistoryEvent" />.</param>
+    /// <param name="calledEvent">The associated <see cref="P.EntityOperationCalledEvent"/>.</param>
+    /// <returns>
+    /// Returns a newly started <see cref="Activity"/> with entity operation metadata.
+    /// </returns>
+    static Activity? StartTraceActivityForSchedulingEntityOperation(
+        string? instanceId,
+        P.HistoryEvent? historyEvent,
+        P.EntityOperationCalledEvent? calledEvent)
+    {
+        if (calledEvent == null)
+        {
+            return null;
+        }
+
+        string targetInstanceId = historyEvent?.EntityOperationCalled?.TargetInstanceId ?? string.Empty;
+
+        // Extract entity name from instance ID (format: "@name@key")
+        string entityName = targetInstanceId;
+        if (targetInstanceId.Length > 1 && targetInstanceId[0] == '@')
+        {
+            int secondAt = targetInstanceId.IndexOf('@', 1);
+            if (secondAt > 1)
+            {
+                entityName = targetInstanceId.Substring(1, secondAt - 1);
+            }
+        }
+        string spanName = string.IsNullOrEmpty(calledEvent.Operation)
+            ? $"{TraceActivityConstants.EntityOperation}:{entityName}"
+            : $"{TraceActivityConstants.EntityOperation}:{entityName}:{calledEvent.Operation}";
+
+        Activity? newActivity = ActivityTraceSource.StartActivity(
+            spanName,
+            kind: ActivityKind.Client,
+            startTime: historyEvent?.Timestamp?.ToDateTimeOffset() ?? default,
+            parentContext: Activity.Current?.Context ?? default);
+
+        if (newActivity == null)
+        {
+            return null;
+        }
+
+        if (calledEvent.ParentTraceContext != null
+            && ActivityContext.TryParse(
+                calledEvent.ParentTraceContext.TraceParent,
+                calledEvent.ParentTraceContext?.TraceState,
+                out ActivityContext parentContext))
+        {
+            newActivity.SetSpanId(parentContext.SpanId.ToString());
+        }
+
+        newActivity.AddTag(Schema.Task.Type, TraceActivityConstants.EntityOperation);
+        newActivity.AddTag(Schema.Task.Name, entityName);
+        newActivity.AddTag(Schema.Task.InstanceId, instanceId);
+
+        if (!string.IsNullOrEmpty(calledEvent.Operation))
+        {
+            newActivity.AddTag(Schema.Task.EntityOperationName, calledEvent.Operation);
+        }
+
+        return newActivity;
+    }
+
     /// <summary>
     /// Starts a new trace activity for sub-orchestrations. Represents the time between enqueuing
     /// the sub-orchestration message and it completing.

src/Worker/Grpc/GrpcDurableTaskWorker.Processor.cs

@@ -502,6 +502,14 @@ async Task OnRunOrchestratorAsync(
                     return taskScheduledEvent;
                 }
 
+                P.HistoryEvent? GetEntityOperationCalledEvent(string requestId)
+                {
+                    return request
+                        .PastEvents
+                        .Where(x => x.EventTypeCase == P.HistoryEvent.EventTypeOneofCase.EntityOperationCalled)
+                        .FirstOrDefault(x => x.EntityOperationCalled.RequestId == requestId);
+                }
+
                 foreach (var newEvent in request.NewEvents)
                 {
                     switch (newEvent.EventTypeCase)
@@ -565,6 +573,33 @@ async Task OnRunOrchestratorAsync(
                                 newEvent.Timestamp.ToDateTime(),
                                 newEvent.TimerFired);
                             break;
+
+                        case P.HistoryEvent.EventTypeOneofCase.EntityOperationCompleted:
+                            {
+                                P.HistoryEvent? entityCalledEvent =
+                                    GetEntityOperationCalledEvent(
+                                        newEvent.EntityOperationCompleted.RequestId);
+
+                                TraceHelper.EmitTraceActivityForEntityOperationCompleted(
+                                    request.InstanceId,
+                                    entityCalledEvent,
+                                    entityCalledEvent?.EntityOperationCalled);
+                                break;
+                            }
+
+                        case P.HistoryEvent.EventTypeOneofCase.EntityOperationFailed:
+                            {
+                                P.HistoryEvent? entityCalledEvent =
+                                    GetEntityOperationCalledEvent(
+                                        newEvent.EntityOperationFailed.RequestId);
+
+                                TraceHelper.EmitTraceActivityForEntityOperationFailed(
+                                    request.InstanceId,
+                                    entityCalledEvent,
+                                    entityCalledEvent?.EntityOperationCalled,
+                                    newEvent.EntityOperationFailed);
+                                break;
+                            }
                     }
                 }
             }
@@ -865,6 +900,20 @@ async Task OnRunEntityBatchAsync(
             var coreEntityId = DTCore.Entities.EntityId.FromString(batchRequest.InstanceId!);
             EntityId entityId = new(coreEntityId.Name, coreEntityId.Key);
 
+            // Start a Server trace span for each entity operation in the batch.
+            // Each operation may come from a different orchestration with its own trace context,
+            // so we create individual spans to preserve correct parent-child relationships.
+            List<Activity> traceActivities = batchRequest.Operations?
+                .Select(op => TraceHelper.StartTraceActivityForEntityOperation(
+                    entityId.Name,
+                    op.Operation,
+                    batchRequest.InstanceId!,
+                    op.TraceContext?.TraceParent,
+                    op.TraceContext?.TraceState))
+                .OfType<Activity>()
+                .ToList()
+                ?? [];
+
             TaskName name = new(entityId.Name);
 
             EntityBatchResult? batchResult;
@@ -885,6 +934,9 @@ async Task OnRunEntityBatchAsync(
                 {
                     // we could not find the entity. This is considered an application error,
                     // so we return a non-retriable error-OperationResult for each operation in the batch.
+                    string errorMessage = $"No entity task named '{name}' was found.";
+                    traceActivities.ForEach(a => a.SetStatus(ActivityStatusCode.Error, errorMessage));
+
                     batchResult = new EntityBatchResult()
                     {
                         Actions = [], // no actions
@@ -894,7 +946,7 @@ async Task OnRunEntityBatchAsync(
                             {
                                 FailureDetails = new FailureDetails(
                                     errorType: "EntityTaskNotFound",
-                                    errorMessage: $"No entity task named '{name}' was found.",
+                                    errorMessage: errorMessage,
                                     stackTrace: null,
                                     innerFailure: null,
                                     isNonRetriable: true),
@@ -913,6 +965,12 @@ async Task OnRunEntityBatchAsync(
                 {
                     FailureDetails = new FailureDetails(frameworkException),
                 };
+
+                traceActivities.ForEach(a => a.SetStatus(ActivityStatusCode.Error, frameworkException.Message));
+            }
+            finally
+            {
+                traceActivities.ForEach(a => a.Dispose());
             }
 
             P.EntityBatchResult response = batchResult.ToEntityBatchResult(