diff --git a/dotnet/src/Microsoft.Agents.AI.A2A/A2AAgent.cs b/dotnet/src/Microsoft.Agents.AI.A2A/A2AAgent.cs index bdb1f72928..0b4cacef85 100644 --- a/dotnet/src/Microsoft.Agents.AI.A2A/A2AAgent.cs +++ b/dotnet/src/Microsoft.Agents.AI.A2A/A2AAgent.cs @@ -13,6 +13,7 @@ using Microsoft.Extensions.Logging; using Microsoft.Extensions.Logging.Abstractions; using Microsoft.Shared.Diagnostics; +using Microsoft.Agents.AI; namespace Microsoft.Agents.AI.A2A; @@ -219,6 +220,13 @@ private A2AAgentThread GetA2AThread(AgentThread? thread, AgentRunOptions? option throw new InvalidOperationException("A thread must be provided when AllowBackgroundResponses is enabled."); } + // New logic: Check for ContextId + string? contextId = options?.ContextId ?? AgentContext.ContextId; + if (thread is null && contextId is not null) + { + thread = this.GetNewThread(contextId); + } + thread ??= this.GetNewThread(); if (thread is not A2AAgentThread typedThread) diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions/AgentContext.cs b/dotnet/src/Microsoft.Agents.AI.Abstractions/AgentContext.cs new file mode 100644 index 0000000000..0d2c72f8ac --- /dev/null +++ b/dotnet/src/Microsoft.Agents.AI.Abstractions/AgentContext.cs @@ -0,0 +1,73 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; +using System.Threading; +using Microsoft.Extensions.AI; + +namespace Microsoft.Agents.AI; + +/// +/// Provides access to the current agent execution context. +/// +public static class AgentContext +{ + private class ContextState + { + public AgentThread? Thread { get; init; } + public string? ContextId { get; init; } + public AdditionalPropertiesDictionary? AdditionalProperties { get; init; } + } + + private static readonly AsyncLocal _current = new(); + + /// + /// Gets the current agent thread for the executing operation. + /// + public static AgentThread? CurrentThread => _current.Value?.Thread; + + /// + /// Gets the current context ID (e.g., session ID, correlation ID). + /// + public static string? ContextId => _current.Value?.ContextId; + + /// + /// Gets the additional properties associated with the current execution context. + /// + public static AdditionalPropertiesDictionary? AdditionalProperties => _current.Value?.AdditionalProperties; + + /// + /// Sets the current agent thread for the scope of the returned disposable. + /// + /// The thread to set as current. + /// A disposable that restores the previous thread when disposed. + public static IDisposable SetCurrentThread(AgentThread? thread) + { + return BeginScope(thread: thread); + } + + /// + /// Begins a new execution scope with the specified context data. + /// + /// The context ID. + /// The agent thread. + /// The additional properties. + /// A disposable object that restores the previous context when disposed. + public static IDisposable BeginScope(string? contextId = null, AgentThread? thread = null, AdditionalPropertiesDictionary? additionalProperties = null) + { + var parent = _current.Value; + _current.Value = new ContextState + { + Thread = thread ?? parent?.Thread, + ContextId = contextId ?? parent?.ContextId, + AdditionalProperties = additionalProperties ?? parent?.AdditionalProperties?.Clone() + }; + return new DisposableAction(() => _current.Value = parent); + } + + private sealed class DisposableAction : IDisposable + { + private readonly Action _action; + public DisposableAction(Action action) => _action = action; + public void Dispose() => _action(); + } +} diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions/AgentRunOptions.cs b/dotnet/src/Microsoft.Agents.AI.Abstractions/AgentRunOptions.cs index 9cd6d51680..891a07c56e 100644 --- a/dotnet/src/Microsoft.Agents.AI.Abstractions/AgentRunOptions.cs +++ b/dotnet/src/Microsoft.Agents.AI.Abstractions/AgentRunOptions.cs @@ -33,9 +33,15 @@ public AgentRunOptions(AgentRunOptions options) _ = Throw.IfNull(options); this.ContinuationToken = options.ContinuationToken; this.AllowBackgroundResponses = options.AllowBackgroundResponses; + this.ContextId = options.ContextId; this.AdditionalProperties = options.AdditionalProperties?.Clone(); } + /// + /// Gets or sets the context identifier (e.g., session ID, correlation ID) for the agent run. + /// + public string? ContextId { get; set; } + /// /// Gets or sets the continuation token for resuming and getting the result of the agent response identified by this token. /// @@ -88,6 +94,14 @@ public AgentRunOptions(AgentRunOptions options) /// Additional properties provide a way to include custom metadata or provider-specific /// information that doesn't fit into the standard options schema. This is useful for /// preserving implementation-specific details or extending the options with custom data. + /// + /// Note on A2A Protocol Usage: + /// When communicating via the A2A protocol, this dictionary maps to the Metadata field of A2A messages. + /// Use this for "out-of-band" information like Trace IDs, Tenant IDs, or routing hints. + /// Do not use this for core message content or data payloads; those should be passed as + /// parts (e.g., TextPart, DataPart) in the message body. + /// Use for session or correlation grouping. + /// /// public AdditionalPropertiesDictionary? AdditionalProperties { get; set; } } diff --git a/dotnet/src/Microsoft.Agents.AI.Hosting.A2A/AIAgentExtensions.cs b/dotnet/src/Microsoft.Agents.AI.Hosting.A2A/AIAgentExtensions.cs index c54af66bb8..081ca2e602 100644 --- a/dotnet/src/Microsoft.Agents.AI.Hosting.A2A/AIAgentExtensions.cs +++ b/dotnet/src/Microsoft.Agents.AI.Hosting.A2A/AIAgentExtensions.cs @@ -44,9 +44,12 @@ async Task OnMessageReceivedAsync(MessageSendParams messageSendPara var contextId = messageSendParams.Message.ContextId ?? Guid.NewGuid().ToString("N"); var thread = await hostAgent.GetOrCreateThreadAsync(contextId, cancellationToken).ConfigureAwait(false); + var runOptions = new AgentRunOptions { ContextId = contextId }; + var response = await hostAgent.RunAsync( messageSendParams.ToChatMessages(), thread: thread, + options: runOptions, cancellationToken: cancellationToken).ConfigureAwait(false); await hostAgent.SaveThreadAsync(contextId, thread, cancellationToken).ConfigureAwait(false); diff --git a/dotnet/src/Microsoft.Agents.AI/AgentExtensions.cs b/dotnet/src/Microsoft.Agents.AI/AgentExtensions.cs index 097b789a84..d29b6fa14e 100644 --- a/dotnet/src/Microsoft.Agents.AI/AgentExtensions.cs +++ b/dotnet/src/Microsoft.Agents.AI/AgentExtensions.cs @@ -73,7 +73,24 @@ async Task InvokeAgentAsync( [Description("Input query to invoke the agent.")] string query, CancellationToken cancellationToken) { - var response = await agent.RunAsync(query, thread: thread, cancellationToken: cancellationToken).ConfigureAwait(false); + // If no thread was provided at creation time, try to use the current ambient thread context. + var effectiveThread = thread ?? AgentContext.CurrentThread; + + // Propagate ambient context (ContextId, AdditionalProperties) + var contextId = AgentContext.ContextId; + var additionalProperties = AgentContext.AdditionalProperties; + + AgentRunOptions? runOptions = null; + if (contextId is not null || additionalProperties is not null) + { + runOptions = new AgentRunOptions + { + ContextId = contextId, + AdditionalProperties = additionalProperties?.Clone() + }; + } + + var response = await agent.RunAsync(query, thread: effectiveThread, options: runOptions, cancellationToken: cancellationToken).ConfigureAwait(false); return response.Text; } diff --git a/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgent.cs b/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgent.cs index df7477241c..10d23b580f 100644 --- a/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgent.cs +++ b/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgent.cs @@ -218,44 +218,25 @@ public override async IAsyncEnumerable RunStreamingAsync IAsyncEnumerator responseUpdatesEnumerator; - try - { - // Using the enumerator to ensure we consider the case where no updates are returned for notification. - responseUpdatesEnumerator = chatClient.GetStreamingResponseAsync(inputMessagesForChatClient, chatOptions, cancellationToken).GetAsyncEnumerator(cancellationToken); - } - catch (Exception ex) - { - await NotifyAIContextProviderOfFailureAsync(safeThread, ex, inputMessages, aiContextProviderMessages, cancellationToken).ConfigureAwait(false); - throw; - } - - this._logger.LogAgentChatClientInvokedStreamingAgent(nameof(RunStreamingAsync), this.Id, loggingAgentName, this._chatClientType); - - bool hasUpdates; - try - { - // Ensure we start the streaming request - hasUpdates = await responseUpdatesEnumerator.MoveNextAsync().ConfigureAwait(false); - } - catch (Exception ex) - { - await NotifyAIContextProviderOfFailureAsync(safeThread, ex, inputMessages, aiContextProviderMessages, cancellationToken).ConfigureAwait(false); - throw; - } - - while (hasUpdates) + using (AgentContext.BeginScope(options?.ContextId, safeThread, options?.AdditionalProperties)) { - var update = responseUpdatesEnumerator.Current; - if (update is not null) + try { - update.AuthorName ??= this.Name; - - responseUpdates.Add(update); - yield return new(update) { AgentId = this.Id }; + // Using the enumerator to ensure we consider the case where no updates are returned for notification. + responseUpdatesEnumerator = chatClient.GetStreamingResponseAsync(inputMessagesForChatClient, chatOptions, cancellationToken).GetAsyncEnumerator(cancellationToken); + } + catch (Exception ex) + { + await NotifyAIContextProviderOfFailureAsync(safeThread, ex, inputMessages, aiContextProviderMessages, cancellationToken).ConfigureAwait(false); + throw; } + this._logger.LogAgentChatClientInvokedStreamingAgent(nameof(RunStreamingAsync), this.Id, loggingAgentName, this._chatClientType); + + bool hasUpdates; try { + // Ensure we start the streaming request hasUpdates = await responseUpdatesEnumerator.MoveNextAsync().ConfigureAwait(false); } catch (Exception ex) @@ -263,6 +244,28 @@ public override async IAsyncEnumerable RunStreamingAsync await NotifyAIContextProviderOfFailureAsync(safeThread, ex, inputMessages, aiContextProviderMessages, cancellationToken).ConfigureAwait(false); throw; } + + while (hasUpdates) + { + var update = responseUpdatesEnumerator.Current; + if (update is not null) + { + update.AuthorName ??= this.Name; + + responseUpdates.Add(update); + yield return new(update) { AgentId = this.Id }; + } + + try + { + hasUpdates = await responseUpdatesEnumerator.MoveNextAsync().ConfigureAwait(false); + } + catch (Exception ex) + { + await NotifyAIContextProviderOfFailureAsync(safeThread, ex, inputMessages, aiContextProviderMessages, cancellationToken).ConfigureAwait(false); + throw; + } + } } var chatResponse = responseUpdates.ToChatResponse(); @@ -394,7 +397,10 @@ private async Task RunCoreAsync