1536 lines · 60.3 KB · Swift
RawDownload
#if canImport(Darwin)import Darwin#elseimport Glibc#endifimport Foundationimport Testing@testable import CodexBarCore private final class FakeAntigravityProcessHandle: AntigravityCLIProcessHandle, @unchecked Sendable {    private let lock = NSLock()    let pid: pid_t    var descendants: [pid_t]    private var running: Bool    private let terminateRootStopsProcess: Bool    private var assignedProcessGroup: pid_t?    private var events: [String] = []    private var drainOutputChunks: [Data] = []     init(pid: pid_t, running: Bool = true, descendants: [pid_t] = [], terminateRootStopsProcess: Bool = true) {        self.pid = pid        self.running = running        self.descendants = descendants        self.terminateRootStopsProcess = terminateRootStopsProcess    }     var isRunning: Bool {        self.lock.lock()        let value = self.running        self.events.append("isRunning:\(value)")        self.lock.unlock()        return value    }     var processGroup: pid_t? {        self.lock.lock()        let value = self.assignedProcessGroup        self.lock.unlock()        return value    }     func assignProcessGroup() -> pid_t? {        self.lock.lock()        self.assignedProcessGroup = self.pid        self.events.append("assignProcessGroup")        self.lock.unlock()        return self.pid    }     func sendExit() throws {        self.append("sendExit")    }     func closePTY() {        self.append("closePTY")    }     func terminateRoot() {        self.lock.lock()        if self.terminateRootStopsProcess {            self.running = false        }        self.events.append("terminateRoot")        self.lock.unlock()    }     func killRoot() {        self.lock.lock()        self.running = false        self.events.append("killRoot")        self.lock.unlock()    }     func descendantPIDs() -> [pid_t] {        self.lock.lock()        let value = self.descendants        self.events.append("descendantPIDs")        self.lock.unlock()        return value    }     func terminateTree(signal: Int32, knownDescendants _: [pid_t]) {        self.lock.lock()        if signal == SIGKILL {            self.running = false        }        self.events.append("terminateTree:\(signal)")        self.lock.unlock()    }     func killDescendants(_ descendants: [pid_t]) {        self.append("killDescendants:\(descendants.map(String.init).joined(separator: ","))")    }     func drainOutput() -> Data {        self.lock.lock()        self.events.append("drainOutput")        let output = self.drainOutputChunks.isEmpty ? Data() : self.drainOutputChunks.removeFirst()        self.lock.unlock()        return output    }     func enqueueDrainOutput(_ output: Data) {        self.lock.lock()        self.drainOutputChunks.append(output)        self.lock.unlock()    }     func snapshotEvents() -> [String] {        self.lock.lock()        let value = self.events        self.lock.unlock()        return value    }     private func append(_ event: String) {        self.lock.lock()        self.events.append(event)        self.lock.unlock()    }} private final class FakeAntigravityProcessLauncher: AntigravityCLIProcessLaunching, @unchecked Sendable {    private let lock = NSLock()    private var nextPID: pid_t    private var launchError: Error?    private var launchedBinaries: [String] = []    private var terminateRootStopsProcess = true    private var handles: [FakeAntigravityProcessHandle] = []     init(nextPID: pid_t = 1) {        self.nextPID = nextPID    }     func launch(binary: String) throws -> any AntigravityCLIProcessHandle {        self.lock.lock()        defer { self.lock.unlock() }        if let launchError {            throw launchError        }        let handle = FakeAntigravityProcessHandle(            pid: self.nextPID,            descendants: [self.nextPID + 100],            terminateRootStopsProcess: self.terminateRootStopsProcess)        self.nextPID += 1        self.launchedBinaries.append(binary)        self.handles.append(handle)        return handle    }     func setLaunchError(_ error: Error?) {        self.lock.lock()        self.launchError = error        self.lock.unlock()    }     func setTerminateRootStopsProcess(_ value: Bool) {        self.lock.lock()        self.terminateRootStopsProcess = value        self.lock.unlock()    }     func launchedBinarySnapshot() -> [String] {        self.lock.lock()        let value = self.launchedBinaries        self.lock.unlock()        return value    }     func handleSnapshot() -> [FakeAntigravityProcessHandle] {        self.lock.lock()        let value = self.handles        self.lock.unlock()        return value    }} private final class FakeAntigravityIdentityProvider: AntigravityCLIProcessIdentityProviding, @unchecked Sendable {    private let lock = NSLock()    private var identities: [pid_t: AntigravityCLIProcessIdentity] = [:]     func setIdentity(pid: pid_t, executablePath: String, startEpoch: TimeInterval) {        self.lock.lock()        self.identities[pid] = AntigravityCLIProcessIdentity(executablePath: executablePath, startEpoch: startEpoch)        self.lock.unlock()    }     func removeIdentity(pid: pid_t) {        self.lock.lock()        self.identities[pid] = nil        self.lock.unlock()    }     func identity(for pid: pid_t) -> AntigravityCLIProcessIdentity? {        self.lock.lock()        let value = self.identities[pid]        self.lock.unlock()        return value    }} private final class MemoryAntigravitySessionRecordStore: AntigravityCLISessionRecordStoring, @unchecked Sendable {    private let lock = NSLock()    private var records: [AntigravityCLISessionRecord]    private let failSaves: Bool    private var saves = 0    private var removes = 0     init(record: AntigravityCLISessionRecord? = nil, failSaves: Bool = false) {        self.records = record.map { [$0] } ?? []        self.failSaves = failSaves    }     func load() throws -> [AntigravityCLISessionRecord] {        self.lock.lock()        let value = self.records        self.lock.unlock()        return value    }     func save(_ record: AntigravityCLISessionRecord) throws {        self.lock.lock()        guard !self.failSaves else {            self.lock.unlock()            throw CocoaError(.fileWriteNoPermission)        }        self.records.removeAll { existing in            if let existingOwnerPID = existing.ownerPID,               let recordOwnerPID = record.ownerPID,               let existingOwnerPath = existing.ownerExecutablePath,               let recordOwnerPath = record.ownerExecutablePath,               let existingOwnerStart = existing.ownerStartEpoch,               let recordOwnerStart = record.ownerStartEpoch            {                return existingOwnerPID == recordOwnerPID &&                    existingOwnerPath == recordOwnerPath &&                    abs(existingOwnerStart - recordOwnerStart) < 0.001            }            return existing.pid == record.pid &&                existing.executablePath == record.executablePath &&                abs(existing.startEpoch - record.startEpoch) < 0.001        }        self.records.append(record)        self.saves += 1        self.lock.unlock()    }     func remove(_ record: AntigravityCLISessionRecord) throws {        self.lock.lock()        self.records.removeAll {            $0.pid == record.pid &&                $0.executablePath == record.executablePath &&                abs($0.startEpoch - record.startEpoch) < 0.001        }        self.removes += 1        self.lock.unlock()    }     func snapshot() -> AntigravityCLISessionRecord? {        self.lock.lock()        let value = self.records.first        self.lock.unlock()        return value    }     func snapshots() -> [AntigravityCLISessionRecord] {        self.lock.lock()        let value = self.records        self.lock.unlock()        return value    }     var saveCount: Int {        self.lock.lock()        let value = self.saves        self.lock.unlock()        return value    }     var removeCount: Int {        self.lock.lock()        let value = self.removes        self.lock.unlock()        return value    }} private final class MemoryAntigravitySessionLaunchLock: AntigravityCLISessionLaunchLocking, @unchecked Sendable {    private let lock = NSLock()     func withLock<T>(_ operation: () throws -> T) throws -> T {        self.lock.lock()        defer { self.lock.unlock() }        return try operation()    }} private struct FailingAntigravitySessionLaunchLock: AntigravityCLISessionLaunchLocking {    func withLock<T>(_: () throws -> T) throws -> T {        throw CocoaError(.fileWriteNoPermission)    }} private final class AntigravitySessionTerminationRecorder: @unchecked Sendable {    private let lock = NSLock()    private var events: [(pid: pid_t, group: pid_t?, signal: Int32, descendants: [pid_t])] = []     func append(pid: pid_t, group: pid_t?, signal: Int32, descendants: [pid_t]) {        self.lock.lock()        self.events.append((pid: pid, group: group, signal: signal, descendants: descendants))        self.lock.unlock()    }     func snapshot() -> [(pid: pid_t, group: pid_t?, signal: Int32, descendants: [pid_t])] {        self.lock.lock()        let value = self.events        self.lock.unlock()        return value    }} private final class AntigravityRegistryRecorder: @unchecked Sendable {    private let lock = NSLock()    private var shouldRegister = true    private var registered: [pid_t] = []    private var unregistered: [pid_t] = []    private var groups: [pid_t: pid_t?] = [:]     func setShouldRegister(_ value: Bool) {        self.lock.lock()        self.shouldRegister = value        self.lock.unlock()    }     func register(pid: pid_t, _: String) -> Bool {        self.lock.lock()        defer { self.lock.unlock() }        guard self.shouldRegister else { return false }        self.registered.append(pid)        return true    }     func update(pid: pid_t, group: pid_t?) {        self.lock.lock()        self.groups[pid] = group        self.lock.unlock()    }     func unregister(pid: pid_t) {        self.lock.lock()        self.unregistered.append(pid)        self.lock.unlock()    }     func registeredSnapshot() -> [pid_t] {        self.lock.lock()        let value = self.registered        self.lock.unlock()        return value    }     func unregisteredSnapshot() -> [pid_t] {        self.lock.lock()        let value = self.unregistered        self.lock.unlock()        return value    }} private final class AntigravityLaunchReservationRecorder: @unchecked Sendable {    private let lock = NSLock()    private var beginCount = 0    private var endCount = 0     func begin() -> Bool {        self.lock.lock()        self.beginCount += 1        self.lock.unlock()        return true    }     func end() {        self.lock.lock()        self.endCount += 1        self.lock.unlock()    }     func counts() -> (begin: Int, end: Int) {        self.lock.lock()        let counts = (begin: self.beginCount, end: self.endCount)        self.lock.unlock()        return counts    }} private final class AntigravityManualSleeper: @unchecked Sendable {    private let lock = NSLock()    private var continuations: [CheckedContinuation<Void, Error>] = []     func sleep(_: UInt64) async throws {        try await withCheckedThrowingContinuation { continuation in            self.lock.lock()            self.continuations.append(continuation)            self.lock.unlock()        }    }     func resumeAll() {        self.lock.lock()        let continuations = self.continuations        self.continuations.removeAll()        self.lock.unlock()         for continuation in continuations {            continuation.resume()        }    }     func waitForSleeps(_ expectedCount: Int) async {        for _ in 0..<200 {            if self.pendingSleepCount >= expectedCount { return }            try? await Task.sleep(nanoseconds: 1_000_000)        }        Issue.record("Timed out waiting for \(expectedCount) sleep continuation(s)")    }     private var pendingSleepCount: Int {        self.lock.lock()        let count = self.continuations.count        self.lock.unlock()        return count    }} struct AntigravityCLISessionTests {    @Test    func `reuses alive process for same binary`() async throws {        let fixture = self.makeFixture()        fixture.identity.setIdentity(pid: 10, executablePath: "/bin/agy", startEpoch: 100)         let firstPID = try await fixture.session.beginProbe(binary: "/bin/agy")        await fixture.session.finishProbe(success: true, resetAfterFetch: false)        let secondPID = try await fixture.session.beginProbe(binary: "/bin/agy")        await fixture.session.finishProbe(success: true, resetAfterFetch: false)         #expect(firstPID == 10)        #expect(secondPID == 10)        #expect(fixture.launcher.launchedBinarySnapshot() == ["/bin/agy"])        #expect(fixture.launchReservations.counts().begin == 1)        #expect(fixture.launchReservations.counts().end == 1)    }     @Test    func `relaunches when binary changes`() async throws {        let fixture = self.makeFixture()        fixture.identity.setIdentity(pid: 10, executablePath: "/bin/agy", startEpoch: 100)        fixture.identity.setIdentity(pid: 11, executablePath: "/new/agy", startEpoch: 101)         _ = try await fixture.session.beginProbe(binary: "/bin/agy")        await fixture.session.finishProbe(success: true, resetAfterFetch: false)        let secondPID = try await fixture.session.beginProbe(binary: "/new/agy")        await fixture.session.finishProbe(success: true, resetAfterFetch: false)         #expect(secondPID == 11)        #expect(fixture.launcher.launchedBinarySnapshot() == ["/bin/agy", "/new/agy"])        #expect(fixture.registry.unregisteredSnapshot() == [10])    }     @Test    func `replacement launch waits for in progress teardown`() async throws {        let fixture = self.makeFixture(            manualSleep: true,            terminationGracePeriod: 1,            terminateRootStopsProcess: false)        fixture.identity.setIdentity(pid: 10, executablePath: "/old/agy", startEpoch: 100)        fixture.identity.setIdentity(pid: 11, executablePath: "/new/agy", startEpoch: 101)         _ = try await fixture.session.beginProbe(binary: "/old/agy")        await fixture.session.finishProbe(success: true, resetAfterFetch: false)         let firstReplacement = Task {            try await fixture.session.beginProbe(binary: "/new/agy")        }        await fixture.sleeper?.waitForSleeps(1)         let secondReplacement = Task {            try await fixture.session.beginProbe(binary: "/new/agy")        }        await Task.yield()        #expect(fixture.launcher.launchedBinarySnapshot() == ["/old/agy"])         fixture.launcher.handleSnapshot().first?.killRoot()        fixture.sleeper?.resumeAll()        let firstPID = try await firstReplacement.value        let secondPID = try await secondReplacement.value        await fixture.session.finishProbe(success: true, resetAfterFetch: false)        await fixture.session.finishProbe(success: true, resetAfterFetch: false)         #expect(firstPID == 11)        #expect(secondPID == 11)

Lines 1–500 of 1536.