Move timed allowance lock handling off the main thread

Sources/AskKeyAppKit/App/AppDelegate+Approval.swift

@@ -15,8 +15,10 @@ extension AppDelegate {
             screenNotifications: DistributedNotificationCenter.default,
             workspaceNotifications: NSWorkspace.shared.notificationCenter,
             screenState: { AgentApprovalScreenSession.current() },
-            updateSession: { active, ending, suspended, sleeping in
-                Vault.shared.updateTimedAllowanceSession(active: active, ending: ending, suspended: suspended, sleeping: sleeping)
+            beginEndingEvent: { Vault.shared.approvalRequests.beginTimedAllowanceSessionEvent(sleeping: $0) },
+            updateSession: { active, ending, suspended, sleeping, sleepOffsetUpperBound in
+                Vault.shared.updateTimedAllowanceSession(active: active, ending: ending, suspended: suspended,
+                    sleeping: sleeping, sleepOffsetUpperBound: sleepOffsetUpperBound)
             }
         )
         Vault.shared.approvalRequests.configureAuthentication { purpose in

Sources/AskKeyAppKit/TimedAllowanceSessionObserver.swift

@@ -1,4 +1,5 @@
 import AppKit
+import AskKeyBroker
 
 /// Listens without changing pending approvals or the management session.
 final class TimedAllowanceSessionObserver: NSObject {
@@ -12,11 +13,14 @@ final class TimedAllowanceSessionObserver: NSObject {
         screenNotifications: NotificationCenter,
         workspaceNotifications: NotificationCenter,
         screenState: @escaping @Sendable () -> AgentApprovalScreenState,
-        updateSession: @escaping @Sendable (Bool, Bool, Bool, Bool?) -> Void
+        queue: DispatchQueue = DispatchQueue(label: "com.sudohg.askkey.timed-allowance-session"),
+        beginEndingEvent: @escaping @Sendable (Bool?) -> BrokerTimedAllowanceSessionEvent? = { _ in nil },
+        updateSession: @escaping @Sendable (Bool, Bool, Bool, Bool?, UInt64?) -> Void
     ) {
         self.screenNotifications = screenNotifications
         self.workspaceNotifications = workspaceNotifications
-        state = TimedAllowanceSessionState(screenState: screenState, updateSession: updateSession)
+        state = TimedAllowanceSessionState(screenState: screenState, queue: queue,
+            beginEndingEvent: beginEndingEvent, updateSession: updateSession)
         super.init()
         let state = self.state
         if let distributed = screenNotifications as? DistributedNotificationCenter {
@@ -48,8 +52,8 @@ final class TimedAllowanceSessionObserver: NSObject {
                 state.update(sessionActive: true, refreshScreen: true)
             }
         ]
-        // Register first, then sample under the event lock to avoid missing a
-        // lock between the initial screen-state read and observer installation.
+        // Register first, then enqueue the initial sample in the same FIFO as
+        // notifications. Sampling and policy work never hold the event lock.
         state.update(refreshScreen: true)
     }
 
@@ -70,25 +74,35 @@ final class TimedAllowanceSessionObserver: NSObject {
 
 private final class TimedAllowanceSessionState: @unchecked Sendable {
     private let lock = NSLock()
+    private let queue: DispatchQueue
     private let screenState: @Sendable () -> AgentApprovalScreenState
-    private let updateSession: @Sendable (Bool, Bool, Bool, Bool?) -> Void
+    private let beginEndingEvent: @Sendable (Bool?) -> BrokerTimedAllowanceSessionEvent?
+    private let updateSession: @Sendable (Bool, Bool, Bool, Bool?, UInt64?) -> Void
     private var screenUnlocked = false
     private var sessionActive = true
 
-    init(screenState: @escaping @Sendable () -> AgentApprovalScreenState,
-         updateSession: @escaping @Sendable (Bool, Bool, Bool, Bool?) -> Void) {
+    init(screenState: @escaping @Sendable () -> AgentApprovalScreenState, queue: DispatchQueue,
+         beginEndingEvent: @escaping @Sendable (Bool?) -> BrokerTimedAllowanceSessionEvent?,
+         updateSession: @escaping @Sendable (Bool, Bool, Bool, Bool?, UInt64?) -> Void) {
         self.screenState = screenState
+        self.queue = queue
+        self.beginEndingEvent = beginEndingEvent
         self.updateSession = updateSession
     }
 
     func update(screenUnlocked: Bool? = nil, sleeping: Bool? = nil, sessionActive: Bool? = nil,
                 refreshScreen: Bool = false, ending: Bool = false) {
         lock.lock()
         defer { lock.unlock() }
-        if let screenUnlocked { self.screenUnlocked = screenUnlocked }
-        if let sessionActive { self.sessionActive = sessionActive }
-        if refreshScreen { self.screenUnlocked = screenState() == .unlocked }
-        // Sleep is a transition owned by the Broker, never an observer level.
-        updateSession(self.screenUnlocked && self.sessionActive, ending, !self.sessionActive, sleeping)
+        let event = ending ? beginEndingEvent(sleeping) : nil
+        queue.async { [self] in
+            defer { event?.finish() }
+            if let screenUnlocked { self.screenUnlocked = screenUnlocked }
+            if let sessionActive { self.sessionActive = sessionActive }
+            if refreshScreen { self.screenUnlocked = screenState() == .unlocked }
+            // Sleep is a transition owned by the Broker, never an observer level.
+            updateSession(self.screenUnlocked && self.sessionActive, ending, !self.sessionActive, sleeping,
+                event?.sleepOffsetUpperBound)
+        }
     }
 }

Sources/AskKeyBroker/BrokerApprovalStateMachine+Allowances.swift

@@ -6,13 +6,14 @@ extension BrokerApprovalStateMachine {
     /// sleeping: true enters sleep, false reports didWake, nil leaves it unchanged.
     @discardableResult
     package func updateTimedAllowanceSession(active: Bool, endOnLock: Bool, suspended: Bool = false,
-                                           sleeping: Bool? = nil) -> Bool {
+                                           sleeping: Bool? = nil, sleepOffsetUpperBound: UInt64? = nil) -> Bool {
         let revoked: ([String], @Sendable (String) -> Void)? = mutate {
+            timedAllowanceSessionRevision = UUID()
             timedAllowanceSessionActive = active
             timedAllowanceSessionSuspended = suspended
             if let sleeping {
                 timedAllowanceSessionSleeping = sleeping
-                timedAllowanceSleepOffsetAtSleep = sleeping ? timedAllowanceSleepOffset().upper : nil
+                timedAllowanceSleepOffsetAtSleep = sleeping ? (sleepOffsetUpperBound ?? timedAllowanceSleepOffset().upper) : nil
             } else {
                 clearTimedAllowanceSleepIfProvenLocked(unlocked: active)
             }

Sources/AskKeyBroker/BrokerApprovalStateMachine+Authentication.swift

@@ -7,12 +7,18 @@ extension BrokerApprovalStateMachine {
         timedAllowanceDuration: TimeInterval,
         endTimedAllowancesOnLock: Bool = false
     ) {
-        let revoked: ([String], @Sendable (String) -> Void)? = mutate {
+        let revoked: ([String], @Sendable (String) -> Void)? = mutate(endOnLock: endTimedAllowancesOnLock) {
+            let duration = Self.validDuration(timedAllowanceDuration) ? timedAllowanceDuration : 30 * 60
+            if self.endTimedAllowancesOnLock != endTimedAllowancesOnLock {
+                timedAllowanceSessionRevision = UUID()
+            }
             self.readAuthenticationEnabled = readAuthenticationEnabled
             self.timedAllowanceEnabled = timedAllowanceEnabled
             self.endTimedAllowancesOnLock = endTimedAllowancesOnLock
-            defaultTimedAllowance = Self.validDuration(timedAllowanceDuration) ? timedAllowanceDuration : 30 * 60
-            if endTimedAllowancesOnLock && !timedAllowanceSessionPermitsGrantLocked() {
+            defaultTimedAllowance = duration
+            if endTimedAllowancesOnLock,
+               !timedAllowanceScreenSampleFailed || timedAllowanceSessionFence.snapshot.pending,
+               !timedAllowanceSessionPermitsGrantLocked() {
                 timedAllowanceSessionGeneration = UUID()
                 return revokeAllTimedAllowancesLocked()
             }

Sources/AskKeyBroker/BrokerApprovalStateMachine+Consumption.swift

@@ -124,6 +124,8 @@ extension BrokerApprovalStateMachine {
             if entry.state == .consumed { throw BrokerApprovalError.alreadyConsumed }
             if entry.committing { throw BrokerApprovalError.commitInProgress }
             guard entry.state == .approved else { throw BrokerApprovalError.invalidDecision }
+            guard !entry.approvedByTimedAllowance || !endTimedAllowancesOnLock
+                    || timedAllowanceSessionPermitsGrantLocked() else { throw BrokerApprovalError.invalidDecision }
             entries.append(entry)
         }
         for var entry in entries {
@@ -169,6 +171,8 @@ extension BrokerApprovalStateMachine {
             if entry.state == .consumed { throw BrokerApprovalError.alreadyConsumed }
             if entry.committing { throw BrokerApprovalError.commitInProgress }
             guard entry.state == .approved else { throw BrokerApprovalError.invalidDecision }
+            guard !entry.approvedByTimedAllowance || !endTimedAllowancesOnLock
+                    || timedAllowanceSessionPermitsGrantLocked() else { throw BrokerApprovalError.invalidDecision }
             let result = try perform()
             entry.state = .consumed
             entriesByOperationID[operationID] = entry
@@ -192,6 +196,8 @@ extension BrokerApprovalStateMachine {
             }
             if entry.state == .consumed { throw BrokerApprovalError.alreadyConsumed }
             guard entry.state == .approved else { throw BrokerApprovalError.invalidDecision }
+            guard !entry.approvedByTimedAllowance || !endTimedAllowancesOnLock
+                    || timedAllowanceSessionPermitsGrantLocked() else { throw BrokerApprovalError.invalidDecision }
             guard !entry.committing else { throw BrokerApprovalError.commitInProgress }
             entry.committing = true
         }

Sources/AskKeyBroker/BrokerApprovalStateMachine+Observers.swift

@@ -32,8 +32,9 @@ extension BrokerApprovalStateMachine {
         observerQueue.sync {}
     }
 
-    func mutate<T>(_ body: () throws -> T) rethrows -> T {
-        lock.lock()
+    func mutate<T>(screenProvider: (@Sendable () -> Bool)? = nil, endOnLock: Bool? = nil,
+                   _ body: () throws -> T) rethrows -> T {
+        lockWithTimedAllowanceScreenSample(screenProvider: screenProvider, endOnLock: endOnLock)
         let schedulingNow = clock()
         let previousEntries = entriesByOperationID.mapValues { ($0.state, $0.request) }
         let screenRevocation = revokeForScreenBackstopLocked()
@@ -44,6 +45,7 @@ extension BrokerApprovalStateMachine {
             let terminalTransitions = terminalTransitionsLocked(from: previousEntries)
             let cleanups = runtimeInvalidationCleanupsLocked(now: schedulingNow)
             rescheduleExpirationLocked(now: schedulingNow)
+            clearTimedAllowanceScreenSampleLocked()
             lock.unlock()
             cleanups.forEach { $0() }
             if let screenRevocation { screenRevocation.0.forEach(screenRevocation.1) }
@@ -56,6 +58,7 @@ extension BrokerApprovalStateMachine {
             let terminalTransitions = terminalTransitionsLocked(from: previousEntries)
             let cleanups = runtimeInvalidationCleanupsLocked(now: schedulingNow)
             rescheduleExpirationLocked(now: schedulingNow)
+            clearTimedAllowanceScreenSampleLocked()
             lock.unlock()
             cleanups.forEach { $0() }
             if let screenRevocation { screenRevocation.0.forEach(screenRevocation.1) }

Sources/AskKeyBroker/BrokerApprovalStateMachine+Requests.swift

@@ -92,6 +92,7 @@ extension BrokerApprovalStateMachine {
                 throw BrokerApprovalError.capacityReached
             }
             let allowanceEnd = request.operation == .read && timedAllowanceEnabled
+                && (!endTimedAllowancesOnLock || timedAllowanceSessionPermitsGrantLocked())
                 ? allowanceLocked(credentialID: request.credentialID, digest: request.readProjectionDigest)?.expiresAt
                 : nil
             let state: BrokerRequestState = allowanceEnd != nil ? .approved : .pending

Sources/AskKeyBroker/BrokerApprovalStateMachine+RuntimeAuthorization.swift

@@ -3,13 +3,17 @@ import Foundation
 extension BrokerApprovalStateMachine {
     func withRuntimeAuthorization<T>(id: UUID, _ body: () throws -> T) throws -> T {
         lock.lock()
-        if endTimedAllowancesOnLock {
+        if endTimedAllowancesOnLock, let authorization = runtimeAuthorizations[id],
+           !authorization.timedCredentialIDs.isEmpty {
             lock.unlock()
-            return try mutate {
-                try runtimeAuthorizationLocked(id: id, body)
-            }
+            lockWithTimedAllowanceScreenSample(screenProvider: nil, endOnLock: nil)
         }
-        defer { lock.unlock() }
+        defer {
+            clearTimedAllowanceScreenSampleLocked()
+            lock.unlock()
+        }
+        // The caller may hold the Vault spawn gate and delivery lease lock.
+        // Refusal here must not publish observers, revoke or run any cleanup.
         return try runtimeAuthorizationLocked(id: id, body)
     }
 
@@ -18,6 +22,10 @@ extension BrokerApprovalStateMachine {
               !paused, clock() < authorization.expiresAt else {
             throw BrokerProviderError.requestRejected
         }
+        if endTimedAllowancesOnLock, !authorization.timedCredentialIDs.isEmpty,
+           !timedAllowanceSessionPermitsGrantLocked() {
+            throw BrokerProviderError.requestRejected
+        }
         // No observer, file IO, or acquisition of the Vault gate while locked.
         return try body()
     }

Sources/AskKeyBroker/BrokerApprovalStateMachine+ScreenSession.swift

@@ -2,16 +2,68 @@ import Foundation
 import Darwin
 
 extension BrokerApprovalStateMachine {
-    /// App-supplied live session check. The provider must be pure, nonblocking,
-    /// and must not call back into the Broker or Vault while the lock is held.
+    /// App-supplied live session check, called outside the approval lock.
+    /// Providers must not recursively start another screen-sampling mutation.
     package func configureTimedAllowanceScreenState(_ unlocked: @escaping @Sendable () -> Bool) {
-        mutate { timedAllowanceScreenUnlocked = unlocked }
+        mutate(screenProvider: unlocked) {
+            timedAllowanceScreenUnlocked = unlocked
+            timedAllowanceSessionRevision = UUID()
+        }
     }
 
     package func configureTimedAllowanceSleepOffset(
         _ offset: @escaping @Sendable () -> (lower: UInt64, upper: UInt64)
     ) {
-        mutate { timedAllowanceSleepOffset = offset }
+        mutate {
+            timedAllowanceSessionRevision = UUID()
+            timedAllowanceSleepOffset = offset
+            timedAllowanceSessionFence.configureSleepOffset(offset)
+        }
+    }
+
+    /// The observer calls this at receipt of each ending event, never for resume.
+    /// Only that event's completion can release its fence.
+    package func beginTimedAllowanceSessionEvent(sleeping: Bool?) -> BrokerTimedAllowanceSessionEvent {
+        timedAllowanceSessionFence.begin(sleeping: sleeping)
+    }
+
+    /// Returns with the approval lock held. An unlocked sample is usable only
+    /// for this mutation and only if no state or event changed during sampling.
+    /// One retry bounds contention; a second race refuses timed authority.
+    func lockWithTimedAllowanceScreenSample(
+        screenProvider: (@Sendable () -> Bool)?, endOnLock: Bool?
+    ) {
+        for attempt in 0...1 {
+            lock.lock()
+            let enabled = endOnLock ?? endTimedAllowancesOnLock
+            let provider = enabled ? (screenProvider ?? timedAllowanceScreenUnlocked) : nil
+            guard let provider else {
+                timedAllowanceScreenSample = nil
+                timedAllowanceScreenSampleFailed = false
+                return
+            }
+            let revision = timedAllowanceSessionRevision
+            let fenceRevision = timedAllowanceSessionFence.snapshot.revision
+            lock.unlock()
+            let sample = provider()
+            lock.lock()
+            if timedAllowanceSessionRevision == revision, timedAllowanceSessionFence.snapshot.revision == fenceRevision {
+                timedAllowanceScreenSample = sample
+                timedAllowanceScreenSampleFailed = false
+                return
+            }
+            if attempt == 1 {
+                timedAllowanceScreenSample = nil
+                timedAllowanceScreenSampleFailed = true
+                return
+            }
+            lock.unlock()
+        }
+    }
+
+    func clearTimedAllowanceScreenSampleLocked() {
+        timedAllowanceScreenSample = nil
+        timedAllowanceScreenSampleFailed = false
     }
 
     static func currentTimedAllowanceSleepOffset() -> (lower: UInt64, upper: UInt64) {
@@ -26,7 +78,10 @@ extension BrokerApprovalStateMachine {
     }
 
     func timedAllowanceSessionPermitsGrantLocked() -> Bool {
-        if let unlocked = timedAllowanceScreenUnlocked?() {
+        guard !timedAllowanceSessionFence.snapshot.pending, !timedAllowanceScreenSampleFailed else { return false }
+        if let unlocked = timedAllowanceScreenSample {
+            // Proven sleep may be cleared here even on the runtime gate path;
+            // this changes no authority records and publishes no callbacks.
             clearTimedAllowanceSleepIfProvenLocked(unlocked: unlocked)
             return unlocked && !timedAllowanceSessionSuspended && !timedAllowanceSessionSleeping
         }
@@ -47,7 +102,10 @@ extension BrokerApprovalStateMachine {
     /// and clear a missed wake after proven sleep, but cannot undo an inactive
     /// console session.
     func revokeForScreenBackstopLocked() -> ([String], @Sendable (String) -> Void)? {
-        guard endTimedAllowancesOnLock, timedAllowanceScreenUnlocked != nil,
+        guard endTimedAllowancesOnLock,
+              !timedAllowanceScreenSampleFailed || timedAllowanceSessionFence.snapshot.pending,
+              timedAllowanceScreenSample != nil || timedAllowanceScreenUnlocked != nil
+                || timedAllowanceSessionFence.snapshot.pending,
               !timedAllowanceSessionPermitsGrantLocked() else { return nil }
         timedAllowanceSessionGeneration = UUID()
         return revokeAllTimedAllowancesLocked()

Sources/AskKeyBroker/BrokerApprovalStateMachine.swift

@@ -54,6 +54,10 @@ public final class BrokerApprovalStateMachine: @unchecked Sendable {
         BrokerApprovalStateMachine.currentTimedAllowanceSleepOffset()
     }
     var timedAllowanceScreenUnlocked: (@Sendable () -> Bool)?
+    var timedAllowanceScreenSample: Bool?
+    var timedAllowanceScreenSampleFailed = false
+    let timedAllowanceSessionFence = TimedAllowanceSessionFence()
+    var timedAllowanceSessionRevision = UUID()
     var timedAllowanceSessionGeneration = UUID()
     var entriesByOperationID: [String: Entry] = [:]
     var operationOrder: [String] = []