diff --git a/TFM/App/AppDependencies.swift b/TFM/App/AppDependencies.swift index c86dac26..22dba657 100644 --- a/TFM/App/AppDependencies.swift +++ b/TFM/App/AppDependencies.swift @@ -7,6 +7,14 @@ import TermindRemoteInfrastructure import struct TermindDomain.Host +@MainActor +struct AppShutdownDependencies { + let stopMCP: @MainActor () async -> Void + let stopSFTP: @MainActor () async -> Void + let closeAgentBackgroundCommands: @MainActor () async -> Void + let closeAgentConnections: @MainActor () async -> Void +} + @MainActor struct AppDependencies { let hostRepository: any HostRepository @@ -21,6 +29,8 @@ struct AppDependencies { let mcpServerSecrets: any SecretStoring let logging: any ApplicationLogManaging let settingsDefaults: UserDefaults + let assistantConversations: AssistantConversationStore? + let shutdown: AppShutdownDependencies? init( hostRepository: any HostRepository, outboxRunner: OutboxRunner?, @@ -33,7 +43,9 @@ struct AppDependencies { metricsPreferences: (any MetricsConfigurationStoring)? = nil, mcpServerSecrets: any SecretStoring = InMemorySecretStore(), logging: any ApplicationLogManaging = DisabledApplicationLogManager.shared, - settingsDefaults: UserDefaults = .standard + settingsDefaults: UserDefaults = .standard, + assistantConversations: AssistantConversationStore? = nil, + shutdown: AppShutdownDependencies? = nil ) { self.hostRepository = hostRepository self.outboxRunner = outboxRunner @@ -47,6 +59,8 @@ struct AppDependencies { self.mcpServerSecrets = mcpServerSecrets self.logging = logging self.settingsDefaults = settingsDefaults + self.assistantConversations = assistantConversations + self.shutdown = shutdown } static func inMemory(hosts: [Host], groups: [HostGroup]) -> Self { diff --git a/TFM/App/AppModel.swift b/TFM/App/AppModel.swift index 04bdbd7c..e8ed3ac0 100644 --- a/TFM/App/AppModel.swift +++ b/TFM/App/AppModel.swift @@ -42,6 +42,7 @@ nonisolated enum HostOpenTarget: Hashable, Sendable { @Observable final class AppModel { private var terminationTask: Task? + private var isTerminating = false private final class WeakWindow { weak var value: WindowModel? init(_ value: WindowModel) { self.value = value } @@ -104,7 +105,7 @@ final class AppModel { let speechTranscription: SpeechTranscriptionCoordinator /// The one set of conversations every window navigates. See /// `AssistantConversationStore`. - private let assistantConversations = AssistantConversationStore() + private let assistantConversations: AssistantConversationStore private let agentRunRepository: any AgentRunPersisting private let agentHosts: AgentHostResolver private let agentConnections: AgentConnectionPool @@ -115,6 +116,7 @@ final class AppModel { private let agentBackgroundCommands: AgentBackgroundCommands private let outboxRunner: OutboxRunner? private let settingsDefaults: UserDefaults + private let shutdownDependencies: AppShutdownDependencies? private var windows: [ObjectIdentifier: WeakWindow] = [:] private var deletedAgentConversationIDs = Set() /// Conversations already handed to the summarizer. Titling is once per @@ -162,6 +164,8 @@ final class AppModel { self.windowSizeStore = windowSizeStore let dependencies = dependencies ?? .inMemory(hosts: hosts.hosts, groups: hosts.groups) settingsDefaults = dependencies.settingsDefaults + assistantConversations = dependencies.assistantConversations ?? AssistantConversationStore() + shutdownDependencies = dependencies.shutdown logging = dependencies.logging let externalAuthentication = ExternalAuthCoordinator( logger: logging.makeLogger(label: "termind.external-mcp")) @@ -393,7 +397,10 @@ final class AppModel { @ObservationIgnored private lazy var assistantExecution = AssistantExecutionController( ai: ai, skills: skills, repository: agentRunRepository, memory: memory, - authorize: { [weak self] in self?.proAccess.require(.assistant) == true }, + authorize: { [weak self] in + guard let self, !self.isTerminating else { return false } + return self.proAccess.require(.assistant) + }, makeRegistry: { [weak self] conversation, model, provider, skills, channel in guard let self else { return AgentToolRegistry([AnyAgentTool]()) } return AssistantToolRegistry.make( @@ -717,9 +724,47 @@ final class AppModel { /// The delegate gives this asynchronous cleanup a bounded termination window. @discardableResult func shutdownForTermination() -> Task { + isTerminating = true + if let terminationTask { return terminationTask } + let conversations = assistantConversations.conversations + let preparationCancellation = assistantExecution.cancelPendingRuns() + for conversation in conversations { conversation.stop() } + let conversationDrains = conversations.map { conversation in + Task { @MainActor in await conversation.stopAndWait() } + } + let task = Task { + [shutdownDependencies, mcpServer, sftpTransfers, agentBackgroundCommands, agentConnections] in + async let stopMCP: Void = { + if let stop = shutdownDependencies?.stopMCP { + await stop() + } else { + await mcpServer.stop() + } + }() + async let stopSFTP: Void = { + if let stop = shutdownDependencies?.stopSFTP { + await stop() + } else { + await sftpTransfers.waitForCancellation() + } + }() + await preparationCancellation.value + for drain in conversationDrains { await drain.value } + _ = await (stopMCP, stopSFTP) + if let close = shutdownDependencies?.closeAgentBackgroundCommands { + await close() + } else { + await agentBackgroundCommands.closeAll() + } + if let close = shutdownDependencies?.closeAgentConnections { + await close() + } else { + await agentConnections.closeAll() + } + } + terminationTask = task memory.suspend() sftpEdits.setMonitoring(false) - if let terminationTask { return terminationTask } serverMetrics.shutdown() speechTranscription.cancelAll() let liveWindows = windows.values.compactMap(\.value) @@ -728,14 +773,6 @@ final class AppModel { sftpTransfers.cancelAll() externalAuthentication.cancelAll() forwarding.stopAll() - let task = Task { [mcpServer, sftpTransfers, agentConnections, agentBackgroundCommands] in - async let stopMCP: Void = mcpServer.stop() - async let stopSFTP: Void = sftpTransfers.waitForCancellation() - _ = await (stopMCP, stopSFTP) - await agentBackgroundCommands.closeAll() - await agentConnections.closeAll() - } - terminationTask = task outboxRunner?.stop() return task } diff --git a/TFM/App/TermindApp.swift b/TFM/App/TermindApp.swift index 14290784..b2410a2c 100644 --- a/TFM/App/TermindApp.swift +++ b/TFM/App/TermindApp.swift @@ -213,6 +213,34 @@ struct TermindApp: App { } #if canImport(AppKit) + @MainActor + final class AppTerminationReplyCoordinator { + typealias Sleep = @Sendable (Duration) async throws -> Void + + private var replied = false + + func begin( + cleanup: Task, timeout: Duration = .seconds(3), + sleep: @escaping Sleep = { try await Task.sleep(for: $0) }, + reply: @escaping @MainActor () -> Void + ) { + Task { @MainActor [weak self] in + await cleanup.value + self?.finish(reply) + } + Task { @MainActor [weak self] in + try? await sleep(timeout) + self?.finish(reply) + } + } + + private func finish(_ reply: @MainActor () -> Void) { + guard !replied else { return } + replied = true + reply() + } + } + /// Owns native application lifecycle; window and pane commands live in `TermindCommands`. final class AppDelegate: NSObject, NSApplicationDelegate { weak var app: AppModel? { @@ -228,7 +256,7 @@ struct TermindApp: App { private let metricsActivity = MetricsActivityObserver() private let metricsMenu = MetricsMenuBarController() private var awaitingTermination = false - private var repliedToTermination = false + private let terminationReply = AppTerminationReplyCoordinator() private var quietLaunchObserver: (any NSObjectProtocol)? private let appearanceController = AppAppearanceController() @@ -381,20 +409,9 @@ struct TermindApp: App { private func beginTermination(_ app: AppModel) { app.mainWindows.saveSizeForTermination() let cleanup = app.shutdownForTermination() - Task { @MainActor [weak self] in - await cleanup.value - self?.finishTermination() + terminationReply.begin(cleanup: cleanup) { + NSApp.reply(toApplicationShouldTerminate: true) } - Task { @MainActor [weak self] in - try? await Task.sleep(for: .seconds(3)) - self?.finishTermination() - } - } - - private func finishTermination() { - guard !repliedToTermination else { return } - repliedToTermination = true - NSApp.reply(toApplicationShouldTerminate: true) } func applicationWillTerminate(_ notification: Notification) { diff --git a/TFMTests/AppTerminationTests.swift b/TFMTests/AppTerminationTests.swift new file mode 100644 index 00000000..f9a00cbe --- /dev/null +++ b/TFMTests/AppTerminationTests.swift @@ -0,0 +1,355 @@ +import AgentKit +import Foundation +import TermindApplication +import TermindDomain +import TermindInfrastructure +import Testing + +@testable import Termind + +@Suite(.timeLimit(.minutes(1))) @MainActor +struct AppTerminationTests { + @Test + func shutdownCancelsEveryStoredConversationBeforeAwaitingAndClosesResourcesAfterDrains() async throws { + let events = TerminationEventLog() + let firstRepository = TerminationRepository( + blocksTerminalSave: true, + terminalSaved: { await events.append("save") }) + let secondRepository = TerminationRepository() + let firstCalls = TerminationModelCalls() + let secondCalls = TerminationModelCalls() + let first = queuedConversation(repository: firstRepository, calls: firstCalls) + let second = queuedConversation(repository: secondRepository, calls: secondCalls) + let store = AssistantConversationStore(conversations: [first, second]) + let app = makeApp( + repository: firstRepository, store: store, + shutdown: AppShutdownDependencies( + stopMCP: { await events.append("mcp") }, + stopSFTP: { await events.append("sftp") }, + closeAgentBackgroundCommands: { await events.append("background") }, + closeAgentConnections: { await events.append("connections") } + )) + let window = WindowModel() + app.registerWindow(window) + #expect(window.assistantLibrary.conversation(id: first.id) === first) + window.openAssistantInSidebar(second.id) + #expect(window.sidebar.assistantConversationID == second.id) + app.unregisterWindow(window) + + let cleanup = app.shutdownForTermination() + let repeatedCleanup = app.shutdownForTermination() + await firstRepository.waitForTerminalSave() + + #expect(first.runState == .cancelled) + #expect(second.runState == .cancelled) + #expect(await firstCalls.count == 0) + #expect(await secondCalls.count == 0) + #expect(!(await events.values).contains("background")) + #expect(!(await events.values).contains("connections")) + + await firstRepository.releaseTerminalSaves() + await cleanup.value + await repeatedCleanup.value + + let values = await events.values + let save = try #require(values.firstIndex(of: "save")) + let background = try #require(values.firstIndex(of: "background")) + let connections = try #require(values.firstIndex(of: "connections")) + #expect(values.filter { $0 == "mcp" }.count == 1) + #expect(values.filter { $0 == "sftp" }.count == 1) + #expect(values.filter { $0 == "background" }.count == 1) + #expect(values.filter { $0 == "connections" }.count == 1) + #expect(save < background) + #expect(background < connections) + #expect(await firstRepository.savedTerminalState == .cancelled) + #expect(await secondRepository.savedTerminalState == .cancelled) + } + + @Test + func shutdownFencesEveryAssistantAdmissionAndInvalidatesAnAlreadyPreparedEdit() async throws { + let repository = TerminationRepository() + let original = AgentTranscriptMessage(role: .user, text: "before", authoredText: "before") + let preparedEdit = ConversationModel( + messages: [ChatMessage(id: original.id, role: .user, text: original.text)], + transcript: [original]) + let draft = ConversationModel() + draft.composer.draft = "draft after shutdown" + let normal = ConversationModel() + let directStart = ConversationModel() + let planConversation = ConversationModel() + let plan = AgentPlan( + title: "Shutdown plan", body: "Must not run", scratchPath: "/tmp/shutdown-plan.md") + planConversation.isPlanning = true + planConversation.handle(.planPresented(plan)) + let postFenceEditMessage = AgentTranscriptMessage(role: .user, text: "edit me", authoredText: "edit me") + let postFenceEdit = ConversationModel( + messages: [ChatMessage(id: postFenceEditMessage.id, role: .user, text: postFenceEditMessage.text)], + transcript: [postFenceEditMessage]) + let store = AssistantConversationStore( + conversations: [preparedEdit, draft, normal, directStart, planConversation, postFenceEdit]) + let app = makeApp(repository: repository, store: store) + let window = WindowModel() + app.registerWindow(window) + + window.assistantCommands?.edit( + preparedEdit.id, original.id, "prepared before shutdown", .none, false) + #expect(preparedEdit.isPreparingAgentRun) + let cleanup = app.shutdownForTermination() + + app.sendAgentDraft(draft) + let newConversation = ConversationModel() + newConversation.composer.draft = "new conversation after shutdown" + #expect(app.sendNewAgentConversation(newConversation, in: window.assistantLibrary) == nil) + window.assistantLibrary.draftConversation.composer.draft = "new after shutdown" + #expect(app.sendNewAgentDraft(in: window.assistantLibrary) == nil) + #expect(!app.sendAgentText("normal after shutdown", in: normal)) + + let planCard = try #require(planConversation.plans.first) + #expect(!app.sendPlanImplementation(planCard, in: planConversation)) + + app.startAgentRun( + directStart, prompt: "explicit start after shutdown", + authoredPrompt: "explicit start after shutdown") + window.assistantCommands?.edit( + postFenceEdit.id, postFenceEditMessage.id, "edit after shutdown", .none, false) + + await cleanup.value + for _ in 0..<10 { await Task.yield() } + + #expect(preparedEdit.transcript == [original]) + #expect(draft.transcript.isEmpty) + #expect(draft.composer.draft == "draft after shutdown") + #expect(normal.transcript.isEmpty) + #expect(directStart.transcript.isEmpty) + #expect(planConversation.plans.first?.outcome == .pending) + #expect(planConversation.isPlanning) + #expect(postFenceEdit.transcript == [postFenceEditMessage]) + #expect(newConversation.transcript.isEmpty) + #expect(await repository.truncationCount == 0) + #expect(await repository.savedSnapshots.map(\.state) == [.interrupted]) + } + + @Test + func shutdownJoinsAnUnstoredCapturedFailureBeforeClosingResources() async { + let events = TerminationEventLog() + let repository = TerminationRepository( + blocksTerminalSave: true, + terminalSaved: { await events.append("save") }) + let app = makeApp( + repository: repository, + shutdown: AppShutdownDependencies( + stopMCP: {}, stopSFTP: {}, + closeAgentBackgroundCommands: { await events.append("background") }, + closeAgentConnections: { await events.append("connections") })) + let conversation = ConversationModel() + + #expect(app.sendAgentText("captured failure", in: conversation)) + await repository.waitForTerminalSave() + #expect(conversation.isPersistingAgentPreparationFailure) + + let cleanup = app.shutdownForTermination() + for _ in 0..<10 { await Task.yield() } + #expect(!(await events.values).contains("background")) + #expect(await repository.savedSnapshots.map(\.state) == [.failed]) + + await repository.releaseTerminalSaves() + await cleanup.value + + #expect(conversation.runState == .failed) + #expect(await repository.savedSnapshots.map(\.state) == [.failed]) + #expect(await events.values == ["save", "background", "connections"]) + } + + @Test + func terminationReplyFallsBackWhenCleanupHangsAndRepliesOnlyOnce() async { + let cleanupGate = TerminationGate() + let fallbackGate = TerminationGate() + let replies = TerminationReplyCount() + let cleanup = Task { await cleanupGate.wait() } + let coordinator = AppTerminationReplyCoordinator() + + coordinator.begin( + cleanup: cleanup, timeout: .seconds(3), + sleep: { duration in + #expect(duration == .seconds(3)) + await fallbackGate.wait() + }, + reply: { replies.value += 1 }) + + await fallbackGate.waitUntilEntered() + #expect(replies.value == 0) + await fallbackGate.release() + await waitUntil { replies.value == 1 } + + await cleanupGate.release() + await cleanup.value + for _ in 0..<10 { await Task.yield() } + #expect(replies.value == 1) + + let lateFallback = TerminationGate() + let cleanupFirstReplies = TerminationReplyCount() + let cleanupFirst = AppTerminationReplyCoordinator() + cleanupFirst.begin( + cleanup: Task {}, + sleep: { _ in await lateFallback.wait() }, + reply: { cleanupFirstReplies.value += 1 }) + await waitUntil { cleanupFirstReplies.value == 1 } + await lateFallback.release() + for _ in 0..<10 { await Task.yield() } + #expect(cleanupFirstReplies.value == 1) + } + + private func queuedConversation( + repository: any AgentRunPersisting, calls: TerminationModelCalls + ) -> ConversationModel { + let conversation = ConversationModel() + let runtime = AgentRuntime( + model: TerminationCountingModel(calls: calls), tools: [], + approval: AgentApprovalBroker(reviewer: nil, manualApproval: { _ in .deny("test") }), + repository: repository) + conversation.start( + runtime: runtime, + request: AgentRunRequest( + conversationID: conversation.id, prompt: "queued", + permissionMode: .askForApproval), + repository: repository) + return conversation + } + + private func makeApp( + repository: any AgentRunPersisting, + store: AssistantConversationStore = AssistantConversationStore(), + shutdown: AppShutdownDependencies? = nil + ) -> AppModel { + let defaults = UserDefaults(suiteName: "termind.tests.termination.\(UUID().uuidString)")! + return AppModel( + hosts: HostLibraryModel(), + dependencies: AppDependencies( + hostRepository: InMemoryHostRepository(), + outboxRunner: nil, + agentRunRepository: repository, + settingsDefaults: defaults, + assistantConversations: store, + shutdown: shutdown)) + } + + private func waitUntil(_ condition: @escaping @MainActor () -> Bool) async { + while !condition() { await Task.yield() } + } +} + +private struct TerminationCountingModel: AgentModelStreaming { + let calls: TerminationModelCalls + + func stream(_ request: AgentModelRequest) -> AsyncThrowingStream { + AsyncThrowingStream { continuation in + Task { + await calls.called() + continuation.yield(.finished(.completed)) + continuation.finish() + } + } + } +} + +private actor TerminationModelCalls { + private(set) var count = 0 + func called() { count += 1 } +} + +private actor TerminationEventLog { + private(set) var values: [String] = [] + func append(_ value: String) { values.append(value) } +} + +private actor TerminationGate { + private var entered = false + private var released = false + private var entryWaiters: [CheckedContinuation] = [] + private var releaseWaiters: [CheckedContinuation] = [] + + func wait() async { + entered = true + let waitingForEntry = entryWaiters + entryWaiters.removeAll() + for waiter in waitingForEntry { waiter.resume() } + guard !released else { return } + await withCheckedContinuation { releaseWaiters.append($0) } + } + + func waitUntilEntered() async { + guard !entered else { return } + await withCheckedContinuation { entryWaiters.append($0) } + } + + func release() { + released = true + let waitingForRelease = releaseWaiters + releaseWaiters.removeAll() + for waiter in waitingForRelease { waiter.resume() } + } +} + +@MainActor +private final class TerminationReplyCount { + var value = 0 +} + +private actor TerminationRepository: AgentRunPersisting { + private let blocksTerminalSave: Bool + private let terminalSaved: (@Sendable () async -> Void)? + private var terminalSaveEntered = false + private var terminalSaveWaiters: [CheckedContinuation] = [] + private var terminalReleaseWaiters: [CheckedContinuation] = [] + private(set) var savedTerminalState: AgentRunState? + private(set) var savedSnapshots: [AgentRunSnapshot] = [] + private(set) var truncationCount = 0 + + init( + blocksTerminalSave: Bool = false, + terminalSaved: (@Sendable () async -> Void)? = nil + ) { + self.blocksTerminalSave = blocksTerminalSave + self.terminalSaved = terminalSaved + } + + func save(_ snapshot: AgentRunSnapshot) async { + savedSnapshots.append(snapshot) + guard snapshot.state.isTerminal else { return } + terminalSaveEntered = true + let waiters = terminalSaveWaiters + terminalSaveWaiters.removeAll() + for waiter in waiters { waiter.resume() } + if blocksTerminalSave { + await withCheckedContinuation { terminalReleaseWaiters.append($0) } + } + await terminalSaved?() + savedTerminalState = snapshot.state + } + + func waitForTerminalSave() async { + guard !terminalSaveEntered else { return } + await withCheckedContinuation { terminalSaveWaiters.append($0) } + } + + func releaseTerminalSaves() { + let waiters = terminalReleaseWaiters + terminalReleaseWaiters.removeAll() + for waiter in waiters { waiter.resume() } + } + + func record(_ event: AgentEvent, runID: UUID) {} + func persistCompaction(_ record: AgentCompactionRecord, conversationID: UUID) {} + func loadConversations() -> [StoredAgentConversation] { [] } + func deleteConversation(_ id: UUID) {} + func setTitle(_ title: String, conversationID: UUID) {} + func setConversationMetadata(_ metadata: StoredAgentConversationMetadata, conversationID: UUID) {} + func reorderConversations(_ updates: [StoredAgentConversationOrder]) {} + func truncateConversation( + _ id: UUID, fromUserMessageID messageID: UUID, + replacingWith message: AgentTranscriptMessage + ) { + truncationCount += 1 + } + func createConversation(_ conversation: StoredAgentConversation) {} +}