fix for stdio proxy

Author: taskbot

cmd/thv/app/proxy.go

@@ -171,10 +171,6 @@ func init() {
 	if err := proxyCmd.MarkFlagRequired("target-uri"); err != nil {
 		logger.Warnf("Warning: Failed to mark flag as required: %v", err)
 	}
-
-	proxyStdioCmd.Flags().StringVar(&stdioWorkloadName, "workload-name", "", "Workload name for the proxy (required)")
-	_ = proxyStdioCmd.MarkFlagRequired("workload-name")
-
 	// Attach the subcommands to the main proxy command
 	proxyCmd.AddCommand(proxyTunnelCmd)
 	proxyCmd.AddCommand(proxyStdioCmd)

cmd/thv/app/proxy_stdio.go

@@ -12,46 +12,34 @@ import (
 	"github.com/stacklok/toolhive/pkg/workloads"
 )
 
-var (
-	stdioWorkloadName string
-)
-
 var proxyStdioCmd = &cobra.Command{
-	Use:   "stdio [flags] SERVER_NAME",
+	Use:   "stdio WORKLOAD-NAME SERVER_NAME",
 	Short: "Create a stdio-based proxy for an MCP server",
 	Long: `Create a stdio-based proxy that connects stdin/stdout to a target MCP server.
 
 Example:
-  thv proxy stdio --workload-name my-server my-server-proxy
-
-Flags:
-  --workload-name  Workload name for the proxy (required)
+  thv proxy stdio my-workload my-server-proxy
 `,
-	Args: cobra.ExactArgs(1),
+	Args: cobra.ExactArgs(2),
 	RunE: proxyStdioCmdFunc,
 }
 
 func proxyStdioCmdFunc(cmd *cobra.Command, args []string) error {
 	ctx, cancel := signal.NotifyContext(cmd.Context(), syscall.SIGINT, syscall.SIGTERM)
 	defer cancel()
 
-	serverName := args[0]
-
-	// validate that workload name exists
-	if stdioWorkloadName == "" {
-		return fmt.Errorf("workload name must be specified with --workload-name")
-	}
+	workloadName := args[0]
+	serverName := args[1]
 
 	workloadManager, err := workloads.NewManager(ctx)
 	if err != nil {
 		return fmt.Errorf("failed to create workload manager: %w", err)
 	}
-	stdioWorkload, err := workloadManager.GetWorkload(ctx, stdioWorkloadName)
+	stdioWorkload, err := workloadManager.GetWorkload(ctx, workloadName)
 	if err != nil {
-		return fmt.Errorf("failed to get workload %q: %w", stdioWorkloadName, err)
+		return fmt.Errorf("failed to get workload %q: %w", workloadName, err)
 	}
-
-	logger.Infof("Starting stdio proxy for server=%q -> %s", serverName, stdioWorkloadName)
+	logger.Infof("Starting stdio proxy for server=%q -> %s", serverName, workloadName)
 
 	bridge, err := transport.NewStdioBridge(stdioWorkload.URL)
 	if err != nil {

pkg/container/images/registry.go

@@ -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("Pulling1 image: %s", imageName)
+	logger.Infof("Pulling image: %s", imageName)
 
 	// Parse the image reference
 	ref, err := name.ParseReference(imageName)

pkg/transport/bridge.go

@@ -73,28 +73,54 @@ func (b *StdioBridge) loop(ctx context.Context) {
 
 // Header include Mcp-Session-Id if needed
 func (b *StdioBridge) detectTransport(ctx context.Context) (string, error) {
-	body, err := buildJSONRPCInitializeRequest()
+	// 1) Probe for legacy SSE without sending any JSON-RPC
+	getReq, err := http.NewRequestWithContext(ctx, "GET", b.baseURL.String(), nil)
 	if err != nil {
-		logger.Errorf("failed to marshal initialization request: %v", err)
 		return "", err
 	}
-	req, err := http.NewRequestWithContext(ctx, "POST", b.baseURL.String(), bytes.NewReader(body))
+	getReq.Header.Set("Accept", "text/event-stream")
+	copyHeaders(getReq.Header, b.headers)
+
+	getResp, err := http.DefaultClient.Do(getReq)
+	if err == nil {
+		ct := getResp.Header.Get("Content-Type")
+		err = getResp.Body.Close() // we’re only peeking at headers here
+		if err != nil {
+			return "", fmt.Errorf("failed to close GET response body: %w", err)
+		}
+		if strings.HasPrefix(ct, "text/event-stream") {
+			// Legacy SSE: runLegacyReader will open the real stream.
+			return "legacy-sse", nil
+		}
+	}
+
+	// 2) Not SSE -> treat as streamable HTTP and do a proper initialize
+	body, err := buildJSONRPCInitializeRequest() // keep if your server expects initialize for streamable
 	if err != nil {
-		logger.Errorf("failed to create HTTP request: %v", err)
 		return "", err
 	}
-	req.Header.Set("Content-Type", "application/json")
-	req.Header.Set("Accept", "application/json, text/event-stream")
-	copyHeaders(req.Header, b.headers)
-	resp, err := http.DefaultClient.Do(req)
+	postReq, err := http.NewRequestWithContext(ctx, "POST", b.baseURL.String(), bytes.NewReader(body))
 	if err != nil {
 		return "", err
 	}
-	defer resp.Body.Close()
-	if resp.StatusCode >= 400 && resp.StatusCode < 500 {
+	postReq.Header.Set("Content-Type", "application/json")
+	postReq.Header.Set("Accept", "application/json, text/event-stream")
+	copyHeaders(postReq.Header, b.headers)
+
+	postResp, err := http.DefaultClient.Do(postReq)
+	if err != nil {
+		return "", err
+	}
+	defer postResp.Body.Close()
+
+	// If a streamable server returns info + Mcp-Session-Id, great.
+	// Some legacy gateways might reply 4xx to POST at base URL.
+	if postResp.StatusCode >= 400 && postResp.StatusCode < 500 {
 		return "legacy-sse", nil
 	}
-	b.handleInitializeResponse(resp)
+
+	// Streamable: capture session headers and emit the response payload.
+	b.handleInitializeResponse(postResp)
 	return "streamable-http", nil
 }
 
@@ -104,7 +130,6 @@ func (b *StdioBridge) handleInitializeResponse(resp *http.Response) {
 			b.headers = make(http.Header)
 		}
 		b.headers.Set("Mcp-Session-Id", sid)
-		logger.Infof("Streamable HTTP session ID: %s", sid)
 	}
 	data, err := io.ReadAll(resp.Body)
 	if err != nil {
@@ -307,15 +332,28 @@ func (b *StdioBridge) updatePostURL(path string) {
 		return
 	}
 	b.postURL = u
-	logger.Infof("POST URL updated to %s", b.postURL)
 }
 
 func buildJSONRPCInitializeRequest() ([]byte, error) {
 	req := map[string]interface{}{
 		"jsonrpc": "2.0",
 		"id":      1,
 		"method":  "initialize",
-		"params":  map[string]interface{}{},
+		"params": map[string]interface{}{
+			"protocolVersion": "2024-02-01",
+			"capabilities": map[string]interface{}{
+				"prompts": true,
+				"tools":   true,
+				"resources": map[string]interface{}{
+					"subscribe":   true,
+					"unsubscribe": true,
+				},
+			},
+			"clientInfo": map[string]interface{}{
+				"name":    "toolhive-stdio-bridge",
+				"version": "0.1.0",
+			},
+		},
 	}
 	return json.Marshal(req)
 }