llmgw/internal/control/adminapi/mcp_server.go

package adminapi

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/http/httptest"
	"net/url"
	"strconv"
)

// The cluster as an MCP server, read only: what `llmgw ctl status`, `nodes`,
// `models`, `replicas`, `explain-placement` and `logs` show, as tools an
// agent can be given. It answers at POST /api/v1/mcp, behind the admin
// token, and in-process to the gateway as the MCP server llmgw:admin, which
// an administrator registers like any other and grants like any other.
// Nothing here changes the cluster.

// AdminMCPAddress is the address that registers this server with the
// gateway: called in-process, with no token and no network.
const AdminMCPAddress = "llmgw:admin"

// adminMCPAnswerBytes bounds one tool's answer: it becomes a prompt.
const adminMCPAnswerBytes = 256 << 10

var adminMCPTools = []map[string]any{
	adminTool("cluster_status", "The whole cluster in one answer: control planes, store, nodes ready, models with the replicas they "+
		"asked for, GPUs in use, and what is unhealthy or idle.", nil, nil),
	adminTool("list_nodes", "Every node: its state, GPUs and their use, CPU and memory load, fabrics, agent version and engines.", nil, nil),
	adminTool("list_models", "Every model in the catalog: its replicas ready against wanted, state, profiles and why one is parked.", nil, nil),
	adminTool("list_replicas", "Every replica, with its shards: the node and GPUs each is on, its state and how long it has run.",
		map[string]any{"model": map[string]any{"type": "string", "description": "Only this model's replicas."}}, nil),
	adminTool("explain_placement", "Why a model is not placed: each node, and what kept the model off it.",
		map[string]any{"model": map[string]any{"type": "string", "description": "The model's name or an alias."}}, []string{"model"}),
	adminTool("node_logs", "The last lines of a node's agent log, or of one of its engines.",
		map[string]any{
			"node":   map[string]any{"type": "string", "description": "The node's id."},
			"lines":  map[string]any{"type": "integer", "description": "How many lines, up to 300; 100 when absent."},
			"source": map[string]any{"type": "string", "description": "agent when absent, or a shard id for that engine's output."},
		}, []string{"node"}),
}

func adminTool(name, desc string, props map[string]any, required []string) map[string]any {
	schema := map[string]any{"type": "object", "properties": map[string]any{}}
	if props != nil {
		schema["properties"] = props
	}
	if len(required) > 0 {
		schema["required"] = required
	}
	return map[string]any{"name": name, "description": desc, "inputSchema": schema,
		"annotations": map[string]any{"readOnlyHint": true, "openWorldHint": false}}
}

// MCP serves the cluster as a read-only MCP server: JSON-RPC over POST, the
// streamable HTTP transport answered as plain JSON, no session.
func (a *API) MCP(w http.ResponseWriter, r *http.Request) {
	var msg struct {
		ID     json.RawMessage `json:"id"`
		Method string          `json:"method"`
		Params json.RawMessage `json:"params"`
	}
	if err := json.NewDecoder(io.LimitReader(r.Body, 1<<20)).Decode(&msg); err != nil || msg.Method == "" {
		writeJSON(w, http.StatusBadRequest, map[string]any{"jsonrpc": "2.0", "id": nil,
			"error": map[string]any{"code": -32600, "message": "expected one JSON-RPC request object"}})
		return
	}
	if len(msg.ID) == 0 {
		w.WriteHeader(http.StatusAccepted)
		return
	}
	answer := func(result any) {
		writeJSON(w, http.StatusOK, map[string]any{"jsonrpc": "2.0", "id": msg.ID, "result": result})
	}
	switch msg.Method {
	case "initialize":
		var p struct {
			ProtocolVersion string `json:"protocolVersion"`
		}
		_ = json.Unmarshal(msg.Params, &p)
		version := p.ProtocolVersion
		if version != "2025-03-26" && version != "2025-06-18" {
			version = "2025-06-18"
		}
		answer(map[string]any{"protocolVersion": version, "capabilities": map[string]any{"tools": map[string]any{}},
			"serverInfo": map[string]any{"name": "llmgw-admin", "version": "1"},
			"instructions": "This llmgw cluster, read only: its control planes, nodes, models and replicas, why a model is " +
				"not placed, and a node's logs. Start with cluster_status."})
	case "ping":
		answer(map[string]any{})
	case "tools/list":
		answer(map[string]any{"tools": adminMCPTools})
	case "tools/call":
		var call struct {
			Name      string `json:"name"`
			Arguments struct {
				Model  string `json:"model"`
				Node   string `json:"node"`
				Source string `json:"source"`
				Lines  int    `json:"lines"`
			} `json:"arguments"`
		}
		_ = json.Unmarshal(msg.Params, &call)
		answer(a.callAdminTool(r, call.Name, call.Arguments.Model, call.Arguments.Node, call.Arguments.Source, call.Arguments.Lines))
	default:
		writeJSON(w, http.StatusOK, map[string]any{"jsonrpc": "2.0", "id": msg.ID,
			"error": map[string]any{"code": -32601, "message": "method not found: " + msg.Method}})
	}
}

// callAdminTool answers a tool from the API's own handler for it, so the
// tool says what `llmgw ctl` would.
func (a *API) callAdminTool(r *http.Request, name, model, node, source string, lines int) map[string]any {
	var h http.HandlerFunc
	q := url.Values{}
	path := ""
	switch name {
	case "cluster_status":
		h, path = a.summary, "/api/v1/summary"
	case "list_nodes":
		h, path = a.listNodes, "/api/v1/nodes"
	case "list_models":
		h, path = a.listModels, "/api/v1/models"
	case "list_replicas":
		h, path = a.listReplicas, "/api/v1/replicas"
		if model != "" {
			q.Set("model", model)
		}
	case "explain_placement":
		if model == "" {
			return adminToolText("explain_placement needs the model", true)
		}
		h, path = a.explainPlacement, "/api/v1/explain-placement"
		q.Set("model", model)
	case "node_logs":
		if node == "" {
			return adminToolText("node_logs needs the node", true)
		}
		if lines <= 0 {
			lines = 100
		}
		h, path = a.nodeLogs, "/api/v1/nodes/"+url.PathEscape(node)+"/logs"
		q.Set("lines", strconv.Itoa(min(lines, 300)))
		if source != "" {
			q.Set("source", source)
		}
	default:
		return adminToolText("no tool "+name, true)
	}
	req := httptest.NewRequestWithContext(r.Context(), http.MethodGet, path+"?"+q.Encode(), nil)
	if name == "node_logs" {
		req.SetPathValue("id", node)
	}
	rec := httptest.NewRecorder()
	h(rec, req)
	body := bytes.TrimSpace(rec.Body.Bytes())
	if rec.Code >= 300 {
		return adminToolText(fmt.Sprintf("%s: %s", name, body), true)
	}
	if len(body) > adminMCPAnswerBytes {
		body = append(body[:adminMCPAnswerBytes], []byte("… (cut: the answer is larger than an agent should read)")...)
	}
	return adminToolText(string(body), false)
}

func adminToolText(text string, isError bool) map[string]any {
	return map[string]any{"content": []any{map[string]any{"type": "text", "text": text}}, "isError": isError}
}

// InProcessURL is the address an in-process client posts to; the transport
// ignores it.
const InProcessURL = "http://llmgw-admin/api/v1/mcp"

// InProcess is a transport that answers every request from h, in this
// process: the gateway's way to the cluster's own MCP server.
func InProcess(h http.Handler) http.RoundTripper { return inProcess{h} }

type inProcess struct{ h http.Handler }

func (t inProcess) RoundTrip(req *http.Request) (*http.Response, error) {
	rec := httptest.NewRecorder()
	t.h.ServeHTTP(rec, req)
	return rec.Result(), nil
}