nucleic/Tests/NucleicCoreTests/OrchestraWorkerSlotTests.swift

import Foundation
import Testing

@testable import NucleicCore

/// A one-shot gate: `wait()` suspends until `open()` (and returns immediately once open).
private actor SlotTestGate {
    private var opened = false
    private var waiters: [CheckedContinuation<Void, Never>] = []
    func wait() async {
        if opened { return }
        await withCheckedContinuation { waiters.append($0) }
    }
    func open() {
        opened = true
        for waiter in waiters { waiter.resume() }
        waiters.removeAll()
    }
}

/// An `AgentBackend` whose run parks mid-turn until `interrupt()` (or `shutdown()`) arrives, then
/// settles as interrupted — the shape of an Orchestra worker whose supervisor dies mid-`join()`:
/// the spawn path must interrupt it (its result has nowhere to go) rather than let it run on.
private final class InterruptParkingBackend: AgentBackend, @unchecked Sendable {
    static let id = BackendID.claudeCode

    let capabilities = BackendCapabilities(
        interactiveApprovals: true,
        allowAlwaysScopes: [.session, .toolName],
        canModifyToolInput: true,
        partialMessageStreaming: false,
        emitsThinking: false,
        emitsFileChangeEvents: false,
        nativeResume: true,
        sandboxModes: [],
        followUpWhileRunning: false)

    private let gate = SlotTestGate()
    /// Set once `interrupt()` was called — the assertion hook for "the worker was stopped".
    let interrupted = LockedBox(false)

    func start(_ run: RunSpec) -> AsyncThrowingStream<AgentEvent, Error> {
        stream(sessionID: run.sessionID, cwd: run.worktree)
    }

    func resume(_ resume: ResumeSpec) -> AsyncThrowingStream<AgentEvent, Error> {
        stream(sessionID: resume.sessionID, cwd: resume.worktree)
    }

    private func stream(sessionID: SessionID, cwd: String) -> AsyncThrowingStream<AgentEvent, Error> {
        AsyncThrowingStream { continuation in
            let gate = self.gate
            let task = Task {
                var seq: UInt64 = 0
                func emit(_ kind: AgentEvent.Kind) {
                    seq += 1
                    continuation.yield(AgentEvent(
                        sessionID: sessionID, seq: seq, at: Date(timeIntervalSince1970: 1_700_000_000),
                        backend: Self.id, nativeType: nil, kind: kind))
                }
                emit(.sessionStarted(SessionStarted(
                    backendSessionID: "be-parked", model: "m", cwd: cwd, toolNames: [])))
                await gate.wait()  // mid-turn park; only interrupt()/shutdown() releases it
                emit(.runFinished(RunFinished(outcome: .interrupted)))
                continuation.finish()
            }
            continuation.onTermination = { _ in task.cancel() }
        }
    }

    func send(_ input: AgentInput) async throws {}
    func respond(to approvalID: ApprovalID, _ decision: Decision, by responder: String) async throws {}
    func interrupt() async {
        interrupted.set(true)
        await gate.open()
    }
    func shutdown() async { await gate.open() }
}

@MainActor
// Serialized + isolated settings for the same reasons as AppStoreTests: @MainActor timing
// assertions, and controlled-project registration touching the container-service settings.
@Suite("Orchestra worker slots — cancellation", .serialized, .isolatedContainerSettings)
struct OrchestraWorkerSlotTests {
    private func makeStore(
        repo: GitTestRepo, backend: @escaping @Sendable (Session) -> any AgentBackend
    ) -> AppStore {
        let database = try! GRDBMetadataStore(path: nil)
        let worktrees = GitWorktreeManager(now: { Date(timeIntervalSince1970: 1_700_000_000) })
        let transcriptsDir = URL(fileURLWithPath: (repo.container as NSString)
            .appendingPathComponent("transcripts"))
        return AppStore(
            database: database, worktrees: worktrees, transcriptsDir: transcriptsDir,
            now: { Date(timeIntervalSince1970: 1_700_000_000) }, backendFactory: backend)
    }

    private func waitFor(_ condition: @escaping () -> Bool, timeoutMs: Int = 3000) async {
        var elapsed = 0
        while !condition() && elapsed < timeoutMs {
            try? await Task.sleep(for: .milliseconds(10))
            elapsed += 10
        }
    }

    // MARK: - The slot gate itself

    /// A queued spawn whose supervisor dies must leave the queue: cancelling the parked task
    /// removes its waiter, `acquire` reports no slot, and the gate's bookkeeping stays intact
    /// for later acquires.
    @Test func cancelledWaiterIsRemovedAndAcquireReturnsFalse() async throws {
        let repo = try await GitTestRepo()
        defer { repo.cleanup() }
        let store = makeStore(repo: repo) { _ in ScriptedBackend { _, _ in } }
        store.defaultOrchestraMaxConcurrentWorkers = 1

        #expect(await store.acquireOrchestraWorkerSlot())  // occupy the only slot
        let parked = Task { @MainActor in await store.acquireOrchestraWorkerSlot() }
        await waitFor { store.orchestraWorkerWaiterCount == 1 }
        #expect(store.orchestraWorkerWaiterCount == 1)

        parked.cancel()
        #expect(await parked.value == false)  // no slot handed to the dead caller
        #expect(store.orchestraWorkerWaiterCount == 0)  // and its waiter is gone, not a ghost
        #expect(store.orchestraWorkersRunning == 1)  // the running count never moved

        // The gate still works: release the held slot, a fresh acquire succeeds immediately.
        store.releaseOrchestraWorkerSlot()
        #expect(await store.acquireOrchestraWorkerSlot())
        store.releaseOrchestraWorkerSlot()
        #expect(store.orchestraWorkersRunning == 0)
    }

    /// A cancelled waiter must not strand the ones queued behind it: after B is cancelled,
    /// releasing A's slot wakes C, which acquires normally.
    @Test func cancelledWaiterDoesNotStrandLaterWaiters() async throws {
        let repo = try await GitTestRepo()
        defer { repo.cleanup() }
        let store = makeStore(repo: repo) { _ in ScriptedBackend { _, _ in } }
        store.defaultOrchestraMaxConcurrentWorkers = 1

        #expect(await store.acquireOrchestraWorkerSlot())  // A holds the slot
        let b = Task { @MainActor in await store.acquireOrchestraWorkerSlot() }
        await waitFor { store.orchestraWorkerWaiterCount == 1 }
        let c = Task { @MainActor in await store.acquireOrchestraWorkerSlot() }
        await waitFor { store.orchestraWorkerWaiterCount == 2 }

        b.cancel()
        #expect(await b.value == false)
        #expect(store.orchestraWorkerWaiterCount == 1)  // only C remains queued

        store.releaseOrchestraWorkerSlot()  // A ends → the freed slot must reach C, not vanish
        #expect(await c.value == true)
        #expect(store.orchestraWorkersRunning == 1)
        store.releaseOrchestraWorkerSlot()
    }

    /// Cancellation is honored even when no parking would occur: a spawn task cancelled before
    /// (or while) acquiring never takes a slot — including under an unlimited (`<= 0`) cap.
    @Test func preCancelledAcquireNeverTakesASlot() async throws {
        let repo = try await GitTestRepo()
        defer { repo.cleanup() }
        let store = makeStore(repo: repo) { _ in ScriptedBackend { _, _ in } }

        for cap in [4, 0] {
            store.defaultOrchestraMaxConcurrentWorkers = cap
            let task = Task { @MainActor in await store.acquireOrchestraWorkerSlot() }
            task.cancel()
            #expect(await task.value == false)
            #expect(store.orchestraWorkersRunning == 0)
        }
    }

    // MARK: - The spawn path

    /// The end-to-end guarantee for a queued spawn: a `nucleic_subagent` call parked on the cap
    /// whose supervisor dies is denied and never creates a worker session — no tokens, no
    /// worktree, no orphaned run for a dead parent.
    @Test func queuedSpawnCancelledBeforeSlotIsDeniedAndCreatesNothing() async throws {
        let repo = try await GitTestRepo()
        defer { repo.cleanup() }
        let store = makeStore(repo: repo) { _ in ScriptedBackend { _, _ in } }
        store.defaultOrchestraMaxConcurrentWorkers = 1

        let project = try #require(
            await store.addProject(name: "demo", rootPath: repo.root, defaultBranch: "main"))
        // An idle parent (empty prompt starts no run) to hang the spawn off.
        let parentID = try await store.createSession(in: project, title: "supervisor", prompt: "")
        let parent = try #require(await store.liveSnapshot(parentID)).session
        let sessionsBefore = store.summaries.count

        #expect(await store.acquireOrchestraWorkerSlot())  // occupy the only slot
        let spawn = Task { @MainActor in
            await store.spawnOrchestraSubagent(
                OrchestraSubagentRequest(task: "scan", prompt: "check it"),
                parent: parent, project: project)
        }
        await waitFor { store.orchestraWorkerWaiterCount == 1 }
        #expect(store.orchestraWorkerWaiterCount == 1)  // the spawn is parked on the cap

        spawn.cancel()  // the supervisor's run ended (Stop/kill → its handler task is cancelled)
        guard case .denied(let message) = await spawn.value else {
            Issue.record("expected .denied for a spawn cancelled while queued")
            return
        }
        #expect(message.contains("cancelled"))
        #expect(store.summaries.count == sessionsBefore)  // no worker session was ever created
        #expect(store.orchestraWorkersRunning == 1)  // only the manual hold; nothing leaked
        store.releaseOrchestraWorkerSlot()
    }

    /// A spawn cancelled mid-`join()` interrupts its in-flight worker (the result has nowhere to
    /// go), reports `.failed`, and releases the concurrency slot.
    @Test func spawnCancelledMidJoinInterruptsWorkerAndFreesSlot() async throws {
        let repo = try await GitTestRepo()
        defer { repo.cleanup() }
        let backends = LockedBox([InterruptParkingBackend]())
        let store = makeStore(repo: repo) { _ in
            let backend = InterruptParkingBackend()
            backends.set(backends.get() + [backend])
            return backend
        }

        let project = try #require(
            await store.addProject(name: "demo", rootPath: repo.root, defaultBranch: "main"))
        let parentID = try await store.createSession(in: project, title: "supervisor", prompt: "")
        let parent = try #require(await store.liveSnapshot(parentID)).session

        let spawn = Task { @MainActor in
            await store.spawnOrchestraSubagent(
                OrchestraSubagentRequest(task: "long job", prompt: "work forever"),
                parent: parent, project: project)
        }
        // Wait until the worker session exists and is mid-turn (its backend is parked).
        await waitFor { store.orchestraWorkersRunning == 1 && store.summaries.count == 2 }
        #expect(store.orchestraWorkersRunning == 1)
        let worker = try #require(backends.get().last)
        #expect(!worker.interrupted.get())

        spawn.cancel()  // the supervisor's run ended while the worker was mid-turn
        guard case .failed(let sessionID, _, _, let message) = await spawn.value else {
            Issue.record("expected .failed for a spawn cancelled mid-join")
            return
        }
        #expect(sessionID != nil)
        #expect(message.contains("interrupted"))
        #expect(worker.interrupted.get())  // the worker was actively stopped, not abandoned
        #expect(store.orchestraWorkersRunning == 0)  // the slot came back
    }
}