feat: add global library search and safe release upgrades (#12)

Plugins/SecondDimensionWatcherReDive.Chat/Tools/ManageDownloadsTool.cs

@@ -1,3 +1,4 @@
+using System.Diagnostics;
 using SecondDimensionWatcherReDive.AI.Models;
 using SecondDimensionWatcherReDive.Framework.AI;
 using SecondDimensionWatcherReDive.Framework.Attributes;
@@ -14,6 +15,12 @@ internal sealed partial class ManageDownloadsTool(
     IFileMappingRepository fileMappingRepository,
     IFileDownloadClientProvider fileDownloadClientProvider) : ITool
 {
+    private static readonly TimeSpan DownloadSubmissionLeaseDuration = TimeSpan.FromMinutes(3);
+    private static readonly TimeSpan DownloadSubmissionRemoteBudget = TimeSpan.FromSeconds(90);
+    private static readonly TimeSpan DownloadCancellationLeaseDuration = TimeSpan.FromMinutes(3);
+    private static readonly TimeSpan DownloadCancellationRemoteBudget = TimeSpan.FromSeconds(90);
+    private static readonly TimeSpan DownloadLeaseSafetyMargin = TimeSpan.FromSeconds(1);
+
     private async Task<IToolResult> ExecuteCoreAsync(
         ManageDownloadsParams param, CancellationToken cancellationToken)
     {
@@ -45,32 +52,71 @@ private async Task<IToolResult> StartDownloadAsync(
             return new ToolFailureResult("Download already tracked");
 
         var downloadAttemptId = Guid.NewGuid();
+        var submissionLeaseId = Guid.NewGuid();
+        var leaseRequestStartedAt = Stopwatch.GetTimestamp();
         var submissionAttempted = false;
         try
         {
-            if (!await animationInfoRepository.TryStartDownloadAsync(
+            var submissionLease = await animationInfoRepository.TryStartDownloadAsync(
                     info.Id,
                     downloadAttemptId,
+                    submissionLeaseId,
+                    DownloadSubmissionLeaseDuration,
                     DateTimeOffset.Now,
                     queuedDisposition: null,
-                    cancellationToken))
+                    cancellationToken);
+            if (submissionLease is null)
                 return new ToolFailureResult("Download already tracked");
 
+            var remainingRemoteBudget = DownloadSubmissionRemoteBudget -
+                                        Stopwatch.GetElapsedTime(leaseRequestStartedAt);
+            if (remainingRemoteBudget <= TimeSpan.Zero)
+            {
+                await CompensateFailedStartAsync(
+                    info,
+                    client,
+                    downloadAttemptId,
+                    submissionLeaseId,
+                    remoteMayHaveAccepted: false);
+                return new ToolFailureResult("Download submission lease expired");
+            }
+
+            using var submissionCancellation = CancellationTokenSource.CreateLinkedTokenSource(
+                cancellationToken);
+            submissionCancellation.CancelAfter(remainingRemoteBudget);
+            submissionCancellation.Token.ThrowIfCancellationRequested();
             submissionAttempted = true;
             if (!await client.SubmitDownloadTaskAsync(
                     info.Id,
                     info.DownloadUrl,
                     info.CachedDownloadData,
                     info.AdditionalDownloadInfo,
-                    cancellationToken))
+                    submissionCancellation.Token))
             {
                 await CompensateFailedStartAsync(
                     info,
                     client,
                     downloadAttemptId,
+                    submissionLeaseId,
                     remoteMayHaveAccepted: false);
                 return new ToolFailureResult("Download client rejected the task");
             }
+
+            using var markCancellation = CreateDownloadSagaTokenSource();
+            if (!await animationInfoRepository.TryMarkDownloadSubmittedAsync(
+                    info.Id,
+                    downloadAttemptId,
+                    submissionLeaseId,
+                    markCancellation.Token))
+            {
+                await CompensateFailedStartAsync(
+                    info,
+                    client,
+                    downloadAttemptId,
+                    submissionLeaseId,
+                    remoteMayHaveAccepted: true);
+                return new ToolFailureResult("Download state changed during submission");
+            }
         }
         catch
         {
@@ -80,6 +126,7 @@ await CompensateFailedStartAsync(
                     info,
                     client,
                     downloadAttemptId,
+                    submissionLeaseId,
                     submissionAttempted);
             }
             catch
@@ -112,35 +159,57 @@ private async Task<IToolResult> CancelDownloadAsync(
         AnimationInfo info, IFileDownloadClient client, bool removeFile, CancellationToken cancellationToken)
     {
         var cancellationAttemptId = info.DownloadCancellationId ?? Guid.NewGuid();
+        var cancellationLeaseId = Guid.NewGuid();
+        var leaseRequestStartedAt = Stopwatch.GetTimestamp();
+        DownloadCancellationLease? cancellationLease;
         cancellationToken.ThrowIfCancellationRequested();
         using (var beginCancellation = CreateDownloadSagaTokenSource())
         {
-            if (!await animationInfoRepository.TryBeginCancelDownloadAsync(
+            cancellationLease = await animationInfoRepository.TryBeginCancelDownloadAsync(
                     info.Id,
                     info.DownloadAttemptId,
                     cancellationAttemptId,
-                    beginCancellation.Token))
+                    cancellationLeaseId,
+                    DownloadCancellationLeaseDuration,
+                    removeFile,
+                    requireUnfinished: false,
+                    SubscriptionAutomationDisposition.DownloadCancelled,
+                    beginCancellation.Token);
+            if (cancellationLease is null)
                 return new ToolFailureResult("Download state changed before cancellation");
         }
 
+        var remainingRemoteBudget = DownloadCancellationRemoteBudget -
+                                    Stopwatch.GetElapsedTime(leaseRequestStartedAt);
+        if (remainingRemoteBudget <= TimeSpan.Zero)
+            return new ToolFailureResult("Download cancellation lease expired");
+        using var remoteCancellation = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
+        remoteCancellation.CancelAfter(remainingRemoteBudget);
+        remoteCancellation.Token.ThrowIfCancellationRequested();
         var result = await client.CancelDownloadTaskAsync(
             info.Id,
             info.DownloadUrl,
             info.CachedDownloadData,
             info.AdditionalDownloadInfo,
-            removeFile,
-            cancellationToken);
+            cancellationLease.RemoveFile,
+            remoteCancellation.Token);
 
         if (!result.IsSuccess)
         {
             return new ToolSuccessResult<bool>(false);
         }
 
-        using var finalizeCancellation = CreateDownloadSagaTokenSource();
+        using var finalizeCancellation = CreateLeaseBoundSagaTokenSource(
+            leaseRequestStartedAt,
+            DownloadCancellationLeaseDuration);
+        if (finalizeCancellation is null)
+            return new ToolFailureResult("Download cancellation lease expired");
         var cancelled = await fileMappingRepository.TryFinalizeDownloadCancellationAsync(
             info.Id,
             info.DownloadAttemptId,
             cancellationAttemptId,
+            cancellationLease.Id,
+            SubscriptionAutomationDisposition.DownloadCancelled,
             finalizeCancellation.Token);
         if (!cancelled)
             return new ToolFailureResult("Download state changed during cancellation");
@@ -151,19 +220,35 @@ private async Task CompensateFailedStartAsync(
         AnimationInfo info,
         IFileDownloadClient client,
         Guid downloadAttemptId,
+        Guid submissionLeaseId,
         bool remoteMayHaveAccepted)
     {
         using var cleanup = CreateDownloadSagaTokenSource();
+        var cancellationAttemptId = Guid.NewGuid();
+        var cancellationLease = await animationInfoRepository.TryBeginCancelDownloadAsync(
+            info.Id,
+            downloadAttemptId,
+            cancellationAttemptId,
+            submissionLeaseId,
+            DownloadCancellationLeaseDuration,
+            removeFile: false,
+            requireUnfinished: true,
+            terminalDisposition: null,
+            cleanup.Token);
+        if (cancellationLease is null)
+            return;
+
         if (remoteMayHaveAccepted)
         {
             try
             {
+                cleanup.Token.ThrowIfCancellationRequested();
                 var remoteCancellation = await client.CancelDownloadTaskAsync(
                     info.Id,
                     info.DownloadUrl,
                     info.CachedDownloadData,
                     info.AdditionalDownloadInfo,
-                    removeFile: false,
+                    cancellationLease.RemoveFile,
                     cleanup.Token);
                 if (!remoteCancellation.IsSuccess)
                 {
@@ -178,9 +263,12 @@ private async Task CompensateFailedStartAsync(
             }
         }
 
-        await animationInfoRepository.TryCancelDownloadAsync(
+        cleanup.Token.ThrowIfCancellationRequested();
+        await fileMappingRepository.TryFinalizeDownloadCancellationAsync(
             info.Id,
             downloadAttemptId,
+            cancellationAttemptId,
+            cancellationLease.Id,
             terminalDisposition: null,
             cleanup.Token);
     }
@@ -207,6 +295,19 @@ await client.SubmitQueryDownloadProgressAsync(
 
     private static CancellationTokenSource CreateDownloadSagaTokenSource() =>
         new(TimeSpan.FromSeconds(10));
+
+    private static CancellationTokenSource? CreateLeaseBoundSagaTokenSource(
+        long leaseRequestStartedAt,
+        TimeSpan leaseDuration)
+    {
+        var remaining = leaseDuration -
+                        Stopwatch.GetElapsedTime(leaseRequestStartedAt) -
+                        DownloadLeaseSafetyMargin;
+        if (remaining <= TimeSpan.Zero)
+            return null;
+        return new CancellationTokenSource(
+            remaining < TimeSpan.FromSeconds(10) ? remaining : TimeSpan.FromSeconds(10));
+    }
 }
 
 internal enum ManageDownloadsAction

Plugins/SecondDimensionWatcherReDive.Inference.AI/Tools/TmdbTool.cs

@@ -178,6 +178,33 @@ public async Task<string> GetSeasonEpisodesAsync(int tmdbId, int seasonNumber, C
         }
     }
 
+    public async Task<int?> GetExpectedEpisodeCountAsync(
+        int tmdbId,
+        int seasonNumber,
+        CancellationToken cancellationToken)
+    {
+        var tmdbClient = GetClient();
+        if (tmdbClient is null || seasonNumber <= 0) return null;
+
+        try
+        {
+            var show = await tmdbClient.GetTvShowAsync(tmdbId, cancellationToken: cancellationToken);
+            var count = show?.Seasons?
+                .FirstOrDefault(season => season.SeasonNumber == seasonNumber)
+                ?.EpisodeCount;
+            return count is > 0 ? count : null;
+        }
+        catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested)
+        {
+            throw;
+        }
+        catch (Exception ex)
+        {
+            LogGetSeasonsFailed(_logger, ex, tmdbId);
+            return null;
+        }
+    }
+
     /// <summary>
     ///     Fetches localized name, original name, and overview for a TV show from TMDB,
     ///     using the server's current culture as the language.