diff --git a/ios/AidenOnTheGo/Features/Remote/AidenChatFeature.swift b/ios/AidenOnTheGo/Features/Remote/AidenChatFeature.swift index 8528499b..2986e452 100644 --- a/ios/AidenOnTheGo/Features/Remote/AidenChatFeature.swift +++ b/ios/AidenOnTheGo/Features/Remote/AidenChatFeature.swift @@ -1729,6 +1729,13 @@ final class AidenChatViewModel { isAgentRosterStale = agentRoster != nil } + /// Await the current progress observer, if any, without cancelling it or + /// starting a new one. A finished observer releases its own handle inside + /// its task body, so this returns only after that release has happened. + func waitForProgressObservation() async { + await progressTask?.value + } + private func loadProgressSnapshot( context: AidenRemoteRequestContext, observationGeneration: UInt64? = nil diff --git a/ios/AidenOnTheGoTests/AidenChatTests.swift b/ios/AidenOnTheGoTests/AidenChatTests.swift index c261554d..fd9f3efd 100644 --- a/ios/AidenOnTheGoTests/AidenChatTests.swift +++ b/ios/AidenOnTheGoTests/AidenChatTests.swift @@ -300,16 +300,18 @@ final class AidenChatTests: XCTestCase { AidenChatProgressLifecycleURLProtocol.reset() } + // Await each observer's own completion instead of polling for a fixed + // wall-clock budget: a stalled host can deliver the denial long after + // the request was counted. model.startProgressObservation() - try await waitForProgressRequestCount(1) - try await waitForProgressObservationToStop(model) + await model.waitForProgressObservation() XCTAssertFalse(model.isProgressObservationRunning) + XCTAssertEqual(AidenChatProgressLifecycleURLProtocol.progressRequestCount, 1) // The first observer exited through a completed task body. A later // activation must be able to create a fresh observer for the same chat. model.startProgressObservation() - try await waitForProgressRequestCount(2) - try await waitForProgressObservationToStop(model) + await model.waitForProgressObservation() XCTAssertFalse(model.isProgressObservationRunning) XCTAssertEqual(AidenChatProgressLifecycleURLProtocol.progressRequestCount, 2) } @@ -1224,8 +1226,8 @@ final class AidenChatTests: XCTestCase { // it while the stream-consumer hold is armed to prove it cannot steal // the intended gate. model.startProgressObservation() - try await waitForProgressRequestCount(1) - try await waitForProgressObservationToStop(model) + await model.waitForProgressObservation() + XCTAssertFalse(model.isProgressObservationRunning) XCTAssertEqual(AidenChatProgressLifecycleURLProtocol.progressRequestCount, 1) await model.send() await fulfillment(of: [arrived], timeout: 2) @@ -3164,15 +3166,6 @@ final class AidenChatTests: XCTestCase { XCTFail("Timed out waiting for progress SSE request (expected).") } - @MainActor - private func waitForProgressObservationToStop(_ model: AidenChatViewModel) async throws { - for _ in 0..<100 { - if !model.isProgressObservationRunning { return } - try await Task.sleep(for: .milliseconds(10)) - } - XCTFail("Timed out waiting for the progress observer to finish.") - } - @MainActor private func waitForAgentRequestCount(_ expected: Int) async throws { for _ in 0..<200 {