From 96b1be10f5592c7e1b4092028efce69d7c1a5d58 Mon Sep 17 00:00:00 2001 From: Vincent Koc Date: Sat, 11 Jul 2026 13:45:27 +0800 Subject: [PATCH] feat(macos): keep MLX TTS model resident (#104200) * test(macos): cover MLX helper packaging * feat(macos): define framed MLX TTS protocol * feat(macos): keep MLX TTS model resident * feat(macos): reuse MLX TTS helper transport * fix(macos): preserve MLX TTS fallback lifecycle * chore(swift): lint MLX TTS protocol sources * fix(macos): keep MLX helper input nonblocking * test(macos): cover forced MLX helper shutdown * fix(macos): fall back after MLX memory eviction * fix(macos): discard canceled MLX audio --- apps/macos-mlx-tts/Package.swift | 22 +- .../Sources/OpenClawMLXTTSHelper/main.swift | 225 +++----- .../MLXTTSHelperService.swift | 174 ++++++ .../MLXTTSHelperServiceTests.swift | 121 ++++ apps/macos/Package.swift | 3 + apps/macos/Sources/OpenClaw/MenuBar.swift | 1 + .../OpenClaw/TalkMLXSpeechSynthesizer.swift | 515 +++++++++++++++--- .../Sources/OpenClaw/TalkModeRuntime.swift | 27 +- .../TalkMLXSpeechSynthesizerTests.swift | 373 +++++++++++++ .../TalkModeRuntimeSpeechTests.swift | 9 + .../OpenClawMLXTTSProtocol/Package.swift | 18 + .../MLXTTSProtocol.swift | 219 ++++++++ .../MLXTTSProtocolTests.swift | 75 +++ scripts/format-swift.sh | 5 +- test/scripts/codesign-mac-app.test.ts | 10 +- test/scripts/package-mac-app.test.ts | 25 + 16 files changed, 1566 insertions(+), 256 deletions(-) create mode 100644 apps/macos-mlx-tts/Sources/OpenClawMLXTTSRuntime/MLXTTSHelperService.swift create mode 100644 apps/macos-mlx-tts/Tests/OpenClawMLXTTSRuntimeTests/MLXTTSHelperServiceTests.swift create mode 100644 apps/macos/Tests/OpenClawIPCTests/TalkMLXSpeechSynthesizerTests.swift create mode 100644 apps/shared/OpenClawMLXTTSProtocol/Package.swift create mode 100644 apps/shared/OpenClawMLXTTSProtocol/Sources/OpenClawMLXTTSProtocol/MLXTTSProtocol.swift create mode 100644 apps/shared/OpenClawMLXTTSProtocol/Tests/OpenClawMLXTTSProtocolTests/MLXTTSProtocolTests.swift diff --git a/apps/macos-mlx-tts/Package.swift b/apps/macos-mlx-tts/Package.swift index aa7703919d5..8074d0b7320 100644 --- a/apps/macos-mlx-tts/Package.swift +++ b/apps/macos-mlx-tts/Package.swift @@ -14,12 +14,32 @@ let package = Package( ], dependencies: [ .package(url: "https://github.com/Blaizzy/mlx-audio-swift", revision: "fc4fe22dc41c053062e647a4e3db9142193670d2"), + .package(path: "../shared/OpenClawMLXTTSProtocol"), ], targets: [ + .target( + name: "OpenClawMLXTTSRuntime", + dependencies: [ + .product(name: "MLXAudioTTS", package: "mlx-audio-swift"), + .product(name: "OpenClawMLXTTSProtocol", package: "OpenClawMLXTTSProtocol"), + ], + swiftSettings: [ + .enableUpcomingFeature("StrictConcurrency"), + ]), .executableTarget( name: "OpenClawMLXTTSHelper", dependencies: [ - .product(name: "MLXAudioTTS", package: "mlx-audio-swift"), + "OpenClawMLXTTSRuntime", + .product(name: "OpenClawMLXTTSProtocol", package: "OpenClawMLXTTSProtocol"), + ], + swiftSettings: [ + .enableUpcomingFeature("StrictConcurrency"), + ]), + .testTarget( + name: "OpenClawMLXTTSRuntimeTests", + dependencies: [ + "OpenClawMLXTTSRuntime", + .product(name: "OpenClawMLXTTSProtocol", package: "OpenClawMLXTTSProtocol"), ], swiftSettings: [ .enableUpcomingFeature("StrictConcurrency"), diff --git a/apps/macos-mlx-tts/Sources/OpenClawMLXTTSHelper/main.swift b/apps/macos-mlx-tts/Sources/OpenClawMLXTTSHelper/main.swift index 5d4a2d254f2..dcc9aed1bb9 100644 --- a/apps/macos-mlx-tts/Sources/OpenClawMLXTTSHelper/main.swift +++ b/apps/macos-mlx-tts/Sources/OpenClawMLXTTSHelper/main.swift @@ -1,184 +1,95 @@ +import Darwin import Foundation -import MLXAudioTTS +import OpenClawMLXTTSProtocol +import OpenClawMLXTTSRuntime -// swiftformat:disable wrap wrapMultilineStatementBraces trailingCommas redundantSelf extensionAccessControl @main enum OpenClawMLXTTSHelper { static func main() async { do { - let options = try Options.parse(CommandLine.arguments.dropFirst()) - let data = try await synthesize(options) - try data.write(to: options.outputURL, options: [.atomic]) + let protocolOutput = try Self.makeProtocolOutput() + let writer = FrameWriter(output: protocolOutput) + let service = MLXTTSHelperService { event in + do { + try await writer.write(event) + } catch { + Self.log("failed to write event: \(error)") + } + } + + try await writer.write(MLXTTSEvent.ready) + try await Self.readRequests(service: service, writer: writer) } catch { - FileHandle.standardError.write(Data("openclaw-mlx-tts: \(error)\n".utf8)) + self.log("\(error)") exit(1) } } - private static func synthesize(_ options: Options) async throws -> Data { - let model = try await TTS.loadModel(modelRepo: options.modelRepo) - let audio = try await UncheckedSpeechModel(raw: model).generateAudio( - text: options.text, - voice: options.voice, - language: options.language) - return makeWavData(samples: audio, sampleRate: Double(model.sampleRate)) + private static func makeProtocolOutput() throws -> FileHandle { + let protocolFD = dup(STDOUT_FILENO) + guard protocolFD >= 0 else { + throw POSIXError(.EBADF) + } + guard dup2(STDERR_FILENO, STDOUT_FILENO) >= 0 else { + close(protocolFD) + throw POSIXError(.EBADF) + } + return FileHandle(fileDescriptor: protocolFD, closeOnDealloc: true) } - private struct Options { - let text: String - let modelRepo: String - let outputURL: URL - let language: String? - let voice: String? + private static func readRequests(service: MLXTTSHelperService, writer: FrameWriter) async throws { + let input = FileHandle.standardInput + let (chunks, continuation) = AsyncStream.makeStream() + input.readabilityHandler = { handle in + let data = handle.availableData + if data.isEmpty { + handle.readabilityHandler = nil + continuation.finish() + } else { + continuation.yield(data) + } + } + defer { + input.readabilityHandler = nil + continuation.finish() + } - static func parse(_ rawArguments: ArraySlice) throws -> Options { - var text: String? - var modelRepo = "mlx-community/Soprano-80M-bf16" - var outputPath: String? - var language: String? - var voice: String? - var iterator = rawArguments.makeIterator() + var decoder = MLXTTSFrameDecoder() + for await chunk in chunks { + for payload in try decoder.append(chunk) { + let request: MLXTTSRequest + do { + request = try MLXTTSFrameCodec.decode(MLXTTSRequest.self, payload: payload) + } catch { + try await writer.write(MLXTTSEvent.error(MLXTTSErrorEvent( + id: nil, + code: .protocolError, + message: "invalid request frame"))) + continue + } - while let argument = iterator.next() { - switch argument { - case "--text", "-t": - text = try nextValue(&iterator, argument) - case "--model": - modelRepo = try nextValue(&iterator, argument) - case "--output", "-o": - outputPath = try nextValue(&iterator, argument) - case "--language": - language = try nextValue(&iterator, argument) - case "--voice", "-v": - voice = try nextValue(&iterator, argument) - case "--help", "-h": - throw Usage.requested - default: - if text == nil, !argument.hasPrefix("-") { - text = argument - } else { - throw Usage.invalid("unknown option \(argument)") - } + if await !(service.handle(request)) { + return } } - - guard let text = text?.trimmingCharacters(in: .whitespacesAndNewlines), !text.isEmpty else { - throw Usage.invalid("missing --text") - } - guard let outputPath, !outputPath.isEmpty else { - throw Usage.invalid("missing --output") - } - - return Options( - text: text, - modelRepo: modelRepo, - outputURL: URL(fileURLWithPath: outputPath), - language: language?.nilIfBlank, - voice: voice?.nilIfBlank) } - private static func nextValue( - _ iterator: inout ArraySlice.Iterator, - _ option: String) throws -> String - { - guard let value = iterator.next(), !value.isEmpty else { - throw Usage.invalid("missing value for \(option)") - } - return value - } + _ = await service.handle(.shutdown) } - private enum Usage: Error, CustomStringConvertible { - case requested - case invalid(String) - - var description: String { - switch self { - case .requested: - "usage: openclaw-mlx-tts --text --output " - + "[--model ] [--language ] [--voice ]" - case let .invalid(message): - "\(message)\nusage: openclaw-mlx-tts --text --output " - + "[--model ] [--language ] [--voice ]" - } - } - } - - private static func makeWavData(samples: [Float], sampleRate: Double) -> Data { - let channels: UInt16 = 1 - let bitsPerSample: UInt16 = 16 - let blockAlign = channels * (bitsPerSample / 8) - let sampleRateInt = UInt32(sampleRate.rounded()) - let byteRate = sampleRateInt * UInt32(blockAlign) - let dataSize = UInt32(samples.count) * UInt32(blockAlign) - - var data = Data(capacity: Int(44 + dataSize)) - data.append(contentsOf: [0x52, 0x49, 0x46, 0x46]) // RIFF - data.appendLEUInt32(36 + dataSize) - data.append(contentsOf: [0x57, 0x41, 0x56, 0x45]) // WAVE - - data.append(contentsOf: [0x66, 0x6D, 0x74, 0x20]) // fmt - data.appendLEUInt32(16) - data.appendLEUInt16(1) - data.appendLEUInt16(channels) - data.appendLEUInt32(sampleRateInt) - data.appendLEUInt32(byteRate) - data.appendLEUInt16(blockAlign) - data.appendLEUInt16(bitsPerSample) - - data.append(contentsOf: [0x64, 0x61, 0x74, 0x61]) // data - data.appendLEUInt32(dataSize) - - for sample in samples { - let clamped = max(-1.0, min(1.0, sample)) - let scaled = Int16((clamped * Float(Int16.max)).rounded()) - data.appendLEInt16(scaled) - } - return data + private static func log(_ message: String) { + FileHandle.standardError.write(Data("openclaw-mlx-tts: \(message)\n".utf8)) } } -private struct UncheckedSpeechModel { - let raw: any SpeechGenerationModel +private actor FrameWriter { + let output: FileHandle - func generateAudio( - text: String, - voice: String?, - language: String?) async throws -> [Float] { - let generatedAudio = try await raw.generate( - text: text, - voice: voice, - refAudio: nil, - refText: nil, - language: language) - return generatedAudio.asArray(Float.self) + init(output: FileHandle) { + self.output = output + } + + func write(_ event: MLXTTSEvent) throws { + try self.output.write(contentsOf: MLXTTSFrameCodec.encode(event)) } } - -extension UncheckedSpeechModel: @unchecked Sendable {} - -private extension String { - var nilIfBlank: String? { - let trimmed = self.trimmingCharacters(in: .whitespacesAndNewlines) - return trimmed.isEmpty ? nil : trimmed - } -} - -private extension Data { - mutating func appendLEUInt16(_ value: UInt16) { - var littleEndian = value.littleEndian - Swift.withUnsafeBytes(of: &littleEndian) { append(contentsOf: $0) } - } - - mutating func appendLEUInt32(_ value: UInt32) { - var littleEndian = value.littleEndian - Swift.withUnsafeBytes(of: &littleEndian) { append(contentsOf: $0) } - } - - mutating func appendLEInt16(_ value: Int16) { - var littleEndian = value.littleEndian - Swift.withUnsafeBytes(of: &littleEndian) { append(contentsOf: $0) } - } -} - -// swiftformat:enable wrap wrapMultilineStatementBraces trailingCommas redundantSelf extensionAccessControl diff --git a/apps/macos-mlx-tts/Sources/OpenClawMLXTTSRuntime/MLXTTSHelperService.swift b/apps/macos-mlx-tts/Sources/OpenClawMLXTTSRuntime/MLXTTSHelperService.swift new file mode 100644 index 00000000000..7b598174c7c --- /dev/null +++ b/apps/macos-mlx-tts/Sources/OpenClawMLXTTSRuntime/MLXTTSHelperService.swift @@ -0,0 +1,174 @@ +import Foundation +import MLXAudioTTS +import OpenClawMLXTTSProtocol + +protocol MLXTTSSpeechModel: AnyObject, Sendable { + var sampleRate: Int { get } + + func generate(text: String, voice: String?, language: String?) async throws -> [Float] +} + +typealias MLXTTSModelLoader = @Sendable (String) async throws -> any MLXTTSSpeechModel + +public actor MLXTTSHelperService { + private struct CachedModel { + let repo: String + let model: any MLXTTSSpeechModel + } + + private let loadModel: MLXTTSModelLoader + private let emit: @Sendable (MLXTTSEvent) async -> Void + private var cachedModel: CachedModel? + private var currentID: String? + private var currentTask: Task? + + public init(eventSink: @Sendable @escaping (MLXTTSEvent) async -> Void) { + self.loadModel = { repo in + let model = try await TTS.loadModel(modelRepo: repo) + return UncheckedSpeechModel(raw: model) + } + self.emit = eventSink + } + + init( + loadModel: @escaping MLXTTSModelLoader, + eventSink: @Sendable @escaping (MLXTTSEvent) async -> Void) + { + self.loadModel = loadModel + self.emit = eventSink + } + + @discardableResult + public func handle(_ request: MLXTTSRequest) async -> Bool { + switch request { + case let .synthesize(synthesize): + guard self.currentTask == nil else { + await self.emit(.error(MLXTTSErrorEvent( + id: synthesize.id, + code: .busy, + message: "another synthesis is already in flight"))) + return true + } + self.currentID = synthesize.id + self.currentTask = Task { await self.run(synthesize) } + return true + + case let .cancel(id): + guard self.currentID == id, let task = self.currentTask else { + await self.emit(.canceled(id: id)) + return true + } + task.cancel() + return true + + case .shutdown: + self.currentTask?.cancel() + await self.currentTask?.value + self.currentTask = nil + self.currentID = nil + self.cachedModel = nil + return false + } + } + + func waitUntilIdle() async { + await self.currentTask?.value + } + + private func run(_ request: MLXTTSSynthesizeRequest) async { + let model: any MLXTTSSpeechModel + do { + model = try await self.model(repo: request.modelRepo) + } catch is CancellationError { + await self.finishCanceled(id: request.id) + return + } catch { + await self.finish( + event: .error(MLXTTSErrorEvent( + id: request.id, + code: .modelLoadFailed, + message: String(describing: error))), + id: request.id) + return + } + + do { + try Task.checkCancellation() + let samples = try await model.generate( + text: request.text, + voice: request.voice, + language: request.language) + try Task.checkCancellation() + let audio = MLXTTSAudio( + id: request.id, + sampleRate: model.sampleRate, + pcm: Self.makePCM16(samples: samples)) + await self.finish(event: .audio(audio), id: request.id) + } catch is CancellationError { + await self.finishCanceled(id: request.id) + } catch { + await self.finish( + event: .error(MLXTTSErrorEvent( + id: request.id, + code: .generationFailed, + message: String(describing: error))), + id: request.id) + } + } + + private func model(repo: String) async throws -> any MLXTTSSpeechModel { + if let cachedModel = self.cachedModel, cachedModel.repo == repo { + return cachedModel.model + } + + // Only one model is retained. Dropping the previous reference before + // loading a new repo avoids holding both sets of MLX weights at once. + self.cachedModel = nil + let model = try await self.loadModel(repo) + self.cachedModel = CachedModel(repo: repo, model: model) + return model + } + + private func finishCanceled(id: String) async { + await self.finish(event: .canceled(id: id), id: id) + } + + private func finish(event: MLXTTSEvent, id: String) async { + guard self.currentID == id else { return } + self.currentID = nil + self.currentTask = nil + await self.emit(event) + } + + static func makePCM16(samples: [Float]) -> Data { + var data = Data(capacity: samples.count * MemoryLayout.size) + for sample in samples { + let clamped = max(-1, min(1, sample)) + var value = Int16((clamped * Float(Int16.max)).rounded()).littleEndian + Swift.withUnsafeBytes(of: &value) { data.append(contentsOf: $0) } + } + return data + } +} + +private final class UncheckedSpeechModel: MLXTTSSpeechModel, @unchecked Sendable { + let raw: any SpeechGenerationModel + + init(raw: any SpeechGenerationModel) { + self.raw = raw + } + + var sampleRate: Int { + self.raw.sampleRate + } + + func generate(text: String, voice: String?, language: String?) async throws -> [Float] { + let generatedAudio = try await self.raw.generate( + text: text, + voice: voice, + refAudio: nil, + refText: nil, + language: language) + return generatedAudio.asArray(Float.self) + } +} diff --git a/apps/macos-mlx-tts/Tests/OpenClawMLXTTSRuntimeTests/MLXTTSHelperServiceTests.swift b/apps/macos-mlx-tts/Tests/OpenClawMLXTTSRuntimeTests/MLXTTSHelperServiceTests.swift new file mode 100644 index 00000000000..ccbec13996b --- /dev/null +++ b/apps/macos-mlx-tts/Tests/OpenClawMLXTTSRuntimeTests/MLXTTSHelperServiceTests.swift @@ -0,0 +1,121 @@ +import Foundation +import OpenClawMLXTTSProtocol +import XCTest +@testable import OpenClawMLXTTSRuntime + +final class MLXTTSHelperServiceTests: XCTestCase { + func testReusesModelForMatchingRepoAndReloadsForChange() async { + let state = TestState() + let service = MLXTTSHelperService( + loadModel: { repo in + await state.loaded(repo) + return TestModel(state: state) + }, + eventSink: { event in await state.emitted(event) }) + + await service.handle(.synthesize(request(id: "one", repo: "repo-a"))) + await service.waitUntilIdle() + await service.handle(.synthesize(request(id: "two", repo: "repo-a"))) + await service.waitUntilIdle() + await service.handle(.synthesize(request(id: "three", repo: "repo-b"))) + await service.waitUntilIdle() + + let loadedRepos = await state.loadedRepos + let generatedTexts = await state.generatedTexts + let events = await state.events + XCTAssertEqual(loadedRepos, ["repo-a", "repo-b"]) + XCTAssertEqual(generatedTexts, ["hello", "hello", "hello"]) + XCTAssertEqual(events.count, 3) + } + + func testRejectsConcurrentSynthesis() async { + let state = TestState() + let service = MLXTTSHelperService( + loadModel: { _ in SlowTestModel() }, + eventSink: { event in await state.emitted(event) }) + + await service.handle(.synthesize(request(id: "one", repo: "repo-a"))) + await service.handle(.synthesize(request(id: "two", repo: "repo-a"))) + _ = await service.handle(.shutdown) + + let events = await state.events + XCTAssertTrue(events.contains(.error(MLXTTSErrorEvent( + id: "two", + code: .busy, + message: "another synthesis is already in flight")))) + XCTAssertTrue(events.contains(.canceled(id: "one"))) + } + + func testCancellationKeepsCachedModelAvailable() async { + let state = TestState() + let model = SlowTestModel() + let service = MLXTTSHelperService( + loadModel: { repo in + await state.loaded(repo) + return model + }, + eventSink: { event in await state.emitted(event) }) + + await service.handle(.synthesize(request(id: "one", repo: "repo-a"))) + await service.handle(.cancel(id: "one")) + await service.waitUntilIdle() + await service.handle(.synthesize(request(id: "two", repo: "repo-a"))) + await service.handle(.cancel(id: "two")) + await service.waitUntilIdle() + + let loadedRepos = await state.loadedRepos + let events = await state.events + XCTAssertEqual(loadedRepos, ["repo-a"]) + XCTAssertEqual(events, [.canceled(id: "one"), .canceled(id: "two")]) + } + + func testConvertsSamplesToLittleEndianPCM16() { + let pcm = MLXTTSHelperService.makePCM16(samples: [-1, 0, 1]) + XCTAssertEqual(pcm, Data([0x01, 0x80, 0x00, 0x00, 0xFF, 0x7F])) + } +} + +private func request(id: String, repo: String) -> MLXTTSSynthesizeRequest { + MLXTTSSynthesizeRequest(id: id, text: "hello", modelRepo: repo, language: nil, voice: nil) +} + +private actor TestState { + private(set) var loadedRepos: [String] = [] + private(set) var generatedTexts: [String] = [] + private(set) var events: [MLXTTSEvent] = [] + + func loaded(_ repo: String) { + self.loadedRepos.append(repo) + } + + func generated(_ text: String) { + self.generatedTexts.append(text) + } + + func emitted(_ event: MLXTTSEvent) { + self.events.append(event) + } +} + +private final class TestModel: MLXTTSSpeechModel, @unchecked Sendable { + let sampleRate = 32000 + let state: TestState + + init(state: TestState) { + self.state = state + } + + func generate(text: String, voice _: String?, language _: String?) async throws -> [Float] { + await self.state.generated(text) + return [-1, 0, 1] + } +} + +private final class SlowTestModel: MLXTTSSpeechModel, @unchecked Sendable { + let sampleRate = 32000 + + func generate(text _: String, voice _: String?, language _: String?) async throws -> [Float] { + try await Task.sleep(for: .seconds(30)) + return [] + } +} diff --git a/apps/macos/Package.swift b/apps/macos/Package.swift index 4d2d22704a8..013f93429ff 100644 --- a/apps/macos/Package.swift +++ b/apps/macos/Package.swift @@ -22,6 +22,7 @@ let package = Package( .package(url: "https://github.com/steipete/Peekaboo.git", exact: "3.5.2"), .package(url: "https://github.com/pointfreeco/swift-concurrency-extras", from: "1.3.1"), .package(path: "../shared/OpenClawKit"), + .package(path: "../shared/OpenClawMLXTTSProtocol"), .package(path: "../swabble"), ], targets: [ @@ -47,6 +48,7 @@ let package = Package( "OpenClawDiscovery", .product(name: "OpenClawKit", package: "OpenClawKit"), .product(name: "OpenClawChatUI", package: "OpenClawKit"), + .product(name: "OpenClawMLXTTSProtocol", package: "OpenClawMLXTTSProtocol"), .product(name: "OpenClawProtocol", package: "OpenClawKit"), .product(name: "SwabbleKit", package: "swabble"), .product(name: "MenuBarExtraAccess", package: "MenuBarExtraAccess"), @@ -88,6 +90,7 @@ let package = Package( "OpenClawMacCLI", "OpenClawDiscovery", .product(name: "OpenClawKit", package: "OpenClawKit"), + .product(name: "OpenClawMLXTTSProtocol", package: "OpenClawMLXTTSProtocol"), .product(name: "OpenClawProtocol", package: "OpenClawKit"), .product(name: "SwabbleKit", package: "swabble"), ], diff --git a/apps/macos/Sources/OpenClaw/MenuBar.swift b/apps/macos/Sources/OpenClaw/MenuBar.swift index 520195b14ab..c75d332b583 100644 --- a/apps/macos/Sources/OpenClaw/MenuBar.swift +++ b/apps/macos/Sources/OpenClaw/MenuBar.swift @@ -293,6 +293,7 @@ final class AppDelegate: NSObject, NSApplicationDelegate { private var terminationCleanupFinished = false private let webChatAutoLogger = Logger(subsystem: "ai.openclaw", category: "Chat") var nodeTerminationCleanup: @MainActor () async -> Void = { + await TalkMLXSpeechSynthesizer.shared.shutdown() await MacNodeModeCoordinator.shared.stopAndWait() } diff --git a/apps/macos/Sources/OpenClaw/TalkMLXSpeechSynthesizer.swift b/apps/macos/Sources/OpenClaw/TalkMLXSpeechSynthesizer.swift index 751226ea79f..eee4231af98 100644 --- a/apps/macos/Sources/OpenClaw/TalkMLXSpeechSynthesizer.swift +++ b/apps/macos/Sources/OpenClaw/TalkMLXSpeechSynthesizer.swift @@ -1,9 +1,17 @@ +import Dispatch import Foundation +import OpenClawMLXTTSProtocol import OSLog -// swiftformat:disable wrap wrapMultilineStatementBraces trailingCommas redundantSelf extensionAccessControl -/// Runtime access stays serialized through `TalkModeRuntime` actor helper methods. -final class TalkMLXSpeechSynthesizer { +protocol MLXTTSTransport: AnyObject, Sendable { + func send(_ request: MLXTTSRequest) async throws + func nextEvent() async throws -> MLXTTSEvent + func close() async +} + +typealias MLXTTSTransportFactory = @Sendable () async throws -> any MLXTTSTransport + +actor TalkMLXSpeechSynthesizer { enum SynthesizeError: Error { case canceled case modelLoadFailed(String) @@ -16,94 +24,268 @@ final class TalkMLXSpeechSynthesizer { static let defaultModelRepo = "mlx-community/Soprano-80M-bf16" private let logger = Logger(subsystem: "ai.openclaw", category: "talk.mlx") - private var currentToken = UUID() - private var currentProcess: Process? + private let transportFactory: MLXTTSTransportFactory + private let idleDuration: Duration + private let cancelGraceDuration: Duration + private let observesMemoryPressure: Bool + private var transport: (any MLXTTSTransport)? + private var activeID: String? + private var cancelRequestedID: String? + private var fallbackRequiredID: String? + private var cancelEscalationTask: Task? + private var idleTask: Task? + private var memoryPressureMonitor: MLXMemoryPressureMonitor? - private init() {} + private init() { + self.transportFactory = { + try ProcessMLXTTSTransport.launch(invocation: TalkMLXSpeechSynthesizer.helperInvocation()) + } + self.idleDuration = .seconds(300) + self.cancelGraceDuration = .seconds(1) + self.observesMemoryPressure = true + } - func stop() { - self.currentToken = UUID() - self.currentProcess?.terminate() - self.currentProcess = nil + init( + transportFactory: @escaping MLXTTSTransportFactory, + idleDuration: Duration = .seconds(300), + cancelGraceDuration: Duration = .seconds(1), + observesMemoryPressure: Bool = false) + { + self.transportFactory = transportFactory + self.idleDuration = idleDuration + self.cancelGraceDuration = cancelGraceDuration + self.observesMemoryPressure = observesMemoryPressure } func synthesize( text: String, modelRepo: String?, language: String?, - voicePreset: String?) async throws -> Data { + voicePreset: String?) async throws -> Data + { + #if !arch(arm64) + throw SynthesizeError.modelLoadFailed("MLX TTS requires Apple silicon") + #else let trimmed = text.trimmingCharacters(in: .whitespacesAndNewlines) guard !trimmed.isEmpty else { return Data() } - - self.stop() - let token = UUID() - self.currentToken = token - - let tempDir = FileManager.default.temporaryDirectory - .appendingPathComponent("openclaw-mlx-tts-\(token.uuidString)", isDirectory: true) - try FileManager.default.createDirectory(at: tempDir, withIntermediateDirectories: true) - defer { try? FileManager.default.removeItem(at: tempDir) } - - let outputURL = tempDir.appendingPathComponent("speech.wav") - let invocation = Self.helperInvocation() - let resolvedRepo = Self.resolvedModelRepo(modelRepo) - var arguments = invocation.argumentPrefix - arguments += [ - "--text", trimmed, - "--model", resolvedRepo, - "--output", outputURL.path, - ] - if let language = language?.trimmingCharacters(in: .whitespacesAndNewlines), !language.isEmpty { - arguments += ["--language", language] - } - if let voicePreset = voicePreset?.trimmingCharacters(in: .whitespacesAndNewlines), !voicePreset.isEmpty { - arguments += ["--voice", voicePreset] - } - - self.logger.info("talk mlx helper start modelRepo=\(resolvedRepo, privacy: .public)") - let process = Process() - process.executableURL = invocation.executableURL - process.arguments = arguments - let stderr = Pipe() - process.standardError = stderr - process.standardOutput = Pipe() - self.currentProcess = process - - let status: Int32 - do { - status = try await Self.run(process) - } catch { - self.currentProcess = nil - self.logger.error("talk mlx helper launch failed: \(error.localizedDescription, privacy: .public)") - throw SynthesizeError.modelLoadFailed(invocation.displayName) - } - self.currentProcess = nil - - guard self.currentToken == token else { - throw SynthesizeError.canceled - } - guard status == 0 else { - let errorText = Self.readPipe(stderr) - self.logger.error( - "talk mlx helper failed status=\(status, privacy: .public): \(errorText, privacy: .public)") + guard self.activeID == nil else { throw SynthesizeError.audioGenerationFailed } + self.ensureMemoryPressureMonitor() + self.idleTask?.cancel() + self.idleTask = nil + + let id = UUID().uuidString + self.activeID = id + let request = MLXTTSRequest.synthesize(MLXTTSSynthesizeRequest( + id: id, + text: trimmed, + modelRepo: Self.resolvedModelRepo(modelRepo), + language: language?.nilIfBlank, + voice: voicePreset?.nilIfBlank)) + + for attempt in 0...1 { + do { + let transport = try await self.ensureTransport() + guard self.activeID == id, self.cancelRequestedID != id else { + await self.discardTransport() + throw SynthesizeError.canceled + } + try await transport.send(request) + let audio = try await self.waitForAudio(id: id, transport: transport) + self.finishRequest(id: id) + return try Self.makeWAV(audio: audio) + } catch let error as SynthesizeError { + let requiresFallback = self.fallbackRequiredID == id + self.finishRequest(id: id) + if requiresFallback { + throw SynthesizeError.audioGenerationFailed + } + throw error + } catch is CancellationError { + try? await self.transport?.send(.cancel(id: id)) + await self.discardTransport() + self.finishRequest(id: id) + throw SynthesizeError.canceled + } catch { + self.logger.error( + "talk mlx helper transport failed attempt=\(attempt + 1, privacy: .public): \(error.localizedDescription, privacy: .public)") + await self.discardTransport() + if self.fallbackRequiredID == id { + self.finishRequest(id: id) + throw SynthesizeError.audioGenerationFailed + } + if self.cancelRequestedID == id { + self.finishRequest(id: id) + throw SynthesizeError.canceled + } + guard self.activeID == id else { + throw SynthesizeError.canceled + } + if attempt == 0 { + continue + } + self.finishRequest(id: id) + throw SynthesizeError.modelLoadFailed(Self.helperInvocation().displayName) + } + } + + self.finishRequest(id: id) + throw SynthesizeError.audioGenerationFailed + #endif + } + + func cancelCurrent() async { + guard let activeID = self.activeID else { return } + self.cancelRequestedID = activeID do { - return try Data(contentsOf: outputURL) + try await self.transport?.send(.cancel(id: activeID)) } catch { - self.logger.error("talk mlx helper output missing: \(error.localizedDescription, privacy: .public)") - throw SynthesizeError.audioGenerationFailed + await self.discardTransport() + } + self.scheduleCancelEscalation(id: activeID) + } + + func shutdown() async { + self.cancelEscalationTask?.cancel() + self.cancelEscalationTask = nil + self.idleTask?.cancel() + self.idleTask = nil + if let activeID = self.activeID { + try? await self.transport?.send(.cancel(id: activeID)) + } + try? await self.transport?.send(.shutdown) + self.activeID = nil + self.cancelRequestedID = nil + await self.discardTransport() + } + + private func ensureTransport() async throws -> any MLXTTSTransport { + if let transport = self.transport { + return transport + } + + let transport = try await self.transportFactory() + // Publish the starting transport before waiting for `ready` so talk + // cancellation and app shutdown can still terminate a wedged startup. + self.transport = transport + do { + guard try await transport.nextEvent() == .ready else { + throw MLXTTSTransportError.unexpectedEvent + } + self.logger.info("talk mlx helper ready") + return transport + } catch { + self.transport = nil + await transport.close() + throw error } } - private struct HelperInvocation { + private func waitForAudio(id: String, transport: any MLXTTSTransport) async throws -> MLXTTSAudio { + while true { + switch try await transport.nextEvent() { + case let .audio(audio) where audio.id == id: + guard self.cancelRequestedID != id else { + throw SynthesizeError.canceled + } + return audio + case let .canceled(canceledID) where canceledID == id: + throw SynthesizeError.canceled + case let .error(error) where error.id == nil || error.id == id: + switch error.code { + case .canceled: + throw SynthesizeError.canceled + case .modelLoadFailed: + throw SynthesizeError.modelLoadFailed(error.message) + case .busy, .generationFailed, .invalidRequest, .protocolError: + throw SynthesizeError.audioGenerationFailed + } + case .ready, .audio, .error, .canceled: + continue + } + } + } + + private func finishRequest(id: String) { + if self.fallbackRequiredID == id { + self.fallbackRequiredID = nil + } + guard self.activeID == id else { return } + self.activeID = nil + self.cancelRequestedID = nil + self.cancelEscalationTask?.cancel() + self.cancelEscalationTask = nil + self.scheduleIdleShutdown() + } + + private func scheduleCancelEscalation(id: String) { + self.cancelEscalationTask?.cancel() + let duration = self.cancelGraceDuration + self.cancelEscalationTask = Task { [weak self] in + do { + try await Task.sleep(for: duration) + } catch { + return + } + await self?.terminateUnresponsiveCancellation(id: id) + } + } + + private func terminateUnresponsiveCancellation(id: String) async { + guard self.activeID == id, self.cancelRequestedID == id else { return } + // Soprano checks cancellation while producing tokens, but its final + // decoder has no cancellation contract. Bound that phase with a kill. + self.logger.info("talk mlx cancel grace expired; terminating helper") + await self.discardTransport() + } + + private func scheduleIdleShutdown() { + self.idleTask?.cancel() + let duration = self.idleDuration + self.idleTask = Task { [weak self] in + do { + try await Task.sleep(for: duration) + } catch { + return + } + await self?.shutdownIfIdle() + } + } + + private func shutdownIfIdle() async { + guard self.activeID == nil else { return } + self.logger.info("talk mlx helper idle shutdown") + await self.shutdown() + } + + private func ensureMemoryPressureMonitor() { + guard self.observesMemoryPressure, self.memoryPressureMonitor == nil else { return } + self.memoryPressureMonitor = MLXMemoryPressureMonitor { [weak self] in + Task { await self?.handleMemoryPressure() } + } + } + + func handleMemoryPressure() async { + self.logger.info("talk mlx helper memory-pressure shutdown") + self.fallbackRequiredID = self.activeID + await self.shutdown() + } + + private func discardTransport() async { + let transport = self.transport + self.transport = nil + await transport?.close() + } + + fileprivate struct HelperInvocation: Sendable { let executableURL: URL let argumentPrefix: [String] let displayName: String } - private static func helperInvocation() -> HelperInvocation { + fileprivate static func helperInvocation() -> HelperInvocation { let fileManager = FileManager.default if let override = ProcessInfo.processInfo.environment["OPENCLAW_MLX_TTS_BIN"], !override.isEmpty { return HelperInvocation( @@ -129,30 +311,189 @@ final class TalkMLXSpeechSynthesizer { } private static func resolvedModelRepo(_ modelRepo: String?) -> String { - let trimmed = modelRepo?.trimmingCharacters(in: .whitespacesAndNewlines) ?? "" - return trimmed.isEmpty ? Self.defaultModelRepo : trimmed + modelRepo?.nilIfBlank ?? self.defaultModelRepo } - private static func run(_ process: Process) async throws -> Int32 { - try await withCheckedThrowingContinuation { continuation in - process.terminationHandler = { process in - continuation.resume(returning: process.terminationStatus) - } - do { - try process.run() - } catch { - continuation.resume(throwing: error) - } + static func makeWAV(audio: MLXTTSAudio) throws -> Data { + guard audio.format == .pcmS16LE, + audio.sampleRate > 0, + audio.sampleRate <= Int(UInt32.max), + audio.channels > 0, + audio.channels <= Int(UInt16.max), + audio.pcm.count <= Int(UInt32.max) - 36, + audio.pcm.count.isMultiple(of: MemoryLayout.size * audio.channels) + else { + throw SynthesizeError.audioGenerationFailed } - } - private static func readPipe(_ pipe: Pipe) -> String { - let data = (try? pipe.fileHandleForReading.readToEnd()) ?? Data() - let text = String(data: data, encoding: .utf8) ?? "" - return text.trimmingCharacters(in: .whitespacesAndNewlines) + let channels = UInt16(audio.channels) + let sampleRate = UInt32(audio.sampleRate) + let bitsPerSample: UInt16 = 16 + let blockAlign = channels * (bitsPerSample / 8) + let byteRate = sampleRate * UInt32(blockAlign) + let dataSize = UInt32(audio.pcm.count) + + var data = Data(capacity: 44 + audio.pcm.count) + data.append(contentsOf: [0x52, 0x49, 0x46, 0x46]) + data.appendLEUInt32(36 + dataSize) + data.append(contentsOf: [0x57, 0x41, 0x56, 0x45]) + data.append(contentsOf: [0x66, 0x6D, 0x74, 0x20]) + data.appendLEUInt32(16) + data.appendLEUInt16(1) + data.appendLEUInt16(channels) + data.appendLEUInt32(sampleRate) + data.appendLEUInt32(byteRate) + data.appendLEUInt16(blockAlign) + data.appendLEUInt16(bitsPerSample) + data.append(contentsOf: [0x64, 0x61, 0x74, 0x61]) + data.appendLEUInt32(dataSize) + data.append(audio.pcm) + return data } } -extension TalkMLXSpeechSynthesizer: @unchecked Sendable {} +private enum MLXTTSTransportError: Error { + case closed + case unexpectedEvent +} -// swiftformat:enable wrap wrapMultilineStatementBraces trailingCommas redundantSelf extensionAccessControl +private actor ProcessMLXTTSTransport: MLXTTSTransport { + private let process: Process + private let input: FileHandle + private let output: FileHandle + private let chunkContinuation: AsyncStream.Continuation + private let chunks: MLXChunkIterator + private var decoder = MLXTTSFrameDecoder() + private var pendingPayloads: [Data] = [] + private var isClosed = false + + private init( + process: Process, + input: FileHandle, + output: FileHandle, + chunks: AsyncStream, + chunkContinuation: AsyncStream.Continuation) + { + self.process = process + self.input = input + self.output = output + self.chunks = MLXChunkIterator(stream: chunks) + self.chunkContinuation = chunkContinuation + } + + static func launch(invocation: TalkMLXSpeechSynthesizer.HelperInvocation) throws -> ProcessMLXTTSTransport { + let process = Process() + let inputPipe = Pipe() + let outputPipe = Pipe() + process.executableURL = invocation.executableURL + process.arguments = invocation.argumentPrefix + process.standardInput = inputPipe + process.standardOutput = outputPipe + process.standardError = FileHandle.standardError + + let output = outputPipe.fileHandleForReading + let (stream, continuation) = AsyncStream.makeStream() + output.readabilityHandler = { handle in + let data = handle.availableData + if data.isEmpty { + handle.readabilityHandler = nil + continuation.finish() + } else { + continuation.yield(data) + } + } + + do { + try process.run() + } catch { + output.readabilityHandler = nil + continuation.finish() + throw error + } + inputPipe.fileHandleForReading.closeFile() + outputPipe.fileHandleForWriting.closeFile() + + return ProcessMLXTTSTransport( + process: process, + input: inputPipe.fileHandleForWriting, + output: output, + chunks: stream, + chunkContinuation: continuation) + } + + func send(_ request: MLXTTSRequest) async throws { + try self.input.write(contentsOf: MLXTTSFrameCodec.encode(request)) + } + + func nextEvent() async throws -> MLXTTSEvent { + while true { + if !self.pendingPayloads.isEmpty { + let payload = self.pendingPayloads.removeFirst() + return try MLXTTSFrameCodec.decode(MLXTTSEvent.self, payload: payload) + } + guard let chunk = await self.chunks.next() else { + throw MLXTTSTransportError.closed + } + try self.pendingPayloads.append(contentsOf: self.decoder.append(chunk)) + } + } + + func close() async { + guard !self.isClosed else { return } + self.isClosed = true + self.output.readabilityHandler = nil + self.chunkContinuation.finish() + self.input.closeFile() + self.output.closeFile() + if self.process.isRunning { + self.process.terminate() + } + } +} + +private final class MLXChunkIterator: @unchecked Sendable { + private var iterator: AsyncStream.Iterator + + init(stream: AsyncStream) { + self.iterator = stream.makeAsyncIterator() + } + + func next() async -> Data? { + await self.iterator.next() + } +} + +private final class MLXMemoryPressureMonitor: @unchecked Sendable { + private let source: DispatchSourceMemoryPressure + + init(handler: @Sendable @escaping () -> Void) { + self.source = DispatchSource.makeMemoryPressureSource( + eventMask: [.warning, .critical], + queue: .global(qos: .utility)) + self.source.setEventHandler(handler: handler) + self.source.resume() + } + + deinit { + self.source.cancel() + } +} + +extension String { + fileprivate var nilIfBlank: String? { + let trimmed = self.trimmingCharacters(in: .whitespacesAndNewlines) + return trimmed.isEmpty ? nil : trimmed + } +} + +extension Data { + fileprivate mutating func appendLEUInt16(_ value: UInt16) { + var littleEndian = value.littleEndian + Swift.withUnsafeBytes(of: &littleEndian) { self.append(contentsOf: $0) } + } + + fileprivate mutating func appendLEUInt32(_ value: UInt32) { + var littleEndian = value.littleEndian + Swift.withUnsafeBytes(of: &littleEndian) { self.append(contentsOf: $0) } + } +} diff --git a/apps/macos/Sources/OpenClaw/TalkModeRuntime.swift b/apps/macos/Sources/OpenClaw/TalkModeRuntime.swift index 24d68a00142..630c4fe4d4e 100644 --- a/apps/macos/Sources/OpenClaw/TalkModeRuntime.swift +++ b/apps/macos/Sources/OpenClaw/TalkModeRuntime.swift @@ -16,6 +16,11 @@ actor TalkModeRuntime { case systemVoiceOnly } + enum MLXFailureDisposition: Equatable { + case canceled + case fallback + } + private let logger = Logger(subsystem: "ai.openclaw", category: "talk.runtime") private let ttsLogger = Logger(subsystem: "ai.openclaw", category: "talk.tts") private static let defaultModelIdFallback = "eleven_v3" @@ -624,10 +629,11 @@ extension TalkModeRuntime { case .mlxThenSystemVoice: do { try await self.playMLX(input: input) - } catch TalkMLXSpeechSynthesizer.SynthesizeError.canceled { - self.ttsLogger.info("talk mlx canceled") - return } catch { + if Self.mlxFailureDisposition(error) == .canceled { + self.ttsLogger.info("talk mlx canceled") + return + } self.ttsLogger .error( "talk MLX failed: \(error.localizedDescription, privacy: .public); " + @@ -668,6 +674,13 @@ extension TalkModeRuntime { } } + static func mlxFailureDisposition(_ error: Error) -> MLXFailureDisposition { + if case TalkMLXSpeechSynthesizer.SynthesizeError.canceled = error { + return .canceled + } + return .fallback + } + private struct TalkPlaybackInput { let generation: Int let provider: String @@ -937,7 +950,7 @@ extension TalkModeRuntime { voicePreset: input.voicePreset) }) } catch TalkMLXSpeechSynthesizer.SynthesizeError.timedOut { - self.stopMLXVoice() + await self.stopMLXVoice() throw TalkMLXSpeechSynthesizer.SynthesizeError.timedOut } let result = await self.playTalkAudio(data: audioData) @@ -1005,7 +1018,7 @@ extension TalkModeRuntime { _ = usePCM ? await self.stopMP3() : await self.stopPCM() let localInterruptedAt = await self.stopTalkAudio() await TalkSystemSpeechSynthesizer.shared.stop() - self.stopMLXVoice() + await self.stopMLXVoice() guard self.phase == .speaking else { return } let interruptedAt = remoteInterruptedAt ?? localInterruptedAt if reason == .speech, let interruptedAt { @@ -1127,8 +1140,8 @@ extension TalkModeRuntime { voicePreset: voicePreset) } - private func stopMLXVoice() { - TalkMLXSpeechSynthesizer.shared.stop() + private func stopMLXVoice() async { + await TalkMLXSpeechSynthesizer.shared.cancelCurrent() } // MARK: - Config diff --git a/apps/macos/Tests/OpenClawIPCTests/TalkMLXSpeechSynthesizerTests.swift b/apps/macos/Tests/OpenClawIPCTests/TalkMLXSpeechSynthesizerTests.swift new file mode 100644 index 00000000000..f8595718495 --- /dev/null +++ b/apps/macos/Tests/OpenClawIPCTests/TalkMLXSpeechSynthesizerTests.swift @@ -0,0 +1,373 @@ +import Foundation +import OpenClawMLXTTSProtocol +import Testing +@testable import OpenClaw + +#if arch(arm64) +@Suite(.serialized) +struct TalkMLXSpeechSynthesizerTests { + @Test + func `reuses resident helper across utterances`() async throws { + let transport = TestMLXTransport(mode: .audio) + let factory = TestMLXTransportFactory([transport]) + let synthesizer = TalkMLXSpeechSynthesizer( + transportFactory: { try await factory.make() }, + idleDuration: .seconds(60)) + + let first = try await synthesizer.synthesize( + text: "first", + modelRepo: nil, + language: nil, + voicePreset: nil) + let second = try await synthesizer.synthesize( + text: "second", + modelRepo: "repo-a", + language: "en", + voicePreset: "voice-a") + + #expect(first.starts(with: Data("RIFF".utf8))) + #expect(second.starts(with: Data("RIFF".utf8))) + #expect(await factory.callCount == 1) + let requests = await transport.sent + #expect(requests.count == 2) + guard case let .synthesize(firstRequest) = requests[0], + case let .synthesize(secondRequest) = requests[1] + else { + Issue.record("expected two synthesis requests") + return + } + #expect(firstRequest.modelRepo == TalkMLXSpeechSynthesizer.defaultModelRepo) + #expect(secondRequest.modelRepo == "repo-a") + #expect(secondRequest.language == "en") + #expect(secondRequest.voice == "voice-a") + } + + @Test + func `retries once after helper crash`() async throws { + let crashed = TestMLXTransport(mode: .crash) + let restarted = TestMLXTransport(mode: .audio) + let factory = TestMLXTransportFactory([crashed, restarted]) + let synthesizer = TalkMLXSpeechSynthesizer( + transportFactory: { try await factory.make() }, + idleDuration: .seconds(60)) + + let data = try await synthesizer.synthesize( + text: "retry me", + modelRepo: nil, + language: nil, + voicePreset: nil) + + #expect(data.starts(with: Data("RIFF".utf8))) + #expect(await factory.callCount == 2) + #expect(await crashed.closeCount == 1) + #expect(await restarted.sent.count == 1) + } + + @Test + func `cancel uses protocol without closing helper`() async throws { + let transport = TestMLXTransport(mode: .waitForCancel) + let factory = TestMLXTransportFactory([transport]) + let synthesizer = TalkMLXSpeechSynthesizer( + transportFactory: { try await factory.make() }, + idleDuration: .seconds(60)) + + let synthesis = Task { + try await synthesizer.synthesize( + text: "cancel me", + modelRepo: nil, + language: nil, + voicePreset: nil) + } + await transport.waitForSynthesisRequest() + await synthesizer.cancelCurrent() + + do { + _ = try await synthesis.value + Issue.record("expected cancellation") + } catch TalkMLXSpeechSynthesizer.SynthesizeError.canceled { + #expect(await transport.closeCount == 0) + #expect(await transport.sent.contains { request in + if case .cancel = request { return true } + return false + }) + } + } + + @Test + func `late audio after cancel is discarded`() async throws { + let transport = TestMLXTransport(mode: .audioAfterCancel) + let factory = TestMLXTransportFactory([transport]) + let synthesizer = TalkMLXSpeechSynthesizer( + transportFactory: { try await factory.make() }, + idleDuration: .seconds(60)) + + let synthesis = Task { + try await synthesizer.synthesize( + text: "discard me", + modelRepo: nil, + language: nil, + voicePreset: nil) + } + await transport.waitForSynthesisRequest() + await synthesizer.cancelCurrent() + + do { + _ = try await synthesis.value + Issue.record("expected cancellation") + } catch TalkMLXSpeechSynthesizer.SynthesizeError.canceled { + #expect(await transport.closeCount == 0) + } + } + + @Test + func `unresponsive cancel terminates helper without retry`() async throws { + let transport = TestMLXTransport(mode: .ignoreCancel) + let factory = TestMLXTransportFactory([transport]) + let synthesizer = TalkMLXSpeechSynthesizer( + transportFactory: { try await factory.make() }, + idleDuration: .seconds(60), + cancelGraceDuration: .milliseconds(10)) + + let synthesis = Task { + try await synthesizer.synthesize( + text: "cancel me hard", + modelRepo: nil, + language: nil, + voicePreset: nil) + } + await transport.waitForSynthesisRequest() + await synthesizer.cancelCurrent() + + do { + _ = try await synthesis.value + Issue.record("expected cancellation") + } catch TalkMLXSpeechSynthesizer.SynthesizeError.canceled { + #expect(await transport.closeCount == 1) + #expect(await factory.callCount == 1) + } + } + + @Test + func `shutdown terminates unresponsive in-flight helper`() async throws { + let transport = TestMLXTransport(mode: .ignoreCancel) + let factory = TestMLXTransportFactory([transport]) + let synthesizer = TalkMLXSpeechSynthesizer( + transportFactory: { try await factory.make() }, + idleDuration: .seconds(60)) + + let synthesis = Task { + try await synthesizer.synthesize( + text: "stop during shutdown", + modelRepo: nil, + language: nil, + voicePreset: nil) + } + await transport.waitForSynthesisRequest() + await synthesizer.shutdown() + + do { + _ = try await synthesis.value + Issue.record("expected cancellation") + } catch TalkMLXSpeechSynthesizer.SynthesizeError.canceled { + #expect(await transport.closeCount == 1) + #expect(await transport.sent.contains(.shutdown)) + } + } + + @Test + func `memory pressure during synthesis preserves fallback`() async throws { + let transport = TestMLXTransport(mode: .ignoreCancel) + let factory = TestMLXTransportFactory([transport]) + let synthesizer = TalkMLXSpeechSynthesizer( + transportFactory: { try await factory.make() }, + idleDuration: .seconds(60)) + + let synthesis = Task { + try await synthesizer.synthesize( + text: "fall back after pressure", + modelRepo: nil, + language: nil, + voicePreset: nil) + } + await transport.waitForSynthesisRequest() + await synthesizer.handleMemoryPressure() + + do { + _ = try await synthesis.value + Issue.record("expected generation failure") + } catch TalkMLXSpeechSynthesizer.SynthesizeError.audioGenerationFailed { + #expect(await transport.closeCount == 1) + #expect(await transport.sent.contains(.shutdown)) + } + } + + @Test + func `cancel can terminate helper before ready`() async throws { + let transport = TestMLXTransport(mode: .startupHang) + let factory = TestMLXTransportFactory([transport]) + let synthesizer = TalkMLXSpeechSynthesizer( + transportFactory: { try await factory.make() }, + idleDuration: .seconds(60), + cancelGraceDuration: .milliseconds(10)) + + let synthesis = Task { + try await synthesizer.synthesize( + text: "never ready", + modelRepo: nil, + language: nil, + voicePreset: nil) + } + await factory.waitForCall() + await synthesizer.cancelCurrent() + + do { + _ = try await synthesis.value + Issue.record("expected cancellation") + } catch TalkMLXSpeechSynthesizer.SynthesizeError.canceled { + #expect(await transport.closeCount >= 1) + #expect(await factory.callCount == 1) + } + } + + @Test + func `idle timeout shuts down resident helper`() async throws { + let transport = TestMLXTransport(mode: .audio) + let factory = TestMLXTransportFactory([transport]) + let synthesizer = TalkMLXSpeechSynthesizer( + transportFactory: { try await factory.make() }, + idleDuration: .milliseconds(10)) + + _ = try await synthesizer.synthesize( + text: "brief", + modelRepo: nil, + language: nil, + voicePreset: nil) + await transport.waitForShutdown() + + #expect(await transport.closeCount == 1) + } + + @Test + func `pcm response becomes playable WAV`() throws { + let wav = try TalkMLXSpeechSynthesizer.makeWAV(audio: MLXTTSAudio( + id: "one", + sampleRate: 32000, + pcm: Data([0x00, 0x00, 0xFF, 0x7F]))) + + #expect(wav.count == 48) + #expect(wav.prefix(4) == Data("RIFF".utf8)) + #expect(wav.subdata(in: 8..<12) == Data("WAVE".utf8)) + #expect(wav.suffix(4) == Data([0x00, 0x00, 0xFF, 0x7F])) + } +} + +private enum TestMLXTransportError: Error { + case closed +} + +private actor TestMLXTransport: MLXTTSTransport { + enum Mode: Equatable, Sendable { + case audio + case audioAfterCancel + case crash + case ignoreCancel + case startupHang + case waitForCancel + } + + let mode: Mode + private(set) var sent: [MLXTTSRequest] = [] + private(set) var closeCount = 0 + private var events: [MLXTTSEvent] = [.ready] + private var closed = false + + init(mode: Mode) { + self.mode = mode + if mode == .startupHang { + self.events = [] + } + } + + func send(_ request: MLXTTSRequest) { + self.sent.append(request) + switch request { + case let .synthesize(synthesize): + switch self.mode { + case .audio: + self.events.append(.audio(MLXTTSAudio( + id: synthesize.id, + sampleRate: 32000, + pcm: Data([0x00, 0x00, 0xFF, 0x7F])))) + case .crash: + self.closed = true + case .audioAfterCancel, .ignoreCancel, .startupHang, .waitForCancel: + break + } + case let .cancel(id): + if self.mode == .audioAfterCancel { + self.events.append(.audio(MLXTTSAudio( + id: id, + sampleRate: 32000, + pcm: Data([0x00, 0x00, 0xFF, 0x7F])))) + } else if self.mode != .ignoreCancel, self.mode != .startupHang { + self.events.append(.canceled(id: id)) + } + case .shutdown: + self.closed = true + } + } + + func nextEvent() async throws -> MLXTTSEvent { + while self.events.isEmpty { + if self.closed { + throw TestMLXTransportError.closed + } + await Task.yield() + } + return self.events.removeFirst() + } + + func close() { + self.closeCount += 1 + self.closed = true + } + + func waitForSynthesisRequest() async { + while !self.sent.contains(where: { + if case .synthesize = $0 { return true } + return false + }) { + await Task.yield() + } + } + + func waitForShutdown() async { + while !self.sent.contains(.shutdown) || self.closeCount == 0 { + await Task.yield() + } + } +} + +private actor TestMLXTransportFactory { + private var transports: [TestMLXTransport] + private(set) var callCount = 0 + + init(_ transports: [TestMLXTransport]) { + self.transports = transports + } + + func make() throws -> any MLXTTSTransport { + self.callCount += 1 + guard !self.transports.isEmpty else { + throw TestMLXTransportError.closed + } + return self.transports.removeFirst() + } + + func waitForCall() async { + while self.callCount == 0 { + await Task.yield() + } + } +} +#endif diff --git a/apps/macos/Tests/OpenClawIPCTests/TalkModeRuntimeSpeechTests.swift b/apps/macos/Tests/OpenClawIPCTests/TalkModeRuntimeSpeechTests.swift index 361ebb1b27d..bb65374eed8 100644 --- a/apps/macos/Tests/OpenClawIPCTests/TalkModeRuntimeSpeechTests.swift +++ b/apps/macos/Tests/OpenClawIPCTests/TalkModeRuntimeSpeechTests.swift @@ -45,6 +45,15 @@ struct TalkModeRuntimeSpeechTests { #expect(systemPlan == .systemVoiceOnly) } + @Test func `mlx cancellation stops while failures preserve system fallback`() { + #expect(TalkModeRuntime.mlxFailureDisposition( + TalkMLXSpeechSynthesizer.SynthesizeError.canceled) == .canceled) + #expect(TalkModeRuntime.mlxFailureDisposition( + TalkMLXSpeechSynthesizer.SynthesizeError.audioGenerationFailed) == .fallback) + #expect(TalkModeRuntime.mlxFailureDisposition( + TalkMLXSpeechSynthesizer.SynthesizeError.modelLoadFailed("missing")) == .fallback) + } + @Test func `talk speak params carry resolved voice and directive overrides`() { let params = TalkModeRuntime.makeTalkSpeakParams( text: "hello", diff --git a/apps/shared/OpenClawMLXTTSProtocol/Package.swift b/apps/shared/OpenClawMLXTTSProtocol/Package.swift new file mode 100644 index 00000000000..fb745d2cdc9 --- /dev/null +++ b/apps/shared/OpenClawMLXTTSProtocol/Package.swift @@ -0,0 +1,18 @@ +// swift-tools-version: 6.2 + +import PackageDescription + +let package = Package( + name: "OpenClawMLXTTSProtocol", + platforms: [ + .macOS(.v15), + ], + products: [ + .library(name: "OpenClawMLXTTSProtocol", targets: ["OpenClawMLXTTSProtocol"]), + ], + targets: [ + .target(name: "OpenClawMLXTTSProtocol"), + .testTarget( + name: "OpenClawMLXTTSProtocolTests", + dependencies: ["OpenClawMLXTTSProtocol"]), + ]) diff --git a/apps/shared/OpenClawMLXTTSProtocol/Sources/OpenClawMLXTTSProtocol/MLXTTSProtocol.swift b/apps/shared/OpenClawMLXTTSProtocol/Sources/OpenClawMLXTTSProtocol/MLXTTSProtocol.swift new file mode 100644 index 00000000000..b849dc8975a --- /dev/null +++ b/apps/shared/OpenClawMLXTTSProtocol/Sources/OpenClawMLXTTSProtocol/MLXTTSProtocol.swift @@ -0,0 +1,219 @@ +import Foundation + +public enum MLXTTSRequest: Codable, Equatable, Sendable { + case synthesize(MLXTTSSynthesizeRequest) + case cancel(id: String) + case shutdown + + private enum CodingKeys: String, CodingKey { + case type + case synthesize + case id + } + + private enum RequestType: String, Codable { + case synthesize + case cancel + case shutdown + } + + public init(from decoder: Decoder) throws { + let container = try decoder.container(keyedBy: CodingKeys.self) + switch try container.decode(RequestType.self, forKey: .type) { + case .synthesize: + self = try .synthesize(container.decode(MLXTTSSynthesizeRequest.self, forKey: .synthesize)) + case .cancel: + self = try .cancel(id: container.decode(String.self, forKey: .id)) + case .shutdown: + self = .shutdown + } + } + + public func encode(to encoder: Encoder) throws { + var container = encoder.container(keyedBy: CodingKeys.self) + switch self { + case let .synthesize(request): + try container.encode(RequestType.synthesize, forKey: .type) + try container.encode(request, forKey: .synthesize) + case let .cancel(id): + try container.encode(RequestType.cancel, forKey: .type) + try container.encode(id, forKey: .id) + case .shutdown: + try container.encode(RequestType.shutdown, forKey: .type) + } + } +} + +public struct MLXTTSSynthesizeRequest: Codable, Equatable, Sendable { + public let id: String + public let text: String + public let modelRepo: String + public let language: String? + public let voice: String? + + public init(id: String, text: String, modelRepo: String, language: String?, voice: String?) { + self.id = id + self.text = text + self.modelRepo = modelRepo + self.language = language + self.voice = voice + } +} + +public enum MLXTTSAudioFormat: String, Codable, Equatable, Sendable { + case pcmS16LE = "pcm_s16le" +} + +public struct MLXTTSAudio: Codable, Equatable, Sendable { + public let id: String + public let format: MLXTTSAudioFormat + public let sampleRate: Int + public let channels: Int + public let pcm: Data + + public init( + id: String, + format: MLXTTSAudioFormat = .pcmS16LE, + sampleRate: Int, + channels: Int = 1, + pcm: Data) + { + self.id = id + self.format = format + self.sampleRate = sampleRate + self.channels = channels + self.pcm = pcm + } +} + +public enum MLXTTSErrorCode: String, Codable, Equatable, Sendable { + case busy + case canceled + case generationFailed = "generation_failed" + case invalidRequest = "invalid_request" + case modelLoadFailed = "model_load_failed" + case protocolError = "protocol_error" +} + +public struct MLXTTSErrorEvent: Codable, Equatable, Sendable { + public let id: String? + public let code: MLXTTSErrorCode + public let message: String + + public init(id: String?, code: MLXTTSErrorCode, message: String) { + self.id = id + self.code = code + self.message = message + } +} + +public enum MLXTTSEvent: Codable, Equatable, Sendable { + case ready + case audio(MLXTTSAudio) + case error(MLXTTSErrorEvent) + case canceled(id: String) + + private enum CodingKeys: String, CodingKey { + case type + case audio + case error + case id + } + + private enum EventType: String, Codable { + case ready + case audio + case error + case canceled + } + + public init(from decoder: Decoder) throws { + let container = try decoder.container(keyedBy: CodingKeys.self) + switch try container.decode(EventType.self, forKey: .type) { + case .ready: + self = .ready + case .audio: + self = try .audio(container.decode(MLXTTSAudio.self, forKey: .audio)) + case .error: + self = try .error(container.decode(MLXTTSErrorEvent.self, forKey: .error)) + case .canceled: + self = try .canceled(id: container.decode(String.self, forKey: .id)) + } + } + + public func encode(to encoder: Encoder) throws { + var container = encoder.container(keyedBy: CodingKeys.self) + switch self { + case .ready: + try container.encode(EventType.ready, forKey: .type) + case let .audio(audio): + try container.encode(EventType.audio, forKey: .type) + try container.encode(audio, forKey: .audio) + case let .error(error): + try container.encode(EventType.error, forKey: .type) + try container.encode(error, forKey: .error) + case let .canceled(id): + try container.encode(EventType.canceled, forKey: .type) + try container.encode(id, forKey: .id) + } + } +} + +public enum MLXTTSFrameError: Error, Equatable, Sendable { + case emptyFrame + case frameTooLarge(Int) +} + +public enum MLXTTSFrameCodec { + public static let maximumPayloadSize = 64 * 1024 * 1024 + + public static func encode(_ value: some Encodable) throws -> Data { + let encoder = JSONEncoder() + encoder.outputFormatting = [.sortedKeys] + let payload = try encoder.encode(value) + guard !payload.isEmpty else { + throw MLXTTSFrameError.emptyFrame + } + guard payload.count <= self.maximumPayloadSize else { + throw MLXTTSFrameError.frameTooLarge(payload.count) + } + + var length = UInt32(payload.count).bigEndian + var frame = Data(bytes: &length, count: MemoryLayout.size) + frame.append(payload) + return frame + } + + public static func decode(_ type: T.Type, payload: Data) throws -> T { + try JSONDecoder().decode(type, from: payload) + } +} + +public struct MLXTTSFrameDecoder: Sendable { + private var buffer = Data() + + public init() {} + + public mutating func append(_ data: Data) throws -> [Data] { + self.buffer.append(data) + var payloads: [Data] = [] + + while self.buffer.count >= MemoryLayout.size { + let length = self.buffer.prefix(4).reduce(0) { ($0 << 8) | Int($1) } + guard length > 0 else { + throw MLXTTSFrameError.emptyFrame + } + guard length <= MLXTTSFrameCodec.maximumPayloadSize else { + throw MLXTTSFrameError.frameTooLarge(length) + } + guard self.buffer.count >= 4 + length else { + break + } + + payloads.append(self.buffer.subdata(in: 4..<(4 + length))) + self.buffer.removeSubrange(0..<(4 + length)) + } + + return payloads + } +} diff --git a/apps/shared/OpenClawMLXTTSProtocol/Tests/OpenClawMLXTTSProtocolTests/MLXTTSProtocolTests.swift b/apps/shared/OpenClawMLXTTSProtocol/Tests/OpenClawMLXTTSProtocolTests/MLXTTSProtocolTests.swift new file mode 100644 index 00000000000..a1e3d6a5815 --- /dev/null +++ b/apps/shared/OpenClawMLXTTSProtocol/Tests/OpenClawMLXTTSProtocolTests/MLXTTSProtocolTests.swift @@ -0,0 +1,75 @@ +import Foundation +import OpenClawMLXTTSProtocol +import XCTest + +final class MLXTTSProtocolTests: XCTestCase { + func testRequestRoundTrip() throws { + let request = MLXTTSRequest.synthesize(MLXTTSSynthesizeRequest( + id: "request-1", + text: "hello", + modelRepo: "mlx-community/Soprano-80M-bf16", + language: "en", + voice: nil)) + + var decoder = MLXTTSFrameDecoder() + let payloads = try decoder.append(MLXTTSFrameCodec.encode(request)) + + XCTAssertEqual(payloads.count, 1) + XCTAssertEqual(try MLXTTSFrameCodec.decode(MLXTTSRequest.self, payload: payloads[0]), request) + } + + func testDecoderAcceptsFragmentedAndCoalescedFrames() throws { + let first = try MLXTTSFrameCodec.encode(MLXTTSRequest.cancel(id: "one")) + let second = try MLXTTSFrameCodec.encode(MLXTTSRequest.shutdown) + let combined = first + second + let split = combined.count / 2 + + var decoder = MLXTTSFrameDecoder() + XCTAssertTrue(try decoder.append(combined.prefix(split)).isEmpty) + let payloads = try decoder.append(combined.suffix(from: split)) + + XCTAssertEqual(payloads.count, 2) + XCTAssertEqual( + try MLXTTSFrameCodec.decode(MLXTTSRequest.self, payload: payloads[0]), + .cancel(id: "one")) + XCTAssertEqual( + try MLXTTSFrameCodec.decode(MLXTTSRequest.self, payload: payloads[1]), + .shutdown) + } + + func testAudioEventCarriesExplicitPCMFormat() throws { + let event = MLXTTSEvent.audio(MLXTTSAudio( + id: "request-2", + sampleRate: 32000, + pcm: Data([0x00, 0x00, 0xFF, 0x7F]))) + let frame = try MLXTTSFrameCodec.encode(event) + + var decoder = MLXTTSFrameDecoder() + let payload = try XCTUnwrap(decoder.append(frame).first) + let decoded = try MLXTTSFrameCodec.decode(MLXTTSEvent.self, payload: payload) + + XCTAssertEqual(decoded, event) + } + + func testDecoderRejectsEmptyAndOversizedFrames() { + var emptyDecoder = MLXTTSFrameDecoder() + XCTAssertThrowsError(try emptyDecoder.append(Data(repeating: 0, count: 4))) { error in + XCTAssertEqual(error as? MLXTTSFrameError, .emptyFrame) + } + + let oversized = UInt32(MLXTTSFrameCodec.maximumPayloadSize + 1).bigEndian + var length = oversized + let header = Data(bytes: &length, count: MemoryLayout.size) + var oversizedDecoder = MLXTTSFrameDecoder() + XCTAssertThrowsError(try oversizedDecoder.append(header)) { error in + XCTAssertEqual( + error as? MLXTTSFrameError, + .frameTooLarge(MLXTTSFrameCodec.maximumPayloadSize + 1)) + } + } + + func testCodecRejectsMalformedJSONPayload() { + XCTAssertThrowsError( + try MLXTTSFrameCodec.decode(MLXTTSRequest.self, payload: Data("not-json".utf8))) + } +} diff --git a/scripts/format-swift.sh b/scripts/format-swift.sh index 2bdc497d68b..39177a4d4b0 100755 --- a/scripts/format-swift.sh +++ b/scripts/format-swift.sh @@ -17,7 +17,10 @@ if [[ "$scope" != "ios" ]]; then swiftformat --lint apps/macos/Sources \ --config config/swiftformat \ --exclude '**/OpenClawProtocol' - swiftformat --lint apps/macos-mlx-tts/Sources apps/swabble/Sources \ + swiftformat --lint \ + apps/macos-mlx-tts/Sources \ + apps/shared/OpenClawMLXTTSProtocol/Sources \ + apps/swabble/Sources \ --config config/swiftformat fi diff --git a/test/scripts/codesign-mac-app.test.ts b/test/scripts/codesign-mac-app.test.ts index 51c91d97401..47ec8536188 100644 --- a/test/scripts/codesign-mac-app.test.ts +++ b/test/scripts/codesign-mac-app.test.ts @@ -166,6 +166,7 @@ describe("codesign-mac-app temp file hygiene", () => { mkdirSync(path.join(app, "Contents", "MacOS"), { recursive: true }); mkdirSync(binDir); mkdirSync(captureDir); + writeFileSync(path.join(app, "Contents", "MacOS", "openclaw-mlx-tts"), "#!/bin/sh\n"); writeFileSync(path.join(app, "Contents", "MacOS", "OpenClaw"), "#!/bin/sh\n"); installFakeCodesign(binDir); @@ -187,9 +188,12 @@ describe("codesign-mac-app temp file hygiene", () => { expect(result.stdout).toContain(`Codesign complete for ${app}`); const signLines = readFileSync(logPath, "utf8").trim().split("\n"); - expect(signLines).toHaveLength(2); - expect(signLines[0]).toContain(`${path.join(app, "Contents", "MacOS", "OpenClaw")}\t`); - expect(signLines[1]).toContain(`${app}\t`); + expect(signLines).toHaveLength(3); + expect(signLines[0]).toContain( + `${path.join(app, "Contents", "MacOS", "openclaw-mlx-tts")}\t`, + ); + expect(signLines[1]).toContain(`${path.join(app, "Contents", "MacOS", "OpenClaw")}\t`); + expect(signLines[2]).toContain(`${app}\t`); for (const line of signLines) { const [, , entitlementPath, copiedEntitlementsPath] = line.split("\t"); const copiedEntitlements = readFileSync(copiedEntitlementsPath, "utf8"); diff --git a/test/scripts/package-mac-app.test.ts b/test/scripts/package-mac-app.test.ts index ff9c39d3757..5bd212019c5 100644 --- a/test/scripts/package-mac-app.test.ts +++ b/test/scripts/package-mac-app.test.ts @@ -331,6 +331,31 @@ describe("package-mac-app plist stamping", () => { expect(installBlock).not.toContain("--no-frozen-lockfile"); }); + it("builds and bundles the MLX TTS helper for every requested architecture", () => { + const script = readFileSync(scriptPath, "utf8"); + const buildLoop = script.slice( + script.indexOf('for arch in "${BUILD_ARCHS[@]}"; do'), + script.indexOf('BIN_PRIMARY="$(bin_for_arch "$PRIMARY_ARCH")"'), + ); + const helperCopy = script.slice( + script.indexOf('echo "🚚 Copying MLX TTS helper"'), + script.indexOf('SPARKLE_FRAMEWORK_PRIMARY='), + ); + + expect(buildLoop).toContain( + 'swift build --package-path "$MLX_TTS_HELPER_ROOT"', + ); + expect(buildLoop).toContain('--product "$MLX_TTS_HELPER_PRODUCT"'); + expect(buildLoop).toContain('--arch "$arch"'); + expect(helperCopy).toContain( + 'cp "$(helper_bin_for_arch "$PRIMARY_ARCH")" "$APP_ROOT/Contents/MacOS/$MLX_TTS_HELPER_PRODUCT"', + ); + expect(helperCopy).toContain('/usr/bin/lipo -create "${HELPER_BIN_INPUTS[@]}"'); + expect(helperCopy).toContain( + 'chmod +x "$APP_ROOT/Contents/MacOS/$MLX_TTS_HELPER_PRODUCT"', + ); + }); + it("falls back to corepack pnpm when the pnpm shim is absent", () => { const helperBlock = getPackageManagerHelperBlock(); const tempRoot = mkdtempSync(path.join(tmpdir(), "openclaw-package-pnpm-root-"));