Add header propagation functionality #143
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change | ||||||
|---|---|---|---|---|---|---|---|---|
| @@ -0,0 +1,78 @@ | ||||||||
| // Copyright (c) Microsoft Corporation. All rights reserved. | ||||||||
| // Licensed under the MIT License. | ||||||||
|
|
||||||||
| using System; | ||||||||
| using System.Collections.Generic; | ||||||||
| using System.Reflection; | ||||||||
|
|
||||||||
| namespace Microsoft.Agents.Core.HeaderPropagation; | ||||||||
|
|
||||||||
| /// <summary> | ||||||||
| /// Attribute to load headers for header propagation. | ||||||||
| /// This attribute should be applied to classes that implement the <see cref="IHeaderPropagationAttribute"/> interface. | ||||||||
| /// </summary> | ||||||||
| [AttributeUsage(AttributeTargets.Class)] | ||||||||
| public class HeaderPropagationAttribute : Attribute | ||||||||
| { | ||||||||
| internal static void LoadHeaders() | ||||||||
| { | ||||||||
| // Init newly loaded assemblies | ||||||||
| AppDomain.CurrentDomain.AssemblyLoad += (s, o) => LoadHeadersAssembly(o.LoadedAssembly); | ||||||||
|
|
||||||||
| // And all the ones we currently have loaded | ||||||||
| foreach (var assembly in AppDomain.CurrentDomain.GetAssemblies()) | ||||||||
| { | ||||||||
| LoadHeadersAssembly(assembly); | ||||||||
| } | ||||||||
| } | ||||||||
|
|
||||||||
| private static void LoadHeadersAssembly(Assembly assembly) | ||||||||
| { | ||||||||
| foreach (var type in GetLoadHeadersTypes(assembly)) | ||||||||
| { | ||||||||
| #if !NETSTANDARD | ||||||||
| if (!typeof(IHeaderPropagationAttribute).IsAssignableFrom(type)) | ||||||||
| { | ||||||||
| throw new InvalidOperationException( | ||||||||
| $"Type '{type.FullName}' is marked with [HeaderPropagation] but does not implement IHeaderPropagationAttribute."); | ||||||||
| } | ||||||||
| #endif | ||||||||
|
|
||||||||
| var loadHeaders = type.GetMethod(nameof(LoadHeaders), BindingFlags.Static | BindingFlags.Public); | ||||||||
|
|
||||||||
| if (loadHeaders == null) | ||||||||
| { | ||||||||
| continue; | ||||||||
| } | ||||||||
|
|
||||||||
| loadHeaders.Invoke(assembly, [HeaderPropagationContext.HeadersToPropagate]); | ||||||||
| } | ||||||||
|
Comment on lines
+48
to
+49
There was a problem hiding this comment. When invoking a static method via reflection, the instance parameter should be 'null' rather than 'assembly'. Use 'loadHeaders.Invoke(null, new object[]{ HeaderPropagationContext.HeadersToPropagate })'.
Suggested change
Copilot uses AI. Check for mistakes. |
||||||||
| } | ||||||||
|
|
||||||||
| private static IEnumerable<Type> GetLoadHeadersTypes(Assembly assembly) | ||||||||
| { | ||||||||
| IList<Type> result = []; | ||||||||
|
|
||||||||
| Type[] types; | ||||||||
|
|
||||||||
| try | ||||||||
| { | ||||||||
| types = assembly.GetTypes(); | ||||||||
| } | ||||||||
| catch (Exception ex) | ||||||||
| { | ||||||||
| System.Diagnostics.Trace.WriteLine($"HeaderPropagationAttribute.GetLoadHeadersTypes: {ex.Message}"); | ||||||||
| return result; | ||||||||
| } | ||||||||
|
|
||||||||
| foreach (Type type in types) | ||||||||
| { | ||||||||
| if (type.GetCustomAttributes(typeof(HeaderPropagationAttribute), true).Length > 0) | ||||||||
| { | ||||||||
| result.Add(type); | ||||||||
| } | ||||||||
| } | ||||||||
|
|
||||||||
| return result; | ||||||||
| } | ||||||||
| } | ||||||||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,106 @@ | ||
| // Copyright (c) Microsoft Corporation. All rights reserved. | ||
| // Licensed under the MIT License. | ||
|
|
||
| using System; | ||
| using System.Collections.Generic; | ||
| using System.Linq; | ||
| using System.Threading; | ||
| using Microsoft.Extensions.Primitives; | ||
|
|
||
| namespace Microsoft.Agents.Core.HeaderPropagation; | ||
|
|
||
| /// <summary> | ||
| /// Shared context to manage request headers that will be used to propagate them in the <see cref="HeaderPropagationExtensions.AddHeaderPropagation"/>. | ||
| /// </summary> | ||
| public class HeaderPropagationContext() | ||
| { | ||
| private static readonly AsyncLocal<IDictionary<string, StringValues>> _headersFromRequest = new(); | ||
| private static HeaderPropagationEntryCollection _headersToPropagate = new(); | ||
|
There was a problem hiding this comment. I believe this is populated by virtue of the attribute handling, yes? There is a question about concurrency. We do something similar with the serializer options init (via attributes). I added a lock there because I wasn't sure about the exact module loading thread safety. But I'm not 100% that is required. There was a problem hiding this comment. We added a lock to the Add operations in HeaderPropagationEntryCollection. Thanks! |
||
|
|
||
| static HeaderPropagationContext() | ||
| { | ||
| HeaderPropagationAttribute.LoadHeaders(); | ||
| } | ||
|
|
||
| /// <summary> | ||
| /// Gets or sets the request headers that will be propagated based on what's inside the <see cref="HeadersToPropagate"/> property. | ||
| /// </summary> | ||
| public static IDictionary<string, StringValues> HeadersFromRequest | ||
| { | ||
| get | ||
| { | ||
| return _headersFromRequest.Value; | ||
| } | ||
| set | ||
| { | ||
| // Create a copy to ensure headers are not modified by the original request. | ||
| #if !NETSTANDARD | ||
| var headers = value?.ToDictionary(StringComparer.InvariantCultureIgnoreCase); | ||
| #else | ||
| var headers = value?.ToDictionary(x => x.Key, x => x.Value, StringComparer.InvariantCultureIgnoreCase); | ||
| #endif | ||
| _headersFromRequest.Value = FilterHeaders(headers); | ||
| } | ||
| } | ||
|
|
||
| /// <summary> | ||
| /// Gets or sets the headers to allow during the propagation. | ||
| /// </summary> | ||
| public static HeaderPropagationEntryCollection HeadersToPropagate | ||
| { | ||
| get | ||
| { | ||
| return _headersToPropagate; | ||
| } | ||
| set | ||
| { | ||
| _headersToPropagate = value ?? new(); | ||
| } | ||
| } | ||
|
|
||
| /// <summary> | ||
| /// Filters the request headers based on the keys provided in <see cref="HeadersToPropagate"/>. | ||
| /// </summary> | ||
| /// <param name="requestHeaders">Headers collection from an Http request.</param> | ||
| /// <returns>Filtered headers.</returns> | ||
| private static Dictionary<string, StringValues> FilterHeaders(Dictionary<string, StringValues> requestHeaders) | ||
| { | ||
| var result = new Dictionary<string, StringValues>(); | ||
|
|
||
| if (requestHeaders == null || requestHeaders.Count == 0) | ||
| { | ||
| return result; | ||
| } | ||
|
|
||
| // Ensure the default headers are always set by overriding the LoadHeaders configuration. | ||
| _headersToPropagate.Propagate("x-ms-correlation-id"); | ||
|
|
||
| foreach (var header in HeadersToPropagate.Entries) | ||
| { | ||
| var headerExists = requestHeaders.TryGetValue(header.Key, out var requestHeader); | ||
|
|
||
| switch (header.Action) | ||
| { | ||
| case HeaderPropagationEntryAction.Add: | ||
| #if !NETSTANDARD | ||
| result.TryAdd(header.Key, header.Value); | ||
| #else | ||
| result.Add(header.Key, header.Value); | ||
| #endif | ||
| break; | ||
| case HeaderPropagationEntryAction.Append when headerExists: | ||
| StringValues newValue = requestHeader.Concat(header.Value).ToArray(); | ||
| result[header.Key] = newValue; | ||
| break; | ||
| case HeaderPropagationEntryAction.Propagate when headerExists: | ||
| result[header.Key] = requestHeader; | ||
| break; | ||
| case HeaderPropagationEntryAction.Override: | ||
| result[header.Key] = header.Value; | ||
| break; | ||
| } | ||
| } | ||
|
|
||
| return result; | ||
| } | ||
| } | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,27 @@ | ||
| // Copyright (c) Microsoft Corporation. All rights reserved. | ||
| // Licensed under the MIT License. | ||
|
|
||
| using Microsoft.Extensions.Primitives; | ||
|
|
||
| namespace Microsoft.Agents.Core.HeaderPropagation; | ||
|
|
||
| /// <summary> | ||
| /// Represents a single header entry used for header propagation. | ||
| /// </summary> | ||
| public class HeaderPropagationEntry | ||
| { | ||
| /// <summary> | ||
| /// Key of the header entry. | ||
| /// </summary> | ||
| public string Key { get; set; } = string.Empty; | ||
|
|
||
| /// <summary> | ||
| /// Value of the header entry. | ||
| /// </summary> | ||
| public StringValues Value { get; set; } = new StringValues(string.Empty); | ||
|
|
||
| /// <summary> | ||
| /// Action of the header entry (Add, Append, etc.). | ||
| /// </summary> | ||
| public HeaderPropagationEntryAction Action; | ||
| } |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,116 @@ | ||
| // Copyright (c) Microsoft Corporation. All rights reserved. | ||
| // Licensed under the MIT License. | ||
|
|
||
| using System.Collections.Generic; | ||
| using System.Linq; | ||
| using Microsoft.Extensions.Primitives; | ||
|
|
||
| namespace Microsoft.Agents.Core.HeaderPropagation; | ||
|
|
||
| /// <summary> | ||
| /// Represents a collection of all the header entries that are to be propagated to the outgoing request. | ||
| /// </summary> | ||
| public class HeaderPropagationEntryCollection | ||
| { | ||
| private readonly Dictionary<string, HeaderPropagationEntry> _entries = []; | ||
|
|
||
| private static readonly object _optionsLock = new object(); | ||
|
|
||
| /// <summary> | ||
| /// Gets the collection of header entries to be propagated to the outgoing request. | ||
| /// </summary> | ||
| public List<HeaderPropagationEntry> Entries | ||
| { | ||
| get => [.. _entries.Select(x => x.Value)]; | ||
| } | ||
|
|
||
| /// <summary> | ||
| /// Attempts to add a new header entry to the collection. | ||
| /// </summary> | ||
| /// <remarks> | ||
| /// If the key already exists in the incoming request headers collection, it will be ignored. | ||
| /// </remarks> | ||
| /// <param name="key">The key of the element to add.</param> | ||
| /// <param name="value">The value to add for the specified key.</param> | ||
| public void Add(string key, StringValues value) | ||
| { | ||
| lock (_optionsLock) | ||
| { | ||
| _entries[key] = new HeaderPropagationEntry | ||
| { | ||
| Key = key, | ||
| Value = value, | ||
| Action = HeaderPropagationEntryAction.Add | ||
| }; | ||
| } | ||
| } | ||
|
|
||
| /// <summary> | ||
| /// Appends a new header value to an existing key. | ||
| /// </summary> | ||
| /// <remarks> | ||
| /// If the key does not exist in the incoming request headers collection, it will be ignored. | ||
| /// </remarks> | ||
| /// <param name="key">The key of the element to add.</param> | ||
| /// <param name="value">The value to add for the specified key.</param> | ||
| public void Append(string key, StringValues value) | ||
| { | ||
| lock (_optionsLock) | ||
| { | ||
| StringValues newValue; | ||
|
|
||
| if (_entries.TryGetValue(key, out var entry)) | ||
| { | ||
| // If the key already exists, append the new value to the existing one. | ||
| newValue = StringValues.Concat(entry.Value, value); | ||
| } | ||
|
|
||
| _entries[key] = new HeaderPropagationEntry | ||
| { | ||
| Key = key, | ||
| Value = !StringValues.IsNullOrEmpty(newValue) ? newValue : value, | ||
| Action = HeaderPropagationEntryAction.Append | ||
| }; | ||
| } | ||
| } | ||
|
|
||
| /// <summary> | ||
| /// Propagates the incoming request header value to the outgoing request. | ||
| /// </summary> | ||
| /// <remarks> | ||
| /// If the key does not exist in the incoming request headers collection, it will be ignored. | ||
| /// </remarks> | ||
| /// <param name="key">The key of the element to add.</param> | ||
| public void Propagate(string key) | ||
| { | ||
| lock (_optionsLock) | ||
| { | ||
| _entries[key] = new HeaderPropagationEntry | ||
| { | ||
| Key = key, | ||
| Action = HeaderPropagationEntryAction.Propagate | ||
| }; | ||
| } | ||
| } | ||
|
|
||
| /// <summary> | ||
| /// Overrides the header value of an existing key. | ||
| /// </summary> | ||
| /// <remarks> | ||
| /// If the key does not exist in the incoming request headers collection, it will add it. | ||
| /// </remarks> | ||
| /// <param name="key">The key of the element to add.</param> | ||
| /// <param name="value">The value to add for the specified key.</param> | ||
| public void Override(string key, StringValues value) | ||
| { | ||
| lock (_optionsLock) | ||
| { | ||
| _entries[key] = new HeaderPropagationEntry | ||
| { | ||
| Key = key, | ||
| Value = value, | ||
| Action = HeaderPropagationEntryAction.Override | ||
| }; | ||
| } | ||
| } | ||
| } |
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Take a look at
SerializationInitAttribute.GetLoadOnInitTypes. We had to add exception handling because some modules failed on this, and the app would crash. See if this is relevant here too.There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thanks Tracy, we applied the changes.