feat(AgentIsolation): raw pty

Sources/AgentIsolation/AgentSession.swift

@@ -1,3 +1,5 @@
+import Synchronization
+
 #if canImport(FoundationEssentials)
   import FoundationEssentials
 #else
@@ -9,6 +11,18 @@ private struct AgentConfigurationSettings: Decodable {
   var additionalMounts: [String]?
 }
 
+/// Errors surfaced by ``AgentSession``.
+public enum AgentSessionError: Error, Sendable {
+  /// ``AgentSession/write(_:)`` or ``AgentSession/resize(cols:rows:)`` was called
+  /// on a session whose ``IsolationConfig/customPTY`` is `false`.
+  case customPTYNotEnabled
+  /// ``AgentSession/wait()``, ``AgentSession/resize(cols:rows:)``, or
+  /// ``AgentSession/write(_:)`` was called before ``AgentSession/start(entrypoint:timeout:)``.
+  case notStarted
+  /// ``AgentSession/start(entrypoint:timeout:)`` was called more than once.
+  case alreadyStarted
+}
+
 /// Orchestrates running an isolated agent container session using a ``ContainerRuntime``.
 ///
 /// `AgentSession` is responsible for:
@@ -17,21 +31,84 @@ private struct AgentConfigurationSettings: Decodable {
 /// - Building container mounts (profile home, workspace, exclude overlays, configurations, additional mounts)
 /// - Configuring and running the container
 /// - Performing necessary cleanups (temp dirs)
-public struct AgentSession<Runtime: ContainerRuntime>: Sendable {
+///
+/// The session is object-oriented: construct once with ``init(config:runtime:)``,
+/// launch with ``start(entrypoint:timeout:)``, then drive I/O via ``rawOut``,
+/// ``write(_:)``, ``resize(cols:rows:)``, and ``wait()``.
+///
+/// When ``IsolationConfig/customPTY`` is `false` (the default), the container
+/// attaches to the current terminal (or standard streams) just like before;
+/// ``rawOut`` finishes immediately on ``start(entrypoint:timeout:)`` and
+/// ``write(_:)``/``resize(cols:rows:)`` throw
+/// ``AgentSessionError/customPTYNotEnabled``.
+public final class AgentSession<Runtime: ContainerRuntime>: Sendable {
   public let config: IsolationConfig
   public let runtime: Runtime
 
+  private let stdinStream: AsyncStream<Data>
+  private let stdinContinuation: AsyncStream<Data>.Continuation
+  private let rawOutStream: AsyncStream<[UInt8]>
+  private let rawOutContinuation: AsyncStream<[UInt8]>.Continuation
+
+  private struct State: ~Copyable {
+    var container: Runtime.Container? = nil
+    var tempDirs: [URL] = []
+    var timeoutInSeconds: Int64? = nil
+    var hasStarted: Bool = false
+    var waited: Bool = false
+  }
+  private let state = Mutex(State())
+
   public init(config: IsolationConfig, runtime: Runtime) {
     self.config = config
     self.runtime = runtime
+    (self.stdinStream, self.stdinContinuation) = AsyncStream<Data>.makeStream(
+      bufferingPolicy: .unbounded)
+    (self.rawOutStream, self.rawOutContinuation) = AsyncStream<[UInt8]>.makeStream(
+      bufferingPolicy: .unbounded)
+  }
+
+  /// A sequence of raw bytes produced by the container's PTY.
+  ///
+  /// When ``IsolationConfig/customPTY`` is `false`, iteration ends as soon as
+  /// ``start(entrypoint:timeout:)`` completes. Otherwise, bytes stream in as
+  /// the container writes to its terminal and the sequence finishes when the
+  /// container's output closes.
+  public var rawOut: some AsyncSequence<[UInt8], Never> {
+    rawOutStream
   }
 
-  /// Run the agent session and return the container process exit code.
+  /// Start the agent session.
   ///
-  /// - Parameter entrypoint: Optional entrypoint override. When non-nil, the bootstrap
-  ///   executes this instead of the last configuration's entrypoint (e.g. `["/bin/bash"]`
-  ///   for an interactive shell, or `["/bin/bash", "-c", "ls -la"]` for a command).
-  public func run(entrypoint entrypointOverride: [String]? = nil) async throws -> Int32 {
+  /// Prepares the runtime, resolves mounts, creates the container, and starts
+  /// it. I/O routing depends on ``IsolationConfig/customPTY``:
+  /// - `false`: attaches to the current terminal (when ``IsolationConfig/allocateTTY``
+  ///   is `true`) or to the parent process's stdio.
+  /// - `true`: allocates a custom PTY wired up to ``rawOut`` /
+  ///   ``write(_:)`` / ``resize(cols:rows:)``.
+  ///
+  /// - Parameters:
+  ///   - entrypointOverride: Optional entrypoint override. When non-nil, the
+  ///     bootstrap executes this instead of the last configuration's entrypoint
+  ///     (e.g. `["/bin/bash"]` for an interactive shell).
+  ///   - timeout: Optional timeout (seconds) forwarded to ``wait()``.
+  public func start(
+    entrypoint entrypointOverride: [String]? = nil,
+    timeout: Int64? = nil
+  ) async throws {
+    try state.withLock { state in
+      guard !state.hasStarted else { throw AgentSessionError.alreadyStarted }
+      state.hasStarted = true
+      state.timeoutInSeconds = timeout
+    }
+
+    if !config.customPTY {
+      // In non-custom mode, nothing will ever be fed through the rawOut/stdin
+      // streams — close them up front so consumers see an immediate EOF.
+      rawOutContinuation.finish()
+      stdinContinuation.finish()
+    }
+
     try await runtime.prepare()
 
     let canonicalWorkspace = AgentIsolationPathUtils.resolveSymlinksWithPlatformConsiderations(
@@ -45,6 +122,7 @@ public struct AgentSession<Runtime: ContainerRuntime>: Sendable {
 
     // Build mounts list
     var mounts: [ContainerConfiguration.Mount] = []
+    var tempDirs: [URL] = []
 
     // Profile home → /home/agent
     mounts.append(
@@ -61,13 +139,6 @@ public struct AgentSession<Runtime: ContainerRuntime>: Sendable {
       ))
 
     // Excluded folders: each gets an empty temp dir mounted as a read-only overlay
-    var tempDirs: [URL] = []
-    defer {
-      for dir in tempDirs {
-        try? FileManager.default.removeItem(at: dir)
-      }
-    }
-
     for rawFolder in config.excludeFolders {
       let folder = rawFolder.trimmingCharacters(in: .init(charactersIn: "/"))
       guard !folder.isEmpty else { continue }
@@ -115,7 +186,8 @@ public struct AgentSession<Runtime: ContainerRuntime>: Sendable {
     // Additional host mounts (from CLI --additional-mount flags)
     for hostMount in config.additionalHostMounts {
       let canonical = AgentIsolationPathUtils.resolveSymlinksWithPlatformConsiderations(hostMount)
-      let containerPath = "/workspace/\(AgentIsolationPathUtils.pathIdentifier(for: canonical.path))"
+      let containerPath =
+        "/workspace/\(AgentIsolationPathUtils.pathIdentifier(for: canonical.path))"
       mounts.append(
         .init(
           hostPath: canonical.path,
@@ -170,8 +242,17 @@ public struct AgentSession<Runtime: ContainerRuntime>: Sendable {
       entrypoint = containerArgs
     }
 
-    let io: ContainerConfiguration.IO =
-      config.allocateTTY ? .currentTerminal : .standardIO
+    let io: ContainerConfiguration.IO
+    if config.customPTY {
+      io = .custom(
+        stdin: AgentSessionStdinReader(inner: stdinStream),
+        stdout: AgentSessionRawOutWriter(continuation: rawOutContinuation),
+        stderr: AgentSessionNullWriter(),
+        isTerminal: true
+      )
+    } else {
+      io = config.allocateTTY ? .currentTerminal : .standardIO
+    }
 
     let containerConfig = ContainerConfiguration(
       entrypoint: entrypoint,
@@ -184,29 +265,135 @@ public struct AgentSession<Runtime: ContainerRuntime>: Sendable {
       memoryLimitMiB: config.memoryLimitMiB
     )
 
-    let container = try await runtime.runContainer(
-      imageRef: config.image,
-      configuration: containerConfig
-    )
-    defer { _runBlocking { try await runtime.removeContainer(container) } }
+    do {
+      let container = try await runtime.runContainer(
+        imageRef: config.image,
+        configuration: containerConfig
+      )
+      state.withLock { state in
+        state.container = container
+        state.tempDirs = tempDirs
+      }
+    } catch {
+      // Container never came up — purge temp dirs eagerly and finish streams.
+      for dir in tempDirs {
+        try? FileManager.default.removeItem(at: dir)
+      }
+      rawOutContinuation.finish()
+      stdinContinuation.finish()
+      throw error
+    }
+  }
+
+  /// Push bytes into the container's PTY input.
+  ///
+  /// Throws ``AgentSessionError/customPTYNotEnabled`` when ``IsolationConfig/customPTY``
+  /// is `false`, or ``AgentSessionError/notStarted`` if called before
+  /// ``start(entrypoint:timeout:)``.
+  public func write(_ data: Data) throws {
+    guard config.customPTY else { throw AgentSessionError.customPTYNotEnabled }
+    let started = state.withLock { $0.hasStarted }
+    guard started else { throw AgentSessionError.notStarted }
+    stdinContinuation.yield(data)
+  }
+
+  /// Resize the container's PTY.
+  ///
+  /// Throws ``AgentSessionError/customPTYNotEnabled`` when ``IsolationConfig/customPTY``
+  /// is `false`, or ``AgentSessionError/notStarted`` if called before
+  /// ``start(entrypoint:timeout:)``.
+  public func resize(cols: Int, rows: Int) async throws {
+    guard config.customPTY else { throw AgentSessionError.customPTYNotEnabled }
+    let container = state.withLock { $0.container }
+    guard let container else { throw AgentSessionError.notStarted }
+    try await container.resize(cols: cols, rows: rows)
+  }
+
+  /// Wait for the container to exit, then clean up temporary resources and
+  /// return the exit code.
+  public func wait() async throws -> Int32 {
+    let (container, timeout, alreadyWaited) = state.withLock {
+      state -> (Runtime.Container?, Int64?, Bool) in
+      let result = (state.container, state.timeoutInSeconds, state.waited)
+      state.waited = true
+      return result
+    }
+    guard !alreadyWaited else {
+      // Idempotent: a second wait just throws `notStarted` if nothing is live.
+      throw AgentSessionError.notStarted
+    }
+    guard let container else {
+      throw AgentSessionError.notStarted
+    }
 
-    let exitCode = try await container.wait(timeoutInSeconds: nil)
+    let exitCode: Int32
+    do {
+      exitCode = try await container.wait(timeoutInSeconds: timeout)
+    } catch {
+      await cleanup(container: container)
+      throw error
+    }
     try await container.stop()
+    await cleanup(container: container)
     return exitCode
   }
 
   // MARK: - Helpers
 
+  private func cleanup(container: Runtime.Container) async {
+    // Signal consumers that no further IO will arrive.
+    rawOutContinuation.finish()
+    stdinContinuation.finish()
+
+    try? await runtime.removeContainer(container)
+
+    let dirs = state.withLock { state -> [URL] in
+      let d = state.tempDirs
+      state.tempDirs = []
+      state.container = nil
+      return d
+    }
+    for dir in dirs {
+      try? FileManager.default.removeItem(at: dir)
+    }
+  }
+
   private func makeTempDir() throws -> URL {
     let dir = URL(fileURLWithPath: "/tmp/agentc-\(UUID().uuidString.lowercased())")
     try FileManager.default.createDirectory(at: dir, withIntermediateDirectories: true)
     return dir
   }
+}
+
+// MARK: - Custom IO plumbing
+
+/// Adapts ``AgentSession``'s internal stdin stream to the runtime's
+/// ``ReaderStream`` protocol. `stream()` must only be called once.
+private struct AgentSessionStdinReader: ReaderStream {
+  let inner: AsyncStream<Data>
 
+  func stream() -> AsyncStream<Data> {
+    inner
+  }
+}
+
+/// A ``Writer`` that pushes bytes into an ``AsyncStream`` continuation so
+/// they surface via ``AgentSession/rawOut``.
+private struct AgentSessionRawOutWriter: Writer {
+  let continuation: AsyncStream<[UInt8]>.Continuation
+
+  func write(_ data: Data) throws {
+    continuation.yield(Array(data))
+  }
+
+  func close() throws {
+    continuation.finish()
+  }
 }
 
-/// Fire-and-forget helper for calling async cleanup in a defer block.
-private func _runBlocking(_ body: @escaping @Sendable () async throws -> Void) {
-  // Best-effort cleanup — runs on a detached task since defer can't be async.
-  Task { try? await body() }
+/// A ``Writer`` that discards everything. Used for the stderr slot in raw-PTY
+/// mode, where a terminal merges stderr into stdout anyway.
+private struct AgentSessionNullWriter: Writer {
+  func write(_ data: Data) throws {}
+  func close() throws {}
 }

Sources/AgentIsolation/ContainerRuntime.swift

@@ -145,4 +145,26 @@ public protocol ContainerRuntimeContainer: Identifiable, Sendable, AnyObject {
   ///
   /// Will be called even after ``wait(timeoutInSeconds:)``.
   func stop() async throws
+
+  /// Resize the container's PTY, in character cells.
+  ///
+  /// Only meaningful when the container was started with a terminal
+  /// (e.g. ``ContainerConfiguration/IO/currentTerminal`` or
+  /// ``ContainerConfiguration/IO/custom(stdin:stdout:stderr:isTerminal:)`` with
+  /// `isTerminal: true`). The default implementation throws
+  /// ``ContainerRuntimeError/resizeNotSupported``; runtimes that support
+  /// resizing should override it.
+  func resize(cols: Int, rows: Int) async throws
+}
+
+extension ContainerRuntimeContainer {
+  public func resize(cols: Int, rows: Int) async throws {
+    throw ContainerRuntimeError.resizeNotSupported
+  }
+}
+
+/// Errors reported by conforming ``ContainerRuntime`` implementations.
+public enum ContainerRuntimeError: Error, Sendable {
+  /// The runtime does not support resizing this container's PTY.
+  case resizeNotSupported
 }

Sources/AgentIsolation/IsolationConfig.swift

@@ -61,6 +61,16 @@ public struct IsolationConfig: Sendable {
   /// prints extra information (e.g. prepare.sh progress).
   public var verbose: Bool
 
+  /// When true, the session allocates a raw PTY whose bytes flow through
+  /// ``AgentSession/rawOut`` and accept input via ``AgentSession/write(_:)``,
+  /// with ``AgentSession/resize(cols:rows:)`` to adjust the terminal size.
+  ///
+  /// When false (the default), the session attaches to the current terminal
+  /// or standard streams per ``allocateTTY``; in that mode ``AgentSession/rawOut``
+  /// finishes immediately on ``AgentSession/start(entrypoint:timeout:)`` and
+  /// ``AgentSession/write(_:)`` / ``AgentSession/resize(cols:rows:)`` throw.
+  public var customPTY: Bool
+
   public init(
     image: String,
     profileHomeDir: URL,
@@ -74,7 +84,8 @@ public struct IsolationConfig: Sendable {
     cpuCount: Int = 1,
     memoryLimitMiB: Int = 1536,
     additionalHostMounts: [URL] = [],
-    verbose: Bool = false
+    verbose: Bool = false,
+    customPTY: Bool = false
   ) {
     self.image = image
     self.profileHomeDir = profileHomeDir
@@ -89,5 +100,6 @@ public struct IsolationConfig: Sendable {
     self.memoryLimitMiB = memoryLimitMiB
     self.additionalHostMounts = additionalHostMounts
     self.verbose = verbose
+    self.customPTY = customPTY
   }
 }

Sources/AgentIsolationAppleContainerRuntime/AppleContainerRuntime.swift

@@ -224,7 +224,8 @@
 
       let firstComponent = name[..<slashIndex]
       // A registry domain contains a dot, a colon (port), or is "localhost"
-      if firstComponent.contains(".") || firstComponent.contains(":") || firstComponent == "localhost"
+      if firstComponent.contains(".") || firstComponent.contains(":")
+        || firstComponent == "localhost"
       {
         return ref
       }
@@ -342,6 +343,12 @@
       terminal?.tryReset()
       try await underlyingContainer.stop()
     }
+
+    public func resize(cols: Int, rows: Int) async throws {
+      try await underlyingContainer.resize(
+        to: ContainerizationOS.Terminal.Size(
+          width: UInt16(cols), height: UInt16(rows)))
+    }
   }
 
   // MARK: - Errors