diff --git a/cleanr/adapters/http.go b/cleanr/adapters/http.go index 952df33..5d518ff 100644 --- a/cleanr/adapters/http.go +++ b/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 { diff --git a/cleanr/adapters/mcp.go b/cleanr/adapters/mcp.go index 4889647..d40dce8 100644 --- a/cleanr/adapters/mcp.go +++ b/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 { diff --git a/internal/mcpserver/protocol.go b/internal/mcpserver/protocol.go index 3942ebe..88c1747 100644 --- a/internal/mcpserver/protocol.go +++ b/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,7 +53,7 @@ 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, @@ -56,7 +61,7 @@ func successResponse(id any, result any) *responseEnvelope { } } -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, diff --git a/internal/mcpserver/server.go b/internal/mcpserver/server.go index 4ee570c..a2c125a 100644 --- a/internal/mcpserver/server.go +++ b/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 } diff --git a/internal/mcpserver/tools/registry.go b/internal/mcpserver/tools/registry.go index a5e0790..abe13a6 100644 --- a/internal/mcpserver/tools/registry.go +++ b/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) } diff --git a/tests/adapters/mcp_adapter_test.go b/tests/adapters/mcp_adapter_test.go index 5f4d044..c667a8f 100644 --- a/tests/adapters/mcp_adapter_test.go +++ b/tests/adapters/mcp_adapter_test.go @@ -10,6 +10,82 @@ import ( "github.com/devr-tools/cleanr/cleanr" ) +// A transient initialize failure must not be cached: the next Invoke retries +// the handshake instead of failing every remaining scenario in the run. +func TestMCPTargetRetriesInitializeAfterTransientFailure(t *testing.T) { + t.Parallel() + + initAttempts := 0 + client := &http.Client{Transport: roundTripperFunc(func(r *http.Request) (*http.Response, error) { + var body map[string]any + if err := json.NewDecoder(r.Body).Decode(&body); err != nil { + t.Fatalf("decode request: %v", err) + } + switch body["method"] { + case "initialize": + initAttempts++ + if initAttempts == 1 { + return jsonResponse(t, http.StatusServiceUnavailable, map[string]any{ + "error": "restarting", + }), nil + } + return jsonResponse(t, http.StatusOK, map[string]any{ + "jsonrpc": "2.0", + "id": body["id"], + "result": map[string]any{"protocolVersion": "2025-06-18"}, + }), nil + case "notifications/initialized": + return &http.Response{ + StatusCode: http.StatusOK, + Status: "200 OK", + Body: http.NoBody, + Header: http.Header{"Content-Type": []string{"application/json"}}, + }, nil + case "tools/call": + return jsonResponse(t, http.StatusOK, map[string]any{ + "jsonrpc": "2.0", + "id": body["id"], + "result": map[string]any{ + "content": []any{map[string]any{"type": "text", "text": "ok"}}, + }, + }), nil + default: + t.Fatalf("unexpected method: %#v", body["method"]) + } + return nil, nil + })} + + mcpTarget := cleanr.NewMCPTarget(cleanr.TargetConfig{ + Type: "mcp", + MCP: cleanr.MCPConfig{ + URL: "https://mcp.test/rpc", + Tool: "lookup_customer", + }, + }, client) + + first := mcpTarget.Invoke(context.Background(), cleanr.BuildScenarioRequest(cleanr.Scenario{ + Name: "first", + Input: "hello", + }, 2*time.Second)) + if first.Err == nil { + t.Fatal("expected first invoke to fail while the server is restarting") + } + + second := mcpTarget.Invoke(context.Background(), cleanr.BuildScenarioRequest(cleanr.Scenario{ + Name: "second", + Input: "hello", + }, 2*time.Second)) + if second.Err != nil { + t.Fatalf("expected second invoke to retry initialization and succeed, got %v", second.Err) + } + if second.Text != "ok" { + t.Fatalf("unexpected mcp response text: %q", second.Text) + } + if initAttempts != 2 { + t.Fatalf("expected initialize to be retried once, got %d attempts", initAttempts) + } +} + func TestMCPTargetInvokesToolOverJSONRPC(t *testing.T) { t.Parallel() diff --git a/tests/adapters/retry_test.go b/tests/adapters/retry_test.go new file mode 100644 index 0000000..1249621 --- /dev/null +++ b/tests/adapters/retry_test.go @@ -0,0 +1,137 @@ +package tests + +import ( + "context" + "errors" + "net/http" + "strings" + "testing" + "time" + + "github.com/devr-tools/cleanr/cleanr" +) + +// When a 429's Retry-After exceeds the remaining context budget, the response +// is returned instead of retried — and its body must still be readable. It +// used to be closed before being handed back, so callers saw "read on closed +// response body" instead of the actual rate-limit response. +func TestHTTPTargetReturnsReadableBodyWhenRetryBudgetExhausted(t *testing.T) { + t.Parallel() + + client := &http.Client{Transport: roundTripperFunc(func(r *http.Request) (*http.Response, error) { + resp := jsonResponse(t, http.StatusTooManyRequests, map[string]any{ + "error": map[string]any{"message": "rate limited"}, + }) + resp.Header.Set("Retry-After", "3600") + return resp, nil + })} + + target := cleanr.NewHTTPTarget(cleanr.TargetConfig{ + Type: "http", + URL: "https://api.test/v1", + Method: http.MethodPost, + }, client) + + resp := target.Invoke(context.Background(), cleanr.BuildScenarioRequest(cleanr.Scenario{ + Name: "rate-limited", + Input: "hello", + }, time.Second)) + + if resp.Err != nil { + t.Fatalf("expected readable rate-limit response, got error: %v", resp.Err) + } + if resp.StatusCode != http.StatusTooManyRequests { + t.Fatalf("expected 429 status, got %d", resp.StatusCode) + } + if !strings.Contains(string(resp.Body), "rate limited") { + t.Fatalf("expected rate-limit body to be readable, got %q", string(resp.Body)) + } +} + +// A transport error on a POST must not be retried: the connection can die +// after the request was fully delivered, and replaying it risks duplicate +// writes or charges. +func TestHTTPTargetDoesNotRetryPOSTOnTransportError(t *testing.T) { + t.Parallel() + + attempts := 0 + client := &http.Client{Transport: roundTripperFunc(func(r *http.Request) (*http.Response, error) { + attempts++ + return nil, errors.New("connection reset by peer") + })} + + target := cleanr.NewHTTPTarget(cleanr.TargetConfig{ + Type: "http", + URL: "https://api.test/v1", + Method: http.MethodPost, + }, client) + + resp := target.Invoke(context.Background(), cleanr.BuildScenarioRequest(cleanr.Scenario{ + Name: "post-reset", + Input: "hello", + }, time.Second)) + + if resp.Err == nil { + t.Fatal("expected transport error to surface") + } + if attempts != 1 { + t.Fatalf("POST must not be replayed after a transport error, got %d attempts", attempts) + } +} + +// Idempotent requests keep retrying transport errors, and 429/503 responses +// keep retrying for every method (the server rejected the request, so a +// replay cannot double-apply it). +func TestHTTPTargetRetriesIdempotentTransportErrorsAndRetryableStatuses(t *testing.T) { + t.Parallel() + + getAttempts := 0 + getClient := &http.Client{Transport: roundTripperFunc(func(r *http.Request) (*http.Response, error) { + getAttempts++ + if getAttempts < 3 { + return nil, errors.New("connection refused") + } + return jsonResponse(t, http.StatusOK, map[string]any{"output": "recovered"}), nil + })} + getTarget := cleanr.NewHTTPTarget(cleanr.TargetConfig{ + Type: "http", + URL: "https://api.test/v1", + Method: http.MethodGet, + }, getClient) + resp := getTarget.Invoke(context.Background(), cleanr.BuildScenarioRequest(cleanr.Scenario{ + Name: "get-flaky", + Input: "hello", + }, 10*time.Second)) + if resp.Err != nil { + t.Fatalf("expected GET to retry transport errors and succeed, got %v", resp.Err) + } + if getAttempts != 3 { + t.Fatalf("expected 3 GET attempts, got %d", getAttempts) + } + + postAttempts := 0 + postClient := &http.Client{Transport: roundTripperFunc(func(r *http.Request) (*http.Response, error) { + postAttempts++ + if postAttempts < 2 { + resp := jsonResponse(t, http.StatusServiceUnavailable, map[string]any{"error": "overloaded"}) + resp.Header.Set("Retry-After", "0") + return resp, nil + } + return jsonResponse(t, http.StatusOK, map[string]any{"output": "recovered"}), nil + })} + postTarget := cleanr.NewHTTPTarget(cleanr.TargetConfig{ + Type: "http", + URL: "https://api.test/v1", + Method: http.MethodPost, + }, postClient) + resp = postTarget.Invoke(context.Background(), cleanr.BuildScenarioRequest(cleanr.Scenario{ + Name: "post-503", + Input: "hello", + }, 10*time.Second)) + if resp.Err != nil || resp.StatusCode != http.StatusOK { + t.Fatalf("expected POST to retry 503 and succeed, got err=%v status=%d", resp.Err, resp.StatusCode) + } + if postAttempts != 2 { + t.Fatalf("expected 2 POST attempts on 503, got %d", postAttempts) + } +} diff --git a/tests/mcp/protocol_edges_test.go b/tests/mcp/protocol_edges_test.go new file mode 100644 index 0000000..2d5b21c --- /dev/null +++ b/tests/mcp/protocol_edges_test.go @@ -0,0 +1,121 @@ +package tests + +import ( + "bytes" + "context" + "encoding/json" + "strings" + "testing" + + "github.com/devr-tools/cleanr/internal/mcpserver" +) + +// A final request without a trailing newline (common for one-shot piped +// clients) must be answered at EOF, not silently dropped. +func TestMCPServerAnswersFinalUnterminatedLine(t *testing.T) { + t.Parallel() + + input := strings.Join([]string{ + `{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-06-18","capabilities":{},"clientInfo":{"name":"t","version":"1"}}}`, + `{"jsonrpc":"2.0","method":"notifications/initialized"}`, + `{"jsonrpc":"2.0","id":2,"method":"ping"}`, // no trailing newline + }, "\n") + + var out bytes.Buffer + if err := mcpserver.New().Serve(context.Background(), strings.NewReader(input), &out); err != nil { + t.Fatalf("serve: %v", err) + } + responses := strings.Split(strings.TrimSpace(out.String()), "\n") + if len(responses) != 2 { + t.Fatalf("expected initialize and ping responses, got %d: %s", len(responses), out.String()) + } + if !strings.Contains(responses[1], `"id":2`) { + t.Fatalf("expected ping response for the unterminated final line, got %s", responses[1]) + } +} + +// JSON-RPC ids must round-trip exactly: id 0 and id "" were previously +// dropped by omitempty, leaving the client unable to correlate the response. +func TestMCPServerEchoesFalsyRequestIDs(t *testing.T) { + server := initializedMCPServer(t) + + resp := mustHandleMCP(t, server, map[string]any{ + "jsonrpc": "2.0", + "id": 0, + "method": "ping", + }) + id, present := resp["id"] + if !present || id != float64(0) { + t.Fatalf("expected id 0 to be echoed, got %#v (present=%v)", id, present) + } + + resp = mustHandleMCP(t, server, map[string]any{ + "jsonrpc": "2.0", + "id": "", + "method": "ping", + }) + if id, present := resp["id"]; !present || id != "" { + t.Fatalf("expected empty-string id to be echoed, got %#v (present=%v)", id, present) + } +} + +// Parse errors must carry "id": null per JSON-RPC, not omit the field. +func TestMCPServerParseErrorCarriesNullID(t *testing.T) { + t.Parallel() + + raw := mcpserver.New().HandleLine(context.Background(), []byte("{not json")) + if raw == nil { + t.Fatal("expected parse-error response") + } + data, err := json.Marshal(raw) + if err != nil { + t.Fatalf("marshal response: %v", err) + } + if !strings.Contains(string(data), `"id":null`) { + t.Fatalf("expected id:null on parse error, got %s", data) + } + if !strings.Contains(string(data), `-32700`) { + t.Fatalf("expected parse-error code, got %s", data) + } +} + +// Ordinary tool execution failures come back as isError results the calling +// model can read; only unknown tools and contained panics are protocol-level +// errors. +func TestMCPServerToolFailureModes(t *testing.T) { + server := initializedMCPServer(t) + + // Execution failure: missing config → isError result, not a JSON-RPC error. + resp := mustHandleMCP(t, server, map[string]any{ + "jsonrpc": "2.0", + "id": 3, + "method": "tools/call", + "params": map[string]any{"name": "cleanr_generate_dataset", "arguments": map[string]any{}}, + }) + if _, hasErr := resp["error"]; hasErr { + t.Fatalf("execution failure must be an isError result, got protocol error: %#v", resp) + } + result := resp["result"].(map[string]any) + if isErr, _ := result["isError"].(bool); !isErr { + t.Fatalf("expected isError result, got %#v", result) + } + content := result["content"].([]any) + if text := content[0].(map[string]any)["text"].(string); !strings.Contains(text, "config") { + t.Fatalf("expected readable failure text, got %q", text) + } + + // Unknown tool → protocol-level invalid params. + resp = mustHandleMCP(t, server, map[string]any{ + "jsonrpc": "2.0", + "id": 4, + "method": "tools/call", + "params": map[string]any{"name": "cleanr_no_such_tool", "arguments": map[string]any{}}, + }) + respErr, ok := resp["error"].(map[string]any) + if !ok { + t.Fatalf("expected protocol error for unknown tool, got %#v", resp) + } + if code, _ := respErr["code"].(float64); code != -32602 { + t.Fatalf("expected -32602 for unknown tool, got %v", respErr["code"]) + } +}