@@ -62,6 +62,13 @@ public interface IExecutionContext : IAgentService, IKnobValueContext
6262 void QueueAttachFile ( string type , string name , string filePath ) ;
6363 ITraceWriter GetTraceWriter ( ) ;
6464
65+ // correlation context for enhanced tracing
66+ void SetCorrelationStep ( string stepId ) ;
67+ void ClearCorrelationStep ( ) ;
68+ void SetCorrelationTask ( string taskId ) ;
69+ void ClearCorrelationTask ( ) ;
70+ string BuildCorrelationId ( ) ;
71+
6572 // timeline record update methods
6673 void Start ( string currentOperation = null ) ;
6774 TaskResult Complete ( TaskResult ? result = null , string currentOperation = null , string resultCode = null ) ;
@@ -93,7 +100,7 @@ public interface IExecutionContext : IAgentService, IKnobValueContext
93100 void PublishTaskRunnerTelemetry ( Dictionary < string , string > taskRunnerData ) ;
94101 }
95102
96- public sealed class ExecutionContext : AgentService , IExecutionContext , IDisposable
103+ public sealed class ExecutionContext : AgentService , IExecutionContext , ICorrelationContext , IDisposable
97104 {
98105 private const int _maxIssueCount = 10 ;
99106
@@ -104,6 +111,8 @@ public sealed class ExecutionContext : AgentService, IExecutionContext, IDisposa
104111 private readonly HashSet < string > _outputvariables = new HashSet < string > ( StringComparer . OrdinalIgnoreCase ) ;
105112 private readonly List < TaskRestrictions > _restrictions = new List < TaskRestrictions > ( ) ;
106113 private readonly string _buildLogsFolderName = "buildlogs" ;
114+ private readonly AsyncLocal < string > _correlationStep = new AsyncLocal < string > ( ) ;
115+ private readonly AsyncLocal < string > _correlationTask = new AsyncLocal < string > ( ) ;
107116 private IAgentLogPlugin _logPlugin ;
108117 private IPagingLogger _logger ;
109118 private IJobServerQueue _jobServerQueue ;
@@ -202,6 +211,9 @@ public override void Initialize(IHostContext hostContext)
202211 }
203212
204213 _jobServerQueue = HostContext . GetService < IJobServerQueue > ( ) ;
214+
215+ // Register this ExecutionContext for enhanced correlation
216+ HostContext . CorrelationContextManager . SetCurrentExecutionContext ( this ) ;
205217 }
206218
207219 public void CancelToken ( )
@@ -996,8 +1008,68 @@ public void PublishTaskRunnerTelemetry(Dictionary<string, string> taskRunnerData
9961008 PublishTelemetry ( taskRunnerData , IsAgentTelemetry : true ) ;
9971009 }
9981010
1011+ // Correlation context methods for enhanced tracing
1012+ public void SetCorrelationStep ( string stepId )
1013+ {
1014+ _correlationStep . Value = stepId ;
1015+ }
1016+
1017+ public void ClearCorrelationStep ( )
1018+ {
1019+ _correlationStep . Value = null ;
1020+ }
1021+
1022+ public void SetCorrelationTask ( string taskId )
1023+ {
1024+ _correlationTask . Value = taskId ;
1025+ }
1026+
1027+ public void ClearCorrelationTask ( )
1028+ {
1029+ _correlationTask . Value = null ;
1030+ }
1031+
1032+ public string BuildCorrelationId ( )
1033+ {
1034+ var step = _correlationStep . Value ;
1035+ var task = _correlationTask . Value ;
1036+
1037+ if ( string . IsNullOrEmpty ( step ) )
1038+ {
1039+ return string . IsNullOrEmpty ( task ) ? string . Empty : $ "TASK-{ ShortenGuid ( task ) } ";
1040+ }
1041+
1042+ return string . IsNullOrEmpty ( task ) ? $ "STEP-{ ShortenGuid ( step ) } " : $ "STEP-{ ShortenGuid ( step ) } |TASK-{ ShortenGuid ( task ) } ";
1043+ }
1044+
1045+ /// <summary>
1046+ /// Shorten a GUID to first 12 characters for more readable logs while maintaining uniqueness
1047+ /// </summary>
1048+ private static string ShortenGuid ( string guid )
1049+ {
1050+ if ( string . IsNullOrEmpty ( guid ) )
1051+ return guid ;
1052+
1053+ // Use first 12 characters total: 8 from first segment + 4 from second segment
1054+ // This ensures consistent output length regardless of input length
1055+ // e.g., "verylongstring-1234..." becomes "verylong1234" (12 chars)
1056+ // e.g., "60cf5508-70a7-..." becomes "60cf550870a7" (12 chars)
1057+ var parts = guid . Split ( '-' ) ;
1058+ if ( parts . Length >= 2 && parts [ 0 ] . Length >= 8 && parts [ 1 ] . Length >= 4 )
1059+ {
1060+ return parts [ 0 ] . Substring ( 0 , 8 ) + parts [ 1 ] . Substring ( 0 , 4 ) ;
1061+ }
1062+
1063+ // Fallback: remove hyphens and take first 12 chars
1064+ var cleaned = guid . Replace ( "-" , "" ) ;
1065+ return cleaned . Length > 12 ? cleaned . Substring ( 0 , 12 ) : cleaned ;
1066+ }
1067+
9991068 public void Dispose ( )
10001069 {
1070+ // Clear the correlation context registration
1071+ HostContext . CorrelationContextManager . ClearCurrentExecutionContext ( ) ;
1072+
10011073 _cancellationTokenSource ? . Dispose ( ) ;
10021074 _forceCompleteCancellationTokenSource ? . Dispose ( ) ;
10031075
0 commit comments