agent/agent-framework
2
0
Fork 0
mirror of https://github.com/microsoft/agent-framework.git synced 2026-06-16 21:04:09 +08:00
Code Issues Packages Projects Releases Wiki Activity

.NET: fix: filter filesystem checkpoint index by session (#6132)

Browse Source 
* fix: filter filesystem checkpoint index by session

* fix: filter checkpoint index by parent

* .NET: preserve legacy checkpoint index discovery
This commit is contained in:
Yufeng He
2026-06-12 06:35:57 +08:00
committed by GitHub
Unverified
parent e7937947d9
commit 4c1b9efa8c
2 changed files with 155 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
dotnet/src/Microsoft.Agents.AI.Workflows/Checkpointing/FileSystemJsonCheckpointStore.cs
+27 -3
View File
@@ -3,6 +3,7 @@
using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Text;
using System.Text.Json;
using System.Text.Json.Serialization.Metadata;
@@ -11,7 +12,11 @@ using System.Threading.Tasks;
namespace Microsoft.Agents.AI.Workflows.Checkpointing;
internal record CheckpointFileIndexEntry(CheckpointInfo CheckpointInfo, string FileName);
internal record CheckpointFileIndexEntry(
CheckpointInfo CheckpointInfo,
string FileName,
string? ParentCheckpointId = null,
bool HasParentMetadata = false);
/// <summary>
/// Provides a file system-based implementation of a JSON checkpoint store that persists checkpoint data and index
@@ -30,6 +35,8 @@ public sealed class FileSystemJsonCheckpointStore : JsonCheckpointStore, IDispos
internal DirectoryInfo Directory { get; }
internal HashSet<CheckpointInfo> CheckpointIndex { get; }
private Dictionary<CheckpointInfo, string?> CheckpointParents { get; } = [];
private HashSet<CheckpointInfo> CheckpointsWithKnownParent { get; } = [];
private static JsonTypeInfo<CheckpointFileIndexEntry> EntryTypeInfo => WorkflowsJsonUtilities.JsonContext.Default.CheckpointFileIndexEntry;
@@ -74,6 +81,11 @@ public sealed class FileSystemJsonCheckpointStore : JsonCheckpointStore, IDispos
// We never actually use the file names from the index entries since they can be derived from the CheckpointInfo, but it is useful to
// have the UrlEncoded file names in the index file for human readability
this.CheckpointIndex.Add(entry.CheckpointInfo);
this.CheckpointParents[entry.CheckpointInfo] = entry.ParentCheckpointId;
if (entry.HasParentMetadata)
{
this.CheckpointsWithKnownParent.Add(entry.CheckpointInfo);
}
}
}
}
@@ -137,7 +149,11 @@ public sealed class FileSystemJsonCheckpointStore : JsonCheckpointStore, IDispos
using Utf8JsonWriter jsonWriter = new(checkpointStream, new JsonWriterOptions() { Indented = false });
value.WriteTo(jsonWriter);
CheckpointFileIndexEntry entry = new(key, fileName);
string? parentCheckpointId = parent?.CheckpointId;
this.CheckpointParents[key] = parentCheckpointId;
this.CheckpointsWithKnownParent.Add(key);
CheckpointFileIndexEntry entry = new(key, fileName, parentCheckpointId, HasParentMetadata: true);
JsonSerializer.Serialize(this._indexFile!, entry, EntryTypeInfo);
byte[] bytes = Encoding.UTF8.GetBytes(Environment.NewLine);
await this._indexFile!.WriteAsync(bytes, 0, bytes.Length, CancellationToken.None).ConfigureAwait(false);
@@ -148,6 +164,8 @@ public sealed class FileSystemJsonCheckpointStore : JsonCheckpointStore, IDispos
catch (Exception ex)
{
this.CheckpointIndex.Remove(key);
this.CheckpointParents.Remove(key);
this.CheckpointsWithKnownParent.Remove(key);
try
{
@@ -184,6 +202,12 @@ public sealed class FileSystemJsonCheckpointStore : JsonCheckpointStore, IDispos
{
this.CheckDisposed();
return new(this.CheckpointIndex);
return new(this.CheckpointIndex
.Where(checkpoint => checkpoint.SessionId == sessionId &&
(withParent is null ||
!this.CheckpointsWithKnownParent.Contains(checkpoint) ||
(this.CheckpointParents.TryGetValue(checkpoint, out string? parentCheckpointId) &&
parentCheckpointId == withParent.CheckpointId)))
.ToArray());
}
}
dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/FileSystemJsonCheckpointStoreTests.cs
+128
View File
@@ -2,6 +2,7 @@
using System;
using System.IO;
using System.Linq;
using System.Text.Json;
using System.Threading.Tasks;
using FluentAssertions;
@@ -197,4 +198,131 @@ public sealed class FileSystemJsonCheckpointStoreTests
retrieved.GetProperty("name").GetString().Should().Be("test");
retrieved.GetProperty("value").GetInt32().Should().Be(42);
}
[Fact]
public async Task RetrieveIndexAsync_ShouldOnlyReturnCheckpointsForRequestedSessionAsync()
{
// Arrange
using TempDirectory tempDirectory = new();
string firstSessionId = Guid.NewGuid().ToString("N");
string secondSessionId = Guid.NewGuid().ToString("N");
CheckpointInfo firstCheckpoint;
CheckpointInfo secondCheckpoint;
using (FileSystemJsonCheckpointStore store = new(tempDirectory))
{
firstCheckpoint = await store.CreateCheckpointAsync(firstSessionId, TestData);
secondCheckpoint = await store.CreateCheckpointAsync(secondSessionId, TestData);
// Act
CheckpointInfo[] firstSessionIndex = (await store.RetrieveIndexAsync(firstSessionId)).ToArray();
// Assert
firstSessionIndex.Should().ContainSingle().Which.Should().Be(firstCheckpoint);
firstSessionIndex.Should().NotContain(secondCheckpoint);
}
using (FileSystemJsonCheckpointStore reopenedStore = new(tempDirectory))
{
CheckpointInfo[] secondSessionIndex = (await reopenedStore.RetrieveIndexAsync(secondSessionId)).ToArray();
secondSessionIndex.Should().ContainSingle().Which.Should().Be(secondCheckpoint);
secondSessionIndex.Should().NotContain(firstCheckpoint);
}
}
[Fact]
public async Task RetrieveIndexAsync_ShouldFilterByParentCheckpointAsync()
{
// Arrange
using TempDirectory tempDirectory = new();
string sessionId = Guid.NewGuid().ToString("N");
CheckpointInfo parentCheckpoint;
CheckpointInfo childCheckpoint;
CheckpointInfo unrelatedCheckpoint;
using (FileSystemJsonCheckpointStore store = new(tempDirectory))
{
parentCheckpoint = await store.CreateCheckpointAsync(sessionId, TestData);
childCheckpoint = await store.CreateCheckpointAsync(sessionId, TestData, parentCheckpoint);
unrelatedCheckpoint = await store.CreateCheckpointAsync(sessionId, TestData);
// Act
CheckpointInfo[] childIndex = (await store.RetrieveIndexAsync(sessionId, parentCheckpoint)).ToArray();
// Assert
childIndex.Should().ContainSingle().Which.Should().Be(childCheckpoint);
childIndex.Should().NotContain(parentCheckpoint);
childIndex.Should().NotContain(unrelatedCheckpoint);
}
using (FileSystemJsonCheckpointStore reopenedStore = new(tempDirectory))
{
CheckpointInfo[] childIndex = (await reopenedStore.RetrieveIndexAsync(sessionId, parentCheckpoint)).ToArray();
childIndex.Should().ContainSingle().Which.Should().Be(childCheckpoint);
childIndex.Should().NotContain(parentCheckpoint);
childIndex.Should().NotContain(unrelatedCheckpoint);
}
}
[Fact]
public async Task RetrieveIndexAsync_ShouldKeepLegacyEntriesDiscoverableWithParentFilterAsync()
{
// Arrange
using TempDirectory tempDirectory = new();
string sessionId = Guid.NewGuid().ToString("N");
CheckpointInfo parentCheckpoint;
CheckpointInfo childCheckpoint;
string childFileName;
using (FileSystemJsonCheckpointStore store = new(tempDirectory))
{
parentCheckpoint = await store.CreateCheckpointAsync(sessionId, TestData);
childCheckpoint = await store.CreateCheckpointAsync(sessionId, TestData, parentCheckpoint);
childFileName = store.GetFileNameForCheckpoint(sessionId, childCheckpoint);
}
string indexPath = Path.Combine(tempDirectory.FullName, "index.jsonl");
string legacyEntry = JsonSerializer.Serialize(new CheckpointFileIndexEntry(childCheckpoint, childFileName));
File.WriteAllText(indexPath, legacyEntry + Environment.NewLine);
// Act
using FileSystemJsonCheckpointStore reopenedStore = new(tempDirectory);
CheckpointInfo[] childIndex = (await reopenedStore.RetrieveIndexAsync(sessionId, parentCheckpoint)).ToArray();
// Assert
childIndex.Should().ContainSingle().Which.Should().Be(childCheckpoint);
}
[Fact]
public async Task RetrieveIndexAsync_ShouldKeepLegacyChildDiscoverableWithUnrelatedParentFilterAsync()
{
// Arrange
using TempDirectory tempDirectory = new();
string sessionId = Guid.NewGuid().ToString("N");
CheckpointInfo parentCheckpoint;
CheckpointInfo childCheckpoint;
CheckpointInfo unrelatedCheckpoint;
string childFileName;
using (FileSystemJsonCheckpointStore store = new(tempDirectory))
{
parentCheckpoint = await store.CreateCheckpointAsync(sessionId, TestData);
childCheckpoint = await store.CreateCheckpointAsync(sessionId, TestData, parentCheckpoint);
unrelatedCheckpoint = await store.CreateCheckpointAsync(sessionId, TestData);
childFileName = store.GetFileNameForCheckpoint(sessionId, childCheckpoint);
}
string indexPath = Path.Combine(tempDirectory.FullName, "index.jsonl");
string legacyEntry = JsonSerializer.Serialize(new CheckpointFileIndexEntry(childCheckpoint, childFileName));
File.WriteAllText(indexPath, legacyEntry + Environment.NewLine);
// Act
using FileSystemJsonCheckpointStore reopenedStore = new(tempDirectory);
CheckpointInfo[] childIndex = (await reopenedStore.RetrieveIndexAsync(sessionId, unrelatedCheckpoint)).ToArray();
// Assert
childIndex.Should().ContainSingle().Which.Should().Be(childCheckpoint);
}
}