diff --git a/Sources/DecartSDK/Realtime/DecartRealtimeManager.swift b/Sources/DecartSDK/Realtime/DecartRealtimeManager.swift index cf7c036..fda617e 100644 --- a/Sources/DecartSDK/Realtime/DecartRealtimeManager.swift +++ b/Sources/DecartSDK/Realtime/DecartRealtimeManager.swift @@ -1,7 +1,7 @@ import Foundation private struct InitialStateRequest: Sendable { - let message: InitialStateMessage + let message: OutgoingWebSocketMessage let ackTarget: InitialStateAckTarget } @@ -61,6 +61,7 @@ public final class DecartRealtimeManager: @unchecked Sendable { private var mediaConnectionStateTask: Task? private var mediaDisconnectTask: Task? private var connectionQualityTask: Task? + private var initialStateAckTask: Task? private var reconnectTask: Task? private let initialStateAckTimeout: TimeInterval = 30 private let promptAckTimeout: TimeInterval = 15 @@ -152,6 +153,7 @@ public final class DecartRealtimeManager: @unchecked Sendable { mediaConnectionStateTask?.cancel() mediaDisconnectTask?.cancel() connectionQualityTask?.cancel() + initialStateAckTask?.cancel() reconnectTask?.cancel() let liveKitMediaChannel = liveKitMediaChannel let webSocketClient = webSocketClient @@ -289,10 +291,12 @@ private extension DecartRealtimeManager { } let initialStateRequest = buildInitialStateRequest() - try await sendMessageThrowing(.liveKitJoin( - initialState: options.connection.bundleInitialStateInJoin ? initialStateRequest?.message : nil, - encodesInitialState: options.connection.bundleInitialStateInJoin - )) + + // Buffer any ack that lands before the out-of-band observer starts. + isWaitingForInitialStateAck = true + try await sendMessageThrowing(.liveKitJoin(passthrough: isPassthrough)) + try await sendMessageThrowing(initialStateRequest.message) + let roomInfo = try await waitForLiveKitRoomInfo(timeout: options.connection.connectionTimeout) let mediaChannel = LiveKitMediaChannel( @@ -304,20 +308,11 @@ private extension DecartRealtimeManager { liveKitMediaChannel = mediaChannel setupMediaListeners(mediaChannel) - let prearmedInitialStateAck = options.connection.bundleInitialStateInJoin && initialStateRequest != nil - if prearmedInitialStateAck { - isWaitingForInitialStateAck = true - } - defer { - if prearmedInitialStateAck { - isWaitingForInitialStateAck = false - clearPendingInitialState() - } - } - - async let initialStateAck: Void = handleInitialStateAfterRoomInfo(initialStateRequest) + // Watch the ack off the critical path: arm its timeout now that room_info + // arrived (a long queue wait can't trip it) and let it run concurrently + // with room connect + publish, surfacing an error only on rejection. + observeInitialStateAck(initialStateRequest) try await mediaChannel.connect(roomInfo: roomInfo) - try await initialStateAck try await mediaChannel.publishLocalTracks(from: localStream) return mediaChannel.currentRemoteStream } @@ -327,6 +322,9 @@ private extension DecartRealtimeManager { // fail runtime waiters here. failAllPendingRuntimeWaiters(DecartError.websocketError("WebSocket disconnected")) clearPendingInitialState() + initialStateAckTask?.cancel() + initialStateAckTask = nil + isWaitingForInitialStateAck = false webSocketListenerTask?.cancel() webSocketListenerTask = nil mediaListenerTask?.cancel() @@ -513,11 +511,6 @@ private extension DecartRealtimeManager { // MARK: - Messaging private extension DecartRealtimeManager { - private func sendMessage(_ message: OutgoingWebSocketMessage) { - guard let webSocketClient else { return } - Task { [webSocketClient] in try? await webSocketClient.send(message) } - } - private func sendMessageThrowing(_ message: OutgoingWebSocketMessage) async throws { guard let webSocketClient else { throw DecartError.websocketError("WebSocket not connected") @@ -525,18 +518,26 @@ private extension DecartRealtimeManager { try await webSocketClient.send(message) } - func buildInitialStateRequest() -> InitialStateRequest? { + // `false` when the user set a real image/prompt (one real frame follows), + // `true` otherwise (a null-bootstrap frame follows). Session config may override. + var isPassthrough: Bool { + if let override = options.connection.passthrough { return override } + let initialPrompt = options.initialPrompt + let hasImage = options.model.hasReferenceImage && initialPrompt.referenceImageData != nil + return !(hasImage || !initialPrompt.text.isEmpty) + } + + func buildInitialStateRequest() -> InitialStateRequest { let initialPrompt = options.initialPrompt if options.model.hasReferenceImage, let base64Image = initialPrompt.referenceImageData?.base64EncodedString() { - let message = SetImageMessage( - imageData: base64Image, - prompt: initialPrompt.text, - enhancePrompt: initialPrompt.enrich - ) return InitialStateRequest( - message: .setImage(message), + message: .setImage(SetImageMessage( + imageData: base64Image, + prompt: initialPrompt.text, + enhancePrompt: initialPrompt.enrich + )), ackTarget: .setImage( failureMessage: "Failed to set initial image", timeoutMessage: "Initial image acknowledgment timed out" @@ -544,56 +545,46 @@ private extension DecartRealtimeManager { ) } - guard !initialPrompt.text.isEmpty else { return nil } - return InitialStateRequest( - message: .prompt(PromptMessage(prompt: initialPrompt.text, enhancePrompt: initialPrompt.enrich)), - ackTarget: .prompt(initialPrompt.text) - ) - } - - func handleInitialStateAfterRoomInfo(_ request: InitialStateRequest?) async throws { - if options.connection.bundleInitialStateInJoin { - try await waitForBundledInitialStateAck(request) - } else { - try await sendInitialState() - } - } - - func waitForBundledInitialStateAck(_ request: InitialStateRequest?) async throws { - guard let request else { return } - - isWaitingForInitialStateAck = true - defer { - isWaitingForInitialStateAck = false - clearPendingInitialState() + if !initialPrompt.text.isEmpty { + return InitialStateRequest( + message: .prompt(PromptMessage(prompt: initialPrompt.text, enhancePrompt: initialPrompt.enrich)), + ackTarget: .prompt(initialPrompt.text) + ) } - switch request.ackTarget { - case .prompt(let prompt): - try await waitForPromptAck(prompt: prompt, timeout: initialStateAckTimeout) - case .setImage(let failureMessage, let timeoutMessage): - try await waitForSetImageAck( - timeout: initialStateAckTimeout, - failureMessage: failureMessage, - timeoutMessage: timeoutMessage + return InitialStateRequest( + message: .setImage(.passthrough()), + ackTarget: .setImage( + failureMessage: "Failed to apply initial passthrough state", + timeoutMessage: "Initial passthrough acknowledgment timed out" ) - } + ) } - func sendInitialState() async throws { - let initialPrompt = options.initialPrompt - if options.model.hasReferenceImage, - let base64Image = initialPrompt.referenceImageData?.base64EncodedString() - { - try await sendInitialImageAndWait( - base64Image, - prompt: initialPrompt.text, - enhance: initialPrompt.enrich - ) - } else if !initialPrompt.text.isEmpty { - try await sendInitialPromptAndWait(initialPrompt) - } else { - try await sendPassthroughAndWait() + func observeInitialStateAck(_ request: InitialStateRequest) { + initialStateAckTask?.cancel() + initialStateAckTask = Task { [weak self] in + guard let self else { return } + defer { + self.isWaitingForInitialStateAck = false + self.clearPendingInitialState() + } + do { + switch request.ackTarget { + case .prompt(let prompt): + try await self.waitForPromptAck(prompt: prompt, timeout: self.initialStateAckTimeout) + case .setImage(let failureMessage, let timeoutMessage): + try await self.waitForSetImageAck( + timeout: self.initialStateAckTimeout, + failureMessage: failureMessage, + timeoutMessage: timeoutMessage + ) + } + } catch is CancellationError { + return + } catch { + DecartLogger.log("Initial-state acknowledgment failed: \(error.localizedDescription)", level: .error) + } } } @@ -623,77 +614,6 @@ private extension DecartRealtimeManager { } } - func sendInitialPromptAndWait(_ prompt: DecartPrompt) async throws { - guard let webSocketClient else { - connectionState = .error - throw DecartError.websocketError("WebSocket not connected") - } - - clearPendingInitialState() - isWaitingForInitialStateAck = true - defer { - isWaitingForInitialStateAck = false - clearPendingInitialState() - } - - let message: OutgoingWebSocketMessage = .prompt(PromptMessage(prompt: prompt.text, enhancePrompt: prompt.enrich)) - try await webSocketClient.send(message) - try await waitForPromptAck(prompt: prompt.text, timeout: initialStateAckTimeout) - } - - func sendInitialImageAndWait( - _ imageBase64: String, - prompt: String, - enhance: Bool - ) async throws { - guard let webSocketClient else { - connectionState = .error - throw DecartError.websocketError("WebSocket not connected") - } - - clearPendingInitialState() - isWaitingForInitialStateAck = true - defer { - isWaitingForInitialStateAck = false - clearPendingInitialState() - } - - let message = SetImageMessage( - imageData: imageBase64, - prompt: prompt, - enhancePrompt: enhance - ) - let outgoing: OutgoingWebSocketMessage = .setImage(message) - try await webSocketClient.send(outgoing) - try await waitForSetImageAck( - timeout: initialStateAckTimeout, - failureMessage: "Failed to set initial image", - timeoutMessage: "Initial image acknowledgment timed out" - ) - } - - func sendPassthroughAndWait() async throws { - guard let webSocketClient else { - connectionState = .error - throw DecartError.websocketError("WebSocket not connected") - } - - clearPendingInitialState() - isWaitingForInitialStateAck = true - defer { - isWaitingForInitialStateAck = false - clearPendingInitialState() - } - - let passthrough: OutgoingWebSocketMessage = .setImage(.passthrough()) - try await webSocketClient.send(passthrough) - try await waitForSetImageAck( - timeout: initialStateAckTimeout, - failureMessage: "Failed to apply initial passthrough state", - timeoutMessage: "Initial passthrough acknowledgment timed out" - ) - } - func waitForPromptAck(prompt: String, timeout: TimeInterval) async throws { let startTime = Date() while true { @@ -982,6 +902,7 @@ extension DecartRealtimeManager { internal func test_setConnectionState(_ state: DecartRealtimeConnectionState) { connectionState = state } + internal var test_isPassthrough: Bool { isPassthrough } internal var test_hasPendingInitialStateAck: Bool { !pendingPromptAcks.isEmpty || !pendingSetImageAcks.isEmpty || pendingInitialStateError != nil } diff --git a/Sources/DecartSDK/Realtime/RealtimeConfiguration.swift b/Sources/DecartSDK/Realtime/RealtimeConfiguration.swift index 0c20824..79547a7 100644 --- a/Sources/DecartSDK/Realtime/RealtimeConfiguration.swift +++ b/Sources/DecartSDK/Realtime/RealtimeConfiguration.swift @@ -51,16 +51,18 @@ public struct RealtimeConfiguration: Sendable { public struct ConnectionConfig: Sendable { public let connectionTimeout: TimeInterval public let reconnectAttempts: Int - public let bundleInitialStateInJoin: Bool + /// Overrides the `passthrough` flag sent on the join. `nil` derives it from + /// the initial prompt/image: `false` when a real reference is set, else `true`. + public let passthrough: Bool? public init( connectionTimeout: TimeInterval = 15, reconnectAttempts: Int = 10, - bundleInitialStateInJoin: Bool = true + passthrough: Bool? = nil ) { self.connectionTimeout = connectionTimeout self.reconnectAttempts = reconnectAttempts - self.bundleInitialStateInJoin = bundleInitialStateInJoin + self.passthrough = passthrough } var connectOptions: ConnectOptions { diff --git a/Sources/DecartSDK/Realtime/Transport/WebSocket/SignalingModel.swift b/Sources/DecartSDK/Realtime/Transport/WebSocket/SignalingModel.swift index cadbe39..f6230d0 100644 --- a/Sources/DecartSDK/Realtime/Transport/WebSocket/SignalingModel.swift +++ b/Sources/DecartSDK/Realtime/Transport/WebSocket/SignalingModel.swift @@ -15,46 +15,13 @@ struct InitializeConnectionMessage: Codable, Sendable { let initialPrompt: String? } -enum InitialStateMessage: Encodable, Sendable { - case prompt(PromptMessage) - case setImage(SetImageMessage) - - func encode(to encoder: Encoder) throws { - switch self { - case .prompt(let message): - try message.encode(to: encoder) - case .setImage(let message): - try message.encode(to: encoder) - } - } -} - struct LiveKitJoinMessage: Encodable, Sendable { let type: String - let initialState: InitialStateMessage? - let encodesInitialState: Bool + let passthrough: Bool - init(initialState: InitialStateMessage?, encodesInitialState: Bool) { + init(passthrough: Bool) { self.type = "livekit_join" - self.initialState = initialState - self.encodesInitialState = encodesInitialState - } - - private enum CodingKeys: String, CodingKey { - case type - case initialState = "initial_state" - } - - func encode(to encoder: Encoder) throws { - var container = encoder.container(keyedBy: CodingKeys.self) - try container.encode(type, forKey: .type) - - guard encodesInitialState else { return } - if let initialState { - try container.encode(initialState, forKey: .initialState) - } else { - try container.encodeNil(forKey: .initialState) - } + self.passthrough = passthrough } } @@ -289,17 +256,14 @@ enum IncomingWebSocketMessage: Codable, Sendable { } enum OutgoingWebSocketMessage: Encodable, Sendable { - case liveKitJoin(initialState: InitialStateMessage?, encodesInitialState: Bool) + case liveKitJoin(passthrough: Bool) case prompt(PromptMessage) case setImage(SetImageMessage) func encode(to encoder: Encoder) throws { switch self { - case .liveKitJoin(let initialState, let encodesInitialState): - try LiveKitJoinMessage( - initialState: initialState, - encodesInitialState: encodesInitialState - ).encode(to: encoder) + case .liveKitJoin(let passthrough): + try LiveKitJoinMessage(passthrough: passthrough).encode(to: encoder) case .prompt(let msg): try msg.encode(to: encoder) case .setImage(let msg): diff --git a/Tests/DecartSDKTests/SignalingModelTests.swift b/Tests/DecartSDKTests/SignalingModelTests.swift index 4dda6e1..a26394e 100644 --- a/Tests/DecartSDKTests/SignalingModelTests.swift +++ b/Tests/DecartSDKTests/SignalingModelTests.swift @@ -2,45 +2,76 @@ import XCTest @testable import DecartSDK final class SignalingModelTests: XCTestCase { - func testEncodesLiveKitJoinMessageWithNullInitialState() throws { - let data = try JSONEncoder().encode(OutgoingWebSocketMessage.liveKitJoin( - initialState: nil, - encodesInitialState: true - )) + func testEncodesLeanLiveKitJoinWithPassthroughFalse() throws { + let data = try JSONEncoder().encode(OutgoingWebSocketMessage.liveKitJoin(passthrough: false)) let json = try XCTUnwrap(JSONSerialization.jsonObject(with: data) as? [String: Any]) XCTAssertEqual(json["type"] as? String, "livekit_join") - XCTAssertTrue(json["initial_state"] is NSNull) + XCTAssertEqual(json["passthrough"] as? Bool, false) + XCTAssertNil(json["initial_state"], "join must be lean — no nested initial state") } - func testEncodesLiveKitJoinMessageWithBundledInitialState() throws { - let data = try JSONEncoder().encode(OutgoingWebSocketMessage.liveKitJoin( - initialState: .setImage(SetImageMessage( - imageData: "base64-image", - prompt: "wear the jacket", - enhancePrompt: true - )), - encodesInitialState: true - )) + func testEncodesLeanLiveKitJoinWithPassthroughTrue() throws { + let data = try JSONEncoder().encode(OutgoingWebSocketMessage.liveKitJoin(passthrough: true)) let json = try XCTUnwrap(JSONSerialization.jsonObject(with: data) as? [String: Any]) - let initialState = try XCTUnwrap(json["initial_state"] as? [String: Any]) XCTAssertEqual(json["type"] as? String, "livekit_join") - XCTAssertEqual(initialState["type"] as? String, "set_image") - XCTAssertEqual(initialState["image_data"] as? String, "base64-image") - XCTAssertEqual(initialState["prompt"] as? String, "wear the jacket") - XCTAssertEqual(initialState["enhance_prompt"] as? Bool, true) + XCTAssertEqual(json["passthrough"] as? Bool, true) + XCTAssertEqual(json.count, 2, "join carries only type + passthrough") } - func testEncodesLegacyLiveKitJoinMessageWithoutInitialStateField() throws { - let data = try JSONEncoder().encode(OutgoingWebSocketMessage.liveKitJoin( - initialState: nil, - encodesInitialState: false - )) - let json = try XCTUnwrap(JSONSerialization.jsonObject(with: data) as? [String: Any]) + // MARK: - passthrough derivation - XCTAssertEqual(json["type"] as? String, "livekit_join") - XCTAssertNil(json["initial_state"]) + private func makeManager( + hasReferenceImage: Bool = false, + initialPrompt: DecartPrompt = .init(text: ""), + passthroughOverride: Bool? = nil + ) -> DecartRealtimeManager { + let model = ModelDefinition( + name: "test-model", + urlPath: "/v1/test", + fps: 24, + width: 512, + height: 512, + hasReferenceImage: hasReferenceImage + ) + return DecartRealtimeManager( + signalingServerURL: URL(string: "wss://example.test")!, + options: RealtimeConfiguration( + model: model, + initialPrompt: initialPrompt, + connection: .init(passthrough: passthroughOverride) + ) + ) + } + + func testPassthroughTrueWhenNoInitialReference() { + XCTAssertTrue(makeManager().test_isPassthrough) + } + + func testPassthroughFalseWhenPromptSet() { + XCTAssertFalse(makeManager(initialPrompt: .init(text: "a city")).test_isPassthrough) + } + + func testPassthroughFalseWhenReferenceImageSet() { + let manager = makeManager( + hasReferenceImage: true, + initialPrompt: .init(text: "", referenceImageData: Data([0x1, 0x2])) + ) + XCTAssertFalse(manager.test_isPassthrough) + } + + func testReferenceImageIgnoredWhenModelHasNoReferenceSupport() { + let manager = makeManager( + hasReferenceImage: false, + initialPrompt: .init(text: "", referenceImageData: Data([0x1, 0x2])) + ) + XCTAssertTrue(manager.test_isPassthrough) + } + + func testExplicitPassthroughOverridesDerivation() { + XCTAssertTrue(makeManager(initialPrompt: .init(text: "a city"), passthroughOverride: true).test_isPassthrough) + XCTAssertFalse(makeManager(passthroughOverride: false).test_isPassthrough) } func testDecodesLiveKitRoomInfoMessage() throws {