diff --git a/FrameworkTests/HermesMemoryTests.swift b/FrameworkTests/HermesMemoryTests.swift index 3d37a4da..ba021eea 100644 --- a/FrameworkTests/HermesMemoryTests.swift +++ b/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]] = [:] private var applyEntryWaiters: [CheckedContinuation] = [] private var applyReleaseWaiters: [CheckedContinuation] = [] 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 { diff --git a/TFM/Application/MCPServer/MCPUploadCoordinator.swift b/TFM/Application/MCPServer/MCPUploadCoordinator.swift index fdc222ab..a3307a80 100644 --- a/TFM/Application/MCPServer/MCPUploadCoordinator.swift +++ b/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( diff --git a/TFMTests/MCPUploadTests.swift b/TFMTests/MCPUploadTests.swift index f6ce6a12..4264ee69 100644 --- a/TFMTests/MCPUploadTests.swift +++ b/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])