Skip to content

Commit f5b35d8

Browse files
adityamandaleekaadityamandaleeka
authored
.NET: CosmosDB Actor State Storage (#262)
* Implement CosmosDB actor state storage. * Fix. * Minor fixes. * Fixes. * Make CosmosDB initialization be lazy. * Remove unnecessary read from write path. * Throw on empty writes. * Add arg validation for read. * Add CosmosIdSanitizer. * Fix. * Fix. * Simplify doc IDs. * Update comment. * fb * Make LazyCosmosContainer internal and add tests. * Make test constants public and remove IVT. * Use source generated JSON context for future nativeAOT support. * Re-add dropped comments.
1 parent 32f1132 commit f5b35d8

27 files changed

+2528
-6
lines changed

‎dotnet/Directory.Packages.props‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,9 @@
1010
<PackageVersion Include="Aspire.Azure.AI.OpenAI" Version="9.3.1-preview.1.25305.6" />
1111
<PackageVersion Include="Aspire.Hosting.AppHost" Version="9.3.1" />
1212
<PackageVersion Include="Aspire.Hosting.Azure.CognitiveServices" Version="9.3.1" />
13+
<PackageVersion Include="Aspire.Hosting.Azure.CosmosDB" Version="9.3.1" />
14+
<PackageVersion Include="Aspire.Microsoft.Azure.Cosmos" Version="9.3.1" />
15+
<PackageVersion Include="Aspire.Hosting.Testing" Version="9.3.1" />
1316
<PackageVersion Include="Azure.AI.Agents.Persistent" Version="1.2.0-beta.1" />
1417
<PackageVersion Include="Azure.AI.OpenAI" Version="2.2.0-beta.5" />
1518
<PackageVersion Include="Azure.Identity" Version="1.14.2" />
@@ -21,6 +24,9 @@
2124
<PackageVersion Include="OpenTelemetry.Instrumentation.AspNetCore" Version="1.12.0" />
2225
<PackageVersion Include="OpenTelemetry.Instrumentation.Http" Version="1.12.0" />
2326
<PackageVersion Include="OpenTelemetry.Instrumentation.Runtime" Version="1.12.0" />
27+
<PackageVersion Include="Microsoft.Azure.Cosmos" Version="3.52.0" />
28+
<!-- Newtonsoft (Required by CosmosClient) -->
29+
<PackageVersion Include="Newtonsoft.Json" Version="13.0.3" />
2430
<!-- System.* -->
2531
<PackageVersion Include="System.Linq.Async" Version="6.0.3" />
2632
<PackageVersion Include="System.Text.Json" Version="9.0.7" />

‎dotnet/agent-framework-dotnet.slnx‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -145,6 +145,9 @@
145145
<Project Path="src/Microsoft.Extensions.AI.Agents.CopilotStudio/Microsoft.Extensions.AI.Agents.CopilotStudio.csproj" />
146146
<Project Path="src/Microsoft.Extensions.AI.Agents.OpenAI/Microsoft.Extensions.AI.Agents.OpenAI.csproj" />
147147
<Project Path="src/Microsoft.Extensions.AI.Agents.Runtime.Abstractions/Microsoft.Extensions.AI.Agents.Runtime.Abstractions.csproj" />
148+
<Project Path="src/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB.csproj">
149+
<BuildType Solution="Publish|*" Project="Release" />
150+
</Project>
148151
<Project Path="src/Microsoft.Extensions.AI.Agents.Runtime/Microsoft.Extensions.AI.Agents.Runtime.csproj" Id="35d72ad5-61e1-45cc-a9ad-fa8490dbd146">
149152
<BuildType Solution="Publish|*" Project="Release" />
150153
</Project>
@@ -163,5 +166,14 @@
163166
<Project Path="tests/Microsoft.Extensions.AI.Agents.UnitTests/Microsoft.Extensions.AI.Agents.UnitTests.csproj">
164167
<BuildType Solution="Publish|*" Project="Debug" />
165168
</Project>
169+
<Project Path="tests/CosmosDB.IntegrationTests/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB.Tests.AppHost/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB.Tests.AppHost.csproj">
170+
<BuildType Solution="Publish|*" Project="Release" />
171+
</Project>
172+
<Project Path="tests/CosmosDB.IntegrationTests/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB.Tests/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB.Tests.csproj">
173+
<BuildType Solution="Publish|*" Project="Release" />
174+
</Project>
175+
<Folder Name="/CosmosDB/">
176+
177+
</Folder>
166178
</Folder>
167179
</Solution>

‎dotnet/samples/HelloHttpApi/HelloHttpApi.ApiService/HelloHttpApi.ApiService.csproj‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,12 +12,14 @@
1212
<ProjectReference Include="..\..\..\src\Microsoft.Extensions.AI.Agents.Runtime.Abstractions\Microsoft.Extensions.AI.Agents.Runtime.Abstractions.csproj" />
1313
<ProjectReference Include="..\..\..\src\Microsoft.Extensions.AI.Agents.Runtime\Microsoft.Extensions.AI.Agents.Runtime.csproj" />
1414
<ProjectReference Include="..\..\..\src\Microsoft.Extensions.AI.Agents\Microsoft.Extensions.AI.Agents.csproj" />
15+
<ProjectReference Include="..\..\..\src\Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB\Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB.csproj" />
1516
<ProjectReference Include="..\HelloHttpApi.ServiceDefaults\HelloHttpApi.ServiceDefaults.csproj" />
1617
</ItemGroup>
1718

1819
<ItemGroup>
1920
<PackageReference Include="Aspire.Azure.AI.OpenAI" />
2021
<PackageReference Include="Aspire.Hosting.Azure.CognitiveServices" />
22+
<PackageReference Include="Aspire.Microsoft.Azure.Cosmos" />
2123
<PackageReference Include="CommunityToolkit.Aspire.OllamaSharp" />
2224
<PackageReference Include="Microsoft.Extensions.AI" />
2325
<PackageReference Include="Microsoft.Extensions.AI.Abstractions" />

‎dotnet/samples/HelloHttpApi/HelloHttpApi.ApiService/HostApplicationBuilderAgentExtensions.cs‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
using Microsoft.Extensions.AI;
55
using Microsoft.Extensions.AI.Agents;
66
using Microsoft.Extensions.AI.Agents.Runtime;
7+
using Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB;
78

89
namespace HelloHttpApi.ApiService;
910

@@ -26,8 +27,12 @@ public static IHostApplicationBuilder AddAIAgent(this IHostApplicationBuilder bu
2627
.Add(triage, customerService, "Hand off to the customer service agent for handling rude customer inquiries.")
2728
.Build("PirateWorkflow");
2829
});
30+
2931
var actorBuilder = builder.AddActorRuntime();
3032

33+
// Add CosmosDB state storage to override default storage
34+
builder.Services.AddCosmosActorStateStorage("actor-state-db", "ActorState");
35+
3136
actorBuilder.AddActorType(
3237
new ActorType(agentKey),
3338
(sp, ctx) => new ChatClientAgentActor(

‎dotnet/samples/HelloHttpApi/HelloHttpApi.ApiService/Program.cs‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,9 @@
88
// Add service defaults & Aspire client integrations.
99
builder.AddServiceDefaults();
1010

11+
// Add CosmosDB client integration
12+
builder.AddAzureCosmosClient("hello-http-api-cosmosdb");
13+
1114
// Add services to the container.
1215
builder.Services.AddProblemDetails();
1316

‎dotnet/samples/HelloHttpApi/HelloHttpApi.AppHost/HelloHttpApi.AppHost.csproj‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
<ItemGroup>
1515
<PackageReference Include="Aspire.Hosting.AppHost" />
1616
<PackageReference Include="Aspire.Hosting.Azure.CognitiveServices" />
17+
<PackageReference Include="Aspire.Hosting.Azure.CosmosDB" />
1718
</ItemGroup>
1819

1920
<ItemGroup>

‎dotnet/samples/HelloHttpApi/HelloHttpApi.AppHost/Program.cs‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,8 +8,15 @@
88
var azOpenAiResourceGroup = builder.AddParameterFromConfiguration("AzureOpenAIResourceGroup", "AzureOpenAI:ResourceGroup");
99
var chatModel = builder.AddAIModel("chat-model").AsAzureOpenAI("gpt-4o", o => o.AsExisting(azOpenAiResource, azOpenAiResourceGroup));
1010

11+
var cosmosDbResource = builder.AddParameterFromConfiguration("CosmosDbName", "CosmosDb:Name");
12+
var cosmosDbResourceGroup = builder.AddParameterFromConfiguration("CosmosDbResourceGroup", "CosmosDb:ResourceGroup");
13+
var cosmos = builder.AddAzureCosmosDB("hello-http-api-cosmosdb").RunAsExisting(cosmosDbResource, cosmosDbResourceGroup);
14+
15+
var stateDb = cosmos.AddCosmosDatabase("actor-state-db");
16+
1117
var apiService = builder.AddProject<Projects.HelloHttpApi_ApiService>("apiservice")
12-
.WithReference(chatModel);
18+
.WithReference(chatModel)
19+
.WithReference(cosmos).WaitFor(cosmos);
1320

1421
builder.AddProject<Projects.HelloHttpApi_Web>("webfrontend")
1522
.WithExternalHttpEndpoints()

‎dotnet/src/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB/ActorDocuments.cs‎

Lines changed: 82 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,82 @@
1+
// Copyright (c) Microsoft. All rights reserved.
2+
3+
using System;
4+
using System.Text.Json;
5+
6+
namespace Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB;
7+
8+
/// <summary>
9+
/// Root document for each actor that provides actor-level ETag semantics.
10+
/// Every write operation updates this document to ensure a single ETag represents
11+
/// the entire actor's state for optimistic concurrency control.
12+
/// This document contains no actor state data. It only serves to track last modified
13+
/// time and provide a single ETag for the actor's state.
14+
///
15+
/// Example structure:
16+
/// {
17+
/// "id": "rootdoc", // Root document ID (constant per actor partition)
18+
/// "actorId": "actor-123", // Partition key (actor ID)
19+
/// "lastModified": "2024-...", // Timestamp
20+
/// }
21+
/// </summary>
22+
public sealed class ActorRootDocument
23+
{
24+
/// <summary>
25+
/// The document ID.
26+
/// </summary>
27+
public string Id { get; set; } = default!;
28+
29+
/// <summary>
30+
/// The actor ID.
31+
/// </summary>
32+
public string ActorId { get; set; } = default!;
33+
34+
/// <summary>
35+
/// The last modified timestamp.
36+
/// </summary>
37+
public DateTimeOffset LastModified { get; set; }
38+
}
39+
40+
/// <summary>
41+
/// Actor state document that represents a single key-value pair in the actor's state.
42+
/// Document Structure (one per actor key):
43+
/// {
44+
/// "id": "state_sanitizedkey", // Unique document ID for the state entry
45+
/// "actorId": "actor-123", // Partition key (actor ID)
46+
/// "key": "foo", // Logical key for the state entry
47+
/// "value": { "bar": 42, "baz": "hello" } // Arbitrary JsonElement payload
48+
/// }
49+
/// </summary>
50+
public sealed class ActorStateDocument
51+
{
52+
/// <summary>
53+
/// The document ID.
54+
/// </summary>
55+
public string Id { get; set; } = default!;
56+
57+
/// <summary>
58+
/// The actor ID.
59+
/// </summary>
60+
public string ActorId { get; set; } = default!;
61+
62+
/// <summary>
63+
/// The logical key for the state entry.
64+
/// </summary>
65+
public string Key { get; set; } = default!;
66+
67+
/// <summary>
68+
/// The value payload.
69+
/// </summary>
70+
public JsonElement Value { get; set; } = default!;
71+
}
72+
73+
/// <summary>
74+
/// Projection class for Cosmos DB queries to retrieve keys.
75+
/// </summary>
76+
public sealed class KeyProjection
77+
{
78+
/// <summary>
79+
/// The key value.
80+
/// </summary>
81+
public string Key { get; set; } = default!;
82+
}

‎dotnet/src/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB/CosmosActorStateJsonContext.cs‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
// Copyright (c) Microsoft. All rights reserved.
2+
3+
using System.Text.Json;
4+
using System.Text.Json.Serialization;
5+
6+
namespace Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB;
7+
8+
/// <summary>
9+
/// Source-generated JSON type information for Cosmos DB actor state documents.
10+
/// </summary>
11+
[JsonSourceGenerationOptions(
12+
JsonSerializerDefaults.Web,
13+
UseStringEnumConverter = true,
14+
DefaultIgnoreCondition = JsonIgnoreCondition.WhenWritingNull,
15+
WriteIndented = false)]
16+
[JsonSerializable(typeof(ActorStateDocument))]
17+
[JsonSerializable(typeof(ActorRootDocument))]
18+
[JsonSerializable(typeof(KeyProjection))]
19+
internal sealed partial class CosmosActorStateJsonContext : JsonSerializerContext;

0 commit comments

Comments
 (0)