When a supervisor's run ended (Stop, kill, dropped connection), its in-flight and queued worker spawns survived it: MCPApprovalServer never cancelled handler tasks on connection close or token unregister, and queued spawns parked in a plain continuation would later create and run full worker sessions for a dead parent, burning tokens and worktrees with the results discarded. Wedged workers also held concurrency slots forever (MCP_TOOL_TIMEOUT is deliberately unbounded), freezing Orchestra at the cap. - MCPApprovalServer tracks in-flight handler tasks by connection and bearer token; connection teardown, unregister(token:), and stop() now cancel them. - acquireOrchestraWorkerSlot is cancellation-aware via withTaskCancellationHandler: a cancelled waiter is removed from the queue (passing any wake-up along) and the spawn returns .denied without creating a session. - A spawn cancelled mid-join interrupts its worker (result has nowhere to go) and reports .failed, releasing the slot; the worker session stays visible/resumable. - UnixSocketByteConn sets SO_NOSIGPIPE: a client hanging up mid-response used to SIGPIPE the whole app instead of surfacing as EPIPE. Co-Authored-By: Claude Fable 5 <[email protected]>
250 lines
11 KiB
Swift
250 lines
11 KiB
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
|
|
}
|
|
}
|