Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
213 changes: 67 additions & 146 deletions Sources/DecartSDK/Realtime/DecartRealtimeManager.swift
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
import Foundation

private struct InitialStateRequest: Sendable {
let message: InitialStateMessage
let message: OutgoingWebSocketMessage
let ackTarget: InitialStateAckTarget
}

Expand Down Expand Up @@ -61,6 +61,7 @@ public final class DecartRealtimeManager: @unchecked Sendable {
private var mediaConnectionStateTask: Task<Void, Never>?
private var mediaDisconnectTask: Task<Void, Never>?
private var connectionQualityTask: Task<Void, Never>?
private var initialStateAckTask: Task<Void, Never>?
private var reconnectTask: Task<Void, Never>?
private let initialStateAckTimeout: TimeInterval = 30
private let promptAckTimeout: TimeInterval = 15
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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(
Expand All @@ -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()
}
}

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Runtime prompts lose acks after connect

High Severity

connect can finish while isWaitingForInitialStateAck is still true. Any setPrompt/setImage in that window routes acks into the initial-state buffer; when the observer completes, clearPendingInitialState() drops unmatched acks, so runtime waiters time out or hang even though the server responded.

Additional Locations (2)
Fix in Cursor Fix in Web

Reviewed by Cursor Bugbot for commit e232327. Configure here.

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
}
Expand All @@ -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()
Expand Down Expand Up @@ -513,87 +511,80 @@ 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")
}
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"
)
)
}

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)
}

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Initial ack timeout marks error

Medium Severity

The out-of-band initial-state observer reuses waitForPromptAck/waitForSetImageAck, which set connectionState to .error on timeout. After a successful connect, a late initial-state ack timeout can flip an already-connected session to error, though the PR describes surfacing errors only on server rejection.

Additional Locations (1)
Fix in Cursor Fix in Web

Reviewed by Cursor Bugbot for commit e232327. Configure here.

}
}

Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -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
}
Expand Down
8 changes: 5 additions & 3 deletions Sources/DecartSDK/Realtime/RealtimeConfiguration.swift
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
}

Expand Down Expand Up @@ -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):
Expand Down
Loading
Loading