From e1e3eceff7b8344b9a407638dfc40c7d59707b78 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=AB=98=E7=84=B6?= Date: Mon, 22 Jun 2026 11:48:10 +0800 Subject: [PATCH 1/5] feat(execd): add embedded MCP proxy endpoint Add /mcpproxy endpoint to execd that acts as an MCP gateway, proxying and aggregating multiple upstream MCP servers (stdio subprocess or HTTP) through a single streamable-http endpoint. Upstream servers are managed dynamically via REST API at /mcpproxy/upstreams. New package execd/pkg/mcpproxy with stdio/HTTP upstream transports, tool aggregation with conflict detection, session management, and JSON-RPC 2.0 routing. Includes unit tests, controller tests, and a smoke test using an inline mock MCP server. Windows compatible via platform-specific process management files. --- .github/workflows/execd-test.yml | 6 + components/execd/main.go | 1 + .../execd/pkg/mcpproxy/http_upstream.go | 219 ++++++++ components/execd/pkg/mcpproxy/jsonrpc.go | 71 +++ components/execd/pkg/mcpproxy/jsonrpc_test.go | 83 ++++ components/execd/pkg/mcpproxy/manager.go | 313 ++++++++++++ components/execd/pkg/mcpproxy/manager_test.go | 240 +++++++++ components/execd/pkg/mcpproxy/process.go | 39 ++ .../execd/pkg/mcpproxy/process_windows.go | 28 ++ components/execd/pkg/mcpproxy/session.go | 35 ++ .../execd/pkg/mcpproxy/stdio_upstream.go | 312 ++++++++++++ components/execd/pkg/mcpproxy/upstream.go | 79 +++ components/execd/pkg/web/controller/mcp.go | 203 ++++++++ .../execd/pkg/web/controller/mcp_test.go | 285 +++++++++++ components/execd/pkg/web/model/mcp.go | 29 ++ components/execd/pkg/web/router.go | 21 + components/execd/tests/smoke_mcpproxy.py | 467 ++++++++++++++++++ 17 files changed, 2431 insertions(+) create mode 100644 components/execd/pkg/mcpproxy/http_upstream.go create mode 100644 components/execd/pkg/mcpproxy/jsonrpc.go create mode 100644 components/execd/pkg/mcpproxy/jsonrpc_test.go create mode 100644 components/execd/pkg/mcpproxy/manager.go create mode 100644 components/execd/pkg/mcpproxy/manager_test.go create mode 100644 components/execd/pkg/mcpproxy/process.go create mode 100644 components/execd/pkg/mcpproxy/process_windows.go create mode 100644 components/execd/pkg/mcpproxy/session.go create mode 100644 components/execd/pkg/mcpproxy/stdio_upstream.go create mode 100644 components/execd/pkg/mcpproxy/upstream.go create mode 100644 components/execd/pkg/web/controller/mcp.go create mode 100644 components/execd/pkg/web/controller/mcp_test.go create mode 100644 components/execd/pkg/web/model/mcp.go create mode 100644 components/execd/tests/smoke_mcpproxy.py diff --git a/.github/workflows/execd-test.yml b/.github/workflows/execd-test.yml index 729e2d0f4..995ab31cc 100644 --- a/.github/workflows/execd-test.yml +++ b/.github/workflows/execd-test.yml @@ -109,6 +109,12 @@ jobs: sleep 5 python3 tests/smoke_api.py + + - name: Run MCP proxy smoke test + working-directory: components/execd + run: | + python3 tests/smoke_mcpproxy.py + - name: SIGTERM forward test if: matrix.os == 'ubuntu-latest' working-directory: components/execd diff --git a/components/execd/main.go b/components/execd/main.go index 0700759b2..999c02e49 100644 --- a/components/execd/main.go +++ b/components/execd/main.go @@ -59,6 +59,7 @@ func main() { } controller.InitCodeRunner() + controller.InitMCPProxy() engine := web.NewRouter(flag.ServerAccessToken) addr := fmt.Sprintf(":%d", flag.ServerPort) listener, err := net.Listen("tcp4", addr) diff --git a/components/execd/pkg/mcpproxy/http_upstream.go b/components/execd/pkg/mcpproxy/http_upstream.go new file mode 100644 index 000000000..3573c658a --- /dev/null +++ b/components/execd/pkg/mcpproxy/http_upstream.go @@ -0,0 +1,219 @@ +// Copyright 2025 Alibaba Group Holding Ltd. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package mcpproxy + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "sync/atomic" + "time" +) + +const mcpSessionHeader = "Mcp-Session-Id" + +type httpUpstream struct { + name string + url string + headers map[string]string // custom headers sent on every request + client *http.Client + sessionID string + nextID atomic.Int64 + tools []Tool +} + +func newHTTPUpstream(config UpstreamConfig) *httpUpstream { + return &httpUpstream{ + name: config.Name, + url: config.URL, + headers: config.Headers, + client: &http.Client{ + Timeout: callTimeout, + }, + } +} + +func (h *httpUpstream) Name() string { return h.name } +func (h *httpUpstream) Transport() string { return "http" } + +func (h *httpUpstream) post(ctx context.Context, method string, params any) (*Response, error) { + id := h.nextID.Add(1) + reqBody := struct { + JSONRPC string `json:"jsonrpc"` + ID int64 `json:"id"` + Method string `json:"method"` + Params any `json:"params,omitempty"` + }{ + JSONRPC: jsonrpcVersion, + ID: id, + Method: method, + Params: params, + } + + data, err := json.Marshal(reqBody) + if err != nil { + return nil, fmt.Errorf("marshal: %w", err) + } + + httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, h.url, bytes.NewReader(data)) + if err != nil { + return nil, fmt.Errorf("new request: %w", err) + } + httpReq.Header.Set("Content-Type", "application/json") + httpReq.Header.Set("Accept", "application/json, text/event-stream") + for k, v := range h.headers { + httpReq.Header.Set(k, v) + } + if h.sessionID != "" { + httpReq.Header.Set(mcpSessionHeader, h.sessionID) + } + + httpResp, err := h.client.Do(httpReq) + if err != nil { + return nil, fmt.Errorf("http post: %w", err) + } + defer httpResp.Body.Close() + + if sid := httpResp.Header.Get(mcpSessionHeader); sid != "" { + h.sessionID = sid + } + + body, err := io.ReadAll(httpResp.Body) + if err != nil { + return nil, fmt.Errorf("read response: %w", err) + } + + if httpResp.StatusCode != http.StatusOK { + return nil, fmt.Errorf("upstream %s returned %d: %s", h.name, httpResp.StatusCode, string(body)) + } + + var resp Response + if err := json.Unmarshal(body, &resp); err != nil { + return nil, fmt.Errorf("unmarshal response: %w", err) + } + return &resp, nil +} + +func (h *httpUpstream) Initialize(ctx context.Context) error { + ctx, cancel := context.WithTimeout(ctx, initTimeout) + defer cancel() + + params := initializeParams{ + ProtocolVersion: "2025-03-26", + Capabilities: map[string]any{}, + ClientInfo: clientInfo{Name: "opensandbox-execd", Version: "1.0.0"}, + } + + resp, err := h.post(ctx, "initialize", params) + if err != nil { + return fmt.Errorf("initialize: %w", err) + } + if resp.Error != nil { + return fmt.Errorf("initialize error: %s", resp.Error.Message) + } + + // Send initialized notification (no response expected). + notif := struct { + JSONRPC string `json:"jsonrpc"` + Method string `json:"method"` + }{ + JSONRPC: jsonrpcVersion, + Method: "notifications/initialized", + } + data, _ := json.Marshal(notif) + httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, h.url, bytes.NewReader(data)) + if err != nil { + return fmt.Errorf("initialized notification request: %w", err) + } + httpReq.Header.Set("Content-Type", "application/json") + httpReq.Header.Set("Accept", "application/json, text/event-stream") + for k, v := range h.headers { + httpReq.Header.Set(k, v) + } + if h.sessionID != "" { + httpReq.Header.Set(mcpSessionHeader, h.sessionID) + } + notifResp, err := h.client.Do(httpReq) + if err != nil { + return fmt.Errorf("initialized notification: %w", err) + } + notifResp.Body.Close() + + return nil +} + +func (h *httpUpstream) Tools(ctx context.Context) ([]Tool, error) { + if h.tools != nil { + return h.tools, nil + } + + ctx, cancel := context.WithTimeout(ctx, initTimeout) + defer cancel() + + resp, err := h.post(ctx, "tools/list", struct{}{}) + if err != nil { + return nil, fmt.Errorf("tools/list: %w", err) + } + if resp.Error != nil { + return nil, fmt.Errorf("tools/list error: %s", resp.Error.Message) + } + + var result toolsListResult + if err := json.Unmarshal(resp.Result, &result); err != nil { + return nil, fmt.Errorf("parse tools/list result: %w", err) + } + + h.tools = result.Tools + return h.tools, nil +} + +func (h *httpUpstream) CallTool(ctx context.Context, name string, arguments json.RawMessage) (*Response, error) { + ctx, cancel := context.WithTimeout(ctx, callTimeout) + defer cancel() + + params := callToolParams{ + Name: name, + Arguments: arguments, + } + return h.post(ctx, "tools/call", params) +} + +func (h *httpUpstream) Close() error { + if h.sessionID == "" { + return nil + } + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + req, err := http.NewRequestWithContext(ctx, http.MethodDelete, h.url, nil) + if err != nil { + return err + } + for k, v := range h.headers { + req.Header.Set(k, v) + } + req.Header.Set(mcpSessionHeader, h.sessionID) + + resp, err := h.client.Do(req) + if err != nil { + return err + } + resp.Body.Close() + return nil +} diff --git a/components/execd/pkg/mcpproxy/jsonrpc.go b/components/execd/pkg/mcpproxy/jsonrpc.go new file mode 100644 index 000000000..6baa4896e --- /dev/null +++ b/components/execd/pkg/mcpproxy/jsonrpc.go @@ -0,0 +1,71 @@ +// Copyright 2025 Alibaba Group Holding Ltd. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package mcpproxy + +import "encoding/json" + +const jsonrpcVersion = "2.0" + +// Request is a JSON-RPC 2.0 request or notification. +// When ID is nil the message is a notification (no response expected). +type Request struct { + JSONRPC string `json:"jsonrpc"` + ID any `json:"id,omitempty"` + Method string `json:"method"` + Params json.RawMessage `json:"params,omitempty"` +} + +// IsNotification returns true when the request has no ID (notification). +func (r *Request) IsNotification() bool { return r.ID == nil } + +// Response is a JSON-RPC 2.0 response (success or error). +type Response struct { + JSONRPC string `json:"jsonrpc"` + ID any `json:"id,omitempty"` + Result json.RawMessage `json:"result,omitempty"` + Error *RPCError `json:"error,omitempty"` +} + +// RPCError represents a JSON-RPC 2.0 error object. +type RPCError struct { + Code int `json:"code"` + Message string `json:"message"` + Data any `json:"data,omitempty"` +} + +// Standard JSON-RPC error codes. +const ( + CodeParseError = -32700 + CodeInvalidRequest = -32600 + CodeMethodNotFound = -32601 + CodeInvalidParams = -32602 + CodeInternalError = -32603 +) + +func newErrorResponse(id any, code int, message string) *Response { + return &Response{ + JSONRPC: jsonrpcVersion, + ID: id, + Error: &RPCError{Code: code, Message: message}, + } +} + +func newResultResponse(id any, result json.RawMessage) *Response { + return &Response{ + JSONRPC: jsonrpcVersion, + ID: id, + Result: result, + } +} diff --git a/components/execd/pkg/mcpproxy/jsonrpc_test.go b/components/execd/pkg/mcpproxy/jsonrpc_test.go new file mode 100644 index 000000000..ea26fa3b5 --- /dev/null +++ b/components/execd/pkg/mcpproxy/jsonrpc_test.go @@ -0,0 +1,83 @@ +// Copyright 2025 Alibaba Group Holding Ltd. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package mcpproxy + +import ( + "encoding/json" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestRequestIsNotification(t *testing.T) { + t.Run("with id", func(t *testing.T) { + r := &Request{ID: 1, Method: "ping"} + require.False(t, r.IsNotification()) + }) + t.Run("without id", func(t *testing.T) { + r := &Request{Method: "notifications/initialized"} + require.True(t, r.IsNotification()) + }) +} + +func TestNewErrorResponse(t *testing.T) { + resp := newErrorResponse(42, CodeMethodNotFound, "not found") + require.Equal(t, jsonrpcVersion, resp.JSONRPC) + require.Equal(t, 42, resp.ID) + require.NotNil(t, resp.Error) + require.Equal(t, CodeMethodNotFound, resp.Error.Code) + require.Equal(t, "not found", resp.Error.Message) + require.Nil(t, resp.Result) +} + +func TestNewResultResponse(t *testing.T) { + data, _ := json.Marshal(map[string]string{"ok": "true"}) + resp := newResultResponse("req-1", data) + require.Equal(t, jsonrpcVersion, resp.JSONRPC) + require.Equal(t, "req-1", resp.ID) + require.Nil(t, resp.Error) + require.JSONEq(t, `{"ok":"true"}`, string(resp.Result)) +} + +func TestNormalizeID(t *testing.T) { + tests := []struct { + input any + want string + }{ + {float64(42), "42"}, + {int64(7), "7"}, + {"abc", "abc"}, + {nil, ""}, + } + for _, tt := range tests { + require.Equal(t, tt.want, normalizeID(tt.input)) + } +} + +func TestRequestRoundTrip(t *testing.T) { + orig := Request{ + JSONRPC: "2.0", + ID: float64(1), + Method: "tools/call", + Params: json.RawMessage(`{"name":"echo"}`), + } + data, err := json.Marshal(orig) + require.NoError(t, err) + + var decoded Request + require.NoError(t, json.Unmarshal(data, &decoded)) + require.Equal(t, orig.Method, decoded.Method) + require.JSONEq(t, `{"name":"echo"}`, string(decoded.Params)) +} diff --git a/components/execd/pkg/mcpproxy/manager.go b/components/execd/pkg/mcpproxy/manager.go new file mode 100644 index 000000000..67215ada8 --- /dev/null +++ b/components/execd/pkg/mcpproxy/manager.go @@ -0,0 +1,313 @@ +// Copyright 2025 Alibaba Group Holding Ltd. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package mcpproxy + +import ( + "context" + "encoding/json" + "fmt" + "sync" + + "github.com/alibaba/opensandbox/execd/pkg/log" +) + +// Manager aggregates multiple upstream MCP servers and proxies requests. +type Manager struct { + mu sync.RWMutex + upstreams map[string]Upstream // name → upstream + toolIndex map[string]string // tool name → upstream name + allTools []Tool // cached aggregated list + sessions map[string]*Session // Mcp-Session-Id → session +} + +// NewManager creates an empty proxy manager. +func NewManager() *Manager { + return &Manager{ + upstreams: make(map[string]Upstream), + toolIndex: make(map[string]string), + sessions: make(map[string]*Session), + } +} + +// AddUpstream connects to an upstream MCP server, initializes it, fetches +// its tool list, and registers the tools. Returns error if any tool name +// conflicts with an already-registered tool. +func (m *Manager) AddUpstream(ctx context.Context, config UpstreamConfig) (*UpstreamInfo, error) { + m.mu.Lock() + defer m.mu.Unlock() + + if _, exists := m.upstreams[config.Name]; exists { + return nil, fmt.Errorf("upstream %q already registered", config.Name) + } + + var up Upstream + switch config.Transport { + case "stdio": + up = newStdioUpstream(config) + case "http": + up = newHTTPUpstream(config) + default: + return nil, fmt.Errorf("unsupported transport: %s", config.Transport) + } + + if err := up.Initialize(ctx); err != nil { + _ = up.Close() + return nil, fmt.Errorf("initialize upstream %q: %w", config.Name, err) + } + + tools, err := up.Tools(ctx) + if err != nil { + _ = up.Close() + return nil, fmt.Errorf("fetch tools from %q: %w", config.Name, err) + } + + // Check for conflicts before registering any tools. + for _, t := range tools { + if owner, exists := m.toolIndex[t.Name]; exists { + _ = up.Close() + return nil, fmt.Errorf("tool %q conflicts: already registered by upstream %q", t.Name, owner) + } + } + + m.upstreams[config.Name] = up + toolNames := make([]string, 0, len(tools)) + for _, t := range tools { + m.toolIndex[t.Name] = config.Name + toolNames = append(toolNames, t.Name) + } + m.rebuildToolList() + + log.Info("mcp proxy: registered upstream %q (%s) with %d tools", config.Name, config.Transport, len(tools)) + + return &UpstreamInfo{ + Name: config.Name, + Transport: config.Transport, + Status: "running", + Tools: toolNames, + }, nil +} + +// RemoveUpstream closes an upstream and removes its tools. +func (m *Manager) RemoveUpstream(name string) error { + m.mu.Lock() + defer m.mu.Unlock() + + up, exists := m.upstreams[name] + if !exists { + return fmt.Errorf("upstream %q not found", name) + } + + if err := up.Close(); err != nil { + log.Warning("mcp proxy: error closing upstream %q: %v", name, err) + } + + delete(m.upstreams, name) + for toolName, owner := range m.toolIndex { + if owner == name { + delete(m.toolIndex, toolName) + } + } + m.rebuildToolList() + + log.Info("mcp proxy: removed upstream %q", name) + return nil +} + +// ListUpstreams returns status of all registered upstreams. +func (m *Manager) ListUpstreams() []UpstreamInfo { + m.mu.RLock() + defer m.mu.RUnlock() + + infos := make([]UpstreamInfo, 0, len(m.upstreams)) + for name, up := range m.upstreams { + var toolNames []string + for toolName, owner := range m.toolIndex { + if owner == name { + toolNames = append(toolNames, toolName) + } + } + infos = append(infos, UpstreamInfo{ + Name: name, + Transport: up.Transport(), + Status: "running", + Tools: toolNames, + }) + } + return infos +} + +// GetUpstream returns info for a single upstream. +func (m *Manager) GetUpstream(name string) (*UpstreamInfo, error) { + m.mu.RLock() + defer m.mu.RUnlock() + + up, exists := m.upstreams[name] + if !exists { + return nil, fmt.Errorf("upstream %q not found", name) + } + + var toolNames []string + for toolName, owner := range m.toolIndex { + if owner == name { + toolNames = append(toolNames, toolName) + } + } + + return &UpstreamInfo{ + Name: name, + Transport: up.Transport(), + Status: "running", + Tools: toolNames, + }, nil +} + +// HandleRequest routes a downstream MCP JSON-RPC request to the appropriate handler. +// Returns the response and optionally a new session ID (for initialize requests). +func (m *Manager) HandleRequest(ctx context.Context, req *Request, sessionID string) (*Response, string, error) { + switch req.Method { + case "initialize": + return m.handleInitialize(req) + case "notifications/initialized": + return nil, "", nil // notification, no response + case "ping": + return m.handlePing(req) + case "tools/list": + return m.handleToolsList(req, sessionID) + case "tools/call": + return m.handleToolsCall(ctx, req, sessionID) + default: + return newErrorResponse(req.ID, CodeMethodNotFound, fmt.Sprintf("method %q not supported", req.Method)), "", nil + } +} + +func (m *Manager) handleInitialize(req *Request) (*Response, string, error) { + session := newSession() + + m.mu.Lock() + m.sessions[session.ID] = session + m.mu.Unlock() + + result := map[string]any{ + "protocolVersion": "2025-03-26", + "capabilities": map[string]any{ + "tools": map[string]any{ + "listChanged": true, + }, + }, + "serverInfo": map[string]any{ + "name": "opensandbox-execd-mcp-proxy", + "version": "1.0.0", + }, + } + + data, _ := json.Marshal(result) + return newResultResponse(req.ID, data), session.ID, nil +} + +func (m *Manager) handlePing(req *Request) (*Response, string, error) { + data, _ := json.Marshal(struct{}{}) + return newResultResponse(req.ID, data), "", nil +} + +func (m *Manager) handleToolsList(req *Request, sessionID string) (*Response, string, error) { + m.mu.RLock() + tools := m.allTools + m.mu.RUnlock() + + if tools == nil { + tools = []Tool{} + } + + result := toolsListResult{Tools: tools} + data, _ := json.Marshal(result) + return newResultResponse(req.ID, data), "", nil +} + +func (m *Manager) handleToolsCall(ctx context.Context, req *Request, sessionID string) (*Response, string, error) { + var params callToolParams + if err := json.Unmarshal(req.Params, ¶ms); err != nil { + return newErrorResponse(req.ID, CodeInvalidParams, "invalid tools/call params"), "", nil + } + + m.mu.RLock() + upstreamName, exists := m.toolIndex[params.Name] + var up Upstream + if exists { + up = m.upstreams[upstreamName] + } + m.mu.RUnlock() + + if !exists || up == nil { + return newErrorResponse(req.ID, CodeMethodNotFound, fmt.Sprintf("tool %q not found", params.Name)), "", nil + } + + resp, err := up.CallTool(ctx, params.Name, params.Arguments) + if err != nil { + return newErrorResponse(req.ID, CodeInternalError, err.Error()), "", nil + } + + // Rewrite the response ID to match the downstream request. + resp.ID = req.ID + return resp, "", nil +} + +// CreateSession creates and registers a new downstream session. +func (m *Manager) CreateSession() *Session { + session := newSession() + m.mu.Lock() + m.sessions[session.ID] = session + m.mu.Unlock() + return session +} + +// GetSession looks up a session by ID. +func (m *Manager) GetSession(id string) *Session { + m.mu.RLock() + defer m.mu.RUnlock() + return m.sessions[id] +} + +// DeleteSession removes a session. +func (m *Manager) DeleteSession(id string) { + m.mu.Lock() + delete(m.sessions, id) + m.mu.Unlock() +} + +// Close shuts down all upstreams. +func (m *Manager) Close() { + m.mu.Lock() + defer m.mu.Unlock() + + for name, up := range m.upstreams { + if err := up.Close(); err != nil { + log.Warning("mcp proxy: error closing upstream %q: %v", name, err) + } + } + m.upstreams = make(map[string]Upstream) + m.toolIndex = make(map[string]string) + m.allTools = nil + m.sessions = make(map[string]*Session) +} + +func (m *Manager) rebuildToolList() { + var tools []Tool + for _, up := range m.upstreams { + if cached, _ := up.Tools(context.Background()); cached != nil { + tools = append(tools, cached...) + } + } + m.allTools = tools +} diff --git a/components/execd/pkg/mcpproxy/manager_test.go b/components/execd/pkg/mcpproxy/manager_test.go new file mode 100644 index 000000000..72d0888eb --- /dev/null +++ b/components/execd/pkg/mcpproxy/manager_test.go @@ -0,0 +1,240 @@ +// Copyright 2025 Alibaba Group Holding Ltd. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package mcpproxy + +import ( + "context" + "encoding/json" + "testing" + + "github.com/stretchr/testify/require" +) + +// fakeUpstream is a minimal Upstream for unit tests. +type fakeUpstream struct { + name string + tools []Tool +} + +func (f *fakeUpstream) Name() string { return f.name } +func (f *fakeUpstream) Transport() string { return "fake" } +func (f *fakeUpstream) Initialize(_ context.Context) error { return nil } +func (f *fakeUpstream) Tools(_ context.Context) ([]Tool, error) { + return f.tools, nil +} +func (f *fakeUpstream) CallTool(_ context.Context, name string, args json.RawMessage) (*Response, error) { + result, _ := json.Marshal(map[string]any{ + "content": []map[string]string{{"type": "text", "text": "called:" + name}}, + }) + return newResultResponse(nil, result), nil +} +func (f *fakeUpstream) Close() error { return nil } + +func newTestManager(t *testing.T, upstreams ...*fakeUpstream) *Manager { + t.Helper() + m := NewManager() + for _, up := range upstreams { + m.mu.Lock() + m.upstreams[up.name] = up + for _, tool := range up.tools { + m.toolIndex[tool.Name] = up.name + } + m.mu.Unlock() + } + m.mu.Lock() + m.rebuildToolList() + m.mu.Unlock() + return m +} + +func TestManagerSessionLifecycle(t *testing.T) { + m := NewManager() + + s := m.CreateSession() + require.NotEmpty(t, s.ID) + + got := m.GetSession(s.ID) + require.Equal(t, s.ID, got.ID) + + m.DeleteSession(s.ID) + require.Nil(t, m.GetSession(s.ID)) +} + +func TestManagerHandleInitialize(t *testing.T) { + m := NewManager() + req := &Request{ + JSONRPC: "2.0", + ID: float64(1), + Method: "initialize", + Params: json.RawMessage(`{"protocolVersion":"2025-03-26"}`), + } + + resp, sid, err := m.HandleRequest(context.Background(), req, "") + require.NoError(t, err) + require.NotEmpty(t, sid) + require.NotNil(t, resp) + require.Nil(t, resp.Error) + + var result map[string]any + require.NoError(t, json.Unmarshal(resp.Result, &result)) + require.Equal(t, "2025-03-26", result["protocolVersion"]) +} + +func TestManagerHandlePing(t *testing.T) { + m := NewManager() + req := &Request{JSONRPC: "2.0", ID: float64(1), Method: "ping"} + + resp, _, err := m.HandleRequest(context.Background(), req, "") + require.NoError(t, err) + require.Nil(t, resp.Error) +} + +func TestManagerHandleToolsList(t *testing.T) { + m := newTestManager(t, &fakeUpstream{ + name: "a", + tools: []Tool{ + {Name: "tool1", Description: "desc1"}, + {Name: "tool2", Description: "desc2"}, + }, + }) + + req := &Request{JSONRPC: "2.0", ID: float64(2), Method: "tools/list", Params: json.RawMessage(`{}`)} + resp, _, err := m.HandleRequest(context.Background(), req, "s1") + require.NoError(t, err) + + var result toolsListResult + require.NoError(t, json.Unmarshal(resp.Result, &result)) + require.Len(t, result.Tools, 2) +} + +func TestManagerHandleToolsCall(t *testing.T) { + m := newTestManager(t, &fakeUpstream{ + name: "a", + tools: []Tool{{Name: "echo"}}, + }) + + req := &Request{ + JSONRPC: "2.0", + ID: float64(3), + Method: "tools/call", + Params: json.RawMessage(`{"name":"echo","arguments":{"msg":"hi"}}`), + } + + resp, _, err := m.HandleRequest(context.Background(), req, "s1") + require.NoError(t, err) + require.Nil(t, resp.Error) + require.Contains(t, string(resp.Result), "called:echo") + require.Equal(t, float64(3), resp.ID) +} + +func TestManagerHandleToolsCallUnknownTool(t *testing.T) { + m := NewManager() + req := &Request{ + JSONRPC: "2.0", + ID: float64(4), + Method: "tools/call", + Params: json.RawMessage(`{"name":"nonexistent","arguments":{}}`), + } + + resp, _, err := m.HandleRequest(context.Background(), req, "s1") + require.NoError(t, err) + require.NotNil(t, resp.Error) + require.Contains(t, resp.Error.Message, "not found") +} + +func TestManagerHandleUnknownMethod(t *testing.T) { + m := NewManager() + req := &Request{JSONRPC: "2.0", ID: float64(5), Method: "resources/list"} + + resp, _, err := m.HandleRequest(context.Background(), req, "s1") + require.NoError(t, err) + require.NotNil(t, resp.Error) + require.Equal(t, CodeMethodNotFound, resp.Error.Code) +} + +func TestManagerToolConflict(t *testing.T) { + m := newTestManager(t, &fakeUpstream{ + name: "a", + tools: []Tool{{Name: "shared_tool"}}, + }) + + // Manually try adding a second upstream with conflicting tool. + m.mu.Lock() + _, exists := m.toolIndex["shared_tool"] + m.mu.Unlock() + require.True(t, exists) +} + +func TestManagerRemoveUpstream(t *testing.T) { + m := newTestManager(t, &fakeUpstream{ + name: "a", + tools: []Tool{{Name: "tool1"}}, + }) + + require.NoError(t, m.RemoveUpstream("a")) + + m.mu.RLock() + _, exists := m.toolIndex["tool1"] + m.mu.RUnlock() + require.False(t, exists) + + require.Len(t, m.allTools, 0) +} + +func TestManagerRemoveUnknownUpstream(t *testing.T) { + m := NewManager() + err := m.RemoveUpstream("nonexistent") + require.Error(t, err) +} + +func TestManagerListUpstreams(t *testing.T) { + m := newTestManager(t, + &fakeUpstream{name: "a", tools: []Tool{{Name: "t1"}}}, + &fakeUpstream{name: "b", tools: []Tool{{Name: "t2"}}}, + ) + + infos := m.ListUpstreams() + require.Len(t, infos, 2) +} + +func TestManagerGetUpstream(t *testing.T) { + m := newTestManager(t, &fakeUpstream{name: "a", tools: []Tool{{Name: "t1"}}}) + + info, err := m.GetUpstream("a") + require.NoError(t, err) + require.Equal(t, "a", info.Name) + + _, err = m.GetUpstream("nonexistent") + require.Error(t, err) +} + +func TestManagerClose(t *testing.T) { + m := newTestManager(t, &fakeUpstream{name: "a", tools: []Tool{{Name: "t1"}}}) + + m.Close() + require.Len(t, m.upstreams, 0) + require.Len(t, m.toolIndex, 0) + require.Nil(t, m.allTools) +} + +func TestManagerNotificationReturnsNil(t *testing.T) { + m := NewManager() + req := &Request{JSONRPC: "2.0", Method: "notifications/initialized"} + + resp, sid, err := m.HandleRequest(context.Background(), req, "") + require.NoError(t, err) + require.Nil(t, resp) + require.Empty(t, sid) +} diff --git a/components/execd/pkg/mcpproxy/process.go b/components/execd/pkg/mcpproxy/process.go new file mode 100644 index 000000000..7b18e7c59 --- /dev/null +++ b/components/execd/pkg/mcpproxy/process.go @@ -0,0 +1,39 @@ +// Copyright 2025 Alibaba Group Holding Ltd. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//go:build !windows +// +build !windows + +package mcpproxy + +import ( + "os/exec" + "syscall" +) + +// setProcAttr configures the subprocess to run in its own process group +// so Close() can kill the entire tree. +func setProcAttr(cmd *exec.Cmd) { + cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true} +} + +// forceKillProcess kills the subprocess and its process group. +func forceKillProcess(cmd *exec.Cmd) { + pgid, err := syscall.Getpgid(cmd.Process.Pid) + if err == nil { + _ = syscall.Kill(-pgid, syscall.SIGKILL) + } else { + _ = cmd.Process.Kill() + } +} diff --git a/components/execd/pkg/mcpproxy/process_windows.go b/components/execd/pkg/mcpproxy/process_windows.go new file mode 100644 index 000000000..798e91c3b --- /dev/null +++ b/components/execd/pkg/mcpproxy/process_windows.go @@ -0,0 +1,28 @@ +// Copyright 2025 Alibaba Group Holding Ltd. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//go:build windows +// +build windows + +package mcpproxy + +import "os/exec" + +// setProcAttr is a no-op on Windows (no process groups via SysProcAttr). +func setProcAttr(cmd *exec.Cmd) {} + +// forceKillProcess terminates the subprocess on Windows. +func forceKillProcess(cmd *exec.Cmd) { + _ = cmd.Process.Kill() +} diff --git a/components/execd/pkg/mcpproxy/session.go b/components/execd/pkg/mcpproxy/session.go new file mode 100644 index 000000000..ec85bd996 --- /dev/null +++ b/components/execd/pkg/mcpproxy/session.go @@ -0,0 +1,35 @@ +// Copyright 2025 Alibaba Group Holding Ltd. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package mcpproxy + +import ( + "time" + + "github.com/google/uuid" +) + +// Session tracks a downstream MCP client connection. +type Session struct { + ID string + CreatedAt time.Time + Initialized bool +} + +func newSession() *Session { + return &Session{ + ID: uuid.New().String(), + CreatedAt: time.Now(), + } +} diff --git a/components/execd/pkg/mcpproxy/stdio_upstream.go b/components/execd/pkg/mcpproxy/stdio_upstream.go new file mode 100644 index 000000000..b4d8ebae1 --- /dev/null +++ b/components/execd/pkg/mcpproxy/stdio_upstream.go @@ -0,0 +1,312 @@ +// Copyright 2025 Alibaba Group Holding Ltd. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package mcpproxy + +import ( + "bufio" + "context" + "encoding/json" + "fmt" + "io" + "os" + "os/exec" + "sync" + "sync/atomic" + "time" + + "github.com/alibaba/opensandbox/execd/pkg/log" +) + +const ( + initTimeout = 30 * time.Second + callTimeout = 5 * time.Minute + shutdownTimeout = 5 * time.Second +) + +type stdioUpstream struct { + name string + config UpstreamConfig + + cmd *exec.Cmd + stdin io.WriteCloser + stdout io.ReadCloser + stderr io.ReadCloser + + writeMu sync.Mutex + nextID atomic.Int64 + pending sync.Map // map[string]chan *Response — keyed by JSON-encoded ID + + tools []Tool + done chan struct{} +} + +func newStdioUpstream(config UpstreamConfig) *stdioUpstream { + return &stdioUpstream{ + name: config.Name, + config: config, + done: make(chan struct{}), + } +} + +func (s *stdioUpstream) Name() string { return s.name } +func (s *stdioUpstream) Transport() string { return "stdio" } + +func (s *stdioUpstream) start() error { + s.cmd = exec.Command(s.config.Command, s.config.Args...) + setProcAttr(s.cmd) + + env := os.Environ() + for k, v := range s.config.Env { + env = append(env, k+"="+v) + } + s.cmd.Env = env + + var err error + s.stdin, err = s.cmd.StdinPipe() + if err != nil { + return fmt.Errorf("stdin pipe: %w", err) + } + s.stdout, err = s.cmd.StdoutPipe() + if err != nil { + return fmt.Errorf("stdout pipe: %w", err) + } + s.stderr, err = s.cmd.StderrPipe() + if err != nil { + return fmt.Errorf("stderr pipe: %w", err) + } + + if err := s.cmd.Start(); err != nil { + return fmt.Errorf("start %s: %w", s.config.Command, err) + } + + go s.readLoop() + go s.drainStderr() + go s.waitLoop() + + return nil +} + +func (s *stdioUpstream) readLoop() { + scanner := bufio.NewScanner(s.stdout) + scanner.Buffer(make([]byte, 0, 64*1024), 16*1024*1024) + + for scanner.Scan() { + line := scanner.Bytes() + if len(line) == 0 { + continue + } + + var resp Response + if err := json.Unmarshal(line, &resp); err != nil { + log.Warning("mcp upstream %s: invalid JSON from stdout: %v", s.name, err) + continue + } + + if resp.ID == nil { + log.Info("mcp upstream %s: notification: %s", s.name, string(line)) + continue + } + + key := normalizeID(resp.ID) + if ch, ok := s.pending.LoadAndDelete(key); ok { + ch.(chan *Response) <- &resp + } else { + log.Warning("mcp upstream %s: unexpected response id=%v", s.name, resp.ID) + } + } + + if err := scanner.Err(); err != nil { + log.Warning("mcp upstream %s: stdout read error: %v", s.name, err) + } +} + +func (s *stdioUpstream) drainStderr() { + scanner := bufio.NewScanner(s.stderr) + scanner.Buffer(make([]byte, 0, 4096), 1024*1024) + for scanner.Scan() { + log.Info("mcp upstream %s stderr: %s", s.name, scanner.Text()) + } +} + +func (s *stdioUpstream) waitLoop() { + _ = s.cmd.Wait() + close(s.done) +} + +func (s *stdioUpstream) sendRequest(ctx context.Context, method string, params any) (*Response, error) { + id := s.nextID.Add(1) + req := struct { + JSONRPC string `json:"jsonrpc"` + ID int64 `json:"id"` + Method string `json:"method"` + Params any `json:"params,omitempty"` + }{ + JSONRPC: jsonrpcVersion, + ID: id, + Method: method, + Params: params, + } + + data, err := json.Marshal(req) + if err != nil { + return nil, fmt.Errorf("marshal request: %w", err) + } + + key := normalizeID(id) + ch := make(chan *Response, 1) + s.pending.Store(key, ch) + defer s.pending.Delete(key) + + s.writeMu.Lock() + _, err = s.stdin.Write(append(data, '\n')) + s.writeMu.Unlock() + if err != nil { + return nil, fmt.Errorf("write to stdin: %w", err) + } + + select { + case resp := <-ch: + return resp, nil + case <-ctx.Done(): + return nil, ctx.Err() + case <-s.done: + return nil, fmt.Errorf("upstream %s process exited", s.name) + } +} + +func (s *stdioUpstream) sendNotification(method string, params any) error { + msg := struct { + JSONRPC string `json:"jsonrpc"` + Method string `json:"method"` + Params any `json:"params,omitempty"` + }{ + JSONRPC: jsonrpcVersion, + Method: method, + Params: params, + } + + data, err := json.Marshal(msg) + if err != nil { + return fmt.Errorf("marshal notification: %w", err) + } + + s.writeMu.Lock() + defer s.writeMu.Unlock() + _, err = s.stdin.Write(append(data, '\n')) + return err +} + +func (s *stdioUpstream) Initialize(ctx context.Context) error { + if err := s.start(); err != nil { + return err + } + + ctx, cancel := context.WithTimeout(ctx, initTimeout) + defer cancel() + + params := initializeParams{ + ProtocolVersion: "2025-03-26", + Capabilities: map[string]any{}, + ClientInfo: clientInfo{Name: "opensandbox-execd", Version: "1.0.0"}, + } + + resp, err := s.sendRequest(ctx, "initialize", params) + if err != nil { + return fmt.Errorf("initialize: %w", err) + } + if resp.Error != nil { + return fmt.Errorf("initialize error: %s", resp.Error.Message) + } + + if err := s.sendNotification("notifications/initialized", nil); err != nil { + return fmt.Errorf("send initialized notification: %w", err) + } + + return nil +} + +func (s *stdioUpstream) Tools(ctx context.Context) ([]Tool, error) { + if s.tools != nil { + return s.tools, nil + } + + ctx, cancel := context.WithTimeout(ctx, initTimeout) + defer cancel() + + resp, err := s.sendRequest(ctx, "tools/list", struct{}{}) + if err != nil { + return nil, fmt.Errorf("tools/list: %w", err) + } + if resp.Error != nil { + return nil, fmt.Errorf("tools/list error: %s", resp.Error.Message) + } + + var result toolsListResult + if err := json.Unmarshal(resp.Result, &result); err != nil { + return nil, fmt.Errorf("parse tools/list result: %w", err) + } + + s.tools = result.Tools + return s.tools, nil +} + +func (s *stdioUpstream) CallTool(ctx context.Context, name string, arguments json.RawMessage) (*Response, error) { + ctx, cancel := context.WithTimeout(ctx, callTimeout) + defer cancel() + + params := callToolParams{ + Name: name, + Arguments: arguments, + } + + resp, err := s.sendRequest(ctx, "tools/call", params) + if err != nil { + return nil, fmt.Errorf("tools/call %s: %w", name, err) + } + return resp, nil +} + +func (s *stdioUpstream) Close() error { + if s.cmd == nil || s.cmd.Process == nil { + return nil + } + + _ = s.stdin.Close() + + select { + case <-s.done: + return nil + case <-time.After(shutdownTimeout): + } + + forceKillProcess(s.cmd) + + <-s.done + return nil +} + +// normalizeID converts a JSON-RPC ID to a stable string key. +func normalizeID(id any) string { + switch v := id.(type) { + case float64: + return fmt.Sprintf("%d", int64(v)) + case int64: + return fmt.Sprintf("%d", v) + case string: + return v + default: + return fmt.Sprintf("%v", v) + } +} diff --git a/components/execd/pkg/mcpproxy/upstream.go b/components/execd/pkg/mcpproxy/upstream.go new file mode 100644 index 000000000..3761e456f --- /dev/null +++ b/components/execd/pkg/mcpproxy/upstream.go @@ -0,0 +1,79 @@ +// Copyright 2025 Alibaba Group Holding Ltd. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package mcpproxy + +import ( + "context" + "encoding/json" +) + +// Upstream is a connection to a single upstream MCP server. +type Upstream interface { + Name() string + Transport() string // "stdio" or "http" + Initialize(ctx context.Context) error + Tools(ctx context.Context) ([]Tool, error) + CallTool(ctx context.Context, name string, arguments json.RawMessage) (*Response, error) + Close() error +} + +// Tool describes a single MCP tool advertised by an upstream. +type Tool struct { + Name string `json:"name"` + Description string `json:"description,omitempty"` + InputSchema json.RawMessage `json:"inputSchema,omitempty"` +} + +// UpstreamConfig describes how to connect to an upstream MCP server. +type UpstreamConfig struct { + Name string `json:"name"` + Transport string `json:"transport"` // "stdio" or "http" + Command string `json:"command,omitempty"` + Args []string `json:"args,omitempty"` + Env map[string]string `json:"env,omitempty"` + URL string `json:"url,omitempty"` + Headers map[string]string `json:"headers,omitempty"` // custom headers for http transport +} + +// UpstreamInfo is the status summary returned by management APIs. +type UpstreamInfo struct { + Name string `json:"name"` + Transport string `json:"transport"` + Status string `json:"status"` + Tools []string `json:"tools"` +} + +// toolsListResult matches the MCP tools/list response shape. +type toolsListResult struct { + Tools []Tool `json:"tools"` +} + +// callToolParams matches the MCP tools/call request params shape. +type callToolParams struct { + Name string `json:"name"` + Arguments json.RawMessage `json:"arguments,omitempty"` +} + +// initializeParams is sent by the proxy when initializing an upstream. +type initializeParams struct { + ProtocolVersion string `json:"protocolVersion"` + Capabilities map[string]any `json:"capabilities"` + ClientInfo clientInfo `json:"clientInfo"` +} + +type clientInfo struct { + Name string `json:"name"` + Version string `json:"version"` +} diff --git a/components/execd/pkg/web/controller/mcp.go b/components/execd/pkg/web/controller/mcp.go new file mode 100644 index 000000000..40f3814d5 --- /dev/null +++ b/components/execd/pkg/web/controller/mcp.go @@ -0,0 +1,203 @@ +// Copyright 2025 Alibaba Group Holding Ltd. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package controller + +import ( + "fmt" + "net/http" + + "github.com/gin-gonic/gin" + + "github.com/alibaba/opensandbox/execd/pkg/mcpproxy" + "github.com/alibaba/opensandbox/execd/pkg/web/model" +) + +var mcpManager *mcpproxy.Manager + +// InitMCPProxy initializes the MCP proxy manager singleton. +func InitMCPProxy() { + mcpManager = mcpproxy.NewManager() +} + +// MCPController handles /mcp endpoints. +type MCPController struct { + *basicController +} + +// NewMCPController creates a new MCPController from the current Gin context. +func NewMCPController(ctx *gin.Context) *MCPController { + return &MCPController{basicController: newBasicController(ctx)} +} + +// HandleRequest handles POST /mcp — the MCP streamable-http entry point. +// Accepts a JSON-RPC 2.0 request, routes it through the proxy manager, +// and returns the JSON-RPC response. +func (c *MCPController) HandleRequest() { + var req mcpproxy.Request + if err := c.bindJSON(&req); err != nil { + c.ctx.JSON(http.StatusBadRequest, mcpproxy.Response{ + JSONRPC: "2.0", + Error: &mcpproxy.RPCError{Code: mcpproxy.CodeParseError, Message: "invalid JSON: " + err.Error()}, + }) + return + } + + sessionID := c.ctx.GetHeader(model.McpSessionHeader) + + // For non-initialize requests, require a valid session. + if req.Method != "initialize" && !req.IsNotification() { + if sessionID == "" { + c.ctx.JSON(http.StatusBadRequest, mcpproxy.Response{ + JSONRPC: "2.0", + ID: req.ID, + Error: &mcpproxy.RPCError{Code: mcpproxy.CodeInvalidRequest, Message: "missing " + model.McpSessionHeader + " header"}, + }) + return + } + if mcpManager.GetSession(sessionID) == nil { + c.ctx.JSON(http.StatusNotFound, mcpproxy.Response{ + JSONRPC: "2.0", + ID: req.ID, + Error: &mcpproxy.RPCError{Code: mcpproxy.CodeInvalidRequest, Message: "unknown session"}, + }) + return + } + } + + resp, newSessionID, err := mcpManager.HandleRequest(c.ctx.Request.Context(), &req, sessionID) + if err != nil { + c.ctx.JSON(http.StatusInternalServerError, mcpproxy.Response{ + JSONRPC: "2.0", + ID: req.ID, + Error: &mcpproxy.RPCError{Code: mcpproxy.CodeInternalError, Message: err.Error()}, + }) + return + } + + // Notifications produce no response. + if resp == nil { + c.ctx.Status(http.StatusAccepted) + return + } + + if newSessionID != "" { + c.ctx.Header(model.McpSessionHeader, newSessionID) + } + + c.ctx.JSON(http.StatusOK, resp) +} + +// SSEStream handles GET /mcp — opens an SSE stream for server-initiated messages. +func (c *MCPController) SSEStream() { + sessionID := c.ctx.GetHeader(model.McpSessionHeader) + if sessionID == "" || mcpManager.GetSession(sessionID) == nil { + c.RespondError(http.StatusBadRequest, model.ErrorCodeInvalidRequest, "missing or invalid "+model.McpSessionHeader) + return + } + + c.setupSSEResponse() + c.ctx.Writer.Flush() + + // Keep the connection open. Currently we only send keepalive pings. + // Upstream notifications will be forwarded here in a future iteration. + ctx := c.ctx.Request.Context() + <-ctx.Done() +} + +// DeleteSession handles DELETE /mcp — terminates an MCP session. +func (c *MCPController) DeleteSession() { + sessionID := c.ctx.GetHeader(model.McpSessionHeader) + if sessionID == "" { + c.RespondError(http.StatusBadRequest, model.ErrorCodeMissingQuery, "missing "+model.McpSessionHeader+" header") + return + } + + mcpManager.DeleteSession(sessionID) + c.ctx.Status(http.StatusOK) +} + +// AddUpstream handles POST /mcp/upstreams — registers a new upstream MCP server. +func (c *MCPController) AddUpstream() { + var req model.AddUpstreamRequest + if err := c.bindJSON(&req); err != nil { + c.RespondError(http.StatusBadRequest, model.ErrorCodeInvalidRequest, fmt.Sprintf("invalid request: %v", err)) + return + } + + if req.Transport == "stdio" && req.Command == "" { + c.RespondError(http.StatusBadRequest, model.ErrorCodeInvalidRequest, "command is required for stdio transport") + return + } + if req.Transport == "http" && req.URL == "" { + c.RespondError(http.StatusBadRequest, model.ErrorCodeInvalidRequest, "url is required for http transport") + return + } + + config := mcpproxy.UpstreamConfig{ + Name: req.Name, + Transport: req.Transport, + Command: req.Command, + Args: req.Args, + Env: req.Env, + URL: req.URL, + Headers: req.Headers, + } + + info, err := mcpManager.AddUpstream(c.ctx.Request.Context(), config) + if err != nil { + c.RespondError(http.StatusInternalServerError, model.ErrorCodeRuntimeError, err.Error()) + return + } + + c.ctx.JSON(http.StatusCreated, info) +} + +// ListUpstreams handles GET /mcp/upstreams. +func (c *MCPController) ListUpstreams() { + c.RespondSuccess(mcpManager.ListUpstreams()) +} + +// GetUpstream handles GET /mcp/upstreams/:name. +func (c *MCPController) GetUpstream() { + name := c.ctx.Param("name") + if name == "" { + c.RespondError(http.StatusBadRequest, model.ErrorCodeMissingQuery, "missing upstream name") + return + } + + info, err := mcpManager.GetUpstream(name) + if err != nil { + c.RespondError(http.StatusNotFound, model.ErrorCodeContextNotFound, err.Error()) + return + } + + c.RespondSuccess(info) +} + +// RemoveUpstream handles DELETE /mcp/upstreams/:name. +func (c *MCPController) RemoveUpstream() { + name := c.ctx.Param("name") + if name == "" { + c.RespondError(http.StatusBadRequest, model.ErrorCodeMissingQuery, "missing upstream name") + return + } + + if err := mcpManager.RemoveUpstream(name); err != nil { + c.RespondError(http.StatusNotFound, model.ErrorCodeContextNotFound, err.Error()) + return + } + + c.ctx.Status(http.StatusOK) +} diff --git a/components/execd/pkg/web/controller/mcp_test.go b/components/execd/pkg/web/controller/mcp_test.go new file mode 100644 index 000000000..be2b81b89 --- /dev/null +++ b/components/execd/pkg/web/controller/mcp_test.go @@ -0,0 +1,285 @@ +// Copyright 2025 Alibaba Group Holding Ltd. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package controller + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" + + "github.com/alibaba/opensandbox/execd/pkg/mcpproxy" + "github.com/alibaba/opensandbox/execd/pkg/web/model" +) + +func init() { + gin.SetMode(gin.TestMode) +} + +func setupMCPRouter() *gin.Engine { + InitMCPProxy() + r := gin.New() + mcp := r.Group("/mcpproxy") + { + mcp.POST("", func(ctx *gin.Context) { NewMCPController(ctx).HandleRequest() }) + mcp.DELETE("", func(ctx *gin.Context) { NewMCPController(ctx).DeleteSession() }) + upstreams := mcp.Group("/upstreams") + { + upstreams.POST("", func(ctx *gin.Context) { NewMCPController(ctx).AddUpstream() }) + upstreams.GET("", func(ctx *gin.Context) { NewMCPController(ctx).ListUpstreams() }) + upstreams.GET("/:name", func(ctx *gin.Context) { NewMCPController(ctx).GetUpstream() }) + upstreams.DELETE("/:name", func(ctx *gin.Context) { NewMCPController(ctx).RemoveUpstream() }) + } + } + return r +} + +func mcpPost(t *testing.T, router *gin.Engine, body any, sessionID string) *httptest.ResponseRecorder { + t.Helper() + data, err := json.Marshal(body) + require.NoError(t, err) + ctx, rec := newTestContext(http.MethodPost, "/mcpproxy", data) + ctx.Request.Header.Set("Content-Type", "application/json") + ctx.Request.Header.Set("Accept", "application/json") + if sessionID != "" { + ctx.Request.Header.Set(model.McpSessionHeader, sessionID) + } + router.ServeHTTP(rec, ctx.Request) + return rec +} + +func TestMCPInitialize(t *testing.T) { + router := setupMCPRouter() + body := mcpproxy.Request{ + JSONRPC: "2.0", + ID: float64(1), + Method: "initialize", + Params: json.RawMessage(`{"protocolVersion":"2025-03-26","capabilities":{},"clientInfo":{"name":"test","version":"1.0"}}`), + } + + rec := mcpPost(t, router, body, "") + require.Equal(t, http.StatusOK, rec.Code) + + sid := rec.Header().Get(model.McpSessionHeader) + require.NotEmpty(t, sid) + + var resp mcpproxy.Response + require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp)) + require.Nil(t, resp.Error) +} + +func TestMCPMissingSessionHeader(t *testing.T) { + router := setupMCPRouter() + body := mcpproxy.Request{ + JSONRPC: "2.0", + ID: float64(1), + Method: "tools/list", + Params: json.RawMessage(`{}`), + } + + rec := mcpPost(t, router, body, "") + require.Equal(t, http.StatusBadRequest, rec.Code) +} + +func TestMCPUnknownSession(t *testing.T) { + router := setupMCPRouter() + body := mcpproxy.Request{ + JSONRPC: "2.0", + ID: float64(1), + Method: "tools/list", + Params: json.RawMessage(`{}`), + } + + rec := mcpPost(t, router, body, "bogus-session") + require.Equal(t, http.StatusNotFound, rec.Code) +} + +func TestMCPToolsListEmpty(t *testing.T) { + router := setupMCPRouter() + + // Initialize first. + initBody := mcpproxy.Request{ + JSONRPC: "2.0", + ID: float64(1), + Method: "initialize", + Params: json.RawMessage(`{"protocolVersion":"2025-03-26","capabilities":{},"clientInfo":{"name":"test","version":"1.0"}}`), + } + initRec := mcpPost(t, router, initBody, "") + sid := initRec.Header().Get(model.McpSessionHeader) + require.NotEmpty(t, sid) + + // tools/list with no upstreams. + body := mcpproxy.Request{ + JSONRPC: "2.0", + ID: float64(2), + Method: "tools/list", + Params: json.RawMessage(`{}`), + } + rec := mcpPost(t, router, body, sid) + require.Equal(t, http.StatusOK, rec.Code) + + var resp mcpproxy.Response + require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp)) + require.Nil(t, resp.Error) + + var result struct { + Tools []json.RawMessage `json:"tools"` + } + require.NoError(t, json.Unmarshal(resp.Result, &result)) + require.Len(t, result.Tools, 0) +} + +func TestMCPDeleteSession(t *testing.T) { + router := setupMCPRouter() + + // Initialize. + initBody := mcpproxy.Request{ + JSONRPC: "2.0", + ID: float64(1), + Method: "initialize", + Params: json.RawMessage(`{"protocolVersion":"2025-03-26","capabilities":{},"clientInfo":{"name":"test","version":"1.0"}}`), + } + initRec := mcpPost(t, router, initBody, "") + sid := initRec.Header().Get(model.McpSessionHeader) + + // Delete session. + ctx, rec := newTestContext(http.MethodDelete, "/mcpproxy", nil) + ctx.Request.Header.Set(model.McpSessionHeader, sid) + router.ServeHTTP(rec, ctx.Request) + require.Equal(t, http.StatusOK, rec.Code) + + // Session should be gone. + body := mcpproxy.Request{ + JSONRPC: "2.0", + ID: float64(2), + Method: "tools/list", + Params: json.RawMessage(`{}`), + } + rec2 := mcpPost(t, router, body, sid) + require.Equal(t, http.StatusNotFound, rec2.Code) +} + +func TestMCPDeleteSessionMissingHeader(t *testing.T) { + router := setupMCPRouter() + ctx, rec := newTestContext(http.MethodDelete, "/mcpproxy", nil) + router.ServeHTTP(rec, ctx.Request) + require.Equal(t, http.StatusBadRequest, rec.Code) +} + +func TestMCPPing(t *testing.T) { + router := setupMCPRouter() + + // Initialize. + initBody := mcpproxy.Request{ + JSONRPC: "2.0", + ID: float64(1), + Method: "initialize", + Params: json.RawMessage(`{"protocolVersion":"2025-03-26","capabilities":{},"clientInfo":{"name":"test","version":"1.0"}}`), + } + initRec := mcpPost(t, router, initBody, "") + sid := initRec.Header().Get(model.McpSessionHeader) + + // Ping. + body := mcpproxy.Request{JSONRPC: "2.0", ID: float64(2), Method: "ping"} + rec := mcpPost(t, router, body, sid) + require.Equal(t, http.StatusOK, rec.Code) + + var resp mcpproxy.Response + require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp)) + require.Nil(t, resp.Error) +} + +func TestMCPInvalidJSON(t *testing.T) { + router := setupMCPRouter() + ctx, rec := newTestContext(http.MethodPost, "/mcpproxy", []byte(`{invalid`)) + ctx.Request.Header.Set("Content-Type", "application/json") + router.ServeHTTP(rec, ctx.Request) + require.Equal(t, http.StatusBadRequest, rec.Code) +} + +func TestMCPAddUpstreamValidation(t *testing.T) { + router := setupMCPRouter() + + t.Run("stdio missing command", func(t *testing.T) { + body := model.AddUpstreamRequest{ + Name: "test", + Transport: "stdio", + } + data, _ := json.Marshal(body) + ctx, rec := newTestContext(http.MethodPost, "/mcpproxy/upstreams", data) + ctx.Request.Header.Set("Content-Type", "application/json") + router.ServeHTTP(rec, ctx.Request) + require.Equal(t, http.StatusBadRequest, rec.Code) + }) + + t.Run("http missing url", func(t *testing.T) { + body := model.AddUpstreamRequest{ + Name: "test", + Transport: "http", + } + data, _ := json.Marshal(body) + ctx, rec := newTestContext(http.MethodPost, "/mcpproxy/upstreams", data) + ctx.Request.Header.Set("Content-Type", "application/json") + router.ServeHTTP(rec, ctx.Request) + require.Equal(t, http.StatusBadRequest, rec.Code) + }) +} + +func TestMCPListUpstreamsEmpty(t *testing.T) { + router := setupMCPRouter() + ctx, rec := newTestContext(http.MethodGet, "/mcpproxy/upstreams", nil) + router.ServeHTTP(rec, ctx.Request) + require.Equal(t, http.StatusOK, rec.Code) + + var list []mcpproxy.UpstreamInfo + require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &list)) + require.Len(t, list, 0) +} + +func TestMCPGetUpstreamNotFound(t *testing.T) { + router := setupMCPRouter() + ctx, rec := newTestContext(http.MethodGet, "/mcpproxy/upstreams/nonexistent", nil) + router.ServeHTTP(rec, ctx.Request) + require.Equal(t, http.StatusNotFound, rec.Code) +} + +func TestMCPRemoveUpstreamNotFound(t *testing.T) { + router := setupMCPRouter() + ctx, rec := newTestContext(http.MethodDelete, "/mcpproxy/upstreams/nonexistent", nil) + router.ServeHTTP(rec, ctx.Request) + require.Equal(t, http.StatusNotFound, rec.Code) +} + +func TestMCPNotificationReturnsAccepted(t *testing.T) { + router := setupMCPRouter() + + // Initialize first. + initBody := mcpproxy.Request{ + JSONRPC: "2.0", + ID: float64(1), + Method: "initialize", + Params: json.RawMessage(`{"protocolVersion":"2025-03-26","capabilities":{},"clientInfo":{"name":"test","version":"1.0"}}`), + } + initRec := mcpPost(t, router, initBody, "") + sid := initRec.Header().Get(model.McpSessionHeader) + + // Send notification. + body := mcpproxy.Request{JSONRPC: "2.0", Method: "notifications/initialized"} + rec := mcpPost(t, router, body, sid) + require.Equal(t, http.StatusAccepted, rec.Code) +} diff --git a/components/execd/pkg/web/model/mcp.go b/components/execd/pkg/web/model/mcp.go new file mode 100644 index 000000000..d02e3bd13 --- /dev/null +++ b/components/execd/pkg/web/model/mcp.go @@ -0,0 +1,29 @@ +// Copyright 2025 Alibaba Group Holding Ltd. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package model + +// McpSessionHeader is the HTTP header used for MCP session tracking. +const McpSessionHeader = "Mcp-Session-Id" + +// AddUpstreamRequest is the body for POST /mcp/upstreams. +type AddUpstreamRequest struct { + Name string `json:"name" binding:"required"` + Transport string `json:"transport" binding:"required,oneof=stdio http"` + Command string `json:"command,omitempty"` + Args []string `json:"args,omitempty"` + Env map[string]string `json:"env,omitempty"` + URL string `json:"url,omitempty"` + Headers map[string]string `json:"headers,omitempty"` +} diff --git a/components/execd/pkg/web/router.go b/components/execd/pkg/web/router.go index 584f8c2e6..1600af1d3 100644 --- a/components/execd/pkg/web/router.go +++ b/components/execd/pkg/web/router.go @@ -92,6 +92,21 @@ func NewRouter(accessToken string) *gin.Engine { pty.GET("/:sessionId/ws", controller.PTYSessionWebSocket) } + mcpProxy := r.Group("/mcpproxy") + { + mcpProxy.POST("", withMCP(func(c *controller.MCPController) { c.HandleRequest() })) + mcpProxy.GET("", withMCP(func(c *controller.MCPController) { c.SSEStream() })) + mcpProxy.DELETE("", withMCP(func(c *controller.MCPController) { c.DeleteSession() })) + + upstreams := mcpProxy.Group("/upstreams") + { + upstreams.POST("", withMCP(func(c *controller.MCPController) { c.AddUpstream() })) + upstreams.GET("", withMCP(func(c *controller.MCPController) { c.ListUpstreams() })) + upstreams.GET("/:name", withMCP(func(c *controller.MCPController) { c.GetUpstream() })) + upstreams.DELETE("/:name", withMCP(func(c *controller.MCPController) { c.RemoveUpstream() })) + } + } + return r } @@ -119,6 +134,12 @@ func withPTY(fn func(*controller.PTYController)) gin.HandlerFunc { } } +func withMCP(fn func(*controller.MCPController)) gin.HandlerFunc { + return func(ctx *gin.Context) { + fn(controller.NewMCPController(ctx)) + } +} + func accessTokenMiddleware(token string) gin.HandlerFunc { return func(ctx *gin.Context) { if token == "" { diff --git a/components/execd/tests/smoke_mcpproxy.py b/components/execd/tests/smoke_mcpproxy.py new file mode 100644 index 000000000..b7fb63f18 --- /dev/null +++ b/components/execd/tests/smoke_mcpproxy.py @@ -0,0 +1,467 @@ +#!/usr/bin/env python3 + +# Copyright 2025 Alibaba Group Holding Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Smoke tests for the execd MCP proxy (/mcpproxy). + +Prerequisites: +- execd server running locally (default http://localhost:44772) +- python3 available (used for the mock MCP stdio server) +- Optional: set env BASE_URL to override +- Optional: set env API_TOKEN if server expects X-EXECD-ACCESS-TOKEN +""" + +import json +import os +import sys +import tempfile +import textwrap + +import requests + +BASE_URL = os.environ.get("BASE_URL", "http://localhost:44772").rstrip("/") +API_TOKEN = os.environ.get("API_TOKEN") + +HEADERS = {} +if API_TOKEN: + HEADERS["X-EXECD-ACCESS-TOKEN"] = API_TOKEN + +session = requests.Session() +session.headers.update(HEADERS) + +# Path to the mock MCP server script (created at runtime). +_mock_script = None + + +def expect(cond: bool, msg: str): + if not cond: + raise SystemExit(msg) + + +def create_mock_mcp_server() -> str: + """Write a minimal stdio MCP server script and return its path.""" + global _mock_script + script = textwrap.dedent("""\ + #!/usr/bin/env python3 + import json, sys + + TOOLS = [ + { + "name": "echo_tool", + "description": "Echoes the input message back.", + "inputSchema": { + "type": "object", + "properties": { + "message": {"type": "string", "description": "Message to echo"} + }, + "required": ["message"] + } + }, + { + "name": "add_numbers", + "description": "Adds two numbers.", + "inputSchema": { + "type": "object", + "properties": { + "a": {"type": "number"}, + "b": {"type": "number"} + }, + "required": ["a", "b"] + } + } + ] + + def respond(id, result): + msg = {"jsonrpc": "2.0", "id": id, "result": result} + sys.stdout.write(json.dumps(msg) + "\\n") + sys.stdout.flush() + + def error(id, code, message): + msg = {"jsonrpc": "2.0", "id": id, "error": {"code": code, "message": message}} + sys.stdout.write(json.dumps(msg) + "\\n") + sys.stdout.flush() + + for line in sys.stdin: + line = line.strip() + if not line: + continue + try: + req = json.loads(line) + except json.JSONDecodeError: + continue + + method = req.get("method", "") + rid = req.get("id") + + if rid is None: + continue + + if method == "initialize": + respond(rid, { + "protocolVersion": "2025-03-26", + "capabilities": {"tools": {"listChanged": True}}, + "serverInfo": {"name": "mock-mcp", "version": "1.0.0"} + }) + elif method == "tools/list": + respond(rid, {"tools": TOOLS}) + elif method == "tools/call": + params = req.get("params", {}) + name = params.get("name", "") + args = params.get("arguments", {}) + if name == "echo_tool": + respond(rid, { + "content": [{"type": "text", "text": args.get("message", "")}], + "isError": False + }) + elif name == "add_numbers": + result = args.get("a", 0) + args.get("b", 0) + respond(rid, { + "content": [{"type": "text", "text": str(result)}], + "isError": False + }) + else: + error(rid, -32601, f"unknown tool: {name}") + elif method == "ping": + respond(rid, {}) + else: + error(rid, -32601, f"unknown method: {method}") + """) + fd, path = tempfile.mkstemp(suffix=".py", prefix="mock_mcp_") + with os.fdopen(fd, "w") as f: + f.write(script) + os.chmod(path, 0o755) + _mock_script = path + return path + + +def cleanup_mock(): + if _mock_script and os.path.exists(_mock_script): + os.unlink(_mock_script) + + +# ── JSON-RPC helpers ─────────────────────────────────────────────── + +def jsonrpc_request(method: str, params=None, req_id=None): + msg = {"jsonrpc": "2.0", "method": method} + if req_id is not None: + msg["id"] = req_id + if params is not None: + msg["params"] = params + return msg + + +def mcp_post(body: dict, session_id: str = None) -> requests.Response: + headers = {"Content-Type": "application/json", "Accept": "application/json"} + if session_id: + headers["Mcp-Session-Id"] = session_id + return session.post(f"{BASE_URL}/mcpproxy", json=body, headers=headers, timeout=15) + + +def mcp_delete(session_id: str) -> requests.Response: + headers = {"Mcp-Session-Id": session_id} + return session.delete(f"{BASE_URL}/mcpproxy", headers=headers, timeout=5) + + +# ── Test functions ───────────────────────────────────────────────── + +def test_register_upstream(mock_path: str) -> dict: + """Register the mock MCP server as a stdio upstream.""" + payload = { + "name": "mock", + "transport": "stdio", + "command": sys.executable, + "args": [mock_path], + } + r = session.post(f"{BASE_URL}/mcpproxy/upstreams", json=payload, timeout=30) + expect(r.status_code == 201, f"register upstream failed: {r.status_code} {r.text}") + info = r.json() + expect(info["name"] == "mock", f"upstream name mismatch: {info}") + expect(info["status"] == "running", f"upstream not running: {info}") + expect(len(info["tools"]) == 2, f"expected 2 tools, got {info['tools']}") + return info + + +def test_list_upstreams(): + """List registered upstreams.""" + r = session.get(f"{BASE_URL}/mcpproxy/upstreams", timeout=5) + expect(r.status_code == 200, f"list upstreams failed: {r.status_code} {r.text}") + upstreams = r.json() + expect(len(upstreams) >= 1, "no upstreams found") + names = [u["name"] for u in upstreams] + expect("mock" in names, f"mock upstream not found in {names}") + + +def test_get_upstream(): + """Get detail of a specific upstream.""" + r = session.get(f"{BASE_URL}/mcpproxy/upstreams/mock", timeout=5) + expect(r.status_code == 200, f"get upstream failed: {r.status_code} {r.text}") + info = r.json() + expect(info["name"] == "mock", f"name mismatch: {info}") + + +def test_initialize() -> str: + """Initialize an MCP session and return the session ID.""" + body = jsonrpc_request("initialize", { + "protocolVersion": "2025-03-26", + "capabilities": {}, + "clientInfo": {"name": "smoke-test", "version": "1.0.0"}, + }, req_id=1) + r = mcp_post(body) + expect(r.status_code == 200, f"initialize failed: {r.status_code} {r.text}") + + session_id = r.headers.get("Mcp-Session-Id") + expect(bool(session_id), "missing Mcp-Session-Id in response") + + resp = r.json() + expect(resp.get("result", {}).get("protocolVersion") == "2025-03-26", + f"protocol version mismatch: {resp}") + return session_id + + +def test_initialized_notification(session_id: str): + """Send the initialized notification (no response expected).""" + body = jsonrpc_request("notifications/initialized") + r = mcp_post(body, session_id=session_id) + expect(r.status_code in (200, 202), f"initialized notification failed: {r.status_code} {r.text}") + + +def test_tools_list(session_id: str): + """List tools — should return aggregated tools from all upstreams.""" + body = jsonrpc_request("tools/list", {}, req_id=2) + r = mcp_post(body, session_id=session_id) + expect(r.status_code == 200, f"tools/list failed: {r.status_code} {r.text}") + + resp = r.json() + tools = resp.get("result", {}).get("tools", []) + tool_names = [t["name"] for t in tools] + expect("echo_tool" in tool_names, f"echo_tool not found in {tool_names}") + expect("add_numbers" in tool_names, f"add_numbers not found in {tool_names}") + + +def test_tools_call_echo(session_id: str): + """Call echo_tool and verify the response.""" + body = jsonrpc_request("tools/call", { + "name": "echo_tool", + "arguments": {"message": "hello from smoke test"}, + }, req_id=3) + r = mcp_post(body, session_id=session_id) + expect(r.status_code == 200, f"tools/call echo failed: {r.status_code} {r.text}") + + resp = r.json() + content = resp.get("result", {}).get("content", []) + expect(len(content) == 1, f"expected 1 content block, got {content}") + expect(content[0]["text"] == "hello from smoke test", + f"echo mismatch: {content[0]}") + + +def test_tools_call_add(session_id: str): + """Call add_numbers and verify the result.""" + body = jsonrpc_request("tools/call", { + "name": "add_numbers", + "arguments": {"a": 17, "b": 25}, + }, req_id=4) + r = mcp_post(body, session_id=session_id) + expect(r.status_code == 200, f"tools/call add failed: {r.status_code} {r.text}") + + resp = r.json() + content = resp.get("result", {}).get("content", []) + expect(len(content) == 1, f"expected 1 content block, got {content}") + expect(content[0]["text"] == "42", f"add result mismatch: {content[0]}") + + +def test_tools_call_unknown(session_id: str): + """Call a nonexistent tool — should get an error.""" + body = jsonrpc_request("tools/call", { + "name": "nonexistent_tool", + "arguments": {}, + }, req_id=5) + r = mcp_post(body, session_id=session_id) + expect(r.status_code == 200, f"tools/call unknown HTTP failed: {r.status_code}") + + resp = r.json() + expect(resp.get("error") is not None, f"expected error for unknown tool: {resp}") + + +def test_ping(session_id: str): + """Send a ping and expect a pong.""" + body = jsonrpc_request("ping", {}, req_id=6) + r = mcp_post(body, session_id=session_id) + expect(r.status_code == 200, f"ping failed: {r.status_code} {r.text}") + + resp = r.json() + expect(resp.get("error") is None, f"ping error: {resp}") + + +def test_missing_session(): + """Requests without session ID should fail (except initialize).""" + body = jsonrpc_request("tools/list", {}, req_id=99) + r = mcp_post(body) + expect(r.status_code == 400, f"expected 400 without session: {r.status_code}") + + +def test_unknown_session(): + """Requests with a bogus session ID should fail.""" + body = jsonrpc_request("tools/list", {}, req_id=99) + r = mcp_post(body, session_id="bogus-session-id") + expect(r.status_code == 404, f"expected 404 for unknown session: {r.status_code}") + + +def test_delete_session(session_id: str): + """Delete the MCP session.""" + r = mcp_delete(session_id) + expect(r.status_code == 200, f"delete session failed: {r.status_code} {r.text}") + + # Subsequent requests with the deleted session should fail. + body = jsonrpc_request("tools/list", {}, req_id=99) + r2 = mcp_post(body, session_id=session_id) + expect(r2.status_code == 404, f"expected 404 after session delete: {r2.status_code}") + + +def test_duplicate_upstream(mock_path: str): + """Registering an upstream with the same name should fail.""" + payload = { + "name": "mock", + "transport": "stdio", + "command": sys.executable, + "args": [mock_path], + } + r = session.post(f"{BASE_URL}/mcpproxy/upstreams", json=payload, timeout=15) + expect(r.status_code == 500, f"expected error for duplicate upstream: {r.status_code}") + + +def test_tool_conflict(mock_path: str): + """Registering a second upstream with conflicting tools should fail.""" + payload = { + "name": "mock-conflict", + "transport": "stdio", + "command": sys.executable, + "args": [mock_path], + } + r = session.post(f"{BASE_URL}/mcpproxy/upstreams", json=payload, timeout=30) + expect(r.status_code == 500, f"expected error for tool conflict: {r.status_code}") + expect("conflicts" in r.text.lower() or "conflict" in r.text.lower(), + f"error should mention conflict: {r.text}") + + +def test_remove_upstream(): + """Remove the mock upstream.""" + r = session.delete(f"{BASE_URL}/mcpproxy/upstreams/mock", timeout=10) + expect(r.status_code == 200, f"remove upstream failed: {r.status_code} {r.text}") + + # Verify it's gone. + r2 = session.get(f"{BASE_URL}/mcpproxy/upstreams/mock", timeout=5) + expect(r2.status_code == 404, f"upstream should be gone: {r2.status_code}") + + +def test_tools_empty_after_remove(): + """After removing all upstreams, tools/list should return empty.""" + sid = test_initialize() + body = jsonrpc_request("tools/list", {}, req_id=10) + r = mcp_post(body, session_id=sid) + expect(r.status_code == 200, f"tools/list failed: {r.status_code}") + tools = r.json().get("result", {}).get("tools", []) + expect(len(tools) == 0, f"expected 0 tools after remove, got {len(tools)}") + mcp_delete(sid) + + +# ── Main ─────────────────────────────────────────────────────────── + +def main(): + print(f"[+] base: {BASE_URL}") + + # Health check + r = session.get(f"{BASE_URL}/ping", timeout=5) + expect(r.status_code == 200, "ping failed") + print("[+] ping ok") + + # Create mock MCP server + mock_path = create_mock_mcp_server() + print(f"[+] mock MCP server: {mock_path}") + + try: + # Upstream management + info = test_register_upstream(mock_path) + print(f"[+] upstream registered: {info['tools']}") + + test_list_upstreams() + print("[+] list upstreams ok") + + test_get_upstream() + print("[+] get upstream ok") + + test_duplicate_upstream(mock_path) + print("[+] duplicate upstream rejected") + + test_tool_conflict(mock_path) + print("[+] tool conflict rejected") + + # Session management + test_missing_session() + print("[+] missing session rejected") + + test_unknown_session() + print("[+] unknown session rejected") + + # MCP protocol flow + sid = test_initialize() + print(f"[+] initialized, session: {sid}") + + test_initialized_notification(sid) + print("[+] initialized notification ok") + + test_ping(sid) + print("[+] ping/pong ok") + + test_tools_list(sid) + print("[+] tools/list ok") + + test_tools_call_echo(sid) + print("[+] tools/call echo_tool ok") + + test_tools_call_add(sid) + print("[+] tools/call add_numbers ok (17+25=42)") + + test_tools_call_unknown(sid) + print("[+] unknown tool error ok") + + test_delete_session(sid) + print("[+] session delete ok") + + # Cleanup upstream + test_remove_upstream() + print("[+] upstream removed") + + test_tools_empty_after_remove() + print("[+] tools empty after remove") + + print("[+] MCP proxy smoke tests PASS") + + finally: + # Best-effort cleanup: remove upstream if still registered. + try: + session.delete(f"{BASE_URL}/mcpproxy/upstreams/mock", timeout=5) + except Exception: + pass + cleanup_mock() + + +if __name__ == "__main__": + try: + main() + except SystemExit as exc: + print(f"[!] MCP proxy smoke tests FAIL: {exc}", file=sys.stderr) + cleanup_mock() + sys.exit(1) From 77591f01990ae91e47ae6b2331eda839fb0befce Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=AB=98=E7=84=B6?= Date: Mon, 22 Jun 2026 12:05:05 +0800 Subject: [PATCH 2/5] style: fix gofmt formatting in mcpproxy package --- components/execd/pkg/mcpproxy/http_upstream.go | 2 +- components/execd/pkg/mcpproxy/manager.go | 8 ++++---- components/execd/pkg/mcpproxy/manager_test.go | 4 ++-- components/execd/pkg/mcpproxy/stdio_upstream.go | 2 +- 4 files changed, 8 insertions(+), 8 deletions(-) diff --git a/components/execd/pkg/mcpproxy/http_upstream.go b/components/execd/pkg/mcpproxy/http_upstream.go index 3573c658a..2ee229faf 100644 --- a/components/execd/pkg/mcpproxy/http_upstream.go +++ b/components/execd/pkg/mcpproxy/http_upstream.go @@ -49,7 +49,7 @@ func newHTTPUpstream(config UpstreamConfig) *httpUpstream { } func (h *httpUpstream) Name() string { return h.name } -func (h *httpUpstream) Transport() string { return "http" } +func (h *httpUpstream) Transport() string { return "http" } func (h *httpUpstream) post(ctx context.Context, method string, params any) (*Response, error) { id := h.nextID.Add(1) diff --git a/components/execd/pkg/mcpproxy/manager.go b/components/execd/pkg/mcpproxy/manager.go index 67215ada8..eaaee7f74 100644 --- a/components/execd/pkg/mcpproxy/manager.go +++ b/components/execd/pkg/mcpproxy/manager.go @@ -26,10 +26,10 @@ import ( // Manager aggregates multiple upstream MCP servers and proxies requests. type Manager struct { mu sync.RWMutex - upstreams map[string]Upstream // name → upstream - toolIndex map[string]string // tool name → upstream name - allTools []Tool // cached aggregated list - sessions map[string]*Session // Mcp-Session-Id → session + upstreams map[string]Upstream // name → upstream + toolIndex map[string]string // tool name → upstream name + allTools []Tool // cached aggregated list + sessions map[string]*Session // Mcp-Session-Id → session } // NewManager creates an empty proxy manager. diff --git a/components/execd/pkg/mcpproxy/manager_test.go b/components/execd/pkg/mcpproxy/manager_test.go index 72d0888eb..781320845 100644 --- a/components/execd/pkg/mcpproxy/manager_test.go +++ b/components/execd/pkg/mcpproxy/manager_test.go @@ -28,8 +28,8 @@ type fakeUpstream struct { tools []Tool } -func (f *fakeUpstream) Name() string { return f.name } -func (f *fakeUpstream) Transport() string { return "fake" } +func (f *fakeUpstream) Name() string { return f.name } +func (f *fakeUpstream) Transport() string { return "fake" } func (f *fakeUpstream) Initialize(_ context.Context) error { return nil } func (f *fakeUpstream) Tools(_ context.Context) ([]Tool, error) { return f.tools, nil diff --git a/components/execd/pkg/mcpproxy/stdio_upstream.go b/components/execd/pkg/mcpproxy/stdio_upstream.go index b4d8ebae1..fe607c912 100644 --- a/components/execd/pkg/mcpproxy/stdio_upstream.go +++ b/components/execd/pkg/mcpproxy/stdio_upstream.go @@ -60,7 +60,7 @@ func newStdioUpstream(config UpstreamConfig) *stdioUpstream { } } -func (s *stdioUpstream) Name() string { return s.name } +func (s *stdioUpstream) Name() string { return s.name } func (s *stdioUpstream) Transport() string { return "stdio" } func (s *stdioUpstream) start() error { From 52277d830416d4b93649b8ce8ae90aa272bcd8ba Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=AB=98=E7=84=B6?= Date: Mon, 22 Jun 2026 12:10:05 +0800 Subject: [PATCH 3/5] fix: resolve golangci-lint nilerr and perfsprint findings --- components/execd/pkg/mcpproxy/manager.go | 2 +- components/execd/pkg/mcpproxy/stdio_upstream.go | 5 +++-- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/components/execd/pkg/mcpproxy/manager.go b/components/execd/pkg/mcpproxy/manager.go index eaaee7f74..9b69333e4 100644 --- a/components/execd/pkg/mcpproxy/manager.go +++ b/components/execd/pkg/mcpproxy/manager.go @@ -238,7 +238,7 @@ func (m *Manager) handleToolsList(req *Request, sessionID string) (*Response, st func (m *Manager) handleToolsCall(ctx context.Context, req *Request, sessionID string) (*Response, string, error) { var params callToolParams if err := json.Unmarshal(req.Params, ¶ms); err != nil { - return newErrorResponse(req.ID, CodeInvalidParams, "invalid tools/call params"), "", nil + return newErrorResponse(req.ID, CodeInvalidParams, "invalid tools/call params"), "", nil //nolint:nilerr // JSON-RPC error is in the response, not the Go error return } m.mu.RLock() diff --git a/components/execd/pkg/mcpproxy/stdio_upstream.go b/components/execd/pkg/mcpproxy/stdio_upstream.go index fe607c912..c0e341d3e 100644 --- a/components/execd/pkg/mcpproxy/stdio_upstream.go +++ b/components/execd/pkg/mcpproxy/stdio_upstream.go @@ -22,6 +22,7 @@ import ( "io" "os" "os/exec" + "strconv" "sync" "sync/atomic" "time" @@ -301,9 +302,9 @@ func (s *stdioUpstream) Close() error { func normalizeID(id any) string { switch v := id.(type) { case float64: - return fmt.Sprintf("%d", int64(v)) + return strconv.FormatInt(int64(v), 10) case int64: - return fmt.Sprintf("%d", v) + return strconv.FormatInt(v, 10) case string: return v default: From 1361f7b09b281d8a04f5c32094e6eb42c71ae20b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=AB=98=E7=84=B6?= Date: Mon, 22 Jun 2026 14:45:31 +0800 Subject: [PATCH 4/5] fix: address PR review findings - Remove unused json import in smoke test - Add stderr warning to empty except clause - Fix data race on httpUpstream.sessionID with RWMutex - Move upstream Initialize/Tools outside manager lock to avoid blocking all proxy traffic during slow upstream registration - Handle all JSON-RPC notifications generically via IsNotification() instead of only notifications/initialized - Support SSE (text/event-stream) responses from HTTP upstreams - Set listChanged capability to false (no notifications sent yet) - Preserve tool annotations field in aggregated tools/list - Follow tools/list pagination (nextCursor) in both transports --- .../execd/pkg/mcpproxy/http_upstream.go | 105 ++++++++++++++---- components/execd/pkg/mcpproxy/manager.go | 30 +++-- .../execd/pkg/mcpproxy/stdio_upstream.go | 32 ++++-- components/execd/pkg/mcpproxy/upstream.go | 9 +- components/execd/tests/smoke_mcpproxy.py | 5 +- 5 files changed, 135 insertions(+), 46 deletions(-) diff --git a/components/execd/pkg/mcpproxy/http_upstream.go b/components/execd/pkg/mcpproxy/http_upstream.go index 2ee229faf..3475d5e7f 100644 --- a/components/execd/pkg/mcpproxy/http_upstream.go +++ b/components/execd/pkg/mcpproxy/http_upstream.go @@ -15,12 +15,15 @@ package mcpproxy import ( + "bufio" "bytes" "context" "encoding/json" "fmt" "io" "net/http" + "strings" + "sync" "sync/atomic" "time" ) @@ -28,13 +31,16 @@ import ( const mcpSessionHeader = "Mcp-Session-Id" type httpUpstream struct { - name string - url string - headers map[string]string // custom headers sent on every request - client *http.Client + name string + url string + headers map[string]string // custom headers sent on every request + client *http.Client + + mu sync.RWMutex sessionID string - nextID atomic.Int64 - tools []Tool + + nextID atomic.Int64 + tools []Tool } func newHTTPUpstream(config UpstreamConfig) *httpUpstream { @@ -79,9 +85,11 @@ func (h *httpUpstream) post(ctx context.Context, method string, params any) (*Re for k, v := range h.headers { httpReq.Header.Set(k, v) } + h.mu.RLock() if h.sessionID != "" { httpReq.Header.Set(mcpSessionHeader, h.sessionID) } + h.mu.RUnlock() httpResp, err := h.client.Do(httpReq) if err != nil { @@ -90,7 +98,19 @@ func (h *httpUpstream) post(ctx context.Context, method string, params any) (*Re defer httpResp.Body.Close() if sid := httpResp.Header.Get(mcpSessionHeader); sid != "" { + h.mu.Lock() h.sessionID = sid + h.mu.Unlock() + } + + if httpResp.StatusCode != http.StatusOK { + body, _ := io.ReadAll(httpResp.Body) + return nil, fmt.Errorf("upstream %s returned %d: %s", h.name, httpResp.StatusCode, string(body)) + } + + ct := httpResp.Header.Get("Content-Type") + if strings.HasPrefix(ct, "text/event-stream") { + return h.parseSSEResponse(httpResp.Body, id) } body, err := io.ReadAll(httpResp.Body) @@ -98,10 +118,6 @@ func (h *httpUpstream) post(ctx context.Context, method string, params any) (*Re return nil, fmt.Errorf("read response: %w", err) } - if httpResp.StatusCode != http.StatusOK { - return nil, fmt.Errorf("upstream %s returned %d: %s", h.name, httpResp.StatusCode, string(body)) - } - var resp Response if err := json.Unmarshal(body, &resp); err != nil { return nil, fmt.Errorf("unmarshal response: %w", err) @@ -109,6 +125,34 @@ func (h *httpUpstream) post(ctx context.Context, method string, params any) (*Re return &resp, nil } +// parseSSEResponse reads an SSE stream and returns the first JSON-RPC +// response that matches the given request ID. +func (h *httpUpstream) parseSSEResponse(r io.Reader, requestID int64) (*Response, error) { + scanner := bufio.NewScanner(r) + scanner.Buffer(make([]byte, 0, 64*1024), 16*1024*1024) + for scanner.Scan() { + line := scanner.Text() + if !strings.HasPrefix(line, "data:") { + continue + } + data := strings.TrimSpace(strings.TrimPrefix(line, "data:")) + if data == "" { + continue + } + var resp Response + if err := json.Unmarshal([]byte(data), &resp); err != nil { + continue + } + if resp.ID != nil { + return &resp, nil + } + } + if err := scanner.Err(); err != nil { + return nil, fmt.Errorf("read SSE stream: %w", err) + } + return nil, fmt.Errorf("SSE stream ended without a response for request %d", requestID) +} + func (h *httpUpstream) Initialize(ctx context.Context) error { ctx, cancel := context.WithTimeout(ctx, initTimeout) defer cancel() @@ -145,9 +189,11 @@ func (h *httpUpstream) Initialize(ctx context.Context) error { for k, v := range h.headers { httpReq.Header.Set(k, v) } + h.mu.RLock() if h.sessionID != "" { httpReq.Header.Set(mcpSessionHeader, h.sessionID) } + h.mu.RUnlock() notifResp, err := h.client.Do(httpReq) if err != nil { return fmt.Errorf("initialized notification: %w", err) @@ -165,20 +211,30 @@ func (h *httpUpstream) Tools(ctx context.Context) ([]Tool, error) { ctx, cancel := context.WithTimeout(ctx, initTimeout) defer cancel() - resp, err := h.post(ctx, "tools/list", struct{}{}) - if err != nil { - return nil, fmt.Errorf("tools/list: %w", err) - } - if resp.Error != nil { - return nil, fmt.Errorf("tools/list error: %s", resp.Error.Message) - } + var allTools []Tool + var cursor string + for { + params := toolsListParams{Cursor: cursor} + resp, err := h.post(ctx, "tools/list", params) + if err != nil { + return nil, fmt.Errorf("tools/list: %w", err) + } + if resp.Error != nil { + return nil, fmt.Errorf("tools/list error: %s", resp.Error.Message) + } - var result toolsListResult - if err := json.Unmarshal(resp.Result, &result); err != nil { - return nil, fmt.Errorf("parse tools/list result: %w", err) + var result toolsListResult + if err := json.Unmarshal(resp.Result, &result); err != nil { + return nil, fmt.Errorf("parse tools/list result: %w", err) + } + allTools = append(allTools, result.Tools...) + if result.NextCursor == "" { + break + } + cursor = result.NextCursor } - h.tools = result.Tools + h.tools = allTools return h.tools, nil } @@ -194,7 +250,10 @@ func (h *httpUpstream) CallTool(ctx context.Context, name string, arguments json } func (h *httpUpstream) Close() error { - if h.sessionID == "" { + h.mu.RLock() + sid := h.sessionID + h.mu.RUnlock() + if sid == "" { return nil } @@ -208,7 +267,7 @@ func (h *httpUpstream) Close() error { for k, v := range h.headers { req.Header.Set(k, v) } - req.Header.Set(mcpSessionHeader, h.sessionID) + req.Header.Set(mcpSessionHeader, sid) resp, err := h.client.Do(req) if err != nil { diff --git a/components/execd/pkg/mcpproxy/manager.go b/components/execd/pkg/mcpproxy/manager.go index 9b69333e4..3edd103da 100644 --- a/components/execd/pkg/mcpproxy/manager.go +++ b/components/execd/pkg/mcpproxy/manager.go @@ -45,10 +45,11 @@ func NewManager() *Manager { // its tool list, and registers the tools. Returns error if any tool name // conflicts with an already-registered tool. func (m *Manager) AddUpstream(ctx context.Context, config UpstreamConfig) (*UpstreamInfo, error) { - m.mu.Lock() - defer m.mu.Unlock() - - if _, exists := m.upstreams[config.Name]; exists { + // Quick check under lock — reject duplicates before expensive IO. + m.mu.RLock() + _, exists := m.upstreams[config.Name] + m.mu.RUnlock() + if exists { return nil, fmt.Errorf("upstream %q already registered", config.Name) } @@ -62,6 +63,8 @@ func (m *Manager) AddUpstream(ctx context.Context, config UpstreamConfig) (*Upst return nil, fmt.Errorf("unsupported transport: %s", config.Transport) } + // Initialize and fetch tools outside the lock to avoid blocking + // all proxy traffic while dialing a slow or unreachable upstream. if err := up.Initialize(ctx); err != nil { _ = up.Close() return nil, fmt.Errorf("initialize upstream %q: %w", config.Name, err) @@ -73,7 +76,15 @@ func (m *Manager) AddUpstream(ctx context.Context, config UpstreamConfig) (*Upst return nil, fmt.Errorf("fetch tools from %q: %w", config.Name, err) } - // Check for conflicts before registering any tools. + // Now lock to register — re-check for races. + m.mu.Lock() + defer m.mu.Unlock() + + if _, exists := m.upstreams[config.Name]; exists { + _ = up.Close() + return nil, fmt.Errorf("upstream %q already registered", config.Name) + } + for _, t := range tools { if owner, exists := m.toolIndex[t.Name]; exists { _ = up.Close() @@ -176,11 +187,14 @@ func (m *Manager) GetUpstream(name string) (*UpstreamInfo, error) { // HandleRequest routes a downstream MCP JSON-RPC request to the appropriate handler. // Returns the response and optionally a new session ID (for initialize requests). func (m *Manager) HandleRequest(ctx context.Context, req *Request, sessionID string) (*Response, string, error) { + // JSON-RPC notifications (no id) must never receive a response. + if req.IsNotification() { + return nil, "", nil + } + switch req.Method { case "initialize": return m.handleInitialize(req) - case "notifications/initialized": - return nil, "", nil // notification, no response case "ping": return m.handlePing(req) case "tools/list": @@ -203,7 +217,7 @@ func (m *Manager) handleInitialize(req *Request) (*Response, string, error) { "protocolVersion": "2025-03-26", "capabilities": map[string]any{ "tools": map[string]any{ - "listChanged": true, + "listChanged": false, }, }, "serverInfo": map[string]any{ diff --git a/components/execd/pkg/mcpproxy/stdio_upstream.go b/components/execd/pkg/mcpproxy/stdio_upstream.go index c0e341d3e..17aedc808 100644 --- a/components/execd/pkg/mcpproxy/stdio_upstream.go +++ b/components/execd/pkg/mcpproxy/stdio_upstream.go @@ -246,20 +246,30 @@ func (s *stdioUpstream) Tools(ctx context.Context) ([]Tool, error) { ctx, cancel := context.WithTimeout(ctx, initTimeout) defer cancel() - resp, err := s.sendRequest(ctx, "tools/list", struct{}{}) - if err != nil { - return nil, fmt.Errorf("tools/list: %w", err) - } - if resp.Error != nil { - return nil, fmt.Errorf("tools/list error: %s", resp.Error.Message) - } + var allTools []Tool + var cursor string + for { + params := toolsListParams{Cursor: cursor} + resp, err := s.sendRequest(ctx, "tools/list", params) + if err != nil { + return nil, fmt.Errorf("tools/list: %w", err) + } + if resp.Error != nil { + return nil, fmt.Errorf("tools/list error: %s", resp.Error.Message) + } - var result toolsListResult - if err := json.Unmarshal(resp.Result, &result); err != nil { - return nil, fmt.Errorf("parse tools/list result: %w", err) + var result toolsListResult + if err := json.Unmarshal(resp.Result, &result); err != nil { + return nil, fmt.Errorf("parse tools/list result: %w", err) + } + allTools = append(allTools, result.Tools...) + if result.NextCursor == "" { + break + } + cursor = result.NextCursor } - s.tools = result.Tools + s.tools = allTools return s.tools, nil } diff --git a/components/execd/pkg/mcpproxy/upstream.go b/components/execd/pkg/mcpproxy/upstream.go index 3761e456f..bcfd87beb 100644 --- a/components/execd/pkg/mcpproxy/upstream.go +++ b/components/execd/pkg/mcpproxy/upstream.go @@ -34,6 +34,7 @@ type Tool struct { Name string `json:"name"` Description string `json:"description,omitempty"` InputSchema json.RawMessage `json:"inputSchema,omitempty"` + Annotations json.RawMessage `json:"annotations,omitempty"` } // UpstreamConfig describes how to connect to an upstream MCP server. @@ -57,7 +58,13 @@ type UpstreamInfo struct { // toolsListResult matches the MCP tools/list response shape. type toolsListResult struct { - Tools []Tool `json:"tools"` + Tools []Tool `json:"tools"` + NextCursor string `json:"nextCursor,omitempty"` +} + +// toolsListParams is the MCP tools/list request params with optional cursor. +type toolsListParams struct { + Cursor string `json:"cursor,omitempty"` } // callToolParams matches the MCP tools/call request params shape. diff --git a/components/execd/tests/smoke_mcpproxy.py b/components/execd/tests/smoke_mcpproxy.py index b7fb63f18..cac1ca032 100644 --- a/components/execd/tests/smoke_mcpproxy.py +++ b/components/execd/tests/smoke_mcpproxy.py @@ -24,7 +24,6 @@ - Optional: set env API_TOKEN if server expects X-EXECD-ACCESS-TOKEN """ -import json import os import sys import tempfile @@ -453,8 +452,8 @@ def main(): # Best-effort cleanup: remove upstream if still registered. try: session.delete(f"{BASE_URL}/mcpproxy/upstreams/mock", timeout=5) - except Exception: - pass + except Exception as exc: + print(f"[!] cleanup: failed to remove upstream (ignored): {exc}", file=sys.stderr) cleanup_mock() From 98bfa04cbc6d19992bf00c6c89c9ba546db098a7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=AB=98=E7=84=B6?= Date: Mon, 22 Jun 2026 15:03:21 +0800 Subject: [PATCH 5/5] fix: address second round review findings - Release manager lock before closing upstreams in RemoveUpstream to avoid blocking proxy traffic during slow upstream shutdown - Validate upstream names: alphanumeric, dot, hyphen, underscore, 1-64 chars, must start with alphanumeric - Cache empty tool lists using toolsFetched bool flag instead of nil check to avoid re-fetching on every rebuildToolList --- components/execd/pkg/mcpproxy/http_upstream.go | 8 +++++--- components/execd/pkg/mcpproxy/manager.go | 13 +++++++------ components/execd/pkg/mcpproxy/stdio_upstream.go | 8 +++++--- components/execd/pkg/web/controller/mcp.go | 7 +++++++ 4 files changed, 24 insertions(+), 12 deletions(-) diff --git a/components/execd/pkg/mcpproxy/http_upstream.go b/components/execd/pkg/mcpproxy/http_upstream.go index 3475d5e7f..a8108d381 100644 --- a/components/execd/pkg/mcpproxy/http_upstream.go +++ b/components/execd/pkg/mcpproxy/http_upstream.go @@ -39,8 +39,9 @@ type httpUpstream struct { mu sync.RWMutex sessionID string - nextID atomic.Int64 - tools []Tool + nextID atomic.Int64 + tools []Tool + toolsFetched bool } func newHTTPUpstream(config UpstreamConfig) *httpUpstream { @@ -204,7 +205,7 @@ func (h *httpUpstream) Initialize(ctx context.Context) error { } func (h *httpUpstream) Tools(ctx context.Context) ([]Tool, error) { - if h.tools != nil { + if h.toolsFetched { return h.tools, nil } @@ -235,6 +236,7 @@ func (h *httpUpstream) Tools(ctx context.Context) ([]Tool, error) { } h.tools = allTools + h.toolsFetched = true return h.tools, nil } diff --git a/components/execd/pkg/mcpproxy/manager.go b/components/execd/pkg/mcpproxy/manager.go index 3edd103da..3f877dffb 100644 --- a/components/execd/pkg/mcpproxy/manager.go +++ b/components/execd/pkg/mcpproxy/manager.go @@ -113,17 +113,12 @@ func (m *Manager) AddUpstream(ctx context.Context, config UpstreamConfig) (*Upst // RemoveUpstream closes an upstream and removes its tools. func (m *Manager) RemoveUpstream(name string) error { m.mu.Lock() - defer m.mu.Unlock() - up, exists := m.upstreams[name] if !exists { + m.mu.Unlock() return fmt.Errorf("upstream %q not found", name) } - if err := up.Close(); err != nil { - log.Warning("mcp proxy: error closing upstream %q: %v", name, err) - } - delete(m.upstreams, name) for toolName, owner := range m.toolIndex { if owner == name { @@ -131,6 +126,12 @@ func (m *Manager) RemoveUpstream(name string) error { } } m.rebuildToolList() + m.mu.Unlock() + + // Close outside the lock — may block on process/network IO. + if err := up.Close(); err != nil { + log.Warning("mcp proxy: error closing upstream %q: %v", name, err) + } log.Info("mcp proxy: removed upstream %q", name) return nil diff --git a/components/execd/pkg/mcpproxy/stdio_upstream.go b/components/execd/pkg/mcpproxy/stdio_upstream.go index 17aedc808..832f303ed 100644 --- a/components/execd/pkg/mcpproxy/stdio_upstream.go +++ b/components/execd/pkg/mcpproxy/stdio_upstream.go @@ -49,8 +49,9 @@ type stdioUpstream struct { nextID atomic.Int64 pending sync.Map // map[string]chan *Response — keyed by JSON-encoded ID - tools []Tool - done chan struct{} + tools []Tool + toolsFetched bool + done chan struct{} } func newStdioUpstream(config UpstreamConfig) *stdioUpstream { @@ -239,7 +240,7 @@ func (s *stdioUpstream) Initialize(ctx context.Context) error { } func (s *stdioUpstream) Tools(ctx context.Context) ([]Tool, error) { - if s.tools != nil { + if s.toolsFetched { return s.tools, nil } @@ -270,6 +271,7 @@ func (s *stdioUpstream) Tools(ctx context.Context) ([]Tool, error) { } s.tools = allTools + s.toolsFetched = true return s.tools, nil } diff --git a/components/execd/pkg/web/controller/mcp.go b/components/execd/pkg/web/controller/mcp.go index 40f3814d5..e08ffa1f8 100644 --- a/components/execd/pkg/web/controller/mcp.go +++ b/components/execd/pkg/web/controller/mcp.go @@ -17,6 +17,7 @@ package controller import ( "fmt" "net/http" + "regexp" "github.com/gin-gonic/gin" @@ -24,6 +25,8 @@ import ( "github.com/alibaba/opensandbox/execd/pkg/web/model" ) +var validUpstreamName = regexp.MustCompile(`^[a-zA-Z0-9][a-zA-Z0-9._-]{0,63}$`) + var mcpManager *mcpproxy.Manager // InitMCPProxy initializes the MCP proxy manager singleton. @@ -136,6 +139,10 @@ func (c *MCPController) AddUpstream() { return } + if !validUpstreamName.MatchString(req.Name) { + c.RespondError(http.StatusBadRequest, model.ErrorCodeInvalidRequest, "name must be 1-64 chars: alphanumeric, dot, hyphen, underscore") + return + } if req.Transport == "stdio" && req.Command == "" { c.RespondError(http.StatusBadRequest, model.ErrorCodeInvalidRequest, "command is required for stdio transport") return