Fix client socket close

client/comfyclientrequests.go

@@ -46,6 +46,7 @@ func (c *ComfyClient) GetSystemStats() (*SystemStats, error) {
 	if err != nil {
 		return nil, err
 	}
+	defer resp.Body.Close()
 
 	body, _ := io.ReadAll(resp.Body)
 	retv := &SystemStats{}
@@ -85,6 +86,7 @@ func (c *ComfyClient) GetPromptHistoryByID() (map[string]PromptHistoryItem, erro
 	if err != nil {
 		return nil, err
 	}
+	defer resp.Body.Close()
 
 	// we need to re-arrange the data into something more coherent
 	// We're going to have to make an adapter that reconstructs an actual prompt
@@ -170,6 +172,7 @@ func (c *ComfyClient) GetViewMetadata(folder string, file string) (string, error
 	if err != nil {
 		return "", err
 	}
+	defer resp.Body.Close()
 
 	body, _ := io.ReadAll(resp.Body)
 	return string(body), nil
@@ -186,6 +189,7 @@ func (c *ComfyClient) GetImage(image_data DataOutput) (*[]byte, error) {
 	if err != nil {
 		return nil, err
 	}
+	defer resp.Body.Close()
 
 	body, _ := io.ReadAll(resp.Body)
 	return &body, nil
@@ -197,6 +201,7 @@ func (c *ComfyClient) GetEmbeddings() ([]string, error) {
 	if err != nil {
 		return nil, err
 	}
+	defer resp.Body.Close()
 
 	body, _ := io.ReadAll(resp.Body)
 	retv := make([]string, 0)
@@ -213,6 +218,7 @@ func (c *ComfyClient) GetQueueExecutionInfo() (*QueueExecInfo, error) {
 	if err != nil {
 		return nil, err
 	}
+	defer resp.Body.Close()
 
 	body, _ := io.ReadAll(resp.Body)
 	queue_exec := &QueueExecInfo{}
@@ -230,6 +236,7 @@ func (c *ComfyClient) GetExtensions() ([]string, error) {
 	if err != nil {
 		return nil, err
 	}
+	defer resp.Body.Close()
 
 	body, _ := io.ReadAll(resp.Body)
 	retv := make([]string, 0)
@@ -247,6 +254,7 @@ func (c *ComfyClient) GetObjectInfos() (*graphapi.NodeObjects, error) {
 	if err != nil {
 		return nil, err
 	}
+	defer resp.Body.Close()
 
 	body, _ := io.ReadAll(resp.Body)
 	result := &graphapi.NodeObjects{}
@@ -296,6 +304,7 @@ func (c *ComfyClient) QueueRawPrompt(graph *graphapi.Graph, prompt *graphapi.Pro
 		ws.Close()
 		return nil, err
 	}
+	defer resp.Body.Close()
 
 	body, _ := io.ReadAll(resp.Body)
 
@@ -304,6 +313,7 @@ func (c *ComfyClient) QueueRawPrompt(graph *graphapi.Graph, prompt *graphapi.Pro
 		Workflow:  graph,
 		Messages:  make(chan PromptMessage),
 		webSocket: ws,
+		done:      make(chan struct{}),
 	}
 
 	err = json.Unmarshal(body, &item)
@@ -357,6 +367,7 @@ func (c *ComfyClient) Interrupt() error {
 	if err != nil {
 		return err
 	}
+	defer resp.Body.Close()
 
 	io.ReadAll(resp.Body)
 	return nil
@@ -369,6 +380,7 @@ func (c *ComfyClient) EraseHistory() error {
 	if err != nil {
 		return err
 	}
+	defer resp.Body.Close()
 
 	io.ReadAll(resp.Body)
 	return nil
@@ -381,6 +393,7 @@ func (c *ComfyClient) EraseHistoryItem(promptID string) error {
 	if err != nil {
 		return err
 	}
+	defer resp.Body.Close()
 
 	io.ReadAll(resp.Body)
 	return nil

client/dataitems.go

@@ -8,7 +8,7 @@ type DataOutput struct {
 	Filename  string `json:"filename"`
 	Subfolder string `json:"subfolder"`
 	Type      string `json:"type"`
-	Text      string `json:"-"` // for "text" type data output
+	Text      string `json:"text"` // for "text" type data output
 }
 
 type SystemStats struct {

client/messagehandler.go

@@ -122,76 +122,86 @@ func (h *MessageHandlers) WithCompleteHandler(fn func()) *MessageHandlers {
 }
 
 // ProcessMessages processes messages from the QueueItem using the provided handlers.
-// This function blocks until execution stops or an error occurs.
-// Returns an error if execution failed, nil if successful.
+// This function blocks until execution stops, the QueueItem is closed, or an error occurs.
+// Returns an error if execution failed or the QueueItem was closed, nil if successful.
 func (qi *QueueItem) ProcessMessages(handlers *MessageHandlers) error {
 	if handlers == nil {
 		handlers = &MessageHandlers{}
 	}
 
 	var executionError error
 
-	// Ensure OnComplete is called when we exit
+	// Ensure OnComplete is called when we exit.
 	if handlers.OnComplete != nil {
 		defer handlers.OnComplete()
 	}
 
 	for {
-		msg := <-qi.Messages
-
-		switch msg.Type {
-		case "started":
-			if handlers.OnStarted != nil {
-				handlers.OnStarted(msg.ToPromptMessageStarted())
+		select {
+		case msg, ok := <-qi.Messages:
+			if !ok {
+				return fmt.Errorf("message channel closed")
 			}
 
-		case "executing":
-			if handlers.OnExecuting != nil {
-				handlers.OnExecuting(msg.ToPromptMessageExecuting())
-			}
+			switch msg.Type {
+			case "started":
+				if handlers.OnStarted != nil {
+					handlers.OnStarted(msg.ToPromptMessageStarted())
+				}
 
-		case "progress":
-			if handlers.OnProgress != nil {
-				handlers.OnProgress(msg.ToPromptMessageProgress())
-			}
+			case "executing":
+				if handlers.OnExecuting != nil {
+					handlers.OnExecuting(msg.ToPromptMessageExecuting())
+				}
 
-		case "progress_state":
-			if handlers.OnProgressState != nil {
-				handlers.OnProgressState(msg.ToPromptMessageProgressState())
-			}
+			case "progress":
+				if handlers.OnProgress != nil {
+					handlers.OnProgress(msg.ToPromptMessageProgress())
+				}
 
-		case "data":
-			if handlers.OnData != nil {
-				handlers.OnData(msg.ToPromptMessageData())
-			}
+			case "progress_state":
+				if handlers.OnProgressState != nil {
+					handlers.OnProgressState(msg.ToPromptMessageProgressState())
+				}
 
-		case "execution_success":
-			if handlers.OnExecutionSuccess != nil {
-				handlers.OnExecutionSuccess(msg.ToPromptMessageExecutionSuccess())
-			}
+			case "data":
+				if handlers.OnData != nil {
+					handlers.OnData(msg.ToPromptMessageData())
+				}
 
-		case "stopped":
-			stopped := msg.ToPromptMessageStopped()
+			case "execution_success":
+				if handlers.OnExecutionSuccess != nil {
+					handlers.OnExecutionSuccess(msg.ToPromptMessageExecutionSuccess())
+				}
 
-			// Handle error first if present
-			if stopped.Exception != nil {
-				if handlers.OnError != nil {
-					handlers.OnError(stopped.Exception)
+			case "stopped":
+				stopped := msg.ToPromptMessageStopped()
+
+				// Handle error first if present.
+				if stopped.Exception != nil {
+					if handlers.OnError != nil {
+						handlers.OnError(stopped.Exception)
+					}
+					executionError = fmt.Errorf(
+						"execution failed: %s - %s",
+						stopped.Exception.ExceptionType,
+						stopped.Exception.ExceptionMessage,
+					)
 				}
-				executionError = fmt.Errorf("execution failed: %s - %s",
-					stopped.Exception.ExceptionType,
-					stopped.Exception.ExceptionMessage)
-			}
 
-			// Then call stopped handler
-			if handlers.OnStopped != nil {
-				handlers.OnStopped(stopped)
-			}
+				// Then call stopped handler.
+				if handlers.OnStopped != nil {
+					handlers.OnStopped(stopped)
+				}
 
-			return executionError
+				return executionError
+
+			default:
+				slog.Warn("Unknown message type received", "type", msg.Type)
+			}
 
-		default:
-			slog.Warn("Unknown message type received", "type", msg.Type)
+		case <-qi.Done():
+			return fmt.Errorf("queue item closed")
 		}
 	}
 }

client/queueitem.go

@@ -1,19 +1,72 @@
 package client
 
-import "github.com/richinsley/comfy2go/graphapi"
+import (
+	"sync"
+
+	"github.com/richinsley/comfy2go/graphapi"
+)
 
 type QueueItem struct {
-	PromptID    string                 `json:"prompt_id"`
-	Number      int                    `json:"number"`
-	NodeErrors  map[string]interface{} `json:"node_errors"`
-	Messages    chan PromptMessage     `json:"-"`
-	Workflow    *graphapi.Graph        `json:"-"`
-	webSocket *WebSocketConnection   `json:"-"`
+	PromptID   string                 `json:"prompt_id"`
+	Number     int                    `json:"number"`
+	NodeErrors map[string]interface{} `json:"node_errors"`
+
+	Messages chan PromptMessage `json:"-"`
+	Workflow *graphapi.Graph    `json:"-"`
+
+	webSocket   *WebSocketConnection `json:"-"`
+	done        chan struct{}        `json:"-"`
+	wsCloseOnce sync.Once            `json:"-"`
+	closeOnce   sync.Once            `json:"-"`
+}
+
+func (qi *QueueItem) Done() <-chan struct{} {
+	// If qi is nil, treat it as already closed.
+	if qi == nil {
+		ch := make(chan struct{})
+		close(ch)
+		return ch
+	}
+	// done should be initialized when the QueueItem is created.
+	// If it's nil (e.g., constructed externally), we treat it as "never closed"
+	// so it won't spuriously win in select and cause message loss.
+	return qi.done
+}
+
+func (qi *QueueItem) send(msg PromptMessage) {
+	if qi == nil || qi.Messages == nil {
+		return
+	}
+	select {
+	case qi.Messages <- msg:
+	case <-qi.Done():
+		// Client/QueueItem is closing; stop delivering messages.
+	}
+}
+
+// CloseWebSocket closes the websocket connection associated with the QueueItem.
+// It does NOT signal Done; callers should send any final messages first.
+func (qi *QueueItem) CloseWebSocket() {
+	if qi == nil {
+		return
+	}
+	qi.wsCloseOnce.Do(func() {
+		if qi.webSocket != nil {
+			qi.webSocket.Close()
+			qi.webSocket = nil
+		}
+	})
 }
 
-// Close closes the websocket connection associated with the QueueItem
+// Close releases QueueItem resources and signals all waiters.
 func (qi *QueueItem) Close() {
-	if qi.webSocket != nil {
-		qi.webSocket.Close()
+	if qi == nil {
+		return
 	}
+	qi.closeOnce.Do(func() {
+		qi.CloseWebSocket()
+		if qi.done != nil {
+			close(qi.done)
+		}
+	})
 }

graphapi/graph.go

@@ -741,6 +741,14 @@ func (t *Graph) GraphToPrompt(clientID string) (Prompt, error) {
 		}
 		p.Nodes = expander.ToPromptNodes()
 	} else {
+		// Pre-compute bypass origin mappings for all bypassed nodes
+		bypassMap := make(map[int]map[int]*Link) // nodeID -> (outputSlot -> upstream Link)
+		for _, node := range t.NodesInExecutionOrder {
+			if node.Mode == NodeModeBypassed {
+				bypassMap[node.ID] = node.GetBypassOrigin()
+			}
+		}
+
 		// Use original logic for backward compatibility
 		for _, node := range t.NodesInExecutionOrder {
 			if node.IsVirtual() {
@@ -749,8 +757,8 @@ func (t *Graph) GraphToPrompt(clientID string) (Prompt, error) {
 				continue
 			}
 
-			if node.Mode == 2 {
-				// Don't serialize muted nodes
+			if node.Mode == NodeModeMuted || node.Mode == NodeModeBypassed {
+				// Don't serialize muted or bypassed nodes
 				continue
 			}
 
@@ -772,6 +780,8 @@ func (t *Graph) GraphToPrompt(clientID string) (Prompt, error) {
 				parent := node.GetNodeForInput(i)
 				if parent != nil {
 					link := t.GetLinkById(slot.Link)
+
+					// Traverse through virtual nodes (PrimitiveNode, Reroute, etc.)
 					for parent != nil && parent.IsVirtual() {
 						link = parent.GetInputLink(link.OriginSlot)
 						if link != nil {
@@ -781,7 +791,23 @@ func (t *Graph) GraphToPrompt(clientID string) (Prompt, error) {
 						}
 					}
 
-					if link != nil {
+					// Traverse through bypassed nodes: follow the bypass pass-through chain
+					for parent != nil && parent.Mode == NodeModeBypassed && link != nil {
+						origins := bypassMap[parent.ID]
+						if origins == nil {
+							link = nil
+							break
+						}
+						upstreamLink, ok := origins[link.OriginSlot]
+						if !ok || upstreamLink == nil {
+							link = nil
+							break
+						}
+						link = upstreamLink
+						parent = t.GetNodeById(link.OriginID)
+					}
+
+					if link != nil && parent != nil {
 						linfo := make([]interface{}, 2)
 						linfo[0] = strconv.Itoa(link.OriginID)
 						linfo[1] = link.OriginSlot