add tests

Author: taskbot

cmd/thv/app/proxy.go

@@ -175,7 +175,8 @@ func init() {
 	proxyStdioCmd.Flags().StringVar(&stdioWorkloadName, "workload-name", "", "Workload name for the proxy (required)")
 	_ = proxyStdioCmd.MarkFlagRequired("workload-name")
 
-	// Attach the subcommand to the main proxy command
+	// Attach the subcommands to the main proxy command
+	proxyCmd.AddCommand(proxyTunnelCmd)
 	proxyCmd.AddCommand(proxyStdioCmd)
 
 }

cmd/thv/app/proxy_stdio.go

@@ -9,7 +9,6 @@ import (
 
 	"github.com/stacklok/toolhive/pkg/logger"
 	"github.com/stacklok/toolhive/pkg/transport"
-	"github.com/stacklok/toolhive/pkg/transport/types"
 	"github.com/stacklok/toolhive/pkg/workloads"
 )
 
@@ -52,12 +51,6 @@ func proxyStdioCmdFunc(cmd *cobra.Command, args []string) error {
 		return fmt.Errorf("failed to get workload %q: %w", stdioWorkloadName, err)
 	}
 
-	// check if workload has http/sse transport
-	if stdioWorkload.TransportType != types.TransportTypeSSE &&
-		stdioWorkload.TransportType != types.TransportTypeStreamableHTTP {
-		return fmt.Errorf("only HTTP/SSE workloads are supported for this proxy")
-	}
-
 	logger.Infof("Starting stdio proxy for server=%q -> %s", serverName, stdioWorkloadName)
 
 	bridge, err := transport.NewStdioBridge(stdioWorkload.URL)

pkg/container/images/registry.go

@@ -9,7 +9,7 @@ import (
 	"github.com/docker/docker/client"
 	"github.com/google/go-containerregistry/pkg/authn"
 	"github.com/google/go-containerregistry/pkg/name"
-	"github.com/google/go-containerregistry/pkg/v1"
+	v1 "github.com/google/go-containerregistry/pkg/v1"
 	"github.com/google/go-containerregistry/pkg/v1/daemon"
 	"github.com/google/go-containerregistry/pkg/v1/remote"
 
@@ -62,7 +62,7 @@ func (r *RegistryImageManager) ImageExists(_ context.Context, imageName string)
 
 // PullImage pulls an image from a registry and saves it to the local daemon
 func (r *RegistryImageManager) PullImage(ctx context.Context, imageName string) error {
-	logger.Infof("Pulling image: %s", imageName)
+	logger.Infof("Pulling1 image: %s", imageName)
 
 	// Parse the image reference
 	ref, err := name.ParseReference(imageName)

pkg/transport/bridge.go

@@ -195,37 +195,67 @@ func (b *StdioBridge) runStreamableWriter(ctx context.Context) {
 
 func (b *StdioBridge) runLegacyReader(ctx context.Context) {
 	defer b.wg.Done()
+
 	req, err := http.NewRequestWithContext(ctx, "GET", b.baseURL.String(), nil)
 	if err != nil {
 		logger.Errorf("Failed to create GET request: %v", err)
 		return
 	}
 	req.Header.Set("Accept", "text/event-stream")
 	copyHeaders(req.Header, b.headers)
+
 	resp, err := http.DefaultClient.Do(req)
 	if err != nil {
 		logger.Errorf("SSE connect error: %v", err)
 		return
 	}
 	defer resp.Body.Close()
+
 	scanner := bufio.NewScanner(resp.Body)
-	var sb strings.Builder
+	var (
+		eventName string
+		dataLines []string
+	)
 	for scanner.Scan() {
 		line := scanner.Text()
+
+		if strings.HasPrefix(line, "event:") {
+			eventName = strings.TrimSpace(strings.TrimPrefix(line, "event:"))
+			continue
+		}
 		if strings.HasPrefix(line, "data:") {
-			sb.WriteString(strings.TrimPrefix(line, "data:"))
+			dataLines = append(dataLines, strings.TrimPrefix(strings.TrimSpace(line), "data:"))
+			continue
 		}
 		if line == "" {
-			raw := strings.TrimSpace(sb.String())
-			sb.Reset()
-			if strings.HasPrefix(raw, "/") && strings.Contains(raw, "sessionId") {
-				b.updatePostURL(raw)
-				b.sendInitialize(ctx)
-			} else {
-				emitJSON(raw)
+			payload := strings.Join(dataLines, "\n")
+			dataLines = dataLines[:0]
+
+			switch eventName {
+			case "endpoint":
+				if payload != "" {
+					b.updatePostURL(payload)
+					b.sendInitialize(ctx)
+				}
+			default:
+				if isJSON(payload) {
+					emitJSON(payload)
+				} else {
+					logger.Debugf("Skipping non-JSON SSE event (%q): %q", eventName, payload)
+				}
 			}
+			eventName = ""
 		}
 	}
+
+	if err := scanner.Err(); err != nil {
+		logger.Errorf("SSE read error: %v", err)
+	}
+}
+
+func isJSON(s string) bool {
+	s = strings.TrimLeft(s, " \r\n\t")
+	return len(s) > 0 && (s[0] == '{' || s[0] == '[')
 }
 
 func (b *StdioBridge) runLegacyWriter(ctx context.Context) {