mcp_server.go

package main

import (
	"crypto/rand"
	"encoding/base64"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"log"
	"math/big"
	"net/http"
	"time"
)

// Configuration
var (
	port        = flag.Int("port", 9001, "Server port")
	errorRate   = flag.Float64("error-rate", 0.05, "Probability of returning a protocol error (0.0-1.0)")
	toolErrorRate = flag.Float64("tool-error-rate", 0.05, "Probability of returning a tool error (0.0-1.0)")
	longRespRate = flag.Float64("long-rate", 0.10, "Probability of returning a large response (0.0-1.0)")
	minLatency  = flag.Duration("min-latency", 5*time.Millisecond, "Minimum processing latency")
	maxLatency  = flag.Duration("max-latency", 50*time.Millisecond, "Maximum processing latency (for normal requests)")
)

// JSON-RPC Types (Simplified for Server)
type JSONRPCRequest struct {
	JSONRPC string      `json:"jsonrpc"`
	Method  string      `json:"method"`
	Params  Params      `json:"params"`
	ID      interface{} `json:"id"` // can be string or number
}

type Params struct {
	Name      string                 `json:"name"`
	Arguments map[string]interface{} `json:"arguments"`
}

type JSONRPCResponse struct {
	JSONRPC string      `json:"jsonrpc"`
	ID      interface{} `json:"id"`
	Result  interface{} `json:"result,omitempty"`
	Error   *JSONRPCError `json:"error,omitempty"`
}

type JSONRPCError struct {
	Code    int    `json:"code"`
	Message string `json:"message"`
}

type ToolResult struct {
	Content []ContentItem `json:"content"`
	IsError bool          `json:"isError"`
}

type ContentItem struct {
	Type string `json:"type"`
	Text string `json:"text"`
}

func main() {
	flag.Parse()

	http.HandleFunc("/mcp", handleMCP)

	addr := fmt.Sprintf(":%d", *port)
	fmt.Printf("Starting MCP Mock Server on %s\n", addr)
	fmt.Printf("  Protocol Error Rate: %.2f\n", *errorRate)
	fmt.Printf("  Tool Error Rate: %.2f\n", *toolErrorRate)
	fmt.Printf("  Long Response Rate: %.2f\n", *longRespRate)
	
	if err := http.ListenAndServe(addr, nil); err != nil {
		log.Fatalf("Server failed: %v", err)
	}
}

func handleMCP(w http.ResponseWriter, r *http.Request) {
	if r.Method != "POST" {
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
		return
	}

	// 1. Parse Request
	body, err := io.ReadAll(r.Body)
	if err != nil {
		http.Error(w, "Failed to read body", http.StatusBadRequest)
		return
	}
	defer r.Body.Close()

	var req JSONRPCRequest
	if err := json.Unmarshal(body, &req); err != nil {
		writeJSONError(w, nil, -32700, "Parse error")
		return
	}

	// 2. Determine Latency
	simulateLatency(req.Params.Name)

	// 3. Prepare Response Headers for SSE
	w.Header().Set("Content-Type", "text/event-stream")
	w.Header().Set("Cache-Control", "no-cache")
	w.Header().Set("Connection", "keep-alive")
	// mcp-headers usually required by spec, echoing standard ones
	w.Header().Set("mcp-protocol-version", "2025-06-18") 

	// 4. Generate Result or Error
	resp := generateResponse(req)

	// 5. Send Response (SSE format)
	// Spec often wraps JSON-RPC in an event, or just sends data. 
	// Common MCP over SSE sends the JSON-RPC message as data.
	
	respBytes, err := json.Marshal(resp)
	if err != nil {
		log.Printf("Failed to marshal response: %v", err)
		return
	}

	// SSE format: "data: <json>\n\n"
	// Ensure implicit chunking by flushing
	fmt.Fprintf(w, "data: %s\n\n", respBytes)
	
	if f, ok := w.(http.Flusher); ok {
		f.Flush()
	}
}

func simulateLatency(toolName string) {
	var duration time.Duration

	switch toolName {
	case "query_database":
		// Slowest: 1s - 2s
		duration = time.Duration(1000+randInt(1000)) * time.Millisecond
	case "resize_image":
		// Slow: 300ms - 600ms
		duration = time.Duration(300+randInt(300)) * time.Millisecond
	default:
		// Normal: min - max
		rangeMs := int64(*maxLatency - *minLatency) / 1e6
		if rangeMs <= 0 {
			duration = *minLatency
		} else {
			duration = *minLatency + time.Duration(randInt(int(rangeMs)))*time.Millisecond
		}
	}
	
	time.Sleep(duration)
}

func generateResponse(req JSONRPCRequest) JSONRPCResponse {
	roll := randFloat()

	// 1. Protocol Error
	if roll < *errorRate {
		return JSONRPCResponse{
			JSONRPC: "2.0",
			ID:      req.ID,
			Error: &JSONRPCError{
				Code:    -32603,
				Message: "Internal error (simulated)",
			},
		}
	}

	// 2. Tool Error (Variant 2)
	// Returns a successful JSON-RPC response, but the result indicates a tool failure.
	if roll < *errorRate+*toolErrorRate {
		return JSONRPCResponse{
			JSONRPC: "2.0",
			ID:      req.ID,
			Result: ToolResult{
				Content: []ContentItem{
					{
						Type: "text",
						Text: fmt.Sprintf("Error executing tool %s: unable to fetch data.", req.Params.Name),
					},
				},
				IsError: true,
			},
		}
	}

	// 3. Success Result
	result := make(map[string]interface{})
	
	// Echo back tool name for verification
	result["tool"] = req.Params.Name
	result["status"] = "success"

	// Determine Size
	if randFloat() < *longRespRate {
		// Large response (> 4KB, e.g., 5KB - 15KB)
		// Simulating an image buffer or large text result
		size := 5000 + randInt(10000)
		result["data"] = generateRandomString(size)
		result["type"] = "large_blob"
	} else {
		// Short response
		result["message"] = "Operation completed successfully"
		result["timestamp"] = time.Now().Unix()
	}

	return JSONRPCResponse{
		JSONRPC: "2.0",
		ID:      req.ID,
		Result:  result,
	}
}

// Helpers
func writeJSONError(w http.ResponseWriter, id interface{}, code int, msg string) {
	w.Header().Set("Content-Type", "application/json") // Or SSE if connection already established?
	// For simplicity, if we fail before SSE upgrade, standard JSON. 
	// But if we are in handleMCP, we prefer consistency.
	// Since the client expects SSE for success/fail mixed, let's try to stick to SSE if possible 
	// OR just standard error if we haven't sent headers. 
	// The prompt implies SSE responses.
	
	resp := JSONRPCResponse{
		JSONRPC: "2.0",
		ID:      id,
		Error: &JSONRPCError{
			Code:    code,
			Message: msg,
		},
	}
	
	bytes, _ := json.Marshal(resp)
	w.WriteHeader(http.StatusOK) // SSE usually returns 200 then stream
	w.Header().Set("Content-Type", "text/event-stream")
	fmt.Fprintf(w, "data: %s\n\n", bytes)
}

func randInt(max int) int {
	n, _ := rand.Int(rand.Reader, big.NewInt(int64(max)))
	return int(n.Int64())
}

func randFloat() float64 {
	n, _ := rand.Int(rand.Reader, big.NewInt(1000000))
	return float64(n.Int64()) / 1000000.0
}

func generateRandomString(n int) string {
	b := make([]byte, n)
	rand.Read(b)
	return base64.StdEncoding.EncodeToString(b)
}