diff --git a/client/comfyclient.go b/client/comfyclient.go
index 78c95c7..bf1d415 100644
--- a/client/comfyclient.go
+++ b/client/comfyclient.go
@@ -206,16 +206,14 @@ func (c *ComfyClient) GetQueuedItem(prompt_id string) *QueueItem {
 }
 
 // OnWindowSocketMessage processes each message received from the websocket connection to ComfyUI.
-// The messages are parsed, and translated into PromptMessage structs and placed into the correct QueuedItem's message channel.
+// The messages are parsed and translated into PromptMessage structs and routed into the QueueItem.
 func (c *ComfyClient) OnWindowSocketMessage(msg string, qi *QueueItem) {
 	message := &WSStatusMessage{}
-	err := json.Unmarshal([]byte(msg), &message)
-	if err != nil {
-		slog.Error("Deserializing Status Message:", "error", err)
+	if err := json.Unmarshal([]byte(msg), &message); err != nil {
+		slog.Error("Deserializing Status Message", "error", err)
+		return
 	}
 
-	// fmt.Println(msg)
-
 	switch message.Type {
 	case "status":
 		s := message.Data.(*WSMessageDataStatus)
@@ -223,221 +221,240 @@ func (c *ComfyClient) OnWindowSocketMessage(msg string, qi *QueueItem) {
 			c.queuecount = s.Status.ExecInfo.QueueRemaining
 			c.callbacks.ClientQueueCountChanged(c, s.Status.ExecInfo.QueueRemaining)
 		}
+
 	case "execution_start":
 		s := message.Data.(*WSMessageDataExecutionStart)
-		// update lastProcessedPromptID to indicate we are processing a new prompt
+		// Update lastProcessedPromptID to indicate we are processing a new prompt.
 		c.lastProcessedPromptID = s.PromptID
-		if qi != nil {
-			if c.callbacks != nil && c.callbacks.QueuedItemStarted != nil {
-				c.callbacks.QueuedItemStarted(c, qi)
-			}
-			m := PromptMessage{
-				Type: "started",
-				Message: &PromptMessageStarted{
-					PromptID: qi.PromptID,
-				},
-			}
-			qi.Messages <- m
+		if qi == nil {
+			return
 		}
+		if c.callbacks != nil && c.callbacks.QueuedItemStarted != nil {
+			c.callbacks.QueuedItemStarted(c, qi)
+		}
+		qi.send(PromptMessage{
+			Type: "started",
+			Message: &PromptMessageStarted{
+				PromptID: qi.PromptID,
+			},
+		})
+
 	case "execution_cached":
-		// this is probably not usefull for us
+		// Intentionally ignored.
+
 	case "executing":
 		s := message.Data.(*WSMessageDataExecuting)
-		if qi != nil {
-			if s.Node == nil {
-				// final node was processed
-				m := PromptMessage{
-					Type: "stopped",
-					Message: &PromptMessageStopped{
-						QueueItem: qi,
-						Exception: nil,
-					},
-				}
-				// remove the Item from our Queue before sending the message
-				// no other messages will be sent to the channel after this
-				if c.callbacks != nil && c.callbacks.QueuedItemStopped != nil {
-					c.callbacks.QueuedItemStopped(c, qi, QueuedItemStoppedReasonFinished)
-				}
-				delete(c.queueditems, qi.PromptID)
-				// qi.Close()
-				qi.Messages <- m
-			} else {
-				// Try to find the node in the workflow
-				// For compound IDs like "57:8", parse the first part
-				var node *graphapi.GraphNode
-				nodeIDStr := *s.Node
-				if nodeID, err := strconv.Atoi(nodeIDStr); err == nil {
-					// Simple integer ID
-					node = qi.Workflow.GetNodeById(nodeID)
-				} else if strings.Contains(nodeIDStr, ":") {
-					// Compound ID like "57:8" - try to get the instance node
-					parts := strings.Split(nodeIDStr, ":")
-					if instanceID, err := strconv.Atoi(parts[0]); err == nil {
-						node = qi.Workflow.GetNodeById(instanceID)
-					}
-				}
-
-				if node != nil {
-					m := PromptMessage{
-						Type: "executing",
-						Message: &PromptMessageExecuting{
-							NodeID: *s.Node,
-							Title:  node.DisplayName,
-						},
-					}
-					qi.Messages <- m
-				} else {
-					m := PromptMessage{
-						Type: "executing",
-						Message: &PromptMessageExecuting{
-							NodeID: *s.Node,
-							Title:  *s.Node,
-						},
-					}
-					qi.Messages <- m
-				}
+		if qi == nil {
+			return
+		}
+		if s.Node == nil {
+			// Final node was processed.
+			if c.callbacks != nil && c.callbacks.QueuedItemStopped != nil {
+				c.callbacks.QueuedItemStopped(c, qi, QueuedItemStoppedReasonFinished)
 			}
+			delete(c.queueditems, qi.PromptID)
+			qi.send(PromptMessage{
+				Type: "stopped",
+				Message: &PromptMessageStopped{
+					QueueItem: qi,
+					Exception: nil,
+				},
+			})
+			// Release websocket resources for this item.
+			qi.CloseWebSocket()
+			return
 		}
+
+		// Try to find the node in the workflow.
+		// For compound IDs like "57:8", parse the first part.
+		var node *graphapi.GraphNode
+		nodeIDStr := *s.Node
+		if nodeID, err := strconv.Atoi(nodeIDStr); err == nil {
+			// Simple integer ID.
+			node = qi.Workflow.GetNodeById(nodeID)
+		} else if strings.Contains(nodeIDStr, ":") {
+			// Compound ID like "57:8" - try to get the instance node.
+			parts := strings.Split(nodeIDStr, ":")
+			if instanceID, err := strconv.Atoi(parts[0]); err == nil {
+				node = qi.Workflow.GetNodeById(instanceID)
+			}
+		}
+
+		title := *s.Node
+		if node != nil {
+			title = node.DisplayName
+		}
+		qi.send(PromptMessage{
+			Type: "executing",
+			Message: &PromptMessageExecuting{
+				NodeID: *s.Node,
+				Title:  title,
+			},
+		})
+
 	case "progress":
 		s := message.Data.(*WSMessageDataProgress)
-		if qi != nil {
-			m := PromptMessage{
-				Type: "progress",
-				Message: &PromptMessageProgress{
-					Value: s.Value,
-					Max:   s.Max,
-				},
-			}
-			qi.Messages <- m
+		if qi == nil {
+			return
 		}
+		qi.send(PromptMessage{
+			Type: "progress",
+			Message: &PromptMessageProgress{
+				Value: s.Value,
+				Max:   s.Max,
+			},
+		})
+
 	case "executed":
 		s := message.Data.(*WSMessageDataExecuted)
-		if qi != nil {
-			// mdata := &PromptMessageData{
-			// 	NodeID: s.Node,
-			// 	Images: *s.Output["images"],
-			// }
-
-			// collect the data from the output
-			mdata := &PromptMessageData{
-				NodeID: s.Node,
-				Data:   make(map[string][]DataOutput),
-			}
-
-			for k, v := range s.Output {
-				mdata.Data[k] = *v
-			}
-
-			m := PromptMessage{
-				Type:    "data",
-				Message: mdata,
-			}
-			if c.callbacks != nil && c.callbacks.QueuedItemDataAvailable != nil {
-				c.callbacks.QueuedItemDataAvailable(c, qi, mdata)
-			}
-			qi.Messages <- m
+		if qi == nil {
+			return
+		}
+		// Collect the data from the output.
+		mdata := &PromptMessageData{
+			NodeID: s.Node,
+			Data:   make(map[string][]DataOutput),
+		}
+		for k, v := range s.Output {
+			mdata.Data[k] = *v
 		}
+		if c.callbacks != nil && c.callbacks.QueuedItemDataAvailable != nil {
+			c.callbacks.QueuedItemDataAvailable(c, qi, mdata)
+		}
+		qi.send(PromptMessage{Type: "data", Message: mdata})
+
 	case "execution_interrupted":
-		if qi != nil {
-			m := PromptMessage{
-				Type: "stopped",
-				Message: &PromptMessageStopped{
-					QueueItem: qi,
-					Exception: nil,
-				},
-			}
-			// remove the Item from our Queue before sending the message
-			// no other messages will be sent to the channel after this
-			if c.callbacks != nil && c.callbacks.QueuedItemStopped != nil {
-				c.callbacks.QueuedItemStopped(c, qi, QueuedItemStoppedReasonInterrupted)
-			}
-			delete(c.queueditems, qi.PromptID)
-			qi.Close()
-			qi.Messages <- m
+		if qi == nil {
+			return
+		}
+		if c.callbacks != nil && c.callbacks.QueuedItemStopped != nil {
+			c.callbacks.QueuedItemStopped(c, qi, QueuedItemStoppedReasonInterrupted)
 		}
+		delete(c.queueditems, qi.PromptID)
+		qi.send(PromptMessage{
+			Type: "stopped",
+			Message: &PromptMessageStopped{
+				QueueItem: qi,
+				Exception: nil,
+			},
+		})
+		qi.CloseWebSocket()
+
 	case "execution_error":
 		s := message.Data.(*WSMessageExecutionError)
-		if qi != nil {
-			// Try to find the node in the workflow
-			var tnode *graphapi.GraphNode
-			if nodeID, err := strconv.Atoi(s.Node); err == nil {
-				tnode = qi.Workflow.GetNodeById(nodeID)
-			} else if strings.Contains(s.Node, ":") {
-				// Compound ID - try to get the instance node
-				parts := strings.Split(s.Node, ":")
-				if instanceID, err := strconv.Atoi(parts[0]); err == nil {
-					tnode = qi.Workflow.GetNodeById(instanceID)
-				}
-			}
+		if qi == nil {
+			return
+		}
 
-			nodeName := s.Node
-			if tnode != nil {
-				nodeName = tnode.Title
+		// Try to find the node in the workflow.
+		var tnode *graphapi.GraphNode
+		if nodeID, err := strconv.Atoi(s.Node); err == nil {
+			tnode = qi.Workflow.GetNodeById(nodeID)
+		} else if strings.Contains(s.Node, ":") {
+			// Compound ID - try to get the instance node.
+			parts := strings.Split(s.Node, ":")
+			if instanceID, err := strconv.Atoi(parts[0]); err == nil {
+				tnode = qi.Workflow.GetNodeById(instanceID)
 			}
+		}
 
-			m := PromptMessage{
-				Type: "stopped",
-				Message: &PromptMessageStopped{
-					QueueItem: qi,
-					Exception: &PromptMessageStoppedException{
-						NodeID:           s.Node,
-						NodeType:         s.NodeType,
-						NodeName:         nodeName,
-						ExceptionMessage: s.ExceptionMessage,
-						ExceptionType:    s.ExceptionType,
-						Traceback:        s.Traceback,
-					},
-				},
-			}
-			// remove the Item from our Queue before sending the message
-			// no other messages will be sent to the channel after this
-			if c.callbacks != nil && c.callbacks.QueuedItemStopped != nil {
-				c.callbacks.QueuedItemStopped(c, qi, QueuedItemStoppedReasonError)
-			}
-			delete(c.queueditems, qi.PromptID)
-			qi.Close()
-			qi.Messages <- m
+		nodeName := s.Node
+		if tnode != nil {
+			nodeName = tnode.Title
 		}
+
+		if c.callbacks != nil && c.callbacks.QueuedItemStopped != nil {
+			c.callbacks.QueuedItemStopped(c, qi, QueuedItemStoppedReasonError)
+		}
+		delete(c.queueditems, qi.PromptID)
+		qi.send(PromptMessage{
+			Type: "stopped",
+			Message: &PromptMessageStopped{
+				QueueItem: qi,
+				Exception: &PromptMessageStoppedException{
+					NodeID:           s.Node,
+					NodeType:         s.NodeType,
+					NodeName:         nodeName,
+					ExceptionMessage: s.ExceptionMessage,
+					ExceptionType:    s.ExceptionType,
+					Traceback:        s.Traceback,
+				},
+			},
+		})
+		qi.CloseWebSocket()
+
 	case "progress_state":
 		s := message.Data.(*WSMessageDataProgressState)
-		if qi != nil {
-			// Convert the map of node progress states to application-level format
-			nodes := make(map[string]NodeProgressInfo)
-			for nodeID, nodeState := range s.Nodes {
-				nodes[nodeID] = NodeProgressInfo{
-					Value:         nodeState.Value,
-					Max:           nodeState.Max,
-					State:         nodeState.State,
-					NodeID:        nodeState.NodeID,
-					DisplayNodeID: nodeState.DisplayNodeID,
-					ParentNodeID:  nodeState.ParentNodeID,
-					RealNodeID:    nodeState.RealNodeID,
-				}
-			}
-			m := PromptMessage{
-				Type: "progress_state",
-				Message: &PromptMessageProgressState{
-					PromptID: s.PromptID,
-					Nodes:    nodes,
-				},
+		if qi == nil {
+			return
+		}
+		// Convert the map of node progress states to application-level format.
+		nodes := make(map[string]NodeProgressInfo)
+		for nodeID, nodeState := range s.Nodes {
+			nodes[nodeID] = NodeProgressInfo{
+				Value:         nodeState.Value,
+				Max:           nodeState.Max,
+				State:         nodeState.State,
+				NodeID:        nodeState.NodeID,
+				DisplayNodeID: nodeState.DisplayNodeID,
+				ParentNodeID:  nodeState.ParentNodeID,
+				RealNodeID:    nodeState.RealNodeID,
 			}
-			qi.Messages <- m
 		}
+		qi.send(PromptMessage{
+			Type: "progress_state",
+			Message: &PromptMessageProgressState{
+				PromptID: s.PromptID,
+				Nodes:    nodes,
+			},
+		})
+
 	case "execution_success":
 		s := message.Data.(*WSMessageDataExecutionSuccess)
-		if qi != nil {
-			m := PromptMessage{
-				Type: "execution_success",
-				Message: &PromptMessageExecutionSuccess{
-					PromptID:  s.PromptID,
-					Timestamp: s.Timestamp,
-				},
-			}
-			qi.Messages <- m
+		if qi == nil {
+			return
 		}
+		qi.send(PromptMessage{
+			Type: "execution_success",
+			Message: &PromptMessageExecutionSuccess{
+				PromptID:  s.PromptID,
+				Timestamp: s.Timestamp,
+			},
+		})
+
 	case "crystools.monitor":
+		// Intentionally ignored.
+
 	default:
-		// Handle unknown data types or return a dedicated error here
-		slog.Warn("Unhandled message type: ", "type", message.Type)
+		slog.Warn("Unhandled message type", "type", message.Type)
+	}
+}
+
+// Close closes all websocket connections and cleans up resources
+func (c *ComfyClient) Close() error {
+	var lastErr error
+
+	// Close all queued items' websocket connections.
+	// Copy items first to avoid holding locks while closing resources.
+	items := make([]*QueueItem, 0, len(c.queueditems))
+	for _, item := range c.queueditems {
+		items = append(items, item)
+	}
+
+	for _, item := range items {
+		if item != nil {
+			item.Close()
+		}
+	}
+
+	// Close idle HTTP keep-alive connections (best effort).
+	if c.httpclient != nil {
+		c.httpclient.CloseIdleConnections()
 	}
+
+	// Clear the queue.
+	c.queueditems = make(map[string]*QueueItem)
+	c.initialized = false
+
+	return lastErr
 }
diff --git a/client/comfyclientrequests.go b/client/comfyclientrequests.go
index f29db16..d788759 100644
--- a/client/comfyclientrequests.go
+++ b/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
diff --git a/client/dataitems.go b/client/dataitems.go
index 6966ce9..01f2abc 100644
--- a/client/dataitems.go
+++ b/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 {
diff --git a/client/messagehandler.go b/client/messagehandler.go
index c4a6e8f..1ab4938 100644
--- a/client/messagehandler.go
+++ b/client/messagehandler.go
@@ -122,8 +122,8 @@ 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{}
@@ -131,67 +131,77 @@ func (qi *QueueItem) ProcessMessages(handlers *MessageHandlers) error {
 
 	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")
 		}
 	}
 }
diff --git a/client/queueitem.go b/client/queueitem.go
index f7de315..6eb888d 100644
--- a/client/queueitem.go
+++ b/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)
+		}
+	})
 }
diff --git a/graphapi/graph.go b/graphapi/graph.go
index ad5f00d..1635833 100644
--- a/graphapi/graph.go
+++ b/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
diff --git a/graphapi/node.go b/graphapi/node.go
index 58b36a0..a96946c 100644
--- a/graphapi/node.go
+++ b/graphapi/node.go
@@ -4,6 +4,13 @@ import (
 	"log/slog"
 )
 
+// NodeMode represents the execution mode of a node in ComfyUI
+const (
+	NodeModeActive  = 0 // Node executes normally
+	NodeModeMuted   = 2 // Node is skipped entirely (not serialized)
+	NodeModeBypassed = 4 // Node is bypassed: pass-through matching input to downstream
+)
+
 // GraphNode represents the encapsulation of an individual functionality within a Graph
 type GraphNode struct {
 	ID                 int                     `json:"id"`
@@ -235,3 +242,60 @@ func (n *GraphNode) ApplyToGraph() {
 		*/
 	}
 }
+
+// SetMode sets the node's execution mode
+func (n *GraphNode) SetMode(mode int) {
+	n.Mode = mode
+}
+
+// SetActive sets the node to active mode (normal execution)
+func (n *GraphNode) SetActive() {
+	n.Mode = NodeModeActive
+}
+
+// Mute sets the node to muted mode (skipped entirely)
+func (n *GraphNode) Mute() {
+	n.Mode = NodeModeMuted
+}
+
+// Bypass sets the node to bypass mode (pass-through)
+func (n *GraphNode) Bypass() {
+	n.Mode = NodeModeBypassed
+}
+
+// IsMuted returns true if the node is muted
+func (n *GraphNode) IsMuted() bool {
+	return n.Mode == NodeModeMuted
+}
+
+// IsBypassed returns true if the node is bypassed
+func (n *GraphNode) IsBypassed() bool {
+	return n.Mode == NodeModeBypassed
+}
+
+// IsActive returns true if the node is in active mode
+func (n *GraphNode) IsActive() bool {
+	return n.Mode == NodeModeActive
+}
+
+// GetBypassOrigin finds the upstream link that should be passed through when this node is bypassed.
+// For each output slot, it tries to find the first input slot whose type matches the output type,
+// following the ComfyUI bypass convention.
+// Returns a map of output slot index -> input Link that should be forwarded.
+func (n *GraphNode) GetBypassOrigin() map[int]*Link {
+	result := make(map[int]*Link)
+
+	for outIdx, outSlot := range n.Outputs {
+		// Find the first input slot with a matching type
+		for inIdx, inSlot := range n.Inputs {
+			if inSlot.Type == outSlot.Type {
+				link := n.GetInputLink(inIdx)
+				if link != nil {
+					result[outIdx] = link
+					break // use first match
+				}
+			}
+		}
+	}
+	return result
+}
diff --git a/graphapi/subgraph.go b/graphapi/subgraph.go
index bc622c0..3d7c415 100644
--- a/graphapi/subgraph.go
+++ b/graphapi/subgraph.go
@@ -172,7 +172,7 @@ func NewSubgraphExpander(g *Graph) *SubgraphExpander {
 // ExpandAll expands all nodes, recursively handling subgraphs
 func (e *SubgraphExpander) ExpandAll() error {
 	for _, node := range e.Graph.NodesInExecutionOrder {
-		if node.IsVirtual() || node.Mode == 2 {
+		if node.IsVirtual() || node.Mode == NodeModeMuted || node.Mode == NodeModeBypassed {
 			continue
 		}
 
@@ -226,8 +226,8 @@ func (e *SubgraphExpander) expandSubgraphNode(
 			continue
 		}
 
-		// Skip muted nodes
-		if internalNode.Mode == 2 {
+		// Skip muted or bypassed nodes
+		if internalNode.Mode == NodeModeMuted || internalNode.Mode == NodeModeBypassed {
 			continue
 		}
 
diff --git a/graphapi/workflow_test.go b/graphapi/workflow_test.go
index bce1a12..90870e1 100644
--- a/graphapi/workflow_test.go
+++ b/graphapi/workflow_test.go
@@ -354,12 +354,12 @@ func TestGraphToPromptWithSubgraphs(t *testing.T) {
 	// Count nodes that should appear in prompt
 	expectedNodes := 0
 	for _, node := range graph.Nodes {
-		if !node.IsVirtual() && node.Mode != 2 {
+		if !node.IsVirtual() && node.Mode != NodeModeMuted && node.Mode != NodeModeBypassed {
 			if node.IsSubgraph {
 				// Count internal non-virtual nodes
 				if node.SubgraphDef != nil {
 					for _, internalNode := range node.SubgraphDef.Nodes {
-						if !internalNode.IsVirtual() && internalNode.Mode != 2 {
+						if !internalNode.IsVirtual() && internalNode.Mode != NodeModeMuted && internalNode.Mode != NodeModeBypassed {
 							expectedNodes++
 						}
 					}