From 1479c5322f61544fd13be65805e6bf0b96d668b3 Mon Sep 17 00:00:00 2001 From: Tamir Dresher Date: Tue, 8 Sep 2026 22:08:13 +0300 Subject: [PATCH] Add pressure-based durable history retention and metrics Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../AgentEntity.cs | 104 +- .../CHANGELOG.md | 1 + .../DurableAgentHistoryRetentionMode.cs | 26 + .../DurableAgentStateRetention.cs | 458 ++++++++ ...bleAgentStateSizeLimitExceededException.cs | 56 + .../DurableAgentTelemetry.cs | 207 ++++ .../DurableAgentsOptions.cs | 41 + .../Microsoft.Agents.AI.DurableTask/Logs.cs | 23 + .../Microsoft.Agents.AI.DurableTask/README.md | 86 +- .../RetentionResult.cs | 24 + .../AgentEntityHistoryTests.cs | 372 ++++++ .../AgentEntityTimeToLiveTests.cs | 48 + .../DurableAgentStateRetentionTests.cs | 1045 +++++++++++++++++ .../DurableAgentTelemetryTests.cs | 474 ++++++++ 14 files changed, 2944 insertions(+), 21 deletions(-) create mode 100644 dotnet/src/Microsoft.Agents.AI.DurableTask/DurableAgentHistoryRetentionMode.cs create mode 100644 dotnet/src/Microsoft.Agents.AI.DurableTask/DurableAgentStateRetention.cs create mode 100644 dotnet/src/Microsoft.Agents.AI.DurableTask/DurableAgentStateSizeLimitExceededException.cs create mode 100644 dotnet/src/Microsoft.Agents.AI.DurableTask/DurableAgentTelemetry.cs create mode 100644 dotnet/src/Microsoft.Agents.AI.DurableTask/RetentionResult.cs create mode 100644 dotnet/tests/Microsoft.Agents.AI.DurableTask.UnitTests/DurableAgentStateRetentionTests.cs create mode 100644 dotnet/tests/Microsoft.Agents.AI.DurableTask.UnitTests/DurableAgentTelemetryTests.cs diff --git a/dotnet/src/Microsoft.Agents.AI.DurableTask/AgentEntity.cs b/dotnet/src/Microsoft.Agents.AI.DurableTask/AgentEntity.cs index 7b1c499..f8e8b08 100644 --- a/dotnet/src/Microsoft.Agents.AI.DurableTask/AgentEntity.cs +++ b/dotnet/src/Microsoft.Agents.AI.DurableTask/AgentEntity.cs @@ -70,7 +70,7 @@ internal class AgentEntity(IServiceProvider services, CancellationToken cancella protected override DurableAgentState InitializeState(TaskEntityOperation entityOperation) { - return this._options.EnableMailboxWrites && + return this.MailboxWritesEnabled && entityOperation.Name is nameof(Run) or nameof(RunAgentAsync) ? new DurableAgentState { @@ -135,16 +135,33 @@ public async Task Run(RunRequest request) // without comparing request content, so callers must not reuse it for another request. // Surface failures before optional migration so failed delivery never writes state. AgentResponse committedResponse = existingOutcome.GetResponse(correlationId); - if (this._options.EnableMailboxWrites && + bool legacyMailboxMigrationRequested = + this.MailboxWritesEnabled && + this.State.SchemaVersion != DurableAgentState.RevisedSchemaVersion; + bool migrationAuthorized = + legacyMailboxMigrationRequested && + this._options.AuthorizeLegacyMigration?.Invoke(this.State) == true; + if (this._options.HistoryRetentionMode == DurableAgentHistoryRetentionMode.Auto && this.State.SchemaVersion != DurableAgentState.RevisedSchemaVersion && - this._options.AuthorizeLegacyMigration?.Invoke(this.State) == true && + !migrationAuthorized) + { + throw new DurableAgentStateCorruptionException( + "Automatic history retention requires schema 2 mailbox state. Legacy terminal transcript " + + "entries must be converted from independently authoritative complete history before delivery."); + } + + if (legacyMailboxMigrationRequested && + migrationAuthorized && existingOutcome.Kind != DurableAgentRunOutcomeKind.CompletedResultUnavailable) { // Legacy evidence is converted without constructing or invoking the agent. DurableAgentState migrated = DurableAgentStateOutcomeResolver.PrepareRevisedWorkingState( this.State, hasAuthoritativeLegacyHistory: true); - ValidateForCommit(migrated); - this.State = migrated; + this.ApplyRetentionAndCommit( + migrated, + sessionId, + logger, + deletionCheckExpiration: null); } return committedResponse; @@ -157,12 +174,12 @@ public async Task Run(RunRequest request) nameof(request)); } - if (this._options.EnableMailboxWrites) + if (this.MailboxWritesEnabled) { DurableAgentStateContract.ValidateIdentifier(correlationId, "correlationId"); } - if (!this._options.EnableMailboxWrites && + if (!this.MailboxWritesEnabled && this.State.SchemaVersion == DurableAgentState.RevisedSchemaVersion) { throw new InvalidOperationException("New mailbox requests require EnableMailboxWrites to be enabled."); @@ -171,13 +188,22 @@ public async Task Run(RunRequest request) this._cancellationToken.ThrowIfCancellationRequested(); // TaskEntity hydrates State with the backend-owned object. Mutate an independent copy so // an exception leaves the hydrated state unchanged. - bool migrateLegacy = this._options.EnableMailboxWrites && + bool migrateLegacy = this.MailboxWritesEnabled && this.State.SchemaVersion != DurableAgentState.RevisedSchemaVersion && this._options.AuthorizeLegacyMigration?.Invoke(this.State) == true; + if (this._options.HistoryRetentionMode == DurableAgentHistoryRetentionMode.Auto && + this.State.SchemaVersion != DurableAgentState.RevisedSchemaVersion && + !migrateLegacy) + { + throw new DurableAgentStateCorruptionException( + "Automatic history retention requires schema 2 mailbox state. Legacy terminal transcript " + + "entries must be converted from independently authoritative complete history before execution."); + } + DurableAgentState workingState = migrateLegacy ? DurableAgentStateOutcomeResolver.PrepareRevisedWorkingState(this.State, hasAuthoritativeLegacyHistory: true) : this.State.Clone(); - if (this._options.EnableMailboxWrites && + if (this.MailboxWritesEnabled && workingState.SchemaVersion == DurableAgentState.RevisedSchemaVersion) { workingState.MailboxWritesAuthorized = true; @@ -502,10 +528,13 @@ persistedHistoryBinding is not null || response.Usage?.TotalTokenCount); } - DateTime? deletionCheckExpiration = this.UpdateExpiration(workingState, sessionId, logger); - this._cancellationToken.ThrowIfCancellationRequested(); - this.CommitWorkingState(workingState, sessionId, logger, deletionCheckExpiration); - + DateTime? deletionCheckExpiration = + this.UpdateExpiration(workingState, sessionId, logger); + this.ApplyRetentionAndCommit( + workingState, + sessionId, + logger, + deletionCheckExpiration); return response; } catch (InvalidOperationException exception) when ( @@ -578,7 +607,7 @@ public void CheckAndExpireResults(AgentEntityResultExpirationCheck? scheduledChe DurableAgentState workingState = DurableAgentStateJsonConverter.DeserializeRevisedContract( DurableAgentStateJsonConverter.SerializeRevisedContract(this.State)); workingState.MailboxWritesAuthorized = true; - this.CommitWorkingState(workingState, sessionId, logger, deletionCheckExpiration: null, + this.ApplyRetentionAndCommit(workingState, sessionId, logger, deletionCheckExpiration: null, previousResultCheckTime: scheduledCheck?.ScheduledTime); } catch (Exception exception) @@ -627,10 +656,37 @@ public void CheckAndDeleteIfExpired(AgentEntityDeletionCheck? scheduledCheck = n if (expirationTime.HasValue) { logger.LogTTLExpirationTimeCleared(sessionId); - DurableAgentState workingState = this.State.Clone(); + bool migrateLegacy = + this.MailboxWritesEnabled && + this._options.HistoryRetentionMode == DurableAgentHistoryRetentionMode.Auto && + this.State.SchemaVersion != DurableAgentState.RevisedSchemaVersion && + this._options.AuthorizeLegacyMigration?.Invoke(this.State) == true; + if (this._options.HistoryRetentionMode == DurableAgentHistoryRetentionMode.Auto && + this.State.SchemaVersion != DurableAgentState.RevisedSchemaVersion && + !migrateLegacy) + { + throw new DurableAgentStateCorruptionException( + "Automatic history retention requires schema 2 mailbox state. Legacy terminal transcript " + + "entries must be converted from independently authoritative complete history before TTL mutation."); + } + + DurableAgentState workingState = migrateLegacy + ? DurableAgentStateOutcomeResolver.PrepareRevisedWorkingState( + this.State, + hasAuthoritativeLegacyHistory: true) + : this.State.Clone(); + if (this.MailboxWritesEnabled && + workingState.SchemaVersion == DurableAgentState.RevisedSchemaVersion) + { + workingState.MailboxWritesAuthorized = true; + } + workingState.Data.ExpirationTimeUtc = null; - ValidateForCommit(workingState); - this.State = workingState; + this.ApplyRetentionAndCommit( + workingState, + sessionId, + logger, + deletionCheckExpiration: null); } return; @@ -673,6 +729,9 @@ state.ExtensionData is null && state.UnknownProperties is null; } + private bool MailboxWritesEnabled => + this._options.EnableMailboxWrites; + private static bool IsPostResponseServiceHistoryFailure( InvalidOperationException exception, ChatClientAgent? chatClientAgent) @@ -838,7 +897,7 @@ ownership is DurableAgentHistoryOwnership.Entity or DurableAgentHistoryOwnership : null; } - private void CommitWorkingState( + private void ApplyRetentionAndCommit( DurableAgentState workingState, AgentSessionId sessionId, ILogger logger, @@ -897,8 +956,17 @@ private void CommitWorkingState( workingState = AgentEntityResultExpirySchedule.Write(workingState, entityId, schedule, pending); } + _ = DurableAgentStateRetention.Enforce( + workingState, + this._options.HistoryRetentionMode, + this._options.MaxStateBytes, + currentTime, + logger, + sessionId); + this._cancellationToken.ThrowIfCancellationRequested(); ValidateForCommit(workingState); + if (deletionCheckExpiration.HasValue) { // Pass the working-copy value explicitly: this.State still refers to the original state diff --git a/dotnet/src/Microsoft.Agents.AI.DurableTask/CHANGELOG.md b/dotnet/src/Microsoft.Agents.AI.DurableTask/CHANGELOG.md index 32e43da..58d15b8 100644 --- a/dotnet/src/Microsoft.Agents.AI.DurableTask/CHANGELOG.md +++ b/dotnet/src/Microsoft.Agents.AI.DurableTask/CHANGELOG.md @@ -2,6 +2,7 @@ ## [Unreleased] +- Added opt-in pressure-based durable transcript retention and low-cardinality operational metrics while protecting schema 2 mailbox and execution-control state ([#97](https://github.com/microsoft/agent-framework-durable-extension/pull/97)) - Hardened durable-agent mailbox delivery, duplicate correlation handling, working-state rollback, token-bounded durable result expiry (including zero-write stale cleanup on legacy state), and stale-safe TTL deletion scheduling; added gated real-backend state/outbox atomicity coverage with locally tested SDK failure handling while keeping schema-2 production activation inaccessible; preserved historical message boundaries, opaque schema-2 history/state profiles, committed failure metadata, and legacy response metadata through mailbox promotion; isolated untrusted response text and malformed activity/child payloads from workflow controls with all-or-nothing typed-message validation, preserved opaque child fallback string routing, and rejected unknown target type hints instead of selecting an unrelated handler ([#94](https://github.com/microsoft/agent-framework-durable-extension/pull/94)) - Fail durable workflows with a `MaxSuperstepsExceededException` when they reach the configurable `MaxSupersteps` limit with work still queued, instead of returning a successful partial result ([#84](https://github.com/microsoft/agent-framework-durable-extension/pull/84)) - Added passive .NET DTO, converter, validation, and source-generation support for the proposed durable agent state 2.0 contract ([tamirdresher/agent-framework-durable-extension#1](https://github.com/tamirdresher/agent-framework-durable-extension/pull/1)) diff --git a/dotnet/src/Microsoft.Agents.AI.DurableTask/DurableAgentHistoryRetentionMode.cs b/dotnet/src/Microsoft.Agents.AI.DurableTask/DurableAgentHistoryRetentionMode.cs new file mode 100644 index 0000000..f491850 --- /dev/null +++ b/dotnet/src/Microsoft.Agents.AI.DurableTask/DurableAgentHistoryRetentionMode.cs @@ -0,0 +1,26 @@ +// Copyright (c) Microsoft. All rights reserved. + +namespace Microsoft.Agents.AI.DurableTask; + +/// +/// Controls how durable agent conversation state is retained. +/// +public enum DurableAgentHistoryRetentionMode +{ + /// + /// Never proactively removes conversation entries. Persistence can still fail when a backend or provider + /// state limit is reached. + /// + KeepAll, + + /// + /// Removes the oldest eligible exchanges when serialized entity state reaches the configured high watermark. + /// + /// + /// Only conversation transcript entries are eligible. Mailbox results, completion receipts, history + /// binding, provider continuation, TTL, and other execution-control state are protected. The newest + /// transcript exchange and system messages are never evicted; if protected state cannot fit below the + /// safe write threshold, the operation fails instead of persisting oversized state. + /// + Auto, +} diff --git a/dotnet/src/Microsoft.Agents.AI.DurableTask/DurableAgentStateRetention.cs b/dotnet/src/Microsoft.Agents.AI.DurableTask/DurableAgentStateRetention.cs new file mode 100644 index 0000000..c9534e9 --- /dev/null +++ b/dotnet/src/Microsoft.Agents.AI.DurableTask/DurableAgentStateRetention.cs @@ -0,0 +1,458 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System.Text.Json; +using Microsoft.Agents.AI.DurableTask.State; +using Microsoft.Extensions.AI; +using Microsoft.Extensions.Logging; + +namespace Microsoft.Agents.AI.DurableTask; + +/// +/// Applies deterministic pressure retention to durable agent state. +/// +internal static class DurableAgentStateRetention +{ + internal const double HighWatermark = 0.85; + internal const double LowWatermark = 0.70; + + internal sealed class ExecutionStatistics + { + public int CandidateGroupingPassCount { get; set; } + + public int SerializedStateMeasurementCount { get; set; } + } + + public static int GetSerializedSize(DurableAgentState state) + { + return JsonSerializer.SerializeToUtf8Bytes( + state, + DurableAgentStateJsonContext.Default.DurableAgentState).Length; + } + + public static int Enforce( + DurableAgentState state, + DurableAgentHistoryRetentionMode mode, + int maxStateBytes, + DateTimeOffset now, + ILogger logger, + AgentSessionId sessionId) => + Enforce(state, mode, maxStateBytes, now, logger, sessionId, statistics: null); + + internal static int Enforce( + DurableAgentState state, + DurableAgentHistoryRetentionMode mode, + int maxStateBytes, + DateTimeOffset now, + ILogger logger, + AgentSessionId sessionId, + ExecutionStatistics? statistics) + { + if (mode == DurableAgentHistoryRetentionMode.KeepAll) + { + return 0; + } + + if (mode != DurableAgentHistoryRetentionMode.Auto) + { + throw new ArgumentOutOfRangeException( + nameof(mode), + mode, + "The durable agent history retention mode is not supported."); + } + + int highWatermark = (int)(maxStateBytes * HighWatermark); + int initialSize = GetSerializedSize(state); + if (statistics is not null) + { + statistics.SerializedStateMeasurementCount++; + } + if (initialSize < highWatermark) + { + DurableAgentTelemetry.RecordNoAction(sessionId.Name); + return 0; + } + + DurableAgentStateSchemaVersion schemaVersion = + DurableAgentStateSchemaVersion.ParseSupported(state.SchemaVersion); + if (schemaVersion.Major != DurableAgentState.RevisedSchemaMajorVersion) + { + throw new DurableAgentStateCorruptionException( + "Automatic history retention requires schema 2 mailbox state. Legacy terminal transcript " + + "entries must be converted to authoritative mailbox results before transcript eviction."); + } + + int lowWatermark = (int)(maxStateBytes * LowWatermark); + List> eligibleGroups = + FindEligibleExchanges(state.Data.ConversationHistory); + if (statistics is not null) + { + statistics.CandidateGroupingPassCount++; + } + + int selectedGroupCount = FindRemovalPrefix( + state, + eligibleGroups, + lowWatermark, + now, + statistics); + int removedEntries = 0; + int removedMessages = 0; + for (int index = 0; index < selectedGroupCount; index++) + { + List group = eligibleGroups[index]; + int removedFromGroup = group.Sum(entry => entry.Messages.Count); + removedEntries += group.Count; + removedMessages += removedFromGroup; + foreach (DurableAgentStateEntry entry in group) + { + _ = state.Data.ConversationHistory.Remove(entry); + } + } + + RecordTruncation(state, removedMessages, now); + int finalSize = selectedGroupCount == 0 + ? initialSize + : GetSerializedSize(state); + if (selectedGroupCount > 0 && statistics is not null) + { + statistics.SerializedStateMeasurementCount++; + } + + bool protectedStateCapacityFailure = finalSize >= highWatermark; + RetentionResult result = new( + removedEntries, + removedMessages, + initialSize, + finalSize, + protectedStateCapacityFailure); + DurableAgentTelemetry.RecordRetentionAttempt(sessionId.Name, result); + + if (removedEntries > 0) + { + logger.LogDurableHistoryTruncated( + sessionId, + initialSize, + maxStateBytes, + removedEntries, + removedMessages, + finalSize); + } + + if (protectedStateCapacityFailure) + { + logger.LogDurableHistoryStillOverBudget( + sessionId, + finalSize, + maxStateBytes); + throw new DurableAgentStateSizeLimitExceededException(finalSize, maxStateBytes); + } + + return result.RemovedMessageCount; + } + + private static int FindRemovalPrefix( + DurableAgentState state, + List> eligibleGroups, + int lowWatermark, + DateTimeOffset now, + ExecutionStatistics? statistics) + { + if (eligibleGroups.Count == 0) + { + return 0; + } + + Dictionary measuredSizes = []; + int Measure(int groupCount) + { + if (measuredSizes.TryGetValue(groupCount, out int measured)) + { + return measured; + } + + int size = MeasureRemovalPrefix(state, eligibleGroups, groupCount, now); + if (statistics is not null) + { + statistics.SerializedStateMeasurementCount++; + } + + measuredSizes[groupCount] = size; + return size; + } + + int firstMessageGroupIndex = eligibleGroups.FindIndex( + static group => group.Any(entry => entry.Messages.Count > 0)); + int zeroMessagePrefixCount = firstMessageGroupIndex < 0 + ? eligibleGroups.Count + : firstMessageGroupIndex; + if (zeroMessagePrefixCount > 0 && + Measure(zeroMessagePrefixCount) <= lowWatermark) + { + return FindFirstPrefixAtOrBelow( + 1, + zeroMessagePrefixCount, + lowWatermark, + Measure); + } + + if (firstMessageGroupIndex < 0) + { + return eligibleGroups.Count; + } + + // Introducing truncation evidence can make the first message-bearing eviction larger than + // the preceding zero-message prefix. After that transition, every additional group removes + // a complete serialized entry while truncation metadata only updates its bounded counters. + int firstPrefixWithTruncation = firstMessageGroupIndex + 1; + if (Measure(firstPrefixWithTruncation) <= lowWatermark) + { + return firstPrefixWithTruncation; + } + + if (Measure(eligibleGroups.Count) > lowWatermark) + { + return eligibleGroups.Count; + } + + return FindFirstPrefixAtOrBelow( + firstPrefixWithTruncation + 1, + eligibleGroups.Count, + lowWatermark, + Measure); + } + + private static int FindFirstPrefixAtOrBelow( + int left, + int right, + int targetSize, + Func measure) + { + while (left < right) + { + int middle = left + ((right - left) / 2); + if (measure(middle) <= targetSize) + { + right = middle; + } + else + { + left = middle + 1; + } + } + + return left; + } + + private static int MeasureRemovalPrefix( + DurableAgentState state, + List> eligibleGroups, + int groupCount, + DateTimeOffset now) + { + List originalHistory = [.. state.Data.ConversationHistory]; + DurableAgentStateTruncation? originalTruncation = state.Data.Truncation; + HashSet removedEntries = []; + int removedMessages = 0; + for (int index = 0; index < groupCount; index++) + { + foreach (DurableAgentStateEntry entry in eligibleGroups[index]) + { + removedEntries.Add(entry); + removedMessages += entry.Messages.Count; + } + } + + try + { + state.Data.ConversationHistory.Clear(); + foreach (DurableAgentStateEntry entry in originalHistory) + { + if (!removedEntries.Contains(entry)) + { + state.Data.ConversationHistory.Add(entry); + } + } + + state.Data.Truncation = ProjectTruncation( + originalTruncation, + removedMessages, + now); + return GetSerializedSize(state); + } + finally + { + state.Data.ConversationHistory.Clear(); + foreach (DurableAgentStateEntry entry in originalHistory) + { + state.Data.ConversationHistory.Add(entry); + } + + state.Data.Truncation = originalTruncation; + } + } + + private static DurableAgentStateTruncation? ProjectTruncation( + DurableAgentStateTruncation? original, + int removedMessages, + DateTimeOffset now) + { + if (removedMessages == 0) + { + return original; + } + + DateTimeOffset effectiveTime = GetEffectiveEvictionTime(original, now); + return new DurableAgentStateTruncation + { + EvictedMessageCount = + (original?.EvictedMessageCount ?? 0) + removedMessages, + FirstEvictedAt = original?.FirstEvictedAt ?? effectiveTime, + LastEvictedAt = effectiveTime, + UnknownProperties = original?.UnknownProperties, + }; + } + + private static List> FindEligibleExchanges( + IList history) + { + List> groups = BuildAtomicGroups(history); + List? newestGroup = history.Count == 0 + ? null + : groups.First(group => group.Contains(history[^1])); + + return groups + .Where(group => + !ReferenceEquals(group, newestGroup) && + !group.Any(entry => + entry.Messages.Any(message => message.Role == ChatRole.System.ToString()))) + .ToList(); + } + + private static List> BuildAtomicGroups( + IList history) + { + int[] parents = Enumerable.Range(0, history.Count).ToArray(); + Dictionary correlationOwners = new(StringComparer.Ordinal); + Dictionary toolCallOwners = new(StringComparer.Ordinal); + + for (int index = 0; index < history.Count; index++) + { + DurableAgentStateEntry entry = history[index]; + if (entry.CorrelationId is not null) + { + UnionWithOwner(correlationOwners, entry.CorrelationId, index); + } + + HashSet entryToolCallIds = new(StringComparer.Ordinal); + foreach (DurableAgentStateContent content in entry.Messages.SelectMany(message => message.Contents)) + { + string? callId = content switch + { + DurableAgentStateFunctionCallContent functionCall => functionCall.CallId, + DurableAgentStateFunctionResultContent functionResult => functionResult.CallId, + _ => null, + }; + + if (!string.IsNullOrWhiteSpace(callId) && entryToolCallIds.Add(callId)) + { + UnionWithOwner(toolCallOwners, callId, index); + } + } + } + + Dictionary> components = []; + List roots = []; + for (int index = 0; index < history.Count; index++) + { + int root = Find(index); + if (!components.TryGetValue(root, out List? component)) + { + component = []; + components[root] = component; + roots.Add(root); + } + + component.Add(history[index]); + } + + return roots.ConvertAll(root => components[root]); + + void UnionWithOwner(Dictionary owners, string key, int index) + { + if (owners.TryGetValue(key, out int owner)) + { + Union(owner, index); + } + else + { + owners[key] = index; + } + } + + int Find(int index) + { + while (parents[index] != index) + { + parents[index] = parents[parents[index]]; + index = parents[index]; + } + + return index; + } + + void Union(int first, int second) + { + int firstRoot = Find(first); + int secondRoot = Find(second); + if (firstRoot == secondRoot) + { + return; + } + + if (firstRoot < secondRoot) + { + parents[secondRoot] = firstRoot; + } + else + { + parents[firstRoot] = secondRoot; + } + } + } + + private static void RecordTruncation( + DurableAgentState state, + int removedMessages, + DateTimeOffset now) + { + if (removedMessages == 0) + { + return; + } + + DurableAgentStateTruncation truncation = state.Data.Truncation ??= new() + { + FirstEvictedAt = now, + }; + + truncation.EvictedMessageCount += removedMessages; + truncation.LastEvictedAt = GetEffectiveEvictionTime(truncation, now); + } + + private static DateTimeOffset GetEffectiveEvictionTime( + DurableAgentStateTruncation? truncation, + DateTimeOffset now) + { + if (truncation is null) + { + return now; + } + + DateTimeOffset effectiveTime = now > truncation.LastEvictedAt + ? now + : truncation.LastEvictedAt; + return effectiveTime > truncation.FirstEvictedAt + ? effectiveTime + : truncation.FirstEvictedAt; + } +} diff --git a/dotnet/src/Microsoft.Agents.AI.DurableTask/DurableAgentStateSizeLimitExceededException.cs b/dotnet/src/Microsoft.Agents.AI.DurableTask/DurableAgentStateSizeLimitExceededException.cs new file mode 100644 index 0000000..4e8144c --- /dev/null +++ b/dotnet/src/Microsoft.Agents.AI.DurableTask/DurableAgentStateSizeLimitExceededException.cs @@ -0,0 +1,56 @@ +// Copyright (c) Microsoft. All rights reserved. + +namespace Microsoft.Agents.AI.DurableTask; + +/// +/// The exception thrown when automatic retention cannot reduce durable agent state below its safe write threshold. +/// +public sealed class DurableAgentStateSizeLimitExceededException : InvalidOperationException +{ + /// + /// Initializes a new instance of the class. + /// + public DurableAgentStateSizeLimitExceededException() + { + } + + /// + /// Initializes a new instance with a specified error message. + /// + public DurableAgentStateSizeLimitExceededException(string? message) + : base(message) + { + } + + /// + /// Initializes a new instance with a specified error message and inner exception. + /// + public DurableAgentStateSizeLimitExceededException(string? message, Exception? innerException) + : base(message, innerException) + { + } + + /// + /// Initializes a new instance of the class. + /// + /// The final serialized state size. + /// The configured state budget. + public DurableAgentStateSizeLimitExceededException(int stateSizeBytes, int maxStateBytes) + : base( + $"Durable agent state has a protected floor of {stateSizeBytes} bytes after all eligible transcript " + + $"eviction and cannot be safely persisted within the configured {maxStateBytes} byte budget.") + { + this.StateSizeBytes = stateSizeBytes; + this.MaxStateBytes = maxStateBytes; + } + + /// + /// Gets the final serialized state size. + /// + public int StateSizeBytes { get; } + + /// + /// Gets the configured state budget. + /// + public int MaxStateBytes { get; } +} diff --git a/dotnet/src/Microsoft.Agents.AI.DurableTask/DurableAgentTelemetry.cs b/dotnet/src/Microsoft.Agents.AI.DurableTask/DurableAgentTelemetry.cs new file mode 100644 index 0000000..c1c66e1 --- /dev/null +++ b/dotnet/src/Microsoft.Agents.AI.DurableTask/DurableAgentTelemetry.cs @@ -0,0 +1,207 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System.Diagnostics; +using System.Diagnostics.CodeAnalysis; +using System.Diagnostics.Metrics; + +namespace Microsoft.Agents.AI.DurableTask; + +/// +/// Provides telemetry identifiers for durable agents. +/// +public static class DurableAgentTelemetry +{ + /// + /// Gets the name of the meter that emits durable-agent metrics. + /// + public const string MeterName = "Microsoft.Agents.AI.DurableTask"; + + internal const string EvictedMessagesInstrumentName = + "durable.agent.history.evicted.messages"; + internal const string EvictedEntriesInstrumentName = + "durable.agent.history.evicted.entries"; + internal const string ReclaimedBytesInstrumentName = + "durable.agent.history.reclaimed.bytes"; + internal const string StateSizeBeforeInstrumentName = + "durable.agent.history.state.size.before"; + internal const string StateSizeAfterInstrumentName = + "durable.agent.history.state.size.after"; + internal const string RetentionOperationsInstrumentName = + "durable.agent.history.retention.operations"; + + internal const string AgentNameTagName = "agent.name"; + internal const string OutcomeTagName = "outcome"; + internal const string ReasonTagName = "reason"; + + internal const string NoActionOutcome = "no_action"; + internal const string TranscriptEvictedOutcome = "transcript_evicted"; + internal const string ProtectedStateCapacityFailureOutcome = + "protected_state_capacity_failure"; + + internal const string TranscriptPressureReason = "transcript_pressure"; + + private static class Instruments + { + internal static readonly Meter Meter = new( + MeterName, + typeof(DurableAgentTelemetry).Assembly.GetName().Version?.ToString()); + internal static readonly Counter EvictedMessages = + Meter.CreateCounter( + EvictedMessagesInstrumentName, + unit: "{message}", + description: "Number of messages removed from durable agent history."); + internal static readonly Counter EvictedEntries = + Meter.CreateCounter( + EvictedEntriesInstrumentName, + unit: "{entry}", + description: "Number of entries removed from durable agent history."); + internal static readonly Counter ReclaimedBytes = + Meter.CreateCounter( + ReclaimedBytesInstrumentName, + unit: "By", + description: "Net serialized durable-state bytes reclaimed by history retention."); + internal static readonly Histogram StateSizeBefore = + Meter.CreateHistogram( + StateSizeBeforeInstrumentName, + unit: "By", + description: "Serialized durable-agent state size before a pressure-retention attempt."); + internal static readonly Histogram StateSizeAfter = + Meter.CreateHistogram( + StateSizeAfterInstrumentName, + unit: "By", + description: "Serialized durable-agent state size after a pressure-retention attempt."); + internal static readonly Counter RetentionOperations = + Meter.CreateCounter( + RetentionOperationsInstrumentName, + unit: "{operation}", + description: "Number of automatic durable-agent history retention checks by outcome."); + } + + [SuppressMessage( + "Design", + "CA1031:Do not catch general exception types", + Justification = "Telemetry must never affect durable agent execution.")] + [SuppressMessage( + "Roslynator", + "RCS1075:Avoid empty catch clause that catches System.Exception", + Justification = "Telemetry must never affect durable agent execution.")] + internal static void RecordNoAction(string agentName) + { + try + { + Counter retentionOperations = Instruments.RetentionOperations; + if (!retentionOperations.Enabled) + { + return; + } + + TagList tags = default; + tags.Add(AgentNameTagName, agentName); + tags.Add(OutcomeTagName, NoActionOutcome); + retentionOperations.Add(1, tags); + } + catch (Exception) + { + // Metrics are best-effort operational telemetry. + } + } + + [SuppressMessage( + "Design", + "CA1031:Do not catch general exception types", + Justification = "Telemetry must never affect durable agent execution.")] + [SuppressMessage( + "Roslynator", + "RCS1075:Avoid empty catch clause that catches System.Exception", + Justification = "Telemetry must never affect durable agent execution.")] + internal static void RecordRetentionAttempt( + string agentName, + RetentionResult result) + { + try + { + Counter evictedMessages = Instruments.EvictedMessages; + Counter evictedEntries = Instruments.EvictedEntries; + Counter reclaimedBytes = Instruments.ReclaimedBytes; + Histogram stateSizeBefore = Instruments.StateSizeBefore; + Histogram stateSizeAfter = Instruments.StateSizeAfter; + Counter retentionOperations = Instruments.RetentionOperations; + if (!evictedMessages.Enabled && + !evictedEntries.Enabled && + !reclaimedBytes.Enabled && + !stateSizeBefore.Enabled && + !stateSizeAfter.Enabled && + !retentionOperations.Enabled) + { + return; + } + + string outcome = result.Outcome switch + { + RetentionOutcome.TranscriptEvicted => TranscriptEvictedOutcome, + RetentionOutcome.ProtectedStateCapacityFailure => + ProtectedStateCapacityFailureOutcome, + _ => NoActionOutcome, + }; + + if (stateSizeBefore.Enabled || stateSizeAfter.Enabled) + { + TagList sizeTags = default; + sizeTags.Add(AgentNameTagName, agentName); + sizeTags.Add(OutcomeTagName, outcome); + stateSizeBefore.Record(result.InitialSizeBytes, sizeTags); + stateSizeAfter.Record(result.FinalSizeBytes, sizeTags); + } + + RecordEviction( + agentName, + TranscriptPressureReason, + result.RemovedEntryCount, + result.RemovedMessageCount, + Math.Max(0, result.InitialSizeBytes - result.FinalSizeBytes)); + + if (retentionOperations.Enabled) + { + TagList operationTags = default; + operationTags.Add(AgentNameTagName, agentName); + operationTags.Add(OutcomeTagName, outcome); + retentionOperations.Add(1, operationTags); + } + } + catch (Exception) + { + // Metrics are best-effort operational telemetry. + } + } + + private static void RecordEviction( + string agentName, + string reason, + int evictedEntries, + int evictedMessages, + int reclaimedBytes) + { + if (evictedEntries <= 0 && evictedMessages <= 0 && reclaimedBytes <= 0) + { + return; + } + + TagList tags = default; + tags.Add(AgentNameTagName, agentName); + tags.Add(ReasonTagName, reason); + if (evictedEntries > 0) + { + Instruments.EvictedEntries.Add(evictedEntries, tags); + } + + if (evictedMessages > 0) + { + Instruments.EvictedMessages.Add(evictedMessages, tags); + } + + if (reclaimedBytes > 0) + { + Instruments.ReclaimedBytes.Add(reclaimedBytes, tags); + } + } +} diff --git a/dotnet/src/Microsoft.Agents.AI.DurableTask/DurableAgentsOptions.cs b/dotnet/src/Microsoft.Agents.AI.DurableTask/DurableAgentsOptions.cs index d6aba4e..5a7abbf 100644 --- a/dotnet/src/Microsoft.Agents.AI.DurableTask/DurableAgentsOptions.cs +++ b/dotnet/src/Microsoft.Agents.AI.DurableTask/DurableAgentsOptions.cs @@ -122,6 +122,47 @@ public TimeSpan MinimumTimeToLiveSignalDelay } } = TimeSpan.FromMinutes(5); + /// + /// Gets or sets how durable agent conversation state is retained. Defaults to + /// . + /// + /// + /// Automatic retention requires mailbox-aware schema 2 state, but selecting + /// does not enable schema 2 production writes. + /// Writer activation remains an internal rollout gate until every participating reader is compatible. + /// Existing legacy sessions require explicitly authorized migration from independently authoritative + /// complete history. + /// + public DurableAgentHistoryRetentionMode HistoryRetentionMode + { + get; + set => field = Enum.IsDefined(value) + ? value + : throw new ArgumentOutOfRangeException( + nameof(value), + value, + "The durable agent history retention mode is not supported."); + } = DurableAgentHistoryRetentionMode.KeepAll; + + /// + /// Gets or sets the extension-controlled serialized state budget used when + /// is . + /// Defaults to 1 MiB. + /// + /// + /// This budget measures the exact JSON payload produced by this extension. Durable Task backends can add + /// envelope bytes outside this payload, so the default retention watermarks intentionally leave headroom. + /// The budget is inactive in mode. + /// Automatic retention fails the operation if protected state cannot fit below the high watermark. + /// + public int MaxStateBytes + { + get; + set => field = value > 0 + ? value + : throw new ArgumentOutOfRangeException(nameof(value), value, "The durable agent state budget must be positive."); + } = 1_048_576; + /// /// Declares that the model service manages history for an agent that enables Agent Framework's /// per-service-call history persistence mode. diff --git a/dotnet/src/Microsoft.Agents.AI.DurableTask/Logs.cs b/dotnet/src/Microsoft.Agents.AI.DurableTask/Logs.cs index 62d9d27..de39633 100644 --- a/dotnet/src/Microsoft.Agents.AI.DurableTask/Logs.cs +++ b/dotnet/src/Microsoft.Agents.AI.DurableTask/Logs.cs @@ -101,6 +101,29 @@ public static partial void LogTTLExpirationTimeCleared( this ILogger logger, AgentSessionId sessionId); + [LoggerMessage( + EventId = 12, + Level = LogLevel.Warning, + Message = "[{SessionId}] Durable state reached {InitialSizeBytes} bytes of a {MaxStateBytes} byte budget. Retention evicted {EvictedEntryCount} transcript entries containing {EvictedMessageCount} message(s), leaving {FinalSizeBytes} bytes.")] + public static partial void LogDurableHistoryTruncated( + this ILogger logger, + AgentSessionId sessionId, + int initialSizeBytes, + int maxStateBytes, + int evictedEntryCount, + int evictedMessageCount, + int finalSizeBytes); + + [LoggerMessage( + EventId = 13, + Level = LogLevel.Error, + Message = "[{SessionId}] Durable state has a protected floor of {ProtectedStateSizeBytes} bytes against a {MaxStateBytes} byte budget after all eligible transcript eviction. Mailbox results, completion receipts, history binding, provider continuation, TTL, execution bookkeeping, system content, and the newest transcript exchange were not removed.")] + public static partial void LogDurableHistoryStillOverBudget( + this ILogger logger, + AgentSessionId sessionId, + int protectedStateSizeBytes, + int maxStateBytes); + [LoggerMessage( EventId = 14, Level = LogLevel.Error, diff --git a/dotnet/src/Microsoft.Agents.AI.DurableTask/README.md b/dotnet/src/Microsoft.Agents.AI.DurableTask/README.md index b7c9ec7..35c02b7 100644 --- a/dotnet/src/Microsoft.Agents.AI.DurableTask/README.md +++ b/dotnet/src/Microsoft.Agents.AI.DurableTask/README.md @@ -86,9 +86,14 @@ and defaults to no payload expiry. Result-payload retention and whole-entity TTL there is no implicit 60-second mailbox expiry. Deleting the entity also deletes its idempotency evidence. Keep schema 2.0 writes disabled until the shared rollout gates are agreed and every participating reader/worker is mailbox-aware or explicitly rejects the new major version. -Producer activation and receipt-deleting entity TTL are internal test gates only, disabled by default; -this draft exposes no public schema-2 activation API. Legacy TTL behavior is preserved, but old deadlines -cannot delete schema-2 receipts without a separately agreed deletion policy. Unknown-field +Producer activation and receipt-deleting entity TTL are internal test gates only, disabled by default. +`HistoryRetentionMode.Auto` does not activate schema 2 production writes. Automatic transcript retention +requires a mailbox-aware state writer, but that rollout remains internal until Python, dashboards, pollers, +and every other participating reader can safely consume or reject schema 2. Auto therefore fails closed when +the internal writer gate is disabled or legacy migration is not explicitly authorized from independently +authoritative complete history. +Legacy TTL behavior is preserved, but old deadlines cannot delete schema-2 receipts without a separately agreed +deletion policy. Unknown-field preservation by an older worker is not sufficient. See [state compatibility](State/README.md). Under the internal mailbox-writer gate, a successful new run also removes already-expired mailbox @@ -238,6 +243,81 @@ decorators. Hidden stateful-compaction pipelines are unsupported but cannot be r side effects without an upstream public discovery hook; this implementation does not use reflection, type-name scanning, guessed session keys, or factory double invocation. +Pressure retention is opt-in. `KeepAll` is the default and performs no proactive history eviction; backend or +provider size limits can still reject a write. Select `Auto` and configure its positive serialized-state budget +when bounded transcript storage is preferred: + +```csharp +services.ConfigureDurableAgents(options => +{ + options.HistoryRetentionMode = DurableAgentHistoryRetentionMode.Auto; + options.MaxStateBytes = 1_048_576; + options.AddAIAgent(agent); +}); +``` + +`MaxStateBytes` is active only in `Auto`. The 85% high watermark starts a retention attempt, which removes the +oldest eligible transcript groups toward the 70% low watermark. The measured payload is the complete JSON state +produced by this extension, including terminal-result mailboxes, completion receipts, fixed history binding, +opaque provider or agent continuation, TTL, ingestion and workflow bookkeeping, truncation evidence, media, and +metadata. Durable Task backends can add envelope bytes outside this measurement. + +Selecting `Auto` configures retention policy only; it does not activate mailbox-aware schema 2 writes. +When the internal rollout gate is enabled, existing legacy sessions are migrated only when the configured +migration authorization confirms independently authoritative complete history. Otherwise the operation fails +before model or provider side effects. + +Only `conversationHistory` transcript entries are eligible for pressure eviction. Mailbox result envelopes, +completion receipts, fixed history binding, serialized continuation, TTL, and other execution controls are +protected. Correlation IDs connect transcript request/response entries, and stable tool-call IDs connect calls +with results even across entries or correlations. Duplicate non-empty tool IDs are conservatively connected; +missing or empty IDs create no cross-entry edge. System-message groups and the newest transcript exchange are +also protected. + +Schema 2 mailbox results remain authoritative after their transcript copies are removed, so duplicate execution +and polling return the same retained result. Legacy state is converted to schema 2 before entity retention once +history ownership can be resolved. Retention itself fails closed if legacy transcript terminals are still the +only completion evidence. + +If all eligible transcript is removed and the protected floor still reaches the high watermark, +`DurableAgentStateSizeLimitExceededException` fails the operation without committing the working state. Auto +does not expire mailbox payloads; delivery expiry is a separate mailbox policy. Large inline image and +tool-result offload is not part of this implementation. + +Retention is separate from model-context compaction: retention destructively removes durable history only under +storage pressure, while compaction changes the context supplied to the model. `Auto` is not +`FollowCompaction`, and stateful compaction remains unsupported. + +### Retention metrics + +The package emits automatic-retention metrics through the +`Microsoft.Agents.AI.DurableTask` meter, with the package assembly version as its instrumentation scope version. +Applications can subscribe by using the public `DurableAgentTelemetry.MeterName` constant. The OpenTelemetry SDK +and exporter remain application choices; the product package depends only on `System.Diagnostics.Metrics`. + +| Instrument | Type | Unit | Tags | Meaning | +| --- | --- | --- | --- | --- | +| `durable.agent.history.evicted.entries` | Counter | `{entry}` | `agent.name`, `reason` | Transcript entries removed, including entries that contain no messages. | +| `durable.agent.history.evicted.messages` | Counter | `{message}` | `agent.name`, `reason` | Transcript messages removed. | +| `durable.agent.history.reclaimed.bytes` | Counter | `By` | `agent.name`, `reason` | Positive net serialized state bytes reclaimed by transcript eviction. | +| `durable.agent.history.state.size.before` | Histogram | `By` | `agent.name`, `outcome` | Exact serialized state size when an `Auto` check reaches the high watermark. | +| `durable.agent.history.state.size.after` | Histogram | `By` | `agent.name`, `outcome` | Exact serialized state size after that pressure-retention attempt. | +| `durable.agent.history.retention.operations` | Counter | `{operation}` | `agent.name`, `outcome` | Automatic retention checks by final outcome. | + +The bounded `outcome` values are `no_action`, `transcript_evicted`, and +`protected_state_capacity_failure`; the bounded `reason` value is `transcript_pressure`. +Removing a zero-message entry increments the entry counter without incrementing the message counter, and +reclaimed bytes are emitted only for a positive net reduction so truncation metadata never creates a negative +measurement. `KeepAll` emits no retention metrics. Session IDs, correlation IDs, message IDs, content, +exception text, and provider paths are never tags. + +These are **attempt-level operational metrics**, not durable-state truth. Retention is evaluated before the +entity operation commits, so a later scheduling, persistence, or retry failure can leave measurements for state +that was not committed; retries can also record an attempt more than once. Exporters can buffer or drop +telemetry. Reload persisted state and inspect model input or mailbox outcomes when validating committed behavior; +do not rely on emitted counters alone or exact-once metric delivery. A metric observation is never evidence that +the corresponding retained state committed. + ## Feedback & Contributing We welcome feedback and contributions in [our GitHub repo](https://github.com/microsoft/agent-framework-durable-extension). diff --git a/dotnet/src/Microsoft.Agents.AI.DurableTask/RetentionResult.cs b/dotnet/src/Microsoft.Agents.AI.DurableTask/RetentionResult.cs new file mode 100644 index 0000000..4521663 --- /dev/null +++ b/dotnet/src/Microsoft.Agents.AI.DurableTask/RetentionResult.cs @@ -0,0 +1,24 @@ +// Copyright (c) Microsoft. All rights reserved. + +namespace Microsoft.Agents.AI.DurableTask; + +internal enum RetentionOutcome +{ + NoAction, + TranscriptEvicted, + ProtectedStateCapacityFailure, +} + +internal sealed record RetentionResult( + int RemovedEntryCount, + int RemovedMessageCount, + int InitialSizeBytes, + int FinalSizeBytes, + bool ProtectedStateCapacityFailure) +{ + public RetentionOutcome Outcome => this.ProtectedStateCapacityFailure + ? RetentionOutcome.ProtectedStateCapacityFailure + : this.RemovedEntryCount > 0 || this.FinalSizeBytes < this.InitialSizeBytes + ? RetentionOutcome.TranscriptEvicted + : RetentionOutcome.NoAction; +} diff --git a/dotnet/tests/Microsoft.Agents.AI.DurableTask.UnitTests/AgentEntityHistoryTests.cs b/dotnet/tests/Microsoft.Agents.AI.DurableTask.UnitTests/AgentEntityHistoryTests.cs index 1c516af..23492fd 100644 --- a/dotnet/tests/Microsoft.Agents.AI.DurableTask.UnitTests/AgentEntityHistoryTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.DurableTask.UnitTests/AgentEntityHistoryTests.cs @@ -1,5 +1,7 @@ // Copyright (c) Microsoft. All rights reserved. +using System.Collections.Concurrent; +using System.Diagnostics.Metrics; using System.Text.Json; using Microsoft.Agents.AI.Compaction; using Microsoft.Agents.AI.DurableTask.State; @@ -1560,6 +1562,361 @@ await Assert.ThrowsAsync( entry => entry.CorrelationId == "new"); } + [Fact] + public async Task AutoRetentionRunsOnCompletedEntityExecutionAsync() + { + RecordingChatClient client = new(); + ChatClientAgent agent = new(client, name: "agent"); + DurableAgentState initialState = CreateLargeState(); + ConcurrentQueue measuredInstruments = new(); + using MeterListener listener = new(); + listener.InstrumentPublished = static (instrument, meterListener) => + { + if (instrument.Meter.Name == DurableAgentTelemetry.MeterName) + { + meterListener.EnableMeasurementEvents(instrument); + } + }; + listener.SetMeasurementEventCallback( + (instrument, _, tags, _) => + { + foreach (KeyValuePair tag in tags) + { + if (tag.Key == DurableAgentTelemetry.AgentNameTagName && + string.Equals(tag.Value as string, "agent", StringComparison.Ordinal)) + { + measuredInstruments.Enqueue(instrument.Name); + break; + } + } + }); + listener.Start(); + + DurableAgentState persisted = await RunEntityAsync( + agent, + initialState, + new RunRequest(new string('n', 500)) { CorrelationId = "new" }, + options => + { + options.HistoryRetentionMode = DurableAgentHistoryRetentionMode.Auto; + options.MaxStateBytes = 5_000; + }); + + Assert.NotNull(persisted.Data.Truncation); + Assert.DoesNotContain(persisted.Data.ConversationHistory, entry => entry.CorrelationId == "oldest"); + Assert.Contains(persisted.Data.ConversationHistory, entry => entry.CorrelationId == "new"); + Assert.Contains("oldest", persisted.Data.TerminalResults!.Keys); + Assert.Contains("oldest", persisted.Data.CompletionReceipts!.Keys); + Assert.Contains( + DurableAgentTelemetry.RetentionOperationsInstrumentName, + measuredInstruments); + + DurableAgentState reloaded = DeserializeState(SerializeState(persisted)); + DurableAgentRunOutcome retainedOutcome = + DurableAgentStateOutcomeResolver.Resolve( + reloaded, + "oldest", + DateTimeOffset.UtcNow); + Assert.Equal(DurableAgentRunOutcomeKind.Succeeded, retainedOutcome.Kind); + Assert.Equal(new string('b', 600), retainedOutcome.Response?.Text); + + RecordingChatClient duplicateClient = new(); + AgentResponse duplicate = await CreateHarness( + new ChatClientAgent(duplicateClient, name: "agent"), + reloaded).RunAsync( + new RunRequest("different request") { CorrelationId = "oldest" }); + Assert.Equal(new string('b', 600), duplicate.Text); + Assert.Equal(0, duplicateClient.InvocationCount); + + RecordingChatClient nextClient = new(); + _ = await RunEntityAsync( + new ChatClientAgent(nextClient, name: "agent"), + DeserializeState(SerializeState(persisted)), + new RunRequest("next request") { CorrelationId = "next" }); + Assert.DoesNotContain( + nextClient.LastMessages, + message => message.Text == new string('a', 600) || + message.Text == new string('b', 600)); + } + + [Fact] + public async Task AutoRetentionRemovesMixedMediaToolGroupFromReloadedModelInputAsync() + { + DateTimeOffset now = DateTimeOffset.UtcNow; + DurableAgentState state = CopyState( + DurableAgentStateOutcomeResolver.PrepareRevisedWorkingState( + new DurableAgentState(), + hasAuthoritativeLegacyHistory: true), + session: null, + historyBinding: DurableAgentHistoryBinding.Create( + DurableAgentHistoryOwnership.Entity, + configuredProviderKey: null)); + using JsonDocument opaqueDocument = JsonDocument.Parse( + """{"$runtimeType":"future-opaque","payload":{"value":42}}"""); + state.Data.ConversationHistory.Add( + new DurableAgentStateRequest + { + CorrelationId = "old-call", + CreatedAt = now.AddMinutes(-10), + Messages = + [ + DurableAgentStateMessage.FromChatMessage( + new ChatMessage(ChatRole.User, "invoke old tool")), + ], + }); + state.Data.ConversationHistory.Add( + new DurableAgentStateResponse + { + CorrelationId = "old-call", + CreatedAt = now.AddMinutes(-10), + Messages = + [ + new DurableAgentStateMessage + { + Role = ChatRole.Assistant.Value, + Contents = + [ + new DurableAgentStateFunctionCallContent + { + CallId = "large-call", + Name = "tool", + Arguments = JsonSerializer.SerializeToElement( + new { payload = new string('a', 4_000) }), + }, + new DurableAgentStateUriContent + { + Uri = new Uri("https://example.test/media"), + MediaType = null, + }, + new DurableAgentStateUnknownContent + { + Content = opaqueDocument.RootElement.Clone(), + }, + ], + }, + ], + }); + state.Data.ConversationHistory.Add( + new DurableAgentStateRequest + { + CorrelationId = "old-result", + CreatedAt = now.AddMinutes(-9), + Messages = + [ + new DurableAgentStateMessage + { + Role = ChatRole.Tool.Value, + Contents = + [ + new DurableAgentStateFunctionResultContent + { + CallId = "large-call", + Result = JsonSerializer.SerializeToElement(new string('r', 8_000)), + }, + ], + }, + ], + }); + state.Data.ConversationHistory.Add( + new DurableAgentStateResponse + { + CorrelationId = "old-result", + CreatedAt = now.AddMinutes(-9), + Messages = + [ + DurableAgentStateMessage.FromChatMessage( + new ChatMessage(ChatRole.Assistant, "old tool complete")), + ], + }); + AddExchange(state, "recent", "recent request", "recent response", now.AddMinutes(-1)); + + DurableAgentState persisted = await RunEntityAsync( + new ChatClientAgent(new RecordingChatClient(), name: "agent"), + state, + new RunRequest("first new request") { CorrelationId = "first-new" }, + options => + { + options.HistoryRetentionMode = DurableAgentHistoryRetentionMode.Auto; + options.MaxStateBytes = 7_000; + }); + + Assert.DoesNotContain( + persisted.Data.ConversationHistory, + entry => entry.CorrelationId is "old-call" or "old-result"); + DurableAgentState reloaded = DeserializeState(SerializeState(persisted)); + RecordingChatClient nextClient = new(); + _ = await RunEntityAsync( + new ChatClientAgent(nextClient, name: "agent"), + reloaded, + new RunRequest("second new request") { CorrelationId = "second-new" }); + + Assert.Contains(nextClient.LastMessages, message => message.Text == "recent request"); + Assert.Contains(nextClient.LastMessages, message => message.Text == "first new request"); + Assert.DoesNotContain( + nextClient.LastMessages.SelectMany(message => message.Contents), + content => + content is FunctionCallContent { CallId: "large-call" } || + content is FunctionResultContent { CallId: "large-call" } || + content is UriContent uri && + uri.Uri == new Uri("https://example.test/media") || + content.RawRepresentation is JsonElement element && + element.ValueKind == JsonValueKind.Object && + element.TryGetProperty("$runtimeType", out _)); + } + + [Fact] + public async Task DefaultKeepAllDoesNotEvictTranscriptUnderConfiguredPressureAsync() + { + RecordingChatClient client = new(); + DurableAgentState initialState = CreateLargeState(); + + DurableAgentState persisted = await RunEntityAsync( + new ChatClientAgent(client, name: "agent"), + initialState, + new RunRequest("new request") { CorrelationId = "new" }, + options => options.MaxStateBytes = 500); + + Assert.Null(persisted.Data.Truncation); + Assert.Contains( + persisted.Data.ConversationHistory, + entry => entry.CorrelationId == "oldest"); + } + + [Fact] + public async Task LegacyAutoWithoutAuthorizedMigrationFailsBeforeModelExecutionAsync() + { + RecordingChatClient client = new(); + DurableAgentState initialState = + CreateStateWithExchange("old", "old request", "old response"); + EntityHarness harness = CreateHarness( + new ChatClientAgent(client, name: "agent"), + initialState, + options => + { + options.HistoryRetentionMode = DurableAgentHistoryRetentionMode.Auto; + options.MaxStateBytes = 500; + options.AuthorizeLegacyMigration = null; + }); + + DurableAgentStateCorruptionException exception = + await Assert.ThrowsAsync( + () => harness.RunAsync( + new RunRequest("new request") { CorrelationId = "new" })); + + Assert.Contains( + "independently authoritative complete history", + exception.Message, + StringComparison.Ordinal); + Assert.Equal(0, client.InvocationCount); + Assert.False(harness.StateWasPersisted); + Assert.Equal(DurableAgentState.CurrentSchemaVersion, initialState.SchemaVersion); + } + + [Fact] + public async Task DuplicateLegacyAutoMigrationStillEnforcesStateBudgetAsync() + { + RecordingChatClient client = new(); + DurableAgentState initialState = CreateStateWithExchange( + "duplicate", + new string('q', 1_000), + new string('a', 2_000)); + EntityHarness harness = CreateHarness( + new ChatClientAgent(client, name: "agent"), + initialState, + options => + { + options.HistoryRetentionMode = DurableAgentHistoryRetentionMode.Auto; + options.MaxStateBytes = 500; + }); + + _ = await Assert.ThrowsAsync( + () => harness.RunAsync( + new RunRequest("different request") { CorrelationId = "duplicate" })); + + Assert.Equal(0, client.InvocationCount); + Assert.False(harness.StateWasPersisted); + Assert.Equal(DurableAgentState.CurrentSchemaVersion, initialState.SchemaVersion); + } + + [Fact] + public async Task DuplicateLegacyAutoWithoutAuthorizedMigrationFailsClosedAsync() + { + RecordingChatClient client = new(); + DurableAgentState initialState = + CreateStateWithExchange("duplicate", "request", "response"); + EntityHarness harness = CreateHarness( + new ChatClientAgent(client, name: "agent"), + initialState, + options => + { + options.HistoryRetentionMode = DurableAgentHistoryRetentionMode.Auto; + options.MaxStateBytes = 10_000; + options.AuthorizeLegacyMigration = null; + }); + + DurableAgentStateCorruptionException exception = + await Assert.ThrowsAsync( + () => harness.RunAsync( + new RunRequest("different request") { CorrelationId = "duplicate" })); + + Assert.Contains( + "independently authoritative complete history", + exception.Message, + StringComparison.Ordinal); + Assert.Equal(0, client.InvocationCount); + Assert.False(harness.StateWasPersisted); + } + + [Fact] + public void AutoDoesNotInitializeMailboxStateWithoutInternalGate() + { + DurableAgentsOptions options = new() + { + EnableMailboxWrites = false, + HistoryRetentionMode = DurableAgentHistoryRetentionMode.Auto, + }; + Dictionary services = new() + { + [typeof(DurableTaskClient)] = new Mock("test").Object, + [typeof(ILoggerFactory)] = new ListLoggerFactory(new ListLoggerProvider()), + [typeof(DurableAgentsOptions)] = options, + }; + TestableAgentEntity entity = new(new DictionaryServiceProvider(services)); + Mock operation = new(); + operation.SetupGet(value => value.Name).Returns(nameof(AgentEntity.Run)); + + DurableAgentState initialized = entity.Initialize(operation.Object); + + Assert.Equal(DurableAgentState.CurrentSchemaVersion, initialized.SchemaVersion); + Assert.False(initialized.MailboxWritesAuthorized); + Assert.Null(initialized.Data.TerminalResults); + Assert.Null(initialized.Data.CompletionReceipts); + } + + [Fact] + public async Task OversizedProtectedStateFailsWithoutPersistenceAsync() + { + RecordingChatClient client = new(); + ChatClientAgent agent = new(client, name: "agent"); + DurableAgentState initialState = new(); + string originalState = SerializeState(initialState); + EntityHarness harness = CreateHarness( + agent, + initialState, + options => + { + options.HistoryRetentionMode = DurableAgentHistoryRetentionMode.Auto; + options.MaxStateBytes = 500; + }); + + await Assert.ThrowsAsync( + () => harness.RunAsync( + new RunRequest(new string('x', 2_000)) { CorrelationId = "new" })); + + Assert.False(harness.StateWasPersisted); + Assert.Equal(originalState, SerializeState(initialState)); + } + [Fact] public async Task ProviderLoadFailureDoesNotInvokeModelOrCommitWorkingStateAsync() { @@ -1920,6 +2277,15 @@ private static DurableAgentState CreateStateWithExchange( return state; } + private static DurableAgentState CreateLargeState() + { + DurableAgentState state = new(); + DateTimeOffset now = DateTimeOffset.UtcNow; + AddExchange(state, "oldest", new string('a', 600), new string('b', 600), now.AddMinutes(-10)); + AddExchange(state, "middle", new string('c', 600), new string('d', 600), now.AddMinutes(-5)); + return state; + } + private static void AddExchange( DurableAgentState state, string correlationId, @@ -2032,6 +2398,12 @@ protected override async IAsyncEnumerable RunCoreStreamingA private sealed class RecordingSession : AgentSession; } + private sealed class TestableAgentEntity(IServiceProvider services) : AgentEntity(services) + { + public DurableAgentState Initialize(TaskEntityOperation operation) => + this.InitializeState(operation); + } + private sealed class FailingSerializationAgent(string name) : AIAgent { public override string? Name => name; diff --git a/dotnet/tests/Microsoft.Agents.AI.DurableTask.UnitTests/AgentEntityTimeToLiveTests.cs b/dotnet/tests/Microsoft.Agents.AI.DurableTask.UnitTests/AgentEntityTimeToLiveTests.cs index 1d42e00..4bdd6a4 100644 --- a/dotnet/tests/Microsoft.Agents.AI.DurableTask.UnitTests/AgentEntityTimeToLiveTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.DurableTask.UnitTests/AgentEntityTimeToLiveTests.cs @@ -156,6 +156,54 @@ await harness.CheckExpirationAsync( Assert.Empty(harness.Signals); } + [Fact] + public async Task AutoRejectsUnauthorizedLegacyTtlMutationAsync() + { + DurableAgentState state = new(); + state.Data.ExpirationTimeUtc = s_startTime.AddMinutes(5).UtcDateTime; + EntityHarness harness = CreateHarness( + TimeSpan.FromMinutes(5), + state, + registerAgent: false); + harness.Options.HistoryRetentionMode = DurableAgentHistoryRetentionMode.Auto; + harness.Options.AuthorizeLegacyMigration = null; + + _ = await Assert.ThrowsAsync( + () => harness.CheckExpirationAsync( + new AgentEntityDeletionCheck(state.Data.ExpirationTimeUtc.Value))); + + Assert.Same(state, harness.State); + Assert.Equal(s_startTime.AddMinutes(5).UtcDateTime, state.Data.ExpirationTimeUtc); + } + + [Fact] + public async Task AutoAppliesProtectedFloorWhenClearingTtlAsync() + { + DurableAgentState state = DurableAgentStateOutcomeResolver.PrepareRevisedWorkingState( + new DurableAgentState(), + hasAuthoritativeLegacyHistory: true); + DurableAgentStateOutcomeResolver.AddSuccessfulResult( + state, + "protected", + new AgentResponse(new ChatMessage(ChatRole.Assistant, new string('x', 2_000))), + s_startTime.AddMinutes(-5)); + state.Data.ExpirationTimeUtc = s_startTime.AddMinutes(5).UtcDateTime; + EntityHarness harness = CreateHarness( + TimeSpan.FromMinutes(5), + state, + registerAgent: false); + harness.Options.HistoryRetentionMode = DurableAgentHistoryRetentionMode.Auto; + harness.Options.MaxStateBytes = 500; + + _ = await Assert.ThrowsAsync( + () => harness.CheckExpirationAsync( + new AgentEntityDeletionCheck(state.Data.ExpirationTimeUtc.Value))); + + Assert.Same(state, harness.State); + Assert.Equal(s_startTime.AddMinutes(5).UtcDateTime, state.Data.ExpirationTimeUtc); + Assert.Single(state.Data.TerminalResults!); + } + [Fact] public async Task StaleLaterCheckDoesNotRescheduleEarlierCurrentExpirationAsync() { diff --git a/dotnet/tests/Microsoft.Agents.AI.DurableTask.UnitTests/DurableAgentStateRetentionTests.cs b/dotnet/tests/Microsoft.Agents.AI.DurableTask.UnitTests/DurableAgentStateRetentionTests.cs new file mode 100644 index 0000000..4045f61 --- /dev/null +++ b/dotnet/tests/Microsoft.Agents.AI.DurableTask.UnitTests/DurableAgentStateRetentionTests.cs @@ -0,0 +1,1045 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System.Text.Json; +using Microsoft.Agents.AI.DurableTask.State; +using Microsoft.Extensions.AI; +using Microsoft.Extensions.Logging.Abstractions; + +namespace Microsoft.Agents.AI.DurableTask.Tests.Unit; + +public sealed class DurableAgentStateRetentionTests +{ + [Fact] + public void KeepAllNeverDeletes() + { + DurableAgentState state = CreateLargeState(); + int originalCount = state.Data.ConversationHistory.Count; + + int removed = DurableAgentStateRetention.Enforce( + state, + DurableAgentHistoryRetentionMode.KeepAll, + 500, + DateTimeOffset.UtcNow, + NullLogger.Instance, + new AgentSessionId("agent", "session")); + + Assert.Equal(0, removed); + Assert.Equal(originalCount, state.Data.ConversationHistory.Count); + } + + [Fact] + public void DefaultRetentionModeIsKeepAll() + { + DurableAgentsOptions options = new(); + + Assert.Equal( + DurableAgentHistoryRetentionMode.KeepAll, + options.HistoryRetentionMode); + Assert.Equal(1_048_576, options.MaxStateBytes); + } + + [Theory] + [InlineData(0)] + [InlineData(-1)] + public void StateBudgetMustBePositive(int value) + { + DurableAgentsOptions options = new(); + + _ = Assert.Throws( + () => options.MaxStateBytes = value); + } + + [Fact] + public void UndefinedRetentionModeIsRejected() + { + DurableAgentsOptions options = new(); + const DurableAgentHistoryRetentionMode Invalid = + (DurableAgentHistoryRetentionMode)42; + + _ = Assert.Throws( + () => options.HistoryRetentionMode = Invalid); + _ = Assert.Throws( + () => DurableAgentStateRetention.Enforce( + CreateRevisedState(), + Invalid, + 1_000, + DateTimeOffset.UtcNow, + NullLogger.Instance, + new AgentSessionId("agent", "session"))); + } + + [Fact] + public void AutoEvictsOldestExchangeAndRecordsBoundedEvidence() + { + DateTimeOffset now = DateTimeOffset.UtcNow; + DurableAgentState state = CreateLargeState(now); + + int removed = DurableAgentStateRetention.Enforce( + state, + DurableAgentHistoryRetentionMode.Auto, + 2_500, + now, + NullLogger.Instance, + new AgentSessionId("agent", "session")); + + Assert.True(removed > 0); + Assert.DoesNotContain(state.Data.ConversationHistory, entry => entry.CorrelationId == "oldest"); + Assert.Contains(state.Data.ConversationHistory, entry => entry.CorrelationId == "newest"); + Assert.NotNull(state.Data.Truncation); + Assert.Equal(removed, state.Data.Truncation.EvictedMessageCount); + Assert.True( + DurableAgentStateRetention.GetSerializedSize(state) < + 2_500 * DurableAgentStateRetention.HighWatermark); + } + + [Fact] + public void AutoDoesNotMoveTruncationEvidenceBackwardWhenClockRegresses() + { + DateTimeOffset now = DateTimeOffset.UtcNow; + DateTimeOffset firstEviction = now.AddMinutes(10); + DateTimeOffset lastEviction = now.AddMinutes(20); + DurableAgentState state = CreateLargeState(now); + state.Data.Truncation = new DurableAgentStateTruncation + { + EvictedMessageCount = 4, + FirstEvictedAt = firstEviction, + LastEvictedAt = lastEviction, + }; + + int removed = DurableAgentStateRetention.Enforce( + state, + DurableAgentHistoryRetentionMode.Auto, + 2_500, + now, + NullLogger.Instance, + new AgentSessionId("agent", "session")); + + Assert.True(removed > 0); + Assert.Equal(firstEviction, state.Data.Truncation?.FirstEvictedAt); + Assert.Equal(lastEviction, state.Data.Truncation?.LastEvictedAt); + Assert.Equal(4 + removed, state.Data.Truncation?.EvictedMessageCount); + } + + [Fact] + public void AutoPreservesSystemExchange() + { + DateTimeOffset now = DateTimeOffset.UtcNow; + DurableAgentState state = CreateLargeState(now); + state.Data.ConversationHistory.Insert( + 0, + CreateRequest("system", ChatRole.System, new string('s', 500), now.AddMinutes(-10))); + + _ = DurableAgentStateRetention.Enforce( + state, + DurableAgentHistoryRetentionMode.Auto, + 3_200, + now, + NullLogger.Instance, + new AgentSessionId("agent", "session")); + + Assert.Contains(state.Data.ConversationHistory, entry => entry.CorrelationId == "system"); + } + + [Fact] + public void AutoPreservesNewestExchange() + { + DateTimeOffset now = DateTimeOffset.UtcNow; + DurableAgentState state = CreateLargeState(now); + + _ = DurableAgentStateRetention.Enforce( + state, + DurableAgentHistoryRetentionMode.Auto, + 2_500, + now, + NullLogger.Instance, + new AgentSessionId("agent", "session")); + + Assert.Contains(state.Data.ConversationHistory, entry => entry.CorrelationId == "newest"); + } + + [Fact] + public void AutoPreservesMailboxResultWhileEvictingTranscript() + { + DateTimeOffset now = DateTimeOffset.UtcNow; + DurableAgentState state = CreateRevisedState(); + AddExchange(state, "old", new string('o', 400), now.AddMinutes(-5)); + AddExchange(state, "completed", new string('a', 400), now.AddSeconds(-30)); + AddExchange(state, "newest", new string('b', 400), now); + AddMailboxResult(state, "completed", "authoritative result", now.AddSeconds(-30)); + + _ = DurableAgentStateRetention.Enforce( + state, + DurableAgentHistoryRetentionMode.Auto, + 2_600, + now, + NullLogger.Instance, + new AgentSessionId("agent", "session")); + + Assert.DoesNotContain(state.Data.ConversationHistory, entry => entry.CorrelationId == "old"); + Assert.DoesNotContain(state.Data.ConversationHistory, entry => entry.CorrelationId == "completed"); + Assert.Contains(state.Data.ConversationHistory, entry => entry.CorrelationId == "newest"); + DurableAgentRunOutcome outcome = + DurableAgentStateOutcomeResolver.Resolve(state, "completed", now); + Assert.Equal(DurableAgentRunOutcomeKind.Succeeded, outcome.Kind); + Assert.Equal("authoritative result", outcome.Response?.Text); + Assert.Contains("completed", state.Data.CompletionReceipts!.Keys); + } + + [Fact] + public void AutoFindsZeroMessagePrefixBeforeTruncationSizeJump() + { + DateTimeOffset now = DateTimeOffset.UtcNow; + DurableAgentState state = CreateRevisedState(); + for (int index = 0; index < 3; index++) + { + state.Data.ConversationHistory.Add( + new DurableAgentStateCompaction + { + CreatedAt = now.AddMinutes(index - 4), + }); + } + + state.Data.ConversationHistory.Add( + new DurableAgentStateCompaction + { + CreatedAt = now.AddMinutes(-1), + Messages = + [ + new DurableAgentStateMessage + { + Role = ChatRole.Assistant.Value, + }, + ], + }); + state.Data.ConversationHistory.Add( + new DurableAgentStateCompaction + { + CreatedAt = now, + Messages = + [ + DurableAgentStateMessage.FromChatMessage( + new ChatMessage(ChatRole.Assistant, new string('p', 500))), + ], + }); + + int sizeAfterTwo = MeasureProjectedPrefix(state, 2, now); + int sizeAfterThree = MeasureProjectedPrefix(state, 3, now); + int sizeAfterFour = MeasureProjectedPrefix(state, 4, now); + Assert.True(sizeAfterThree < sizeAfterTwo); + Assert.True(sizeAfterThree < sizeAfterFour); + int maxStateBytes = (int)Math.Ceiling( + sizeAfterThree / DurableAgentStateRetention.LowWatermark); + int lowWatermark = (int)( + maxStateBytes * DurableAgentStateRetention.LowWatermark); + Assert.InRange(lowWatermark, sizeAfterThree, Math.Min(sizeAfterTwo, sizeAfterFour) - 1); + + int removedMessages = DurableAgentStateRetention.Enforce( + state, + DurableAgentHistoryRetentionMode.Auto, + maxStateBytes, + now, + NullLogger.Instance, + new AgentSessionId("agent", "session")); + + Assert.Equal(0, removedMessages); + Assert.Equal(2, state.Data.ConversationHistory.Count); + Assert.Null(state.Data.Truncation); + Assert.Equal(sizeAfterThree, DurableAgentStateRetention.GetSerializedSize(state)); + } + + [Fact] + public void AutoFailsWhenProtectedMailboxAndNewestTranscriptAloneExceedBudget() + { + DateTimeOffset now = DateTimeOffset.UtcNow; + DurableAgentState state = CreateRevisedState(); + AddMailboxResult(state, "completed", new string('m', 2_000), now.AddMinutes(-5)); + AddExchange(state, "newest", new string('c', 500), now); + int protectedFloorSize = DurableAgentStateRetention.GetSerializedSize(state); + const int HighWatermark = + (int)(1_500 * DurableAgentStateRetention.HighWatermark); + + Assert.True(protectedFloorSize >= HighWatermark); + + DurableAgentStateSizeLimitExceededException exception = + Assert.Throws( + () => DurableAgentStateRetention.Enforce( + state, + DurableAgentHistoryRetentionMode.Auto, + 1_500, + now, + NullLogger.Instance, + new AgentSessionId("agent", "session"))); + + Assert.Equal(protectedFloorSize, exception.StateSizeBytes); + Assert.Contains("completed", state.Data.TerminalResults!.Keys); + Assert.Contains("completed", state.Data.CompletionReceipts!.Keys); + Assert.Contains(state.Data.ConversationHistory, entry => entry.CorrelationId == "newest"); + Assert.Equal(2, state.Data.ConversationHistory.Count); + } + + [Fact] + public void AutoFailsRatherThanPersistProtectedStateOverSafeThreshold() + { + DateTimeOffset now = DateTimeOffset.UtcNow; + DurableAgentState state = CreateRevisedState(); + AddExchange(state, "newest", new string('x', 2_000), now); + + DurableAgentStateSizeLimitExceededException exception = + Assert.Throws( + () => DurableAgentStateRetention.Enforce( + state, + DurableAgentHistoryRetentionMode.Auto, + 500, + now, + NullLogger.Instance, + new AgentSessionId("agent", "session"))); + + Assert.True(exception.StateSizeBytes >= 500 * DurableAgentStateRetention.HighWatermark); + Assert.Equal(500, exception.MaxStateBytes); + Assert.Equal(2, state.Data.ConversationHistory.Count); + } + + [Fact] + public void AutoRemovesToolCallAndResultAtomicallyWithExchange() + { + DateTimeOffset now = DateTimeOffset.UtcNow; + DurableAgentState state = CreateRevisedState(); + DurableAgentStateRequest request = CreateRequest("tools", ChatRole.User, new string('a', 400), now.AddMinutes(-10)); + DurableAgentStateResponse response = DurableAgentStateResponse.FromResponse( + "tools", + new AgentResponse( + new ChatMessage( + ChatRole.Assistant, + [ + new FunctionCallContent("call", "tool"), + new FunctionResultContent("call", "result"), + ]) + { + CreatedAt = now.AddMinutes(-10), + })); + state.Data.ConversationHistory.Add(request); + state.Data.ConversationHistory.Add(response); + AddExchange(state, "newest", new string('b', 400), now); + + _ = DurableAgentStateRetention.Enforce( + state, + DurableAgentHistoryRetentionMode.Auto, + 2_000, + now, + NullLogger.Instance, + new AgentSessionId("agent", "session")); + + Assert.DoesNotContain(state.Data.ConversationHistory, entry => entry.CorrelationId == "tools"); + Assert.DoesNotContain( + state.Data.ConversationHistory.SelectMany(entry => entry.Messages).SelectMany(message => message.Contents), + content => content is DurableAgentStateFunctionCallContent or DurableAgentStateFunctionResultContent); + } + + [Fact] + public void AutoEvictsToolCallAndResultAcrossDifferentCorrelations() + { + DateTimeOffset now = DateTimeOffset.UtcNow; + DurableAgentState state = CreateRevisedState(); + state.Data.ConversationHistory.Add( + CreateRequest("call", ChatRole.User, "invoke", now.AddMinutes(-10))); + state.Data.ConversationHistory.Add( + CreateToolCallResponse("call", "shared-call", new string('a', 2_000), now.AddMinutes(-10))); + state.Data.ConversationHistory.Add( + CreateToolResultRequest("result", "shared-call", new string('b', 2_000), now.AddMinutes(-9))); + state.Data.ConversationHistory.Add( + CreateResponse("result", "after tool", now.AddMinutes(-9))); + AddExchange(state, "newest", new string('c', 400), now); + + _ = DurableAgentStateRetention.Enforce( + state, + DurableAgentHistoryRetentionMode.Auto, + 5_000, + now, + NullLogger.Instance, + new AgentSessionId("agent", "session")); + + Assert.DoesNotContain( + state.Data.ConversationHistory, + entry => entry.CorrelationId is "call" or "result"); + Assert.False(ContainsToolCallId(state, "shared-call")); + Assert.Contains(state.Data.ConversationHistory, entry => entry.CorrelationId == "newest"); + Assert.True( + DurableAgentStateRetention.GetSerializedSize(state) < + 5_000 * DurableAgentStateRetention.HighWatermark); + } + + [Fact] + public void AutoTreatsInterleavedToolCallsAsOneConnectedComponent() + { + DateTimeOffset now = DateTimeOffset.UtcNow; + DurableAgentState state = CreateRevisedState(); + state.Data.ConversationHistory.Add( + CreateRequest("calls", ChatRole.User, "invoke", now.AddMinutes(-10))); + state.Data.ConversationHistory.Add( + CreateToolCallResponse( + "calls", + now.AddMinutes(-10), + ("call-a", new string('a', 1_000)), + ("call-b", new string('b', 1_000)))); + state.Data.ConversationHistory.Add( + CreateToolResultRequest("result-a", "call-a", new string('c', 1_000), now.AddMinutes(-9))); + state.Data.ConversationHistory.Add( + CreateResponse("result-a", "after a", now.AddMinutes(-9))); + state.Data.ConversationHistory.Add( + CreateToolResultRequest("result-b", "call-b", new string('d', 1_000), now.AddMinutes(-8))); + state.Data.ConversationHistory.Add( + CreateResponse("result-b", "after b", now.AddMinutes(-8))); + AddExchange(state, "newest", new string('e', 400), now); + + _ = DurableAgentStateRetention.Enforce( + state, + DurableAgentHistoryRetentionMode.Auto, + 5_000, + now, + NullLogger.Instance, + new AgentSessionId("agent", "session")); + + Assert.DoesNotContain( + state.Data.ConversationHistory, + entry => entry.CorrelationId is "calls" or "result-a" or "result-b"); + Assert.False(ContainsToolCallId(state, "call-a")); + Assert.False(ContainsToolCallId(state, "call-b")); + Assert.Contains(state.Data.ConversationHistory, entry => entry.CorrelationId == "newest"); + } + + [Fact] + public void AutoProtectsWholeToolComponentWhenResultIsInNewestExchange() + { + DateTimeOffset now = DateTimeOffset.UtcNow; + DurableAgentState state = CreateRevisedState(); + AddExchange(state, "filler", new string('f', 7_000), now.AddMinutes(-20)); + state.Data.ConversationHistory.Add( + CreateRequest("call", ChatRole.User, "invoke", now.AddMinutes(-10))); + state.Data.ConversationHistory.Add( + CreateToolCallResponse("call", "protected-call", new string('a', 500), now.AddMinutes(-10))); + state.Data.ConversationHistory.Add( + CreateToolResultRequest("newest", "protected-call", new string('b', 500), now)); + state.Data.ConversationHistory.Add( + CreateResponse("newest", "final", now)); + + _ = DurableAgentStateRetention.Enforce( + state, + DurableAgentHistoryRetentionMode.Auto, + 5_000, + now, + NullLogger.Instance, + new AgentSessionId("agent", "session")); + + Assert.DoesNotContain(state.Data.ConversationHistory, entry => entry.CorrelationId == "filler"); + Assert.Contains(state.Data.ConversationHistory, entry => entry.CorrelationId == "call"); + Assert.Contains(state.Data.ConversationHistory, entry => entry.CorrelationId == "newest"); + Assert.True(ContainsToolCallId(state, "protected-call")); + } + + [Fact] + public void AutoTreatsDuplicateToolIdsAsOneConservativeGroup() + { + DateTimeOffset now = DateTimeOffset.UtcNow; + DurableAgentState state = CreateRevisedState(); + AddExchange(state, "filler", new string('f', 7_000), now.AddMinutes(-20)); + state.Data.ConversationHistory.Add( + CreateToolCallResponse("first", "duplicate", new string('a', 500), now.AddMinutes(-10))); + state.Data.ConversationHistory.Add( + CreateToolCallResponse("second", "duplicate", new string('b', 500), now.AddMinutes(-5))); + state.Data.ConversationHistory.Add( + CreateToolResultRequest("newest", "duplicate", "result", now)); + state.Data.ConversationHistory.Add( + CreateResponse("newest", "final", now)); + + _ = DurableAgentStateRetention.Enforce( + state, + DurableAgentHistoryRetentionMode.Auto, + 5_000, + now, + NullLogger.Instance, + new AgentSessionId("agent", "session")); + + Assert.DoesNotContain(state.Data.ConversationHistory, entry => entry.CorrelationId == "filler"); + Assert.Contains(state.Data.ConversationHistory, entry => entry.CorrelationId == "first"); + Assert.Contains(state.Data.ConversationHistory, entry => entry.CorrelationId == "second"); + Assert.Contains(state.Data.ConversationHistory, entry => entry.CorrelationId == "newest"); + } + + [Fact] + public void AutoDoesNotConnectOrphanedToolContentWithDifferentIds() + { + DateTimeOffset now = DateTimeOffset.UtcNow; + DurableAgentState state = CreateRevisedState(); + state.Data.ConversationHistory.Add( + CreateToolCallResponse("orphan-call", "call-only", new string('a', 5_000), now.AddMinutes(-10))); + state.Data.ConversationHistory.Add( + CreateToolResultRequest("orphan-result", "result-only", new string('b', 500), now.AddMinutes(-5))); + state.Data.ConversationHistory.Add( + CreateResponse("orphan-result", "after orphan", now.AddMinutes(-5))); + AddExchange(state, "newest", new string('c', 400), now); + + _ = DurableAgentStateRetention.Enforce( + state, + DurableAgentHistoryRetentionMode.Auto, + 4_000, + now, + NullLogger.Instance, + new AgentSessionId("agent", "session")); + + Assert.DoesNotContain(state.Data.ConversationHistory, entry => entry.CorrelationId == "orphan-call"); + Assert.Contains(state.Data.ConversationHistory, entry => entry.CorrelationId == "orphan-result"); + Assert.True(ContainsToolCallId(state, "result-only")); + } + + [Fact] + public void AutoDoesNotConnectToolContentWithMissingIds() + { + DateTimeOffset now = DateTimeOffset.UtcNow; + DurableAgentState state = CreateRevisedState(); + state.Data.ConversationHistory.Add( + CreateToolCallResponse("missing-call", string.Empty, new string('a', 5_000), now.AddMinutes(-10))); + state.Data.ConversationHistory.Add( + CreateToolResultRequest("missing-result", string.Empty, new string('b', 500), now.AddMinutes(-5))); + state.Data.ConversationHistory.Add( + CreateResponse("missing-result", "after orphan", now.AddMinutes(-5))); + AddExchange(state, "newest", new string('c', 400), now); + + _ = DurableAgentStateRetention.Enforce( + state, + DurableAgentHistoryRetentionMode.Auto, + 4_000, + now, + NullLogger.Instance, + new AgentSessionId("agent", "session")); + + Assert.DoesNotContain(state.Data.ConversationHistory, entry => entry.CorrelationId == "missing-call"); + Assert.Contains(state.Data.ConversationHistory, entry => entry.CorrelationId == "missing-result"); + } + + [Fact] + public void SerializedSizeIncludesSessionAndTruncation() + { + DurableAgentState state = CreateRevisedState(); + int emptySize = DurableAgentStateRetention.GetSerializedSize(state); + state.Data.Session = JsonSerializer.SerializeToElement(new { conversationId = new string('c', 100) }); + state.Data.Truncation = new DurableAgentStateTruncation + { + EvictedMessageCount = 2, + FirstEvictedAt = DateTimeOffset.UtcNow, + LastEvictedAt = DateTimeOffset.UtcNow, + }; + + int completeSize = DurableAgentStateRetention.GetSerializedSize(state); + + Assert.True(completeSize > emptySize + 100); + } + + [Fact] + public void AutoPreservesMailboxContinuationBindingTtlAndBookkeepingFloor() + { + DateTimeOffset now = DateTimeOffset.UtcNow; + DurableAgentState state = CreateRevisedState( + DurableAgentStateHistoryBinding.HistoryProviderOwner, + "external-history.v1"); + AddMailboxResult(state, "completed", new string('r', 2_000), now.AddMinutes(-5)); + state.Data.Session = JsonSerializer.SerializeToElement( + new { continuation = new string('s', 2_000) }); + state.Data.ExpirationTimeUtc = now.AddDays(1).UtcDateTime; + state.Data.IngestedPositions = new Dictionary + { + ["workflow"] = 42, + }; + state.Data.UnknownProperties = new Dictionary + { + ["control"] = JsonSerializer.SerializeToElement(new string('e', 1_000)), + }; + + DurableAgentStateSizeLimitExceededException exception = + Assert.Throws( + () => DurableAgentStateRetention.Enforce( + state, + DurableAgentHistoryRetentionMode.Auto, + 1_000, + now, + NullLogger.Instance, + new AgentSessionId("agent", "session"))); + + Assert.Empty(state.Data.ConversationHistory); + Assert.Contains("completed", state.Data.TerminalResults!.Keys); + Assert.Contains("completed", state.Data.CompletionReceipts!.Keys); + Assert.Equal( + "external-history.v1", + DurableAgentHistoryBinding.Parse(state.Data.HistoryBinding)?.ProviderKey); + Assert.NotNull(state.Data.Session); + Assert.NotNull(state.Data.ExpirationTimeUtc); + Assert.Equal(42, state.Data.IngestedPositions?["workflow"]); + Assert.True(exception.StateSizeBytes > exception.MaxStateBytes); + } + + [Fact] + public void AutoPreservesLosslessStructuredResultAndUnavailableReceipt() + { + DateTimeOffset now = DateTimeOffset.UtcNow; + DurableAgentState state = CreateRevisedState(); + DurableAgentStateOutcomeResolver.AddSuccessfulResult( + state, + "lossless", + new AgentResponse(new ChatMessage(ChatRole.Assistant, "mailbox")), + now.AddMinutes(-5), + structuredValue: JsonSerializer.SerializeToElement( + new { count = 3, label = "retained" })); + DurableAgentStateOutcomeResolver.AddSuccessfulResult( + state, + "unavailable", + new AgentResponse(new ChatMessage(ChatRole.Assistant, "expired")), + now.AddMinutes(-5), + resultExpiresAt: now.AddMinutes(-1)); + Assert.True( + DurableAgentStateOutcomeResolver.MarkExpiredResultUnavailable( + state, + "unavailable", + now)); + AddExchange(state, "old", new string('x', 4_000), now.AddMinutes(-10)); + AddExchange(state, "newest", "newest", now); + + _ = DurableAgentStateRetention.Enforce( + state, + DurableAgentHistoryRetentionMode.Auto, + 4_000, + now, + NullLogger.Instance, + new AgentSessionId("agent", "session")); + + DurableAgentRunOutcome lossless = + DurableAgentStateOutcomeResolver.Resolve(state, "lossless", now); + Assert.Equal(DurableAgentRunOutcomeKind.Succeeded, lossless.Kind); + Assert.Equal(3, lossless.Value.GetProperty("count").GetInt32()); + Assert.Equal("retained", lossless.Value.GetProperty("label").GetString()); + DurableAgentRunOutcome unavailable = + DurableAgentStateOutcomeResolver.Resolve(state, "unavailable", now); + Assert.Equal( + DurableAgentRunOutcomeKind.CompletedResultUnavailable, + unavailable.Kind); + Assert.Equal( + DurableAgentStateCompletionReceipt.UnavailableResult, + unavailable.Receipt?.ResultState); + Assert.DoesNotContain("unavailable", state.Data.TerminalResults!.Keys); + } + + [Fact] + public void AutoAccountsForMixedTextMediaAndMetadataWhenEvictingTranscript() + { + DateTimeOffset now = DateTimeOffset.UtcNow; + DurableAgentState state = CreateRevisedState(); + ChatMessage oldMessage = new( + ChatRole.User, + [ + new TextContent(new string('t', 2_000)), + new DataContent( + "data:application/octet-stream;base64," + + Convert.ToBase64String(new byte[4_000]), + mediaType: null), + ]) + { + CreatedAt = now.AddMinutes(-5), + AdditionalProperties = new() + { + ["metadata"] = new string('m', 2_000), + }, + }; + state.Data.ConversationHistory.Add( + new DurableAgentStateRequest + { + CorrelationId = "mixed", + CreatedAt = now.AddMinutes(-5), + Messages = [DurableAgentStateMessage.FromChatMessage(oldMessage)], + }); + state.Data.ConversationHistory.Add( + CreateResponse("mixed", "old response", now.AddMinutes(-5))); + AddExchange(state, "newest", "newest", now); + int initialSize = DurableAgentStateRetention.GetSerializedSize(state); + + int removed = DurableAgentStateRetention.Enforce( + state, + DurableAgentHistoryRetentionMode.Auto, + 4_000, + now, + NullLogger.Instance, + new AgentSessionId("agent", "session")); + + Assert.True(removed > 0); + Assert.DoesNotContain(state.Data.ConversationHistory, entry => entry.CorrelationId == "mixed"); + Assert.True(DurableAgentStateRetention.GetSerializedSize(state) < initialSize); + } + + [Fact] + public void PublicRetentionModesAreOnlyKeepAllAndAuto() + { + Assert.Equal( + [nameof(DurableAgentHistoryRetentionMode.KeepAll), nameof(DurableAgentHistoryRetentionMode.Auto)], + Enum.GetNames()); + } + + [Fact] + public void AutoRejectsLegacyLayoutBeforeRemovingTerminalEvidence() + { + DateTimeOffset now = DateTimeOffset.UtcNow; + DurableAgentState state = new(); + AddExchange(state, "legacy", new string('x', 2_000), now.AddMinutes(-5)); + AddExchange(state, "newest", "newest", now); + string original = JsonSerializer.Serialize( + state, + DurableAgentStateJsonContext.Default.DurableAgentState); + + DurableAgentStateCorruptionException exception = + Assert.Throws( + () => DurableAgentStateRetention.Enforce( + state, + DurableAgentHistoryRetentionMode.Auto, + 1_000, + now, + NullLogger.Instance, + new AgentSessionId("agent", "session"))); + + Assert.Contains("schema 2 mailbox", exception.Message, StringComparison.Ordinal); + Assert.Equal( + original, + JsonSerializer.Serialize( + state, + DurableAgentStateJsonContext.Default.DurableAgentState)); + } + + [Fact] + public void AutoCanEvictOlderCorrelationlessCompaction() + { + DateTimeOffset now = DateTimeOffset.UtcNow; + DurableAgentState state = CreateRevisedState(); + state.Data.ConversationHistory.Add( + new DurableAgentStateCompaction + { + CreatedAt = now.AddMinutes(-5), + Messages = + [ + DurableAgentStateMessage.FromChatMessage( + new ChatMessage(ChatRole.Assistant, new string('a', 1_000))), + ], + }); + state.Data.ConversationHistory.Add( + new DurableAgentStateCompaction + { + CreatedAt = now, + Messages = + [ + DurableAgentStateMessage.FromChatMessage( + new ChatMessage(ChatRole.Assistant, "newest")), + ], + }); + + int removed = DurableAgentStateRetention.Enforce( + state, + DurableAgentHistoryRetentionMode.Auto, + 1_200, + now, + NullLogger.Instance, + new AgentSessionId("agent", "session")); + + Assert.Equal(1, removed); + DurableAgentStateCompaction remaining = + Assert.IsType(Assert.Single(state.Data.ConversationHistory)); + Assert.Equal("newest", remaining.Messages[0].ToChatMessage().Text); + } + + [Fact] + public void AutoProtectsActualNewestCorrelationlessTranscriptComponent() + { + DateTimeOffset now = DateTimeOffset.UtcNow; + DurableAgentState state = CreateRevisedState(); + AddExchange(state, "older", new string('o', 4_000), now.AddMinutes(-5)); + state.Data.ConversationHistory.Add( + new DurableAgentStateCompaction + { + CreatedAt = now, + Messages = + [ + DurableAgentStateMessage.FromChatMessage( + new ChatMessage(ChatRole.Assistant, new string('n', 500))), + ], + }); + + _ = DurableAgentStateRetention.Enforce( + state, + DurableAgentHistoryRetentionMode.Auto, + 2_500, + now, + NullLogger.Instance, + new AgentSessionId("agent", "session")); + + DurableAgentStateCompaction newest = + Assert.Single(state.Data.ConversationHistory.OfType()); + Assert.Equal(new string('n', 500), newest.Messages[0].ToChatMessage().Text); + Assert.DoesNotContain( + state.Data.ConversationHistory, + entry => entry.CorrelationId == "older"); + } + + [Fact] + public void AutoUsesBoundedExactMeasurementsForManyIndependentExchanges() + { + DateTimeOffset now = DateTimeOffset.UtcNow; + DurableAgentState state = CreateRevisedState(); + const int ExchangeCount = 120; + for (int index = 0; index < ExchangeCount; index++) + { + AddExchange( + state, + $"exchange-{index:D3}", + new string((char)('a' + (index % 26)), 200), + now.AddMinutes(index - ExchangeCount)); + } + + AddExchange(state, "newest", "protected newest", now); + DurableAgentStateRetention.ExecutionStatistics statistics = new(); + + int removedMessages = DurableAgentStateRetention.Enforce( + state, + DurableAgentHistoryRetentionMode.Auto, + 8_000, + now, + NullLogger.Instance, + new AgentSessionId("agent", "session"), + statistics); + + Assert.True(removedMessages > 100); + Assert.Equal(1, statistics.CandidateGroupingPassCount); + Assert.InRange(statistics.SerializedStateMeasurementCount, 3, 20); + Assert.True( + statistics.SerializedStateMeasurementCount < + (removedMessages / 4)); + Assert.Contains( + state.Data.ConversationHistory, + entry => entry.CorrelationId == "newest"); + Assert.True( + DurableAgentStateRetention.GetSerializedSize(state) <= + 8_000 * DurableAgentStateRetention.LowWatermark); + } + + private static DurableAgentState CreateLargeState(DateTimeOffset? now = null) + { + DateTimeOffset current = now ?? DateTimeOffset.UtcNow; + DurableAgentState state = CreateRevisedState(); + AddExchange(state, "oldest", new string('a', 500), current.AddMinutes(-10)); + AddExchange(state, "middle", new string('b', 500), current.AddMinutes(-5)); + AddExchange(state, "newest", new string('c', 500), current); + return state; + } + + private static DurableAgentState CreateRevisedState( + string ownerKind = DurableAgentStateHistoryBinding.DurableStateOwner, + string providerKey = DurableAgentHistoryBinding.DurableStateProviderKey) + { + DurableAgentHistoryOwnership ownership = ownerKind switch + { + DurableAgentStateHistoryBinding.DurableStateOwner => + DurableAgentHistoryOwnership.Entity, + DurableAgentStateHistoryBinding.HistoryProviderOwner => + DurableAgentHistoryOwnership.ExternalProvider, + DurableAgentStateHistoryBinding.ModelServiceOwner => + DurableAgentHistoryOwnership.Service, + _ => throw new ArgumentOutOfRangeException(nameof(ownerKind)), + }; + return new DurableAgentState + { + SchemaVersion = DurableAgentState.RevisedSchemaVersion, + MailboxWritesAuthorized = true, + Data = new DurableAgentStateData + { + TerminalResults = new Dictionary( + StringComparer.Ordinal), + CompletionReceipts = new Dictionary( + StringComparer.Ordinal), + HistoryBinding = DurableAgentHistoryBinding.ToJson( + DurableAgentHistoryBinding.Create( + ownership, + ownership == DurableAgentHistoryOwnership.Entity + ? null + : providerKey)), + }, + }; + } + + private static void AddMailboxResult( + DurableAgentState state, + string correlationId, + string content, + DateTimeOffset completedAt) + { + DurableAgentStateOutcomeResolver.AddSuccessfulResult( + state, + correlationId, + new AgentResponse(new ChatMessage(ChatRole.Assistant, content)), + completedAt); + } + + private static int MeasureProjectedPrefix( + DurableAgentState state, + int entryCount, + DateTimeOffset now) + { + List originalHistory = + [.. state.Data.ConversationHistory]; + DurableAgentStateTruncation? originalTruncation = state.Data.Truncation; + int removedMessages = originalHistory + .Take(entryCount) + .Sum(entry => entry.Messages.Count); + try + { + state.Data.ConversationHistory.Clear(); + foreach (DurableAgentStateEntry entry in originalHistory.Skip(entryCount)) + { + state.Data.ConversationHistory.Add(entry); + } + + state.Data.Truncation = removedMessages == 0 + ? null + : new DurableAgentStateTruncation + { + EvictedMessageCount = removedMessages, + FirstEvictedAt = now, + LastEvictedAt = now, + }; + return DurableAgentStateRetention.GetSerializedSize(state); + } + finally + { + state.Data.ConversationHistory.Clear(); + foreach (DurableAgentStateEntry entry in originalHistory) + { + state.Data.ConversationHistory.Add(entry); + } + + state.Data.Truncation = originalTruncation; + } + } + + private static void AddExchange( + DurableAgentState state, + string correlationId, + string content, + DateTimeOffset createdAt) + { + state.Data.ConversationHistory.Add( + CreateRequest(correlationId, ChatRole.User, content, createdAt)); + state.Data.ConversationHistory.Add( + new DurableAgentStateResponse + { + CorrelationId = correlationId, + CreatedAt = createdAt, + Messages = + [ + DurableAgentStateMessage.FromChatMessage( + new ChatMessage(ChatRole.Assistant, content) { CreatedAt = createdAt }), + ], + }); + } + + private static DurableAgentStateRequest CreateRequest( + string correlationId, + ChatRole role, + string content, + DateTimeOffset createdAt) + { + return new DurableAgentStateRequest + { + CorrelationId = correlationId, + CreatedAt = createdAt, + Messages = + [ + DurableAgentStateMessage.FromChatMessage( + new ChatMessage(role, content) { CreatedAt = createdAt }), + ], + }; + } + + private static DurableAgentStateResponse CreateResponse( + string correlationId, + string content, + DateTimeOffset createdAt) + { + return new DurableAgentStateResponse + { + CorrelationId = correlationId, + CreatedAt = createdAt, + Messages = + [ + DurableAgentStateMessage.FromChatMessage( + new ChatMessage(ChatRole.Assistant, content) { CreatedAt = createdAt }), + ], + }; + } + + private static DurableAgentStateResponse CreateToolCallResponse( + string correlationId, + string callId, + string payload, + DateTimeOffset createdAt) + => CreateToolCallResponse(correlationId, createdAt, (callId, payload)); + + private static DurableAgentStateResponse CreateToolCallResponse( + string correlationId, + DateTimeOffset createdAt, + params (string CallId, string Payload)[] calls) + { + List contents = calls + .Select(call => (AIContent)new FunctionCallContent( + call.CallId, + "tool", + new Dictionary { ["payload"] = call.Payload })) + .ToList(); + return new DurableAgentStateResponse + { + CorrelationId = correlationId, + CreatedAt = createdAt, + Messages = + [ + DurableAgentStateMessage.FromChatMessage( + new ChatMessage(ChatRole.Assistant, contents) { CreatedAt = createdAt }), + ], + }; + } + + private static DurableAgentStateRequest CreateToolResultRequest( + string correlationId, + string callId, + object result, + DateTimeOffset createdAt) + { + return new DurableAgentStateRequest + { + CorrelationId = correlationId, + CreatedAt = createdAt, + Messages = + [ + DurableAgentStateMessage.FromChatMessage( + new ChatMessage( + ChatRole.Tool, + [new FunctionResultContent(callId, result)]) + { + CreatedAt = createdAt, + }), + ], + }; + } + + private static bool ContainsToolCallId(DurableAgentState state, string callId) + { + return state.Data.ConversationHistory + .SelectMany(entry => entry.Messages) + .SelectMany(message => message.Contents) + .Any(content => content switch + { + DurableAgentStateFunctionCallContent functionCall => functionCall.CallId == callId, + DurableAgentStateFunctionResultContent functionResult => functionResult.CallId == callId, + _ => false, + }); + } +} diff --git a/dotnet/tests/Microsoft.Agents.AI.DurableTask.UnitTests/DurableAgentTelemetryTests.cs b/dotnet/tests/Microsoft.Agents.AI.DurableTask.UnitTests/DurableAgentTelemetryTests.cs new file mode 100644 index 0000000..750225a --- /dev/null +++ b/dotnet/tests/Microsoft.Agents.AI.DurableTask.UnitTests/DurableAgentTelemetryTests.cs @@ -0,0 +1,474 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System.Collections.Concurrent; +using System.Diagnostics.Metrics; +using System.Text.Json; +using Microsoft.Agents.AI.DurableTask.State; +using Microsoft.Extensions.AI; +using Microsoft.Extensions.Logging.Abstractions; + +namespace Microsoft.Agents.AI.DurableTask.Tests.Unit; + +public sealed class DurableAgentTelemetryTests +{ + private static readonly string[] s_allowedTagNames = ["agent.name", "outcome", "reason"]; + + [Fact] + public void NormalEvictionRecordsCountsBytesSizesAndBoundedTags() + { + const string AgentName = "metric-normal"; + const string SessionKey = "do-not-export-session"; + const string Content = "do-not-export-content"; + DateTimeOffset now = DateTimeOffset.UtcNow; + DurableAgentState state = CreateLargeState(now, Content); + int initialSize = DurableAgentStateRetention.GetSerializedSize(state); + using RetentionMetricListener listener = new(AgentName); + + int removed = DurableAgentStateRetention.Enforce( + state, + DurableAgentHistoryRetentionMode.Auto, + 2_500, + now, + NullLogger.Instance, + new AgentSessionId(AgentName, SessionKey)); + + int finalSize = DurableAgentStateRetention.GetSerializedSize(state); + MetricMeasurement operation = listener.Single( + DurableAgentTelemetry.RetentionOperationsInstrumentName); + MetricMeasurement evictedEntries = listener.Single( + DurableAgentTelemetry.EvictedEntriesInstrumentName); + MetricMeasurement evicted = listener.Single( + DurableAgentTelemetry.EvictedMessagesInstrumentName); + MetricMeasurement reclaimed = listener.Single( + DurableAgentTelemetry.ReclaimedBytesInstrumentName); + MetricMeasurement before = listener.Single( + DurableAgentTelemetry.StateSizeBeforeInstrumentName); + MetricMeasurement after = listener.Single( + DurableAgentTelemetry.StateSizeAfterInstrumentName); + + Assert.Equal(DurableAgentTelemetry.TranscriptEvictedOutcome, operation.Tags["outcome"]); + Assert.Equal(DurableAgentTelemetry.TranscriptPressureReason, evicted.Tags["reason"]); + Assert.Equal("{operation}", operation.Unit); + Assert.Equal("{entry}", evictedEntries.Unit); + Assert.Equal("{message}", evicted.Unit); + Assert.Equal("By", reclaimed.Unit); + Assert.Equal("By", before.Unit); + Assert.Equal("By", after.Unit); + Assert.Equal(removed, evicted.Value); + Assert.Equal(removed, evictedEntries.Value); + Assert.Equal(initialSize - finalSize, reclaimed.Value); + Assert.Equal(initialSize, before.Value); + Assert.Equal(finalSize, after.Value); + Assert.All( + listener.Measurements, + measurement => + { + Assert.Equal(AgentName, measurement.Tags["agent.name"]); + Assert.DoesNotContain(SessionKey, measurement.Tags.Values); + Assert.DoesNotContain(Content, measurement.Tags.Values); + Assert.All( + measurement.Tags.Keys, + key => Assert.Contains(key, s_allowedTagNames)); + }); + } + + [Fact] + public void ZeroMessageEntryEvictionRecordsEntryOutcomeAndReclaimedBytes() + { + const string AgentName = "metric-empty-entry"; + DateTimeOffset now = DateTimeOffset.UtcNow; + DurableAgentState state = CreateRevisedState(); + state.Data.ConversationHistory.Add( + new DurableAgentStateCompaction + { + CreatedAt = now.AddMinutes(-5), + ExtensionData = new Dictionary + { + ["padding"] = JsonSerializer.SerializeToElement(new string('x', 2_000)), + }, + }); + state.Data.ConversationHistory.Add( + new DurableAgentStateCompaction + { + CreatedAt = now, + Messages = + [ + DurableAgentStateMessage.FromChatMessage( + new ChatMessage(ChatRole.Assistant, "newest")), + ], + }); + int initialSize = DurableAgentStateRetention.GetSerializedSize(state); + using RetentionMetricListener listener = new(AgentName); + + int removedMessages = DurableAgentStateRetention.Enforce( + state, + DurableAgentHistoryRetentionMode.Auto, + 2_000, + now, + NullLogger.Instance, + new AgentSessionId(AgentName, "session")); + + int finalSize = DurableAgentStateRetention.GetSerializedSize(state); + MetricMeasurement operation = listener.Single( + DurableAgentTelemetry.RetentionOperationsInstrumentName); + MetricMeasurement evictedEntries = listener.Single( + DurableAgentTelemetry.EvictedEntriesInstrumentName); + MetricMeasurement reclaimed = listener.Single( + DurableAgentTelemetry.ReclaimedBytesInstrumentName); + + Assert.Equal(0, removedMessages); + Assert.Equal(DurableAgentTelemetry.TranscriptEvictedOutcome, operation.Tags["outcome"]); + Assert.Equal(1, evictedEntries.Value); + Assert.Equal(DurableAgentTelemetry.TranscriptPressureReason, evictedEntries.Tags["reason"]); + Assert.Empty(listener.Find(DurableAgentTelemetry.EvictedMessagesInstrumentName)); + Assert.Equal(initialSize - finalSize, reclaimed.Value); + Assert.True(reclaimed.Value > 0); + } + + [Fact] + public void ZeroMessageEvictionDoesNotCreateInvalidTruncationEvidence() + { + const string AgentName = "metric-truncation-offset"; + DateTimeOffset now = DateTimeOffset.UtcNow; + DurableAgentState state = CreateRevisedState(); + state.Data.ConversationHistory.Add( + new DurableAgentStateCompaction + { + CreatedAt = now.AddMinutes(-5), + }); + state.Data.ConversationHistory.Add( + new DurableAgentStateCompaction + { + CreatedAt = now, + Messages = + [ + DurableAgentStateMessage.FromChatMessage( + new ChatMessage(ChatRole.Assistant, new string('x', 2_000))), + ], + }); + int initialSize = DurableAgentStateRetention.GetSerializedSize(state); + using RetentionMetricListener listener = new(AgentName); + + _ = Assert.Throws( + () => DurableAgentStateRetention.Enforce( + state, + DurableAgentHistoryRetentionMode.Auto, + initialSize, + now, + NullLogger.Instance, + new AgentSessionId(AgentName, "session"))); + + int finalSize = DurableAgentStateRetention.GetSerializedSize(state); + MetricMeasurement operation = listener.Single( + DurableAgentTelemetry.RetentionOperationsInstrumentName); + MetricMeasurement evictedEntries = listener.Single( + DurableAgentTelemetry.EvictedEntriesInstrumentName); + MetricMeasurement reclaimed = listener.Single( + DurableAgentTelemetry.ReclaimedBytesInstrumentName); + + Assert.True(finalSize < initialSize); + Assert.Equal( + DurableAgentTelemetry.ProtectedStateCapacityFailureOutcome, + operation.Tags["outcome"]); + Assert.Equal(1, evictedEntries.Value); + Assert.Equal(DurableAgentTelemetry.TranscriptPressureReason, evictedEntries.Tags["reason"]); + Assert.Empty(listener.Find(DurableAgentTelemetry.EvictedMessagesInstrumentName)); + Assert.Equal(initialSize - finalSize, reclaimed.Value); + Assert.Null(state.Data.Truncation); + } + + [Fact] + public void ProtectedStateFailureRecordsFailedOutcomeAndSizes() + { + const string AgentName = "metric-failure"; + DateTimeOffset now = DateTimeOffset.UtcNow; + DurableAgentState state = CreateRevisedState(); + AddExchange(state, "newest", new string('x', 2_000), now); + int initialSize = DurableAgentStateRetention.GetSerializedSize(state); + using RetentionMetricListener listener = new(AgentName); + + _ = Assert.Throws( + () => DurableAgentStateRetention.Enforce( + state, + DurableAgentHistoryRetentionMode.Auto, + 500, + now, + NullLogger.Instance, + new AgentSessionId(AgentName, "session"))); + + MetricMeasurement operation = listener.Single( + DurableAgentTelemetry.RetentionOperationsInstrumentName); + MetricMeasurement before = listener.Single( + DurableAgentTelemetry.StateSizeBeforeInstrumentName); + MetricMeasurement after = listener.Single( + DurableAgentTelemetry.StateSizeAfterInstrumentName); + + Assert.Equal( + DurableAgentTelemetry.ProtectedStateCapacityFailureOutcome, + operation.Tags["outcome"]); + Assert.Equal(initialSize, before.Value); + Assert.Equal(DurableAgentStateRetention.GetSerializedSize(state), after.Value); + Assert.Empty(listener.Find(DurableAgentTelemetry.EvictedMessagesInstrumentName)); + Assert.Empty(listener.Find(DurableAgentTelemetry.EvictedEntriesInstrumentName)); + Assert.Empty(listener.Find(DurableAgentTelemetry.ReclaimedBytesInstrumentName)); + } + + [Fact] + public void BelowHighWatermarkRecordsOnlyNoActionOperation() + { + const string AgentName = "metric-no-action"; + DateTimeOffset now = DateTimeOffset.UtcNow; + DurableAgentState state = CreateRevisedState(); + AddExchange(state, "newest", "small", now); + using RetentionMetricListener listener = new(AgentName); + + int removed = DurableAgentStateRetention.Enforce( + state, + DurableAgentHistoryRetentionMode.Auto, + 100_000, + now, + NullLogger.Instance, + new AgentSessionId(AgentName, "session")); + + Assert.Equal(0, removed); + MetricMeasurement operation = listener.Single( + DurableAgentTelemetry.RetentionOperationsInstrumentName); + Assert.Equal(DurableAgentTelemetry.NoActionOutcome, operation.Tags["outcome"]); + Assert.Single(listener.Measurements); + } + + [Fact] + public void KeepAllDoesNotEmitRetentionMetrics() + { + const string AgentName = "metric-keep-all"; + DateTimeOffset now = DateTimeOffset.UtcNow; + DurableAgentState state = CreateLargeState(now, "large"); + using RetentionMetricListener listener = new(AgentName); + + int removed = DurableAgentStateRetention.Enforce( + state, + DurableAgentHistoryRetentionMode.KeepAll, + 500, + now, + NullLogger.Instance, + new AgentSessionId(AgentName, "session")); + + Assert.Equal(0, removed); + Assert.Empty(listener.Measurements); + } + + [Fact] + public void ConcurrentRetentionCallsRecordIndependently() + { + const string AgentName = "metric-concurrent"; + const int AttemptCount = 32; + DateTimeOffset now = DateTimeOffset.UtcNow; + using RetentionMetricListener listener = new(AgentName); + + Parallel.For( + 0, + AttemptCount, + _ => + { + DurableAgentState state = CreateLargeState(now, "payload"); + int removed = DurableAgentStateRetention.Enforce( + state, + DurableAgentHistoryRetentionMode.Auto, + 2_500, + now, + NullLogger.Instance, + new AgentSessionId(AgentName, "session")); + Assert.True(removed > 0); + }); + + Assert.Equal( + AttemptCount, + listener.Find(DurableAgentTelemetry.RetentionOperationsInstrumentName).Count); + Assert.Equal( + AttemptCount, + listener.Find(DurableAgentTelemetry.StateSizeBeforeInstrumentName).Count); + Assert.Equal( + AttemptCount, + listener.Find(DurableAgentTelemetry.StateSizeAfterInstrumentName).Count); + } + + [Fact] + public void ListenerAbsenceDoesNotChangeRetention() + { + DateTimeOffset now = DateTimeOffset.UtcNow; + DurableAgentState state = CreateLargeState(now, "payload"); + + int removed = DurableAgentStateRetention.Enforce( + state, + DurableAgentHistoryRetentionMode.Auto, + 2_500, + now, + NullLogger.Instance, + new AgentSessionId("metric-no-listener", "session")); + + Assert.True(removed > 0); + Assert.DoesNotContain( + state.Data.ConversationHistory, + entry => entry.CorrelationId == "oldest"); + Assert.Contains( + state.Data.ConversationHistory, + entry => entry.CorrelationId == "newest"); + } + + [Fact] + public void ThrowingListenerCannotAffectRetention() + { + DateTimeOffset now = DateTimeOffset.UtcNow; + DurableAgentState state = CreateLargeState(now, "payload"); + using MeterListener listener = new(); + listener.InstrumentPublished = static (instrument, meterListener) => + { + if (instrument.Meter.Name == DurableAgentTelemetry.MeterName) + { + meterListener.EnableMeasurementEvents(instrument); + } + }; + listener.SetMeasurementEventCallback( + static (_, _, _, _) => throw new InvalidOperationException("listener failure")); + listener.Start(); + + int removed = DurableAgentStateRetention.Enforce( + state, + DurableAgentHistoryRetentionMode.Auto, + 2_500, + now, + NullLogger.Instance, + new AgentSessionId("metric-throwing-listener", "session")); + + Assert.True(removed > 0); + Assert.Contains( + state.Data.ConversationHistory, + entry => entry.CorrelationId == "newest"); + } + + private static DurableAgentState CreateLargeState( + DateTimeOffset now, + string content) + { + DurableAgentState state = CreateRevisedState(); + AddExchange(state, "oldest", new string('a', 500) + content, now.AddMinutes(-10)); + AddExchange(state, "middle", new string('b', 500), now.AddMinutes(-5)); + AddExchange(state, "newest", new string('c', 500), now); + return state; + } + + private static DurableAgentState CreateRevisedState() + { + return new DurableAgentState + { + SchemaVersion = DurableAgentState.RevisedSchemaVersion, + MailboxWritesAuthorized = true, + Data = new DurableAgentStateData + { + TerminalResults = new Dictionary( + StringComparer.Ordinal), + CompletionReceipts = new Dictionary( + StringComparer.Ordinal), + HistoryBinding = DurableAgentHistoryBinding.ToJson( + DurableAgentHistoryBinding.Create( + DurableAgentHistoryOwnership.Entity, + configuredProviderKey: null)), + }, + }; + } + + private static void AddExchange( + DurableAgentState state, + string correlationId, + string content, + DateTimeOffset createdAt) + { + state.Data.ConversationHistory.Add( + new DurableAgentStateRequest + { + CorrelationId = correlationId, + CreatedAt = createdAt, + Messages = + [ + DurableAgentStateMessage.FromChatMessage( + new ChatMessage(ChatRole.User, content) { CreatedAt = createdAt }), + ], + }); + state.Data.ConversationHistory.Add( + new DurableAgentStateResponse + { + CorrelationId = correlationId, + CreatedAt = createdAt, + Messages = + [ + DurableAgentStateMessage.FromChatMessage( + new ChatMessage(ChatRole.Assistant, content) { CreatedAt = createdAt }), + ], + }); + } + + private sealed record MetricMeasurement( + string InstrumentName, + string? Unit, + long Value, + IReadOnlyDictionary Tags); + + private sealed class RetentionMetricListener : IDisposable + { + private readonly string _agentName; + private readonly ConcurrentQueue _measurements = new(); + private readonly MeterListener _listener = new(); + + public RetentionMetricListener(string agentName) + { + this._agentName = agentName; + this._listener.InstrumentPublished = static (instrument, listener) => + { + if (instrument.Meter.Name == DurableAgentTelemetry.MeterName) + { + listener.EnableMeasurementEvents(instrument); + } + }; + this._listener.SetMeasurementEventCallback(this.Record); + this._listener.Start(); + } + + public IReadOnlyList Measurements => [.. this._measurements]; + + public List Find(string instrumentName) => + this.Measurements + .Where(measurement => measurement.InstrumentName == instrumentName) + .ToList(); + + public MetricMeasurement Single(string instrumentName) => + Assert.Single(this.Find(instrumentName)); + + public void Dispose() => this._listener.Dispose(); + + private void Record( + Instrument instrument, + long measurement, + ReadOnlySpan> tags, + object? state) + { + Dictionary copiedTags = new(StringComparer.Ordinal); + foreach (KeyValuePair tag in tags) + { + copiedTags[tag.Key] = tag.Value; + } + + if (copiedTags.TryGetValue( + DurableAgentTelemetry.AgentNameTagName, + out object? agentName) && + string.Equals(agentName as string, this._agentName, StringComparison.Ordinal)) + { + this._measurements.Enqueue( + new MetricMeasurement( + instrument.Name, + instrument.Unit, + measurement, + copiedTags)); + } + } + } +}