functions/private/New-KlippyWebSocketClient.ps1
|
function New-KlippyWebSocketClient { <# .SYNOPSIS Creates a new WebSocket client for Moonraker real-time communication. .DESCRIPTION Establishes a WebSocket connection to Moonraker for subscribing to printer object updates and receiving real-time notifications. .PARAMETER Printer The printer object. .PARAMETER ConnectTimeout Connection timeout in seconds. Default is 10. .EXAMPLE $ws = New-KlippyWebSocketClient -Printer $printer $ws.Subscribe(@{print_stats = $null; heater_bed = $null}) $ws.Close() .OUTPUTS KlippyWebSocketClient object. #> [CmdletBinding()] param( [Parameter(Mandatory = $true)] [PSCustomObject]$Printer, [Parameter()] [ValidateRange(1, 60)] [int]$ConnectTimeout = 10 ) # Build WebSocket URI $httpUri = [System.Uri]$Printer.Uri $wsScheme = if ($httpUri.Scheme -eq 'https') { 'wss' } else { 'ws' } $wsUri = "$wsScheme`://$($httpUri.Host):$($httpUri.Port)/websocket" Write-Verbose "[$($Printer.PrinterName)] Connecting to WebSocket: $wsUri" try { $webSocket = [System.Net.WebSockets.ClientWebSocket]::new() # Add API key if present if ($Printer.ApiKey) { $webSocket.Options.SetRequestHeader('X-Api-Key', $Printer.ApiKey) } # Connect with timeout $connectTask = $webSocket.ConnectAsync([Uri]$wsUri, [System.Threading.CancellationToken]::None) $completed = $connectTask.Wait([TimeSpan]::FromSeconds($ConnectTimeout)) if (-not $completed) { $webSocket.Dispose() throw "Connection timed out after $ConnectTimeout seconds" } if ($connectTask.IsFaulted) { $webSocket.Dispose() throw $connectTask.Exception.InnerException } Write-Verbose "[$($Printer.PrinterName)] WebSocket connected" # Create client wrapper object $client = [PSCustomObject]@{ PSTypeName = 'KlippyCLI.WebSocketClient' Printer = $Printer WebSocket = $webSocket Uri = $wsUri IsConnected = $true RequestId = 0 Subscriptions = @{} MessageBuffer = [System.Collections.Concurrent.ConcurrentQueue[PSCustomObject]]::new() CancelSource = [System.Threading.CancellationTokenSource]::new() } # Add methods to the client object $client | Add-Member -MemberType ScriptMethod -Name 'GetNextId' -Value { $this.RequestId++ return $this.RequestId } $client | Add-Member -MemberType ScriptMethod -Name 'Send' -Value { param([hashtable]$Message) if (-not $this.IsConnected -or $this.WebSocket.State -ne [System.Net.WebSockets.WebSocketState]::Open) { throw "WebSocket is not connected" } $json = $Message | ConvertTo-Json -Depth 10 -Compress $bytes = [System.Text.Encoding]::UTF8.GetBytes($json) $segment = [System.ArraySegment[byte]]::new($bytes) $sendTask = $this.WebSocket.SendAsync( $segment, [System.Net.WebSockets.WebSocketMessageType]::Text, $true, [System.Threading.CancellationToken]::None ) $sendTask.Wait() Write-Verbose "Sent: $json" } $client | Add-Member -MemberType ScriptMethod -Name 'Receive' -Value { param([int]$TimeoutMs = 5000) if (-not $this.IsConnected -or $this.WebSocket.State -ne [System.Net.WebSockets.WebSocketState]::Open) { return $null } $buffer = [byte[]]::new(65536) $segment = [System.ArraySegment[byte]]::new($buffer) $result = [System.Text.StringBuilder]::new() try { do { # Start receive WITHOUT cancellation token - this prevents WebSocket abort on timeout $receiveTask = $this.WebSocket.ReceiveAsync($segment, [System.Threading.CancellationToken]::None) # Wait with timeout (doesn't cancel the underlying operation) $completed = $receiveTask.Wait($TimeoutMs) if (-not $completed) { # Timeout - but WebSocket stays Open (not aborted) return $null } $receiveResult = $receiveTask.Result if ($receiveResult.MessageType -eq [System.Net.WebSockets.WebSocketMessageType]::Close) { $this.IsConnected = $false return $null } $received = [System.Text.Encoding]::UTF8.GetString($buffer, 0, $receiveResult.Count) $result.Append($received) | Out-Null } while (-not $receiveResult.EndOfMessage) $json = $result.ToString() Write-Verbose "Received: $json" return $json | ConvertFrom-Json } catch { Write-Verbose "Receive error: $_" return $null } } $client | Add-Member -MemberType ScriptMethod -Name 'SendRpc' -Value { param([string]$Method, [hashtable]$Params) $id = $this.GetNextId() $request = @{ jsonrpc = "2.0" method = $Method id = $id } if ($Params) { $request['params'] = $Params } $this.Send($request) return $id } $client | Add-Member -MemberType ScriptMethod -Name 'Subscribe' -Value { param([hashtable]$Objects) # Subscribe to printer objects # Objects format: @{ print_stats = $null; heater_bed = @("temperature", "target") } $this.Subscriptions = $Objects $params = @{ objects = $Objects } $id = $this.SendRpc("printer.objects.subscribe", $params) # Wait for response $response = $this.Receive(10000) if ($response.error) { throw "Subscription failed: $($response.error.message)" } return $response.result } $client | Add-Member -MemberType ScriptMethod -Name 'WaitForMessage' -Value { param( [scriptblock]$Condition, [int]$TimeoutSeconds = 60, [scriptblock]$OnProgress ) $deadline = [datetime]::UtcNow.AddSeconds($TimeoutSeconds) while ([datetime]::UtcNow -lt $deadline) { $message = $this.Receive(1000) if ($null -ne $message) { # Check if it's a notification with status update if ($message.method -eq 'notify_status_update' -and $message.params) { $status = $message.params[0] if ($OnProgress) { & $OnProgress $status } if ($Condition -and (& $Condition $status)) { return $status } } # Handle other message types elseif ($message.result -or $message.error) { # RPC response if ($message.error) { Write-Warning "RPC Error: $($message.error.message)" } } } # Check if connection is still valid if (-not $this.IsConnected -or $this.WebSocket.State -ne [System.Net.WebSockets.WebSocketState]::Open) { throw "WebSocket connection lost" } } throw "Timeout waiting for condition after $TimeoutSeconds seconds" } $client | Add-Member -MemberType ScriptMethod -Name 'Close' -Value { if ($this.WebSocket.State -eq [System.Net.WebSockets.WebSocketState]::Open) { try { $closeTask = $this.WebSocket.CloseAsync( [System.Net.WebSockets.WebSocketCloseStatus]::NormalClosure, "Closing", [System.Threading.CancellationToken]::None ) $closeTask.Wait(5000) | Out-Null } catch { Write-Verbose "Close error (ignored): $_" } } $this.CancelSource.Cancel() $this.WebSocket.Dispose() $this.IsConnected = $false Write-Verbose "[$($this.Printer.PrinterName)] WebSocket closed" } return $client } catch { throw "Failed to connect WebSocket to '$($Printer.PrinterName)': $_" } } |