.NET: Add orchestration ID to durable agent entity state (#2137)

* Propagate orchestration ID (if any).

* Add integration test for orchestration ID in entity state.

* Update schema.

* Fixup formatting issues.

* Fix more formatting issues.

dotnet/src/Microsoft.Agents.AI.DurableTask/DurableAIAgent.cs

@@ -99,6 +99,8 @@ public sealed class DurableAIAgent : AIAgent
         }
 
         RunRequest request = new([.. messages], responseFormat, enableToolCalls, enableToolNames);
+        request.OrchestrationId = this._context.InstanceId;
+
         try
         {
             return await this._context.Entities.CallEntityAsync<AgentRunResponse>(

dotnet/src/Microsoft.Agents.AI.DurableTask/Microsoft.Agents.AI.DurableTask.csproj

@@ -28,6 +28,7 @@
   </ItemGroup>
 
   <ItemGroup>
+    <InternalsVisibleTo Include="Microsoft.Agents.AI.DurableTask.IntegrationTests" />
     <InternalsVisibleTo Include="Microsoft.Agents.AI.DurableTask.UnitTests" />
     <InternalsVisibleTo Include="Microsoft.Agents.AI.Hosting.AzureFunctions" />
   </ItemGroup>

dotnet/src/Microsoft.Agents.AI.DurableTask/RunRequest.cs

@@ -36,6 +36,13 @@ public record RunRequest
     [JsonInclude]
     internal string CorrelationId { get; set; } = Guid.NewGuid().ToString("N");
 
+    /// <summary>
+    /// Gets or sets the ID of the orchestration that initiated this request (if any).
+    /// </summary>
+    [JsonInclude]
+    [JsonIgnore(Condition = JsonIgnoreCondition.WhenWritingNull)]
+    internal string? OrchestrationId { get; set; }
+
     /// <summary>
     /// Initializes a new instance of the <see cref="RunRequest"/> class for a single message.
     /// </summary>

dotnet/src/Microsoft.Agents.AI.DurableTask/State/DurableAgentState.cs

@@ -23,5 +23,5 @@ internal sealed class DurableAgentState
     /// The version is specified in semver (i.e. "major.minor.patch") format.
     /// </remarks>
     [JsonPropertyName("schemaVersion")]
-    public string SchemaVersion { get; init; } = "1.0.0";
+    public string SchemaVersion { get; init; } = "1.1.0";
 }

dotnet/src/Microsoft.Agents.AI.DurableTask/State/DurableAgentStateRequest.cs

@@ -11,6 +11,12 @@ namespace Microsoft.Agents.AI.DurableTask.State;
 /// </summary>
 internal sealed class DurableAgentStateRequest : DurableAgentStateEntry
 {
+    /// <summary>
+    /// Gets the ID of the orchestration that initiated this request (if any).
+    /// </summary>
+    [JsonPropertyName("orchestrationId")]
+    public string? OrchestrationId { get; init; }
+
     /// <summary>
     /// Gets the expected response type for this request (e.g. "json" or "text").
     /// </summary>
@@ -41,6 +47,7 @@ internal sealed class DurableAgentStateRequest : DurableAgentStateEntry
         return new DurableAgentStateRequest()
         {
             CorrelationId = request.CorrelationId,
+            OrchestrationId = request.OrchestrationId,
             Messages = request.Messages.Select(DurableAgentStateMessage.FromChatMessage).ToList(),
             CreatedAt = request.Messages.Min(m => m.CreatedAt) ?? DateTimeOffset.UtcNow,
             ResponseType = request.ResponseFormat is ChatResponseFormatJson ? "json" : "text",

dotnet/tests/Microsoft.Agents.AI.DurableTask.IntegrationTests/AgentEntityTests.cs

@@ -2,6 +2,8 @@
 
 using System.Diagnostics;
 using System.Reflection;
+using Microsoft.Agents.AI.DurableTask.State;
+using Microsoft.DurableTask;
 using Microsoft.DurableTask.Client;
 using Microsoft.DurableTask.Client.Entities;
 using Microsoft.DurableTask.Entities;
@@ -67,8 +69,72 @@ public sealed class AgentEntityTests(ITestOutputHelper outputHelper) : IDisposab
             cancellationToken: this.TestTimeoutToken);
 
         // Assert: verify the agent state was stored with the correct entity name prefix
-        entity = await client.Entities.GetEntityAsync(expectedEntityId, false, this.TestTimeoutToken);
+        entity = await client.Entities.GetEntityAsync(expectedEntityId, true, this.TestTimeoutToken);
 
         Assert.NotNull(entity);
+        Assert.True(entity.IncludesState);
+
+        DurableAgentState state = entity.State.ReadAs<DurableAgentState>();
+
+        DurableAgentStateRequest request = Assert.Single(state.Data.ConversationHistory.OfType<DurableAgentStateRequest>());
+
+        Assert.Null(request.OrchestrationId);
+    }
+
+    [Fact]
+    public async Task OrchestrationIdSetDuringOrchestrationAsync()
+    {
+        // Arrange
+        AIAgent simpleAgent = TestHelper.GetAzureOpenAIChatClient(s_configuration).CreateAIAgent(
+            name: "TestAgent",
+            instructions: "You are a helpful assistant that always responds with a friendly greeting."
+        );
+
+        using TestHelper testHelper = TestHelper.Start(
+            [simpleAgent],
+            this._outputHelper,
+            registry => registry.AddOrchestrator<TestOrchestrator>());
+
+        DurableTaskClient client = testHelper.GetClient();
+
+        // Act
+        string orchestrationId = await client.ScheduleNewOrchestrationInstanceAsync(nameof(TestOrchestrator), "What is the capital of Maine?");
+
+        OrchestrationMetadata? status = await client.WaitForInstanceCompletionAsync(
+            orchestrationId,
+            true,
+            this.TestTimeoutToken);
+
+        // Assert
+        EntityInstanceId expectedEntityId = AgentSessionId.Parse(status.ReadOutputAs<string>()!);
+
+        EntityMetadata? entity = await client.Entities.GetEntityAsync(expectedEntityId, true, this.TestTimeoutToken);
+
+        Assert.NotNull(entity);
+        Assert.True(entity.IncludesState);
+
+        DurableAgentState state = entity.State.ReadAs<DurableAgentState>();
+
+        DurableAgentStateRequest request = Assert.Single(state.Data.ConversationHistory.OfType<DurableAgentStateRequest>());
+
+        Assert.Equal(orchestrationId, request.OrchestrationId);
+    }
+
+    [System.Diagnostics.CodeAnalysis.SuppressMessage("Performance", "CA1812:Avoid uninstantiated internal classes", Justification = "Constructed via reflection.")]
+    private sealed class TestOrchestrator : TaskOrchestrator<string, string>
+    {
+        public override async Task<string> RunAsync(TaskOrchestrationContext context, string input)
+        {
+            DurableAIAgent writer = context.GetAgent("TestAgent");
+            AgentThread writerThread = writer.GetNewThread();
+
+            await writer.RunAsync(
+                message: context.GetInput<string>()!,
+                thread: writerThread);
+
+            AgentSessionId sessionId = writerThread.GetService<AgentSessionId>();
+
+            return sessionId.ToString();
+        }
     }
 }

dotnet/tests/Microsoft.Agents.AI.DurableTask.UnitTests/State/DurableAgentStateRequestTests.cs

@@ -0,0 +1,34 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+using System.Text.Json;
+using Microsoft.Agents.AI.DurableTask.State;
+
+namespace Microsoft.Agents.AI.DurableTask.Tests.Unit.State;
+
+public sealed class DurableAgentStateRequestTests
+{
+    [Fact]
+    public void RequestSerializationDeserialization()
+    {
+        // Arrange
+        RunRequest originalRequest = new("Hello, world!")
+        {
+            OrchestrationId = "orch-456"
+        };
+        DurableAgentStateRequest originalDurableRequest = DurableAgentStateRequest.FromRunRequest(originalRequest);
+
+        // Act
+        string jsonContent = JsonSerializer.Serialize(
+            originalDurableRequest,
+            DurableAgentStateJsonContext.Default.GetTypeInfo(typeof(DurableAgentStateRequest))!);
+
+        DurableAgentStateRequest? convertedJsonContent = (DurableAgentStateRequest?)JsonSerializer.Deserialize(
+            jsonContent,
+            DurableAgentStateJsonContext.Default.GetTypeInfo(typeof(DurableAgentStateRequest))!);
+
+        // Assert
+        Assert.NotNull(convertedJsonContent);
+        Assert.Equal(originalRequest.CorrelationId, convertedJsonContent.CorrelationId);
+        Assert.Equal(originalRequest.OrchestrationId, convertedJsonContent.OrchestrationId);
+    }
+}

schemas/durable-agent-entity-state.json

@@ -166,6 +166,10 @@
       "description": "The request (i.e. prompt) sent to the agent.",
       "properties": {
         "$type": { "type": "string", "const": "request" },
+        "orchestrationId": {
+          "type": "string",
+          "description": "The identifier of the orchestration that initiated this agent request (if any)."
+        },
         "responseSchema": {
           "type": "object",
           "description": "If the expected response type is JSON, this schema defines the expected structure of the response."