From e898ce6297e057ff33c3e87fdd6c853fc17ae376 Mon Sep 17 00:00:00 2001
From: Yibo Zhuang <yzhuang@openai.com>
Date: Thu, 13 Aug 2026 19:01:19 -0700
Subject: [PATCH] Fail VM startup when its control socket cannot bind

---
 Sources/tart/Commands/Run.swift          |  4 +-
 Sources/tart/ControlSocket.swift         | 32 +++++++++-----
 Tests/TartTests/ControlSocketTests.swift | 55 ++++++++++++++++++++++++
 3 files changed, 79 insertions(+), 12 deletions(-)
 create mode 100644 Tests/TartTests/ControlSocketTests.swift

diff --git a/Sources/tart/Commands/Run.swift b/Sources/tart/Commands/Run.swift
index e39a716..f227680 100644
--- a/Sources/tart/Commands/Run.swift
+++ b/Sources/tart/Commands/Run.swift
@@ -570,8 +570,10 @@ struct Run: AsyncParsableCommand {
         }
 
         if #available(macOS 14, *) {
+          let controlSocket = try await ControlSocket(vmDir.controlSocketURL)
+
           ErrorReportingTask("Failed to run control socket") {
-            try await ControlSocket(vmDir.controlSocketURL).run()
+            try await controlSocket.run()
           }
         }
 
diff --git a/Sources/tart/ControlSocket.swift b/Sources/tart/ControlSocket.swift
index fbf1ce4..a8fbc59 100644
--- a/Sources/tart/ControlSocket.swift
+++ b/Sources/tart/ControlSocket.swift
@@ -6,17 +6,20 @@ import NIOPosix
 
 @available(macOS 14, *)
 class ControlSocket {
+  typealias ServerChannel = NIOAsyncChannel<NIOAsyncChannel<ByteBuffer, ByteBuffer>, Never>
+
   let controlSocketURL: URL
   let vmPort: UInt32
-  let eventLoopGroup = MultiThreadedEventLoopGroup(numberOfThreads: 1)
+  let eventLoopGroup: MultiThreadedEventLoopGroup
+  let serverChannel: ServerChannel
   let logger: os.Logger = os.Logger(subsystem: "org.cirruslabs.tart.control-socket", category: "network")
 
-  init(_ controlSocketURL: URL, vmPort: UInt32 = 8080) {
+  init(_ controlSocketURL: URL, vmPort: UInt32 = 8080) async throws {
     self.controlSocketURL = controlSocketURL
     self.vmPort = vmPort
-  }
+    let eventLoopGroup = MultiThreadedEventLoopGroup(numberOfThreads: 1)
+    self.eventLoopGroup = eventLoopGroup
 
-  func run() async throws {
     // Remove control socket file from previous "tart run" invocations,
     // if any, otherwise we may get the "address already in use" error
     try? FileManager.default.removeItem(atPath: controlSocketURL.path())
@@ -29,15 +32,22 @@ class ControlSocket {
       FileManager.default.changeCurrentDirectoryPath(baseURL.path())
     }
 
-    let serverChannel = try await ServerBootstrap(group: eventLoopGroup)
-      .bind(unixDomainSocketPath: controlSocketURL.relativePath) { childChannel in
-        childChannel.eventLoop.makeCompletedFuture {
-          return try NIOAsyncChannel<ByteBuffer, ByteBuffer>(
-            wrappingChannelSynchronously: childChannel
-          )
+    do {
+      self.serverChannel = try await ServerBootstrap(group: eventLoopGroup)
+        .bind(unixDomainSocketPath: controlSocketURL.relativePath) { childChannel in
+          childChannel.eventLoop.makeCompletedFuture {
+            return try NIOAsyncChannel<ByteBuffer, ByteBuffer>(
+              wrappingChannelSynchronously: childChannel
+            )
+          }
         }
-      }
+    } catch {
+      try? await eventLoopGroup.shutdownGracefully()
+      throw error
+    }
+  }
 
+  func run() async throws {
     try await withThrowingDiscardingTaskGroup { group in
       try await serverChannel.executeThenClose { serverInbound in
         for try await clientChannel in serverInbound {
diff --git a/Tests/TartTests/ControlSocketTests.swift b/Tests/TartTests/ControlSocketTests.swift
new file mode 100644
index 0000000..3e2a5fd
--- /dev/null
+++ b/Tests/TartTests/ControlSocketTests.swift
@@ -0,0 +1,55 @@
+import XCTest
+@testable import tart
+
+@available(macOS 14, *)
+final class ControlSocketTests: XCTestCase {
+  func testInitializerCreatesControlSocketBeforeReturning() async throws {
+    let temporaryDirectory = try makeTemporaryDirectory()
+    let originalDirectory = FileManager.default.currentDirectoryPath
+    defer {
+      FileManager.default.changeCurrentDirectoryPath(originalDirectory)
+      try? FileManager.default.removeItem(at: temporaryDirectory)
+    }
+
+    let socketURL = URL(fileURLWithPath: "control.sock", relativeTo: temporaryDirectory)
+    var controlSocket: ControlSocket? = try await ControlSocket(socketURL)
+    let eventLoopGroup = try XCTUnwrap(controlSocket?.eventLoopGroup)
+
+    do {
+      let serverChannel = try XCTUnwrap(controlSocket?.serverChannel)
+      XCTAssertTrue(FileManager.default.fileExists(atPath: socketURL.path))
+
+      try await serverChannel.executeThenClose { _ in }
+    }
+
+    controlSocket = nil
+    try await eventLoopGroup.shutdownGracefully()
+  }
+
+  func testInitializerPropagatesControlSocketCreationFailure() async throws {
+    let temporaryDirectory = try makeTemporaryDirectory()
+    let originalDirectory = FileManager.default.currentDirectoryPath
+    defer {
+      FileManager.default.changeCurrentDirectoryPath(originalDirectory)
+      try? FileManager.default.removeItem(at: temporaryDirectory)
+    }
+
+    let socketURL = URL(fileURLWithPath: "missing/control.sock", relativeTo: temporaryDirectory)
+
+    do {
+      _ = try await ControlSocket(socketURL)
+      XCTFail("Binding should fail when the socket's parent directory does not exist")
+    } catch {
+      XCTAssertFalse(FileManager.default.fileExists(atPath: socketURL.path))
+    }
+  }
+
+  private func makeTemporaryDirectory() throws -> URL {
+    let directory = FileManager.default.temporaryDirectory.appendingPathComponent(
+      UUID().uuidString,
+      isDirectory: true
+    )
+    try FileManager.default.createDirectory(at: directory, withIntermediateDirectories: false)
+    return directory
+  }
+}