Follow-up hardening: MCP init retry, safe HTTP retry semantics, MCP server protocol compliance

cleanr/adapters/http.go

@@ -30,12 +30,16 @@ func NewHTTP(cfg core.TargetConfig, client *http.Client) *HTTP {
 // (initial attempt plus retries).
 const maxRequestAttempts = 3
 
-// doWithRetry issues the request built by build for each attempt, retrying on
-// transport errors and HTTP 429/503 with exponential backoff plus jitter. It
-// honors a Retry-After header when present and never sleeps past the request
-// context deadline: if the next backoff would exceed the remaining budget it
-// returns the last result instead of waiting. build must return a fresh
-// *http.Request per call so the body can be re-read on retry.
+// doWithRetry issues the request built by build for each attempt, retrying
+// HTTP 429/503 (for any method — the server rejected the request, so a replay
+// cannot double-apply it) and transport errors (for idempotent methods only —
+// a connection can die after the request was fully delivered, and replaying a
+// POST there risks duplicate writes or charges) with exponential backoff plus
+// jitter. It honors a Retry-After header when present and never sleeps past
+// the request context deadline: if the next backoff would exceed the
+// remaining budget it returns the last result, with its body still readable.
+// build must return a fresh *http.Request per call so the body can be re-read
+// on retry.
 func doWithRetry(ctx context.Context, client *http.Client, build func() (*http.Request, error)) (*http.Response, error) {
 	var lastResp *http.Response
 	var lastErr error
@@ -47,21 +51,29 @@ func doWithRetry(ctx context.Context, client *http.Client, build func() (*http.R
 		resp, err := client.Do(httpReq)
 		lastResp, lastErr = resp, err
 
-		if err == nil && !retryableStatus(resp.StatusCode) {
+		if err != nil {
+			if !idempotentMethod(httpReq.Method) {
+				return nil, err
+			}
+		} else if !retryableStatus(resp.StatusCode) {
 			return resp, nil
 		}
 		if attempt == maxRequestAttempts-1 {
 			break
 		}
-		// Drain and close the retryable response before the next attempt.
 		wait := backoffDelay(attempt, resp)
-		if err == nil && resp != nil {
-			resp.Body.Close()
-		}
 		if !sleepWithin(ctx, wait) {
-			// Not enough budget left to retry; surface the last result.
+			// Not enough budget left to retry; surface the last result. The
+			// body is intentionally NOT closed here — the caller reads it.
 			return lastResp, lastErr
 		}
+		// Committed to another attempt: drain and close the retryable
+		// response so its keep-alive connection can be reused instead of
+		// forcing a new TCP+TLS handshake against an already-degraded server.
+		if err == nil && resp != nil {
+			_, _ = io.Copy(io.Discard, io.LimitReader(resp.Body, 256<<10))
+			_ = resp.Body.Close()
+		}
 	}
 	return lastResp, lastErr
 }
@@ -70,6 +82,18 @@ func retryableStatus(status int) bool {
 	return status == http.StatusTooManyRequests || status == http.StatusServiceUnavailable
 }
 
+// idempotentMethod reports whether a request may be safely replayed after a
+// transport error, when it is unknowable whether the server already processed
+// the request.
+func idempotentMethod(method string) bool {
+	switch strings.ToUpper(method) {
+	case http.MethodGet, http.MethodHead, http.MethodOptions:
+		return true
+	default:
+		return false
+	}
+}
+
 func backoffDelay(attempt int, resp *http.Response) time.Duration {
 	if resp != nil {
 		if ra := parseRetryAfter(resp.Header.Get("Retry-After")); ra > 0 {

cleanr/adapters/mcp.go

@@ -17,8 +17,8 @@ import (
 type MCP struct {
 	cfg          core.TargetConfig
 	client       *http.Client
-	initOnce     sync.Once
-	initErr      error
+	initMu       sync.Mutex
+	initialized  bool
 	requestIDSeq int64
 	mu           sync.Mutex
 }
@@ -72,34 +72,43 @@ func (t *MCP) Invoke(ctx context.Context, req core.Request) core.Response {
 	}
 }
 
+// initialize performs the MCP handshake at most once per adapter. A failed
+// attempt is NOT cached: the next Invoke retries, so a transient failure (a
+// restarting server, the first scenario's deadline expiring mid-handshake)
+// fails that one scenario instead of poisoning every remaining scenario in
+// the run. The lock serializes concurrent first invokes, matching the
+// blocking behavior sync.Once had.
 func (t *MCP) initialize(ctx context.Context) error {
-	t.initOnce.Do(func() {
-		initReq := map[string]any{
-			"jsonrpc": "2.0",
-			"id":      t.nextRequestID(),
-			"method":  "initialize",
-			"params": map[string]any{
-				"protocolVersion": "2025-06-18",
-				"capabilities":    map[string]any{},
-				"clientInfo": map[string]any{
-					"name":    "cleanr",
-					"version": "v1alpha1",
-				},
+	t.initMu.Lock()
+	defer t.initMu.Unlock()
+	if t.initialized {
+		return nil
+	}
+	initReq := map[string]any{
+		"jsonrpc": "2.0",
+		"id":      t.nextRequestID(),
+		"method":  "initialize",
+		"params": map[string]any{
+			"protocolVersion": "2025-06-18",
+			"capabilities":    map[string]any{},
+			"clientInfo": map[string]any{
+				"name":    "cleanr",
+				"version": "v1alpha1",
 			},
-		}
-		if _, _, err := t.postJSONRPC(ctx, initReq); err != nil {
-			t.initErr = fmt.Errorf("initialize mcp target: %w", err)
-			return
-		}
-		notify := map[string]any{
-			"jsonrpc": "2.0",
-			"method":  "notifications/initialized",
-		}
-		if _, _, err := t.postJSONRPC(ctx, notify); err != nil {
-			t.initErr = fmt.Errorf("notify initialized mcp target: %w", err)
-		}
-	})
-	return t.initErr
+		},
+	}
+	if _, _, err := t.postJSONRPC(ctx, initReq); err != nil {
+		return fmt.Errorf("initialize mcp target: %w", err)
+	}
+	notify := map[string]any{
+		"jsonrpc": "2.0",
+		"method":  "notifications/initialized",
+	}
+	if _, _, err := t.postJSONRPC(ctx, notify); err != nil {
+		return fmt.Errorf("notify initialized mcp target: %w", err)
+	}
+	t.initialized = true
+	return nil
 }
 
 func (t *MCP) nextRequestID() int64 {

internal/mcpserver/protocol.go

@@ -17,18 +17,23 @@ const (
 	jsonRPCServerError    = -32000
 )
 
+// The request/response id is kept as raw JSON so it round-trips exactly:
+// decoding into `any` and re-encoding with omitempty dropped falsy ids
+// (id: 0, id: ""), leaving clients unable to correlate the response. A nil
+// RawMessage marshals as null, which is what JSON-RPC requires on responses
+// to unparseable requests.
 type requestEnvelope struct {
 	JSONRPC string          `json:"jsonrpc"`
-	ID      any             `json:"id,omitempty"`
+	ID      json.RawMessage `json:"id,omitempty"`
 	Method  string          `json:"method,omitempty"`
 	Params  json.RawMessage `json:"params,omitempty"`
 }
 
 type responseEnvelope struct {
-	JSONRPC string         `json:"jsonrpc"`
-	ID      any            `json:"id,omitempty"`
-	Result  any            `json:"result,omitempty"`
-	Error   *errorEnvelope `json:"error,omitempty"`
+	JSONRPC string          `json:"jsonrpc"`
+	ID      json.RawMessage `json:"id"`
+	Result  any             `json:"result,omitempty"`
+	Error   *errorEnvelope  `json:"error,omitempty"`
 }
 
 type errorEnvelope struct {
@@ -48,15 +53,15 @@ type toolCallParams struct {
 	Arguments map[string]any `json:"arguments"`
 }
 
-func successResponse(id any, result any) *responseEnvelope {
+func successResponse(id json.RawMessage, result any) *responseEnvelope {
 	return &responseEnvelope{
 		JSONRPC: "2.0",
 		ID:      id,
 		Result:  result,
 	}
 }
 
-func errorResponse(id any, code int, message string, data any) *responseEnvelope {
+func errorResponse(id json.RawMessage, code int, message string, data any) *responseEnvelope {
 	return &responseEnvelope{
 		JSONRPC: "2.0",
 		ID:      id,

internal/mcpserver/server.go

@@ -29,25 +29,24 @@ func New() *Server {
 func (s *Server) Serve(ctx context.Context, in io.Reader, out io.Writer) error {
 	reader := bufio.NewReader(in)
 	for {
-		line, err := reader.ReadBytes('\n')
-		if err != nil {
-			if errors.Is(err, io.EOF) {
-				return nil
-			}
-			return err
-		}
-
+		// ReadBytes can return data together with io.EOF when the final
+		// request has no trailing newline (common for one-shot piped
+		// clients); process the line before acting on the error so that
+		// request is answered instead of silently dropped.
+		line, readErr := reader.ReadBytes('\n')
 		line = bytes.TrimSpace(line)
-		if len(line) == 0 {
-			continue
-		}
-
-		resp := s.HandleLine(ctx, line)
-		if resp == nil {
-			continue
+		if len(line) > 0 {
+			if resp := s.HandleLine(ctx, line); resp != nil {
+				if err := writeMessage(out, resp); err != nil {
+					return err
+				}
+			}
 		}
-		if err := writeMessage(out, resp); err != nil {
-			return err
+		if readErr != nil {
+			if errors.Is(readErr, io.EOF) {
+				return nil
+			}
+			return readErr
 		}
 	}
 }
@@ -134,23 +133,42 @@ func (s *Server) handleToolCall(ctx context.Context, req requestEnvelope) *respo
 
 	result, err := safeToolCall(ctx, params.Name, params.Arguments)
 	if err != nil {
-		return errorResponse(req.ID, jsonRPCInternalError, err.Error(), nil)
+		// Per the MCP spec: an unknown tool is a protocol-level invalid-params
+		// error and a panic is a genuine server fault, but ordinary execution
+		// failures (bad config path, invalid arguments) are returned as
+		// isError results so the calling model can read them and self-correct
+		// instead of the client treating the server as broken.
+		switch {
+		case errors.Is(err, mcptools.ErrUnknownTool):
+			return errorResponse(req.ID, jsonRPCInvalidParams, err.Error(), nil)
+		case errors.Is(err, errToolPanicked):
+			return errorResponse(req.ID, jsonRPCInternalError, err.Error(), nil)
+		default:
+			return successResponse(req.ID, mcptools.Result{
+				Content: []mcptools.Content{{Type: "text", Text: err.Error()}},
+				IsError: true,
+			})
+		}
 	}
 	return successResponse(req.ID, result)
 }
 
+// errToolPanicked marks a contained tool-handler panic so it surfaces as an
+// internal JSON-RPC error rather than an isError tool result.
+var errToolPanicked = errors.New("tool handler panicked")
+
 // safeToolCall contains a panicking tool handler: the server speaks stdio, so
 // an uncaught panic would kill the whole process and every session with it.
-func safeToolCall(ctx context.Context, name string, args map[string]any) (result any, err error) {
+func safeToolCall(ctx context.Context, name string, args map[string]any) (result mcptools.Result, err error) {
 	defer func() {
 		if r := recover(); r != nil {
-			err = fmt.Errorf("tool %s panicked: %v", name, r)
+			err = fmt.Errorf("%w: tool %s panicked: %v", errToolPanicked, name, r)
 		}
 	}()
 	return mcptools.Call(ctx, name, args)
 }
 
-func (s *Server) requireInitialized(id any) *responseEnvelope {
+func (s *Server) requireInitialized(id json.RawMessage) *responseEnvelope {
 	if s.initialized {
 		return nil
 	}

internal/mcpserver/tools/registry.go

@@ -2,6 +2,7 @@ package tools
 
 import (
 	"context"
+	"errors"
 	"fmt"
 
 	"github.com/devr-tools/cleanr/internal/mcpserver/catalog"
@@ -11,6 +12,12 @@ import (
 
 type Definition = toolkit.Definition
 type Result = toolkit.Result
+type Content = toolkit.Content
+
+// ErrUnknownTool marks a tools/call for a name that is not registered, so the
+// server can answer with a protocol-level invalid-params error instead of a
+// tool-execution failure.
+var ErrUnknownTool = errors.New("unknown tool")
 
 type handler func(context.Context, map[string]any) (toolkit.Result, error)
 
@@ -49,7 +56,7 @@ func Definitions() []Definition {
 func Call(ctx context.Context, name string, args map[string]any) (Result, error) {
 	h, ok := handlers[name]
 	if !ok {
-		return Result{}, fmt.Errorf("unknown tool: %s", name)
+		return Result{}, fmt.Errorf("%w: %s", ErrUnknownTool, name)
 	}
 	return h(ctx, args)
 }