mirror of
https://github.com/coder/coder.git
synced 2026-10-31 19:36:41 +08:00
Michael Suchacz
c8e8b21a88
## Stack Context PR 5 of 6 in a stack that migrates every Coder MCP surface from the archived `github.com/mark3labs/mcp-go` library to the official `github.com/modelcontextprotocol/go-sdk` v1.7.0. Stack: #28056 -> #28057 -> #28058 -> #28059 -> #28060 -> #28061 ## Why The aibridge injected-MCP proxy now owns an official `*mcp.Client`, `*mcp.StreamableClientTransport`, and `*mcp.ClientSession`. - The proxy constructor accepts an optional `*http.Client` instead of mark3labs options; the header-injecting wrapper shallow-copies a supplied client so its Timeout, Jar, and redirect policy survive. - Manual protocol version negotiation and the mark3labs five-second close workaround are removed; the SDK negotiates during `Connect` and fails when no mutually supported version exists. - Repeated `Init` closes the previous session, and a failed tool fetch closes the just-created session so transports do not leak. - Tool and intercept types use the official pointer content types; embedded resource blobs are re-encoded to base64 for model-facing text because the SDK decodes them into raw bytes. - `aibridge/mcpmock` is regenerated, and its stale `go:generate` source path is corrected. > Mux created this PR on Mike's behalf.
131 lines
2.9 KiB
Go
131 lines
2.9 KiB
Go
package mcp
|
|
|
|
import (
|
|
"context"
|
|
"slices"
|
|
"strings"
|
|
"sync"
|
|
|
|
"github.com/modelcontextprotocol/go-sdk/mcp"
|
|
"go.opentelemetry.io/otel/trace"
|
|
"golang.org/x/xerrors"
|
|
|
|
"github.com/coder/coder/v2/aibridge/tracing"
|
|
"github.com/coder/coder/v2/aibridge/utils"
|
|
)
|
|
|
|
var _ ServerProxier = &ServerProxyManager{}
|
|
|
|
// ServerProxyManager can act on behalf of multiple [ServerProxier]s.
|
|
// It aggregates all server resources (currently just tools) across all MCP servers
|
|
// for the purpose of injection into bridged requests and invocation.
|
|
type ServerProxyManager struct {
|
|
proxiers map[string]ServerProxier
|
|
tracer trace.Tracer
|
|
|
|
// Protects access to the tools map.
|
|
toolsMu sync.RWMutex
|
|
tools map[string]*Tool
|
|
}
|
|
|
|
func NewServerProxyManager(proxiers map[string]ServerProxier, tracer trace.Tracer) *ServerProxyManager {
|
|
return &ServerProxyManager{
|
|
proxiers: proxiers,
|
|
tracer: tracer,
|
|
}
|
|
}
|
|
|
|
func (s *ServerProxyManager) addTools(tools []*Tool) {
|
|
s.toolsMu.Lock()
|
|
defer s.toolsMu.Unlock()
|
|
|
|
if s.tools == nil {
|
|
s.tools = make(map[string]*Tool, len(tools))
|
|
}
|
|
|
|
for _, tool := range tools {
|
|
s.tools[tool.ID] = tool
|
|
}
|
|
}
|
|
|
|
// Init concurrently initializes all of its [ServerProxier]s.
|
|
func (s *ServerProxyManager) Init(ctx context.Context) (outErr error) {
|
|
ctx, span := s.tracer.Start(ctx, "ServerProxyManager.Init")
|
|
defer tracing.EndSpanErr(span, &outErr)
|
|
|
|
cg := utils.NewConcurrentGroup()
|
|
for _, proxy := range s.proxiers {
|
|
cg.Go(func() error {
|
|
return proxy.Init(ctx)
|
|
})
|
|
}
|
|
|
|
// Wait for all servers to initialize and load their tools.
|
|
err := cg.Wait()
|
|
|
|
// Aggregate all proxiers' tools.
|
|
for _, proxy := range s.proxiers {
|
|
s.addTools(proxy.ListTools())
|
|
}
|
|
|
|
return err
|
|
}
|
|
|
|
func (s *ServerProxyManager) GetTool(name string) *Tool {
|
|
s.toolsMu.RLock()
|
|
defer s.toolsMu.RUnlock()
|
|
|
|
if s.tools == nil {
|
|
return nil
|
|
}
|
|
|
|
return s.tools[name]
|
|
}
|
|
|
|
func (s *ServerProxyManager) ListTools() []*Tool {
|
|
s.toolsMu.RLock()
|
|
defer s.toolsMu.RUnlock()
|
|
|
|
if s.tools == nil {
|
|
return nil
|
|
}
|
|
|
|
var out []*Tool
|
|
for _, tool := range s.tools {
|
|
out = append(out, tool)
|
|
}
|
|
|
|
slices.SortStableFunc(out, func(a, b *Tool) int {
|
|
return strings.Compare(a.ID, b.ID)
|
|
})
|
|
|
|
return out
|
|
}
|
|
|
|
// CallTool locates the proxier to which the requested tool is associated and
|
|
// delegates the tool call to it.
|
|
func (s *ServerProxyManager) CallTool(ctx context.Context, name string, input any) (*mcp.CallToolResult, error) {
|
|
tool := s.GetTool(name)
|
|
if tool == nil {
|
|
return nil, xerrors.Errorf("%q tool not known", name)
|
|
}
|
|
|
|
proxy, ok := s.proxiers[tool.ServerName]
|
|
if !ok {
|
|
return nil, xerrors.Errorf("%q server not known", tool.ServerName)
|
|
}
|
|
|
|
return proxy.CallTool(ctx, name, input)
|
|
}
|
|
|
|
// Shutdown concurrently shuts down all known proxiers and waits for them *all* to complete.
|
|
func (s *ServerProxyManager) Shutdown(ctx context.Context) error {
|
|
cg := utils.NewConcurrentGroup()
|
|
for _, proxy := range s.proxiers {
|
|
cg.Go(func() error {
|
|
return proxy.Shutdown(ctx)
|
|
})
|
|
}
|
|
return cg.Wait()
|
|
}
|