Fix undo event stream

Sources/SQLiteUndo/UndoCoordinator.swift

@@ -17,11 +17,9 @@ final class UndoCoordinator: Sendable {
   private let untrackedTables: Set<String>
   private let state = LockIsolated(State())
 
-  let events: AsyncStream<UndoEvent>
-  private let eventsContinuation: AsyncStream<UndoEvent>.Continuation
-
   private struct State {
     var openBarriers: [UUID: OpenBarrier] = [:]
+    var subscribers: [UUID: AsyncStream<UndoEvent>.Continuation] = [:]
 
     /// Tracks current seq range for each barrier.
     ///
@@ -63,7 +61,28 @@ final class UndoCoordinator: Sendable {
     self.database = database ?? defaultDatabase
     self.registeredTables = registeredTables
     self.untrackedTables = untrackedTables
-    (self.events, self.eventsContinuation) = AsyncStream.makeStream()
+  }
+
+  /// Create a new stream of undo/redo events.
+  ///
+  /// Each call creates an independent subscription that receives events emitted from
+  /// this point on; earlier events are not replayed.
+  func events() -> AsyncStream<UndoEvent> {
+    let id = UUID()
+    let (stream, continuation) = AsyncStream<UndoEvent>.makeStream()
+    state.withValue { $0.subscribers[id] = continuation }
+    continuation.onTermination = { [state] _ in
+      state.withValue { _ = $0.subscribers.removeValue(forKey: id) }
+    }
+    return stream
+  }
+
+  /// Broadcast an event to all active subscribers.
+  private func emit(_ event: UndoEvent) {
+    // Copy out before yielding so `onTermination` can't re-enter the lock.
+    for continuation in state.withValue({ Array($0.subscribers.values) }) {
+      continuation.yield(event)
+    }
   }
 
   /// Begin recording changes for a new undoable action.
@@ -201,9 +220,7 @@ final class UndoCoordinator: Sendable {
       state.withValue {
         $0.barrierSeqRanges[barrier.id] = result.seqRange
       }
-      eventsContinuation.yield(
-        UndoEvent(kind: .undo, name: barrier.name, affectedItems: result.affectedItems)
-      )
+      emit(UndoEvent(kind: .undo, name: barrier.name, affectedItems: result.affectedItems))
     }
   }
 
@@ -228,9 +245,7 @@ final class UndoCoordinator: Sendable {
       state.withValue {
         $0.barrierSeqRanges[barrier.id] = result.seqRange
       }
-      eventsContinuation.yield(
-        UndoEvent(kind: .redo, name: barrier.name, affectedItems: result.affectedItems)
-      )
+      emit(UndoEvent(kind: .redo, name: barrier.name, affectedItems: result.affectedItems))
     }
   }
 }

Sources/SQLiteUndo/UndoEngine.swift

@@ -68,6 +68,10 @@ public struct UndoEngine: Sendable {
   public var cancelBarrier: @Sendable (_ id: UUID) throws -> Void
 
   /// Stream of events emitted after each undo/redo operation.
+  ///
+  /// Each call returns an independent subscription delivering events from that point
+  /// on; earlier events are not replayed. Cancelling one subscription leaves the others
+  /// unaffected, so callers may freely resubscribe.
   public var events: @Sendable () -> AsyncStream<UndoEvent> = { .finished }
 }
 
@@ -229,7 +233,7 @@ extension UndoEngine: DependencyKey {
         try coordinator.cancelBarrier(id)
       },
       events: {
-        coordinator.events
+        coordinator.events()
       }
     )
   }

Tests/SQLiteUndoTCATests/UndoableEffectTests.swift

@@ -68,6 +68,36 @@ struct UndoableEffectTests {
       #expect(countAfter == 0)
     }
   }
+  /// `.task(id: undoManager)` re-fires whenever the environment's UndoManager changes
+  /// identity, so `.set` arrives more than once and resubscribes to the event stream.
+  @Test
+  func resettingUndoManagerKeepsEventsFlowing() async throws {
+    let testUndoManager = UndoManager()
+
+    await withDependencies {
+      let database = try! makeTestDatabase()
+      $0.defaultDatabase = database
+      $0.defaultUndoStack = .live(testUndoManager)
+      $0.defaultUndoEngine = try! UndoEngine(for: database, tables: TestRecord.self)
+    } operation: {
+      let store = TestStore(initialState: TestFeature.State()) {
+        TestFeature()
+      }
+
+      let first = await store.send(.undoManager(.set(testUndoManager)))
+      let second = await store.send(.undoManager(.set(testUndoManager)))
+
+      await store.send(.insertItem)
+      await store.receive(\.itemInserted)
+
+      testUndoManager.undo()
+
+      await store.receive(\.undoManager.event)
+
+      await first.cancel()
+      await second.cancel()
+    }
+  }
 }
 
 // MARK: - Test Feature
@@ -77,14 +107,16 @@ private struct TestFeature {
   @ObservableState
   struct State: Equatable {}
 
-  enum Action {
+  enum Action: UndoManageableAction {
     case insertItem
     case itemInserted
+    case undoManager(UndoManagingAction)
   }
 
   @Dependency(\.defaultDatabase) var database
 
   var body: some ReducerOf<Self> {
+    UndoManagingReducer()
     Reduce { state, action in
       switch action {
       case .insertItem:
@@ -99,6 +131,9 @@ private struct TestFeature {
 
       case .itemInserted:
         return .none
+
+      case .undoManager:
+        return .none
       }
     }
   }

Tests/SQLiteUndoTests/UndoEngineTests.swift

@@ -848,9 +848,9 @@ enum UndoEngineTests {
       }
       let barrier = try coordinator.endBarrier(barrierId)!
 
+      var iterator = coordinator.events().makeAsyncIterator()
       try coordinator.performUndo(barrier: barrier)
 
-      var iterator = coordinator.events.makeAsyncIterator()
       let event = await iterator.next()
       expectNoDifference(
         event,
@@ -871,10 +871,10 @@ enum UndoEngineTests {
       }
       let barrier = try coordinator.endBarrier(barrierId)!
 
+      var iterator = coordinator.events().makeAsyncIterator()
       try coordinator.performUndo(barrier: barrier)
       try coordinator.performRedo(barrier: barrier)
 
-      var iterator = coordinator.events.makeAsyncIterator()
       let undoEvent = await iterator.next()
       #expect(undoEvent?.kind == .undo)
       let redoEvent = await iterator.next()
@@ -899,9 +899,9 @@ enum UndoEngineTests {
       }
       let barrier = try coordinator.endBarrier(barrierId)!
 
+      var iterator = coordinator.events().makeAsyncIterator()
       try coordinator.performUndo(barrier: barrier)
 
-      var iterator = coordinator.events.makeAsyncIterator()
       let event = await iterator.next()
       expectNoDifference(
         event,
@@ -916,6 +916,72 @@ enum UndoEngineTests {
         ))
     }
 
+    @Test
+    func eventsBroadcastToAllSubscribers() async throws {
+      let (database, coordinator) = try makeTestDatabaseWithUndo()
+
+      var first = coordinator.events().makeAsyncIterator()
+      var second = coordinator.events().makeAsyncIterator()
+
+      let barrierId = try coordinator.beginBarrier("Insert Item")
+      try await database.write { db in
+        try TestRecord.insert { TestRecord(id: 1, name: "Test") }.execute(db)
+      }
+      let barrier = try coordinator.endBarrier(barrierId)!
+      try coordinator.performUndo(barrier: barrier)
+
+      let firstEvent = await first.next()
+      let secondEvent = await second.next()
+      #expect(firstEvent?.kind == .undo)
+      expectNoDifference(firstEvent, secondEvent)
+    }
+
+    @Test
+    func resubscribingAfterCancellationStillReceivesEvents() async throws {
+      let (database, coordinator) = try makeTestDatabaseWithUndo()
+
+      let abandoned = coordinator.events()
+      let task = Task { for await _ in abandoned {} }
+      task.cancel()
+      await task.value
+
+      var iterator = coordinator.events().makeAsyncIterator()
+
+      let barrierId = try coordinator.beginBarrier("Insert Item")
+      try await database.write { db in
+        try TestRecord.insert { TestRecord(id: 1, name: "Test") }.execute(db)
+      }
+      let barrier = try coordinator.endBarrier(barrierId)!
+      try coordinator.performUndo(barrier: barrier)
+
+      let event = await iterator.next()
+      #expect(event?.kind == .undo)
+    }
+
+    @Test
+    func eventsAreNotReplayedToLaterSubscribers() async throws {
+      let (database, coordinator) = try makeTestDatabaseWithUndo()
+
+      let firstId = try coordinator.beginBarrier("First")
+      try await database.write { db in
+        try TestRecord.insert { TestRecord(id: 1, name: "Test") }.execute(db)
+      }
+      let first = try coordinator.endBarrier(firstId)!
+      try coordinator.performUndo(barrier: first)
+
+      var iterator = coordinator.events().makeAsyncIterator()
+
+      let secondId = try coordinator.beginBarrier("Second")
+      try await database.write { db in
+        try TestRecord.insert { TestRecord(id: 2, name: "Test") }.execute(db)
+      }
+      let second = try coordinator.endBarrier(secondId)!
+      try coordinator.performUndo(barrier: second)
+
+      let event = await iterator.next()
+      #expect(event?.name == "Second")
+    }
+
     @Test
     func affectedItemIdAs() {
       let item = AffectedItem(table: TestRecord.self, rowid: 42)