diff --git a/Sources/Writing/NextSentence/NextSentenceClaude.swift b/Sources/Writing/NextSentence/NextSentenceClaude.swift index 791531db0..fff4e12a6 100644 --- a/Sources/Writing/NextSentence/NextSentenceClaude.swift +++ b/Sources/Writing/NextSentence/NextSentenceClaude.swift @@ -42,8 +42,15 @@ final class NextSentenceClaude: @unchecked Sendable { } private let pool: NextSentenceWarmPool + private let scheduleTimeout: (TimeInterval, NextSentenceCallbackBox) -> Void - init() { + init( + spawn: ((Kind, String?) -> NextSentenceLineProcess?)? = nil, + background: @escaping (@escaping @Sendable () -> Void) -> Void = { DispatchQueue.global(qos: .utility).async(execute: $0) }, + scheduleTimeout: @escaping (TimeInterval, NextSentenceCallbackBox) -> Void = { delay, callback in + DispatchQueue.global().asyncAfter(deadline: .now() + delay) { callback.call() } + } + ) { var json: [Kind: String] = [:] for kind in Kind.allCases { guard let schema = kind.schema, @@ -51,7 +58,11 @@ final class NextSentenceClaude: @unchecked Sendable { json[kind] = String(decoding: data, as: UTF8.self) } let schemas = json - pool = NextSentenceWarmPool(spawn: { Self.spawn($0, schema: schemas[$0]) }) + self.scheduleTimeout = scheduleTimeout + pool = NextSentenceWarmPool(spawn: { kind in + if let spawn { return spawn(kind, schemas[kind]) } + return Self.spawn(kind, schema: schemas[kind]) + }, background: background) } static func binary(userDefaults: UserDefaults = .standard) -> URL? { @@ -79,9 +90,11 @@ final class NextSentenceClaude: @unchecked Sendable { let finish: (Result) -> Void = { result in guard once.claim() else { return } process.onLine(nil) + process.onOutputEnd(nil) self.pool.finished(process) completion(result) } + let ended = { finish(.failure(Failure.failed("Claude stopped before returning an answer."))) } process.onLine { line in switch line["type"] as? String { case "stream_event": @@ -99,7 +112,9 @@ final class NextSentenceClaude: @unchecked Sendable { default: break } } - send(prompt: prompt, image: image, to: process, timeout: 60, onTimeout: { finish(.failure(Failure.failed("Claude took too long."))) }) + process.onOutputEnd(ended) + send(prompt: prompt, image: image, to: process, timeout: 60, onSendFailure: ended, + onTimeout: { finish(.failure(Failure.failed("Claude took too long."))) }) } /// A JSON answer matching `kind`'s schema. `completion` runs on a background queue. @@ -110,9 +125,11 @@ final class NextSentenceClaude: @unchecked Sendable { let finish: (Result<[String: Any], Error>) -> Void = { result in guard once.claim() else { return } process.onLine(nil) + process.onOutputEnd(nil) self.pool.finished(process) completion(result) } + let ended = { finish(.failure(Failure.failed("Claude stopped before returning an answer."))) } process.onLine { line in // Structured output arrives as a StructuredOutput tool call, then a result line. if line["type"] as? String == "assistant", @@ -127,7 +144,9 @@ final class NextSentenceClaude: @unchecked Sendable { finish(.failure(Failure.failed("Claude didn't return an answer."))) } } - send(prompt: prompt, image: image, to: process, timeout: timeout, onTimeout: { finish(.failure(Failure.failed("Claude took too long."))) }) + process.onOutputEnd(ended) + send(prompt: prompt, image: image, to: process, timeout: timeout, onSendFailure: ended, + onTimeout: { finish(.failure(Failure.failed("Claude took too long."))) }) } // MARK: Processes @@ -150,7 +169,8 @@ final class NextSentenceClaude: @unchecked Sendable { return process } - private func send(prompt: String, image: Data?, to process: NextSentenceLineProcess, timeout: TimeInterval, onTimeout: @escaping () -> Void) { + private func send(prompt: String, image: Data?, to process: NextSentenceLineProcess, timeout: TimeInterval, + onSendFailure: () -> Void, onTimeout: @escaping () -> Void) { var content: [[String: Any]] = [] if let image { content.append(["type": "image", "source": ["type": "base64", "media_type": "image/jpeg", "data": image.base64EncodedString()]]) @@ -159,10 +179,10 @@ final class NextSentenceClaude: @unchecked Sendable { do { try process.send(["type": "user", "message": ["role": "user", "content": content]]) } catch { - return onTimeout() + return onSendFailure() } let timeoutBox = NextSentenceCallbackBox(onTimeout) - DispatchQueue.global().asyncAfter(deadline: .now() + timeout) { timeoutBox.call() } + scheduleTimeout(timeout, timeoutBox) } private static func failure(_ line: [String: Any], _ process: NextSentenceLineProcess) -> Failure { diff --git a/Sources/Writing/NextSentence/NextSentenceLineProcess.swift b/Sources/Writing/NextSentence/NextSentenceLineProcess.swift index 8ee9af798..e9ac1311d 100644 --- a/Sources/Writing/NextSentence/NextSentenceLineProcess.swift +++ b/Sources/Writing/NextSentence/NextSentenceLineProcess.swift @@ -13,6 +13,12 @@ final class NextSentenceLineProcess: @unchecked Sendable { private var buffer = Data() private var errorTail = Data() private var lineHandler: (([String: Any]) -> Void)? + // The following state belongs to queue, alongside the stdout buffer. + private var outputEnded = false + private var endHandler: EndCallback? + private var endWaiters: [CheckedContinuation] = [] + private var exited = false + private var exitWaiters: [CheckedContinuation] = [] init(executable: URL, arguments: [String], environment: [String: String] = [:], workingDirectory: URL? = nil, label: String) { queue = DispatchQueue(label: "transcripted.next-sentence.\(label)") @@ -28,8 +34,18 @@ final class NextSentenceLineProcess: @unchecked Sendable { var isRunning: Bool { process.isRunning } func start() throws { + process.terminationHandler = { [weak self] _ in + guard let self else { return } + self.queue.async { + self.exited = true + let waiters = self.exitWaiters + self.exitWaiters.removeAll() + for waiter in waiters { waiter.resume() } + } + } output.fileHandleForReading.readabilityHandler = { [weak self] handle in let data = handle.availableData + if data.isEmpty { handle.readabilityHandler = nil } guard let self else { return } self.queue.async { self.receive(data) } } @@ -49,6 +65,37 @@ final class NextSentenceLineProcess: @unchecked Sendable { lock.lock(); lineHandler = handler; lock.unlock() } + /// Runs after stdout's final JSON frame, including when the child ended + /// before this handler was installed. No child output is included. + func onOutputEnd(_ handler: (() -> Void)?) { + let callback = handler.map(EndCallback.init) + queue.async { + self.endHandler = callback + if self.outputEnded { callback?.call() } + } + } + + /// Waits until buffered output and its end callback have been delivered. + func waitForOutputEnd() async { + await withCheckedContinuation { continuation in + queue.async { + if self.outputEnded { continuation.resume() } + else { self.endWaiters.append(continuation) } + } + } + } + + /// Waits for a started child to exit, then drains its output callbacks. + func waitForExit() async { + await withCheckedContinuation { continuation in + queue.async { + if self.exited { continuation.resume() } + else { self.exitWaiters.append(continuation) } + } + } + await waitForOutputEnd() + } + func send(_ object: [String: Any]) throws { let data = try JSONSerialization.data(withJSONObject: object) try input.fileHandleForWriting.write(contentsOf: data + Data([0x0A])) @@ -63,17 +110,44 @@ final class NextSentenceLineProcess: @unchecked Sendable { output.fileHandleForReading.readabilityHandler = nil errors.fileHandleForReading.readabilityHandler = nil if process.isRunning { process.terminate() } + queue.async { self.finishOutput() } } private func receive(_ data: Data) { - guard !data.isEmpty else { return } + guard !outputEnded else { return } + if data.isEmpty { + // A final complete JSON object need not have a trailing newline. + if !buffer.isEmpty { emit(buffer); buffer.removeAll() } + finishOutput() + return + } buffer.append(data) while let newline = buffer.firstIndex(of: 0x0A) { let line = buffer[buffer.startIndex.. Void + init(_ handler: @escaping () -> Void) { self.handler = handler } + func call() { handler() } + } } diff --git a/Tests/NextSentenceExitTests.swift b/Tests/NextSentenceExitTests.swift new file mode 100644 index 000000000..1b0ecc346 --- /dev/null +++ b/Tests/NextSentenceExitTests.swift @@ -0,0 +1,161 @@ +import Foundation + +func testNextSentenceExit() async { + for structured in [false, true] { + await runSuite(structured ? "Structured Claude child exit completes before timeout" : "Streaming Claude child exit completes before timeout") { + let fixture = NextSentenceExitFixture(script: "IFS= read -r request\nexit 7\n") + let result = NextSentenceExitResult() + assertEqual(fixture.started, false, "constructing the client must not start a child") + if structured { + fixture.client.structured(.alternatives, prompt: "synthetic request", image: nil) { + result.record($0.map { _ in true }) + } + } else { + fixture.client.stream(prompt: "synthetic request", image: nil, onText: { _ in }) { + result.record($0.map { _ in true }) + } + } + await fixture.waitForExit() + assertEqual(result.outcomes, ["ended"], "exit must finish the request before the timeout callback") + fixture.timeouts.fire() + assertEqual(result.outcomes, ["ended"], "the later timeout must not complete the request again") + await fixture.cleanup() + } + } + + for structured in [false, true] { + await runSuite(structured ? "Structured final JSON wins over EOF without a trailing newline" : "Streaming final result wins over EOF") { + let payload = structured + ? #"{"type":"result","is_error":false,"structured_output":{"answer":"synthetic"}}"# + : #"{"type":"stream_event","event":{"type":"content_block_delta","delta":{"type":"text_delta","text":"synthetic"}}}"# + + "\n" + #"{"type":"result","is_error":false}"# + "\n" + let fixture = NextSentenceExitFixture(script: "IFS= read -r request\nprintf '%s' '\(payload)'\n") + let result = NextSentenceExitResult() + if structured { + fixture.client.structured(.alternatives, prompt: "synthetic request", image: nil) { + result.record($0.map { $0["answer"] as? String == "synthetic" }) + } + } else { + fixture.client.stream(prompt: "synthetic request", image: nil, onText: { _ in }) { + result.record($0.map { $0 == "synthetic" }) + } + } + await fixture.waitForExit() + assertEqual(result.outcomes, ["success"], "the final parsed answer must win before output ends") + fixture.timeouts.fire() + assertEqual(result.outcomes, ["success"], "EOF and timeout must not replace a successful answer") + await fixture.cleanup() + } + } + + await runSuite("A late output-end observer still sees the exited child") { + let fixture = NextSentenceExitFixture(script: "exit 0\n") + let process = fixture.start() + await fixture.waitForExit() + let seen = NextSentenceExitResult() + process?.onOutputEnd { seen.markEnd() } + await process?.waitForOutputEnd() + assertEqual(seen.outcomes, ["ended"], "output closure must be replayed after handler registration") + await fixture.cleanup() + } +} + +private final class NextSentenceExitResult: @unchecked Sendable { + private let lock = NSLock() + private var values: [String] = [] + var outcomes: [String] { lock.lock(); defer { lock.unlock() }; return values } + func markEnd() { lock.lock(); values.append("ended"); lock.unlock() } + func record(_ result: Result) { + let outcome: String + switch result { + case .success(true): outcome = "success" + case .success(false): outcome = "unexpected_success" + case .failure(let error): + if case NextSentenceClaude.Failure.failed(let message) = error, + message == "Claude stopped before returning an answer." { outcome = "ended" } + else { outcome = "other_failure" } + } + lock.lock(); values.append(outcome); lock.unlock() + } +} + +private final class NextSentenceExitTimeouts: @unchecked Sendable { + private let lock = NSLock() + private var callbacks: [NextSentenceCallbackBox] = [] + func store(_ callback: NextSentenceCallbackBox) { lock.lock(); callbacks.append(callback); lock.unlock() } + func fire() { + lock.lock(); let pending = callbacks; callbacks.removeAll(); lock.unlock() + for callback in pending { callback.call() } + } +} + +private final class NextSentenceExitFixture { + let root = FileManager.default.temporaryDirectory + .appendingPathComponent("NextSentenceExit-\(UUID().uuidString)", isDirectory: true) + let timeouts = NextSentenceExitTimeouts() + private let process: NextSentenceLineProcess + private let startState = NextSentenceExitStartState() + private let prepared: Bool + let client: NextSentenceClaude + var started: Bool { startState.started } + + init(script: String) { + let child = root.appendingPathComponent("child.sh") + var ready = false + do { + try FileManager.default.createDirectory(at: root, withIntermediateDirectories: true) + try script.write(to: child, atomically: true, encoding: .utf8) + ready = true + } catch { recordFailure("could not prepare the owned synthetic child") } + prepared = ready + let isPrepared = ready + let process = NextSentenceLineProcess( + executable: URL(fileURLWithPath: "/bin/sh"), arguments: [child.path], + environment: ["HOME": root.path, "PATH": "/usr/bin:/bin", "ENV": "/dev/null", "BASH_ENV": "/dev/null"], + workingDirectory: root, label: "synthetic-exit" + ) + self.process = process + let state = startState + let timeouts = timeouts + client = NextSentenceClaude(spawn: { _, _ in + guard isPrepared else { return nil } + do { try process.start(); state.didStart(); return process } + catch { return nil } + }, background: { _ in }, scheduleTimeout: { _, callback in timeouts.store(callback) }) + } + + func start() -> NextSentenceLineProcess? { + guard prepared else { return nil } + do { try process.start(); startState.didStart(); return process } + catch { recordFailure("could not start the owned synthetic child"); return nil } + } + + func waitForExit() async { + guard started else { recordFailure("the synthetic child was not started"); return } + await process.waitForExit() + } + + func cleanup() async { + client.stop() + if started { process.terminate(); await process.waitForExit() } + let resolved = root.resolvingSymlinksInPath() + let parent = FileManager.default.temporaryDirectory.resolvingSymlinksInPath() + let prefix = "NextSentenceExit-" + guard root.lastPathComponent.hasPrefix(prefix), + UUID(uuidString: String(root.lastPathComponent.dropFirst(prefix.count))) != nil, + resolved.deletingLastPathComponent() == parent, + resolved.lastPathComponent == root.lastPathComponent else { + return recordFailure("refusing to remove a synthetic child outside its owned root") + } + guard FileManager.default.fileExists(atPath: root.path) else { return } + do { try FileManager.default.removeItem(at: root) } + catch { recordFailure("synthetic child fixture cleanup failed") } + } +} + +private final class NextSentenceExitStartState: @unchecked Sendable { + private let lock = NSLock() + private var value = false + var started: Bool { lock.lock(); defer { lock.unlock() }; return value } + func didStart() { lock.lock(); value = true; lock.unlock() } +} diff --git a/scripts/entrypoints/run-tests.sh b/scripts/entrypoints/run-tests.sh index e4c366a4e..7fa74c949 100755 --- a/scripts/entrypoints/run-tests.sh +++ b/scripts/entrypoints/run-tests.sh @@ -485,6 +485,8 @@ APP_SOURCES=( "Sources/Writing/NextSentence/NextSentencePrivacy.swift" "Sources/Writing/NextSentence/NextSentenceWarmPool.swift" "Sources/Writing/NextSentence/NextSentenceTetheredProcess.swift" + "Sources/Writing/NextSentence/NextSentenceLineProcess.swift" + "Sources/Writing/NextSentence/NextSentenceClaude.swift" "Sources/Writing/NextSentence/NextSentenceSemanticIndex.swift" "Sources/TranscriptedWriting/Core/ScreenMemory/DefaultExcludedApps.swift" "Sources/TranscriptedWriting/Core/Text/SecretRules.swift"