Fix upload disconnect cancellation race

FrameworkTests/HermesMemoryTests.swift

@@ -119,7 +119,8 @@ struct HermesMemoryTests {
         let history = InMemoryAgentRunRepository()
         let repository = MemoryIndexedAgentRepository(base: history, memory: model)
         await model.setEnabled(true)
-        try await waitUntil { await store.historyReplacementCount == 1 && !model.isIndexing }
+        await store.waitForHistoryReplacementCount(1)
+        await waitForIndexingToFinish(model)
 
         model.suspend()
         let conversationID = UUID()
@@ -132,7 +133,8 @@ struct HermesMemoryTests {
         #expect(await store.historyReplacementCount == 1)
 
         model.resume()
-        try await waitUntil { await store.historyReplacementCount == 2 && !model.isIndexing }
+        await store.waitForHistoryReplacementCount(2)
+        await waitForIndexingToFinish(model)
         #expect(await store.lastHistory?.contains(where: { $0.id == conversationID }) == true)
     }
 
@@ -274,6 +276,7 @@ private actor MemoryLifecycleStore: MemoryPersisting {
     private let blocksApply: Bool
     private var applyEntered = false
     private var applyReleased = false
+    private var historyReplacementWaiters: [Int: [CheckedContinuation<Void, Never>]] = [:]
     private var applyEntryWaiters: [CheckedContinuation<Void, Never>] = []
     private var applyReleaseWaiters: [CheckedContinuation<Void, Never>] = []
     private(set) var reviewApplyCount = 0
@@ -284,6 +287,11 @@ private actor MemoryLifecycleStore: MemoryPersisting {
 
     init(blocksApply: Bool = false) { self.blocksApply = blocksApply }
 
+    func waitForHistoryReplacementCount(_ count: Int) async {
+        guard historyReplacements.count < count else { return }
+        await withCheckedContinuation { historyReplacementWaiters[count, default: []].append($0) }
+    }
+
     func waitUntilApplyEntered() async {
         if applyEntered { return }
         await withCheckedContinuation { applyEntryWaiters.append($0) }
@@ -315,7 +323,12 @@ private actor MemoryLifecycleStore: MemoryPersisting {
         return memory
     }
     func snapshot(for id: UUID) -> String { memory.prompt }
-    func replaceHistory(_ conversations: [StoredAgentConversation]) { historyReplacements.append(conversations) }
+    func replaceHistory(_ conversations: [StoredAgentConversation]) {
+        historyReplacements.append(conversations)
+        let readyCounts = historyReplacementWaiters.keys.filter { $0 <= historyReplacements.count }
+        let waiters = readyCounts.flatMap { historyReplacementWaiters.removeValue(forKey: $0) ?? [] }
+        for waiter in waiters { waiter.resume() }
+    }
     func searchSessions(_ request: AgentSessionSearchRequest) -> AgentJSONValue { .array([]) }
     func lastReviewedMessage(_ id: UUID) -> UUID? { nil }
     func markReviewed(_ messageID: UUID, conversationID: UUID) {}
@@ -367,7 +380,7 @@ private actor BlockingMemoryReviewModel: AgentModelCompleting {
 
 @MainActor
 private func waitUntil(
-    timeout: Duration = .seconds(2), _ condition: @escaping @MainActor () async -> Bool
+    timeout: Duration = .seconds(10), _ condition: @escaping @MainActor () async -> Bool
 ) async throws {
     let clock = ContinuousClock()
     let deadline = clock.now.advanced(by: timeout)
@@ -377,6 +390,11 @@ private func waitUntil(
     }
 }
 
+@MainActor
+private func waitForIndexingToFinish(_ model: MemoryModel) async {
+    while model.isIndexing { await Task.yield() }
+}
+
 private struct MemoryLifecycleTestTimeout: Error {}
 
 private struct FailingDestroyStore: MemoryPersisting {

TFM/Application/MCPServer/MCPUploadCoordinator.swift

@@ -177,13 +177,20 @@ final class MCPUploadCoordinator {
             runID: ticket.info.runID, clientName: ticket.info.clientName, hostID: ticket.request.hostID,
             remotePath: ticket.request.remotePath, overwrite: ticket.request.overwrite
         )
+        let cancellation = ticket.info.cancellation
+        // The transfer cancels its own token on any failure, so track HTTP disconnect provenance separately.
+        let disconnection = TransferCancellationToken()
         tickets[id] = ticket
         return .accept(
             { [weak self] read in
                 guard let self else { return Self.error(503, "upload_unavailable") }
-                return await self.receive(id: id, attempt: attempt, read: read)
+                return await self.receive(id: id, attempt: attempt, disconnection: disconnection, read: read)
             },
-            disconnected: { [weak self] in Task { @MainActor in self?.interrupt(id, attempt: attempt) } }
+            disconnected: { [weak self] in
+                disconnection.cancel()
+                cancellation.cancel()
+                Task { @MainActor in self?.interrupt(id, attempt: attempt) }
+            }
         )
     }
 
@@ -220,9 +227,10 @@ final class MCPUploadCoordinator {
         tickets.removeAll()
     }
 
-    private func receive(id: UUID, attempt: UUID, read: @escaping SFTPStreamUploadOperation.ReadChunk) async
-        -> HTTPResponse
-    {
+    private func receive(
+        id: UUID, attempt: UUID, disconnection: TransferCancellationToken,
+        read: @escaping SFTPStreamUploadOperation.ReadChunk
+    ) async -> HTTPResponse {
         guard let ticket = tickets[id], ticket.attemptID == attempt, ticket.info.state == "connecting",
             !ticket.info.cancellation.isCancelled
         else {
@@ -244,8 +252,14 @@ final class MCPUploadCoordinator {
                 return Self.json(200, payload(completed))
             } catch {
                 let cleanup = error as? SFTPTemporaryFileCleanupError
+                let cancelRequested = tickets[id]?.info.cancelRequested == true
+                let interrupted =
+                    tickets[id]?.interrupted == true
+                    || disconnection.isCancelled
+                if interrupted { tickets[id]?.interrupted = true }
                 let (status, code) = Self.failure(
-                    cleanup?.cause ?? error, cancelled: tickets[id]?.info.cancelRequested == true)
+                    cleanup?.cause ?? error,
+                    cancelled: cancelRequested || disconnection.isCancelled)
                 tickets[id]?.temporaryPath = cleanup?.path
                 if let failure = (cleanup?.cause ?? error) as? SFTPStreamUploadOperation.Failure,
                     case .digestMismatch(let actual) = failure
@@ -253,9 +267,9 @@ final class MCPUploadCoordinator {
                     tickets[id]?.actualSHA256 = actual
                 }
                 let canResume =
-                    tickets[id]?.info.cancelRequested != true
+                    !cancelRequested
                     && !Self.integrityFailure(cleanup?.cause ?? error)
-                    && (tickets[id]?.interrupted == true || Self.recoverable(cleanup?.cause ?? error))
+                    && (interrupted || Self.recoverable(cleanup?.cause ?? error))
                     && ticket.info.expiresAt > now()
                 if !canResume { await cleanupTicket(id) }
                 finish(

TFMTests/MCPUploadTests.swift

@@ -359,10 +359,7 @@ struct MCPUploadTests {
                 return Data([1])
             }
         }
-        for _ in 0..<100 {
-            if await fixture.client.started { break }
-            await Task.yield()
-        }
+        #expect(await waitUntil { await fixture.client.started })
         switch reason {
         case "disabled": fixture.uploads.setEnabled(false)
         case "session": await fixture.uploads.endSession(runID: fixture.runID)
@@ -375,6 +372,19 @@ struct MCPUploadTests {
         #expect(await fixture.client.committed == nil)
         #expect(!fixture.uploads.hasActiveUpload(runID: fixture.runID))
         #expect(fixture.history.entries(tool: "http_sftp_upload").first?.status == .error)
+        if reason == "disconnect" {
+            #expect(
+                try fixture.uploads.status(id: fixture.id(ticket), runID: fixture.runID).objectValue?["status"]
+                    == .string("paused"))
+        }
+    }
+
+    private func waitUntil(_ condition: @escaping () async -> Bool) async -> Bool {
+        for _ in 0..<200 {
+            if await condition() { return true }
+            try? await Task.sleep(for: .milliseconds(10))
+        }
+        return false
     }
 
     @Test(arguments: [true, false])