From d82f4041660bad462b74f766eb6c7d7c81dcc98c Mon Sep 17 00:00:00 2001
From: Jamie <2119834+jamieQ@users.noreply.github.com>
Date: Sat, 6 Dec 2025 17:27:45 -0600
Subject: [PATCH] [fix]: eliminate race condition in SocketController (#769)

* [fix]: eliminate race condition in SocketController

* [refactor]: use sequence- rather than iterator-based iteration

* [fix]: remove now-superfluous await

---------

Co-authored-by: Max Goedjen <max.goedjen@gmail.com>
---
 .../SecretAgentKit/SocketController.swift     | 28 ++++++++++++-------
 Sources/SecretAgent/AppDelegate.swift         |  2 +-
 2 files changed, 19 insertions(+), 11 deletions(-)

diff --git a/Sources/Packages/Sources/SecretAgentKit/SocketController.swift b/Sources/Packages/Sources/SecretAgentKit/SocketController.swift
index f8aa7cb..7839037 100644
--- a/Sources/Packages/Sources/SecretAgentKit/SocketController.swift
+++ b/Sources/Packages/Sources/SecretAgentKit/SocketController.swift
@@ -37,7 +37,13 @@ public struct SocketController {
         port = SocketPort(path: path)
         fileHandle = FileHandle(fileDescriptor: port.socket, closeOnDealloc: true)
         Task { @MainActor [fileHandle, sessionsContinuation, logger] in
-            for await notification in NotificationCenter.default.notifications(named: .NSFileHandleConnectionAccepted) {
+            // Create the sequence before triggering the notification to
+            // ensure it will not be missed.
+            let connectionAcceptedNotifications = NotificationCenter.default.notifications(named: .NSFileHandleConnectionAccepted)
+
+            fileHandle.acceptConnectionInBackgroundAndNotify()
+
+            for await notification in connectionAcceptedNotifications {
                 logger.debug("Socket controller accepted connection")
                 guard let new = notification.userInfo?[NSFileHandleNotificationFileHandleItem] as? FileHandle else { continue }
                 let session = Session(fileHandle: new)
@@ -45,7 +51,6 @@ public struct SocketController {
                 fileHandle.acceptConnectionInBackgroundAndNotify()
             }
         }
-        fileHandle.acceptConnectionInBackgroundAndNotify()
         logger.debug("Socket listening at \(path)")
     }
 
@@ -77,8 +82,14 @@ extension SocketController {
             self.fileHandle = fileHandle
             provenance = SigningRequestTracer().provenance(from: fileHandle)
             (messages, messagesContinuation) = AsyncStream.makeStream()
-            Task { [messagesContinuation, logger] in
-                for await _ in NotificationCenter.default.notifications(named: .NSFileHandleDataAvailable, object: fileHandle) {
+            Task { @MainActor [messagesContinuation, logger] in
+                // Create the sequence before triggering the notification to
+                // ensure it will not be missed.
+                let dataAvailableNotifications = NotificationCenter.default.notifications(named: .NSFileHandleDataAvailable, object: fileHandle)
+
+                fileHandle.waitForDataInBackgroundAndNotify()
+
+                for await _ in dataAvailableNotifications {
                     let data = fileHandle.availableData
                     guard !data.isEmpty else {
                         logger.debug("Socket controller received empty data, ending continuation.")
@@ -90,16 +101,13 @@ extension SocketController {
                     logger.debug("Socket controller yielded data.")
                 }
             }
-            fileHandle.waitForDataInBackgroundAndNotify()
         }
         
         /// Writes new data to the socket.
         /// - Parameter data: The data to write.
-        public func write(_ data: Data) async throws {
-            try fileHandle.write(contentsOf: data)
-            await MainActor.run {
-                fileHandle.waitForDataInBackgroundAndNotify()
-            }
+        @MainActor public func write(_ data: Data) throws {
+          try fileHandle.write(contentsOf: data)
+          fileHandle.waitForDataInBackgroundAndNotify()
         }
         
         /// Closes the socket and cleans up resources.
diff --git a/Sources/SecretAgent/AppDelegate.swift b/Sources/SecretAgent/AppDelegate.swift
index 660830b..b49cb81 100644
--- a/Sources/SecretAgent/AppDelegate.swift
+++ b/Sources/SecretAgent/AppDelegate.swift
@@ -42,7 +42,7 @@ class AppDelegate: NSObject, NSApplicationDelegate {
                         for await message in session.messages {
                             let request = try await inputParser.parse(data: message)
                             let agentResponse = await agent.handle(request: request, provenance: session.provenance)
-                            try await session.write(agentResponse)
+                            try session.write(agentResponse)
                         }
                     } catch {
                         try session.close()