From e74d79249cf450af960e5f8a0430cc2768f8fee1 Mon Sep 17 00:00:00 2001
From: Sambit Biswas <sambit.biswas2@gmail.com>
Date: Wed, 30 Sep 2026 16:18:25 -0400
Subject: [PATCH] fix(ios): await progress observer completion instead of a 1s
 poll

testProgressObservationReleasesCompletedHandleAndCanRestart polled
isProgressObservationRunning for 100 x 10ms after the progress request
was counted. On a stalled hosted runner, delivering the 403 denial back
to the observer took longer than that budget, so the observer was still
running when the poll gave up, even though the handle is released
correctly.

Add waitForProgressObservation(), which awaits the current observer
task (the same pattern as waitForDraftPersistence), and use it in place
of the wall-clock poll. The test now asserts the exact request count
after each phase, so a restart that never creates a fresh observer
still fails.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 .../Features/Remote/AidenChatFeature.swift    |  7 ++++++
 ios/AidenOnTheGoTests/AidenChatTests.swift    | 23 +++++++------------
 2 files changed, 15 insertions(+), 15 deletions(-)

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 {