Added Invocation ID to SetFunctionInvocationContextCommand and added a call to Get-CurrentActivityForInvocation

Author: bachuv

src/DurableSDK/Commands/SetFunctionInvocationContextCommand.cs

@@ -18,6 +18,7 @@ public class SetFunctionInvocationContextCommand : PSCmdlet
     {
         internal const string ContextKey = "OrchestrationContext";
         private const string DurableClientKey = "DurableClient";
+        private const string InvocationIdKey = "InvocationId";
 
         [Parameter(Mandatory = true, ParameterSetName = ContextKey)]
         public OrchestrationContext OrchestrationContext { get; set; }
@@ -28,6 +29,12 @@ public class SetFunctionInvocationContextCommand : PSCmdlet
         [Parameter(Mandatory = true, ParameterSetName = DurableClientKey)]
         public object DurableClient { get; set; }
 
+        /// <summary>
+        /// The invocation id.
+        /// </summary>
+        [Parameter(Mandatory = true, ParameterSetName = InvocationIdKey)]
+        public string InvocationId { get; set; }
+
         [Parameter(Mandatory = true, ParameterSetName = "Clear")]
         public SwitchParameter Clear { get; set; }
 
@@ -44,11 +51,16 @@ protected override void EndProcessing()
                     privateData[DurableClientKey] = DurableClient;
                     break;
 
+                case InvocationIdKey:
+                    privateData[InvocationIdKey] = InvocationId;
+                    break;
+
                 default:
                     if (Clear.IsPresent)
                     {
                         privateData.Remove(ContextKey);
                         privateData.Remove(DurableClientKey);
+                        privateData.Remove(InvocationIdKey);
                     }
                     break;
             }

src/DurableSDK/IPowerShellServices.cs

@@ -21,6 +21,8 @@ internal interface IPowerShellServices
 
         void SetDurableClient(object durableClient);
 
+        void SetInvocationId(string invocationId);
+
         OrchestrationBindingInfo SetOrchestrationContext(ParameterBinding context, out IExternalOrchestrationInvoker externalInvoker);
 
         void ClearOrchestrationContext();

src/DurableSDK/PowerShellServices.cs

@@ -72,6 +72,13 @@ public void SetDurableClient(object durableClient)
             _hasInitializedDurableFunctions = true;
         }
 
+        public void SetInvocationId(string invocationId)
+        {
+            _pwsh.AddCommand(SetFunctionInvocationContextCommand)
+                .AddParameter("InvocationId", invocationId)
+                .InvokeAndClearCommands();
+        }
+
         public OrchestrationBindingInfo SetOrchestrationContext(
             ParameterBinding context,
             out IExternalOrchestrationInvoker externalInvoker)

src/DurableWorker/DurableController.cs

@@ -85,7 +85,7 @@ private void tryEnablingExternalSDK()
             }
         }
 
-        public void InitializeBindings(IList<ParameterBinding> inputData, out bool hasExternalSDK)
+        public void InitializeBindings(IList<ParameterBinding> inputData, out bool hasExternalSDK, string invocationId = "")
         {
             this.tryEnablingExternalSDK();
 
@@ -121,6 +121,7 @@ public void InitializeBindings(IList<ParameterBinding> inputData, out bool hasEx
                 _orchestrationInvoker.SetExternalInvoker(externalInvoker);
             }
             hasExternalSDK = _powerShellServices.HasExternalDurableSDK();
+            _powerShellServices.SetInvocationId(invocationId);
         }
 
         public void AfterFunctionInvocation()

src/Modules/Microsoft.Azure.Functions.PowerShellWorker/Microsoft.Azure.Functions.PowerShellWorker.psm1

@@ -19,6 +19,16 @@ function GetDurableClientFromModulePrivateData {
     }
 }
 
+function GetInvocationIdFromModulePrivateData {
+    $PrivateData = $PSCmdlet.MyInvocation.MyCommand.Module.PrivateData
+    if ($null -eq $PrivateData -or $null -eq $PrivateData['InvocationId']) {
+        return $null
+    }
+    else {
+        return $PrivateData['InvocationId']
+    }
+}
+
 function Get-DurableStatus {
     [CmdletBinding()]
     param(
@@ -121,6 +131,9 @@ function Start-DurableOrchestration {
         $InstanceId = (New-Guid).Guid
     }
 
+    $invocationId = GetInvocationIdFromModulePrivateData
+    $headers = Get-TraceHeaders -InvocationId $invocationId
+
     $Uri =
         if ($DurableClient.rpcBaseUrl) {
             # Fast local RPC path
@@ -133,44 +146,47 @@ function Start-DurableOrchestration {
 
     $Body = $InputObject | ConvertTo-Json -Compress
 
-    $invokeParams = @{
-        Uri         = $Uri
-        Method      = 'POST'
-        ContentType = 'application/json'
-        Body        = $Body
-    }
+    $null = Invoke-RestMethod -Uri $Uri -Method 'POST' -ContentType 'application/json' -Body $Body -Headers $headers
+
+    return $instanceId
+}
 
-    try {
-        $activity = Start-FunctionsOpenTelemetrySpan -ActivityName "Starting orchestration"
+function Get-TraceHeaders {
+    param(
+        [string] $InvocationId
+    )
+
+    if ($null -eq $InvocationId -or $InvocationId -eq "") {
+        return @{} # Return an empty headers object
+    }
 
-        $traceId = $activity.activity.TraceId
-        $spanId = $activity.activity.SpanId
+    # Check if Get-CurrentActivityForInvocation is available
+    if (-not (Get-Command -Name Get-CurrentActivityForInvocation -ErrorAction SilentlyContinue)) {
+        Write-Warning "Get-CurrentActivityForInvocation is not available. Skipping call."
+        return @{} # Return an empty headers object
+    }
 
-        # Construct the traceparent header
-        $traceParent = "00-$traceId-$spanId-01"
+    $activityResponse = Get-CurrentActivityForInvocation -InvocationId $invocationId
+    $activity = $activityResponse.activity
 
-        $traceState = $activity.activity.TraceState
+    $traceId = $activity.TraceId
+    $spanId = $activity.SpanId
+    $traceFlags = $activity.TraceFlags
+    $traceState = $activity.TraceStateString
 
-        $invokeParams.Headers = @{
-            'traceparent' = $traceParent
-            'tracestate'  = $traceState
-        }
-        # Do whatever you need to do with the trace information using the activity here
-    } catch {
-        # Do something better - correctly handle errors when the OTel SDK is not available
-        # Output errors from this command in a reasonable way
-        # Detect if calling Stop-FunctionsOpenTelemetrySpan is necessary and change that logic too
+    $flag = "00"
+    if ($null -ne $traceFlags -and $traceFlags -eq "Recorded") {
+        $flag = "01"
     }
-              
-    $null = Invoke-RestMethod @invokeParams
 
-    try {
-        Stop-FunctionsOpenTelemetrySpan -Activity $activity
-    } catch {
-        # Do something better here
+    $traceparent = "00-$traceId-$spanId-$flag"
+
+    $headers = @{
+        "traceparent" = $traceparent
+        "tracestate"  = $traceState
     }
 
-    return $instanceId
+    return $headers
 }
 
 function Stop-DurableOrchestration {

src/PowerShell/PowerShellManager.cs

@@ -215,7 +215,7 @@ public Hashtable InvokeFunction(
 
             try
             {
-                durableFunctionsUtils.InitializeBindings(inputData, out bool hasExternalDFsdk);
+                durableFunctionsUtils.InitializeBindings(inputData: inputData, invocationId: otelContext.InvocationId, hasExternalSDK: out bool hasExternalDFsdk);
                 Logger.Log(isUserOnlyLog: false, LogLevel.Trace, String.Format(PowerShellWorkerStrings.UtilizingExternalDurableSDK, hasExternalDFsdk));
 
                 ImportEntryPointModule(functionInfo);