Shrink ORT arena per window and abort separation on memory warnings

App/Continuity/Views/RootView.swift

@@ -93,6 +93,12 @@ struct RootView: View {
             // per-playlist lastSyncedAt staleness (+ failure backoff) is the real rate limiter.
             // Main-runloop timers pause while suspended, but guard on scenePhase anyway for the
             // coalesced fire a resume delivers before the app is active again.
+            // Memory warning precedes a jetsam kill: abort any in-flight stem separation and
+            // free the ONNX session. Losing stems for one track beats losing the process.
+            .onReceive(NotificationCenter.default.publisher(
+                for: UIApplication.didReceiveMemoryWarningNotification)) { _ in
+                prepQueue.handleMemoryPressure()
+            }
             .onReceive(syncTick) { _ in
                 guard scenePhase == .active else { return }
                 prepQueue.autoSyncIfNeeded(in: modelContext)

Packages/ContinuityKit/Sources/Ingest/PreparationQueue+Stems.swift

@@ -30,6 +30,14 @@ extension PreparationQueue {
         }
     }
 
+    /// System memory warning: iOS warns before it jetsams. Abort any in-flight stem separation
+    /// (it deletes its partial output; the track plays without stems and retries on a later
+    /// `ensureStems` pass) and drop the cached ORT session — staying alive beats finishing stems.
+    public func handleMemoryPressure() {
+        OnnxStemSeparator.requestAbort()
+        OnnxStemSeparator.releaseSession()
+    }
+
     /// Brings a track's stem links in line with the disk: links cached stems that exist unlinked
     /// (e.g. a re-added video), clears links whose files were evicted so `hasStems` tells the
     /// truth and the track becomes eligible for re-separation.

Packages/ContinuityKit/Sources/Ingest/StemSeparator.swift

@@ -84,6 +84,25 @@ final class OnnxStemSeparator: StemSeparating {
         cachedSession = nil
     }
 
+    /// Set on a system memory warning; the streaming loop checks it before each window and
+    /// bails, deleting the partial stems. The track keeps playing without stems and re-enters
+    /// separation on a later `ensureStems` pass. Cleared when the next separation starts.
+    nonisolated(unsafe) private static var abortRequested = false
+    private static let abortLock = NSLock()
+
+    static var isAbortRequested: Bool {
+        abortLock.lock(); defer { abortLock.unlock() }
+        return abortRequested
+    }
+
+    static func requestAbort() {
+        abortLock.lock(); abortRequested = true; abortLock.unlock()
+    }
+
+    static func clearAbort() {
+        abortLock.lock(); abortRequested = false; abortLock.unlock()
+    }
+
     private static func sharedSession(modelURL: URL) throws -> ORTSession {
         sessionLock.lock()
         defer { sessionLock.unlock() }
@@ -109,6 +128,7 @@ final class OnnxStemSeparator: StemSeparating {
     }
 
     func separate(inputURL: URL, vocalsOut: URL, accompanimentOut: URL) throws -> StemPaths {
+        Self.clearAbort()   // a stale warning must not kill this fresh attempt
         // ONNX Runtime session, CPU EP everywhere. The CoreML EP jetsammed real devices:
         // converting/compiling the 85M-param HT-Demucs transformer at session creation ballooned
         // the app past the per-process limit (~3.4 GB observed on iPhone 17 Pro — JetsamEvent
@@ -230,13 +250,22 @@ final class OnnxStemSeparator: StemSeparating {
     /// returns just the vocals source (index 3) as planar (2, segment).
     private static func ortInference(session: ORTSession, segment: Int,
                                      channels: Int, vocalsIndex: Int) -> WindowInference {
-        { input in
+        // Shrink the CPU arena back after every Run: the transformer's attention intermediates
+        // spike the arena to its high-water mark, and by default ORT never returns that memory
+        // between windows — it sits on top of live playback for the whole track (jetsam margin).
+        // Re-growing per window costs a little CPU; an offline cache-once job can afford it.
+        let runOptions = try? ORTRunOptions()
+        try? runOptions?.addConfigEntry(withKey: "memory.enable_memory_arena_shrinkage", value: "cpu:0")
+        return { input in
             let inputData = input.withUnsafeBytes { NSMutableData(bytes: $0.baseAddress, length: $0.count) }
             do {
+                if Self.isAbortRequested {
+                    throw StemSeparationError.inference("aborted under memory pressure")
+                }
                 let inputValue = try ORTValue(tensorData: inputData, elementType: .float,
                                               shape: [1, NSNumber(value: channels), NSNumber(value: segment)])
                 let outputs = try session.run(withInputs: ["mix": inputValue],
-                                              outputNames: ["stems"], runOptions: nil)
+                                              outputNames: ["stems"], runOptions: runOptions)
                 guard let out = outputs["stems"] else { throw StemSeparationError.inference("no 'stems' output") }
                 let outData = try out.tensorData()
                 let vocBase = (vocalsIndex * channels) * segment