mcp-gateway-nexus/internal/router/router.go

package router

import (
	"context"
	"encoding/json"
	"fmt"
	"log/slog"
	"sync"

	"gitea.thetempleofdoom.com/drjones/mcp-gateway-nexus/internal/mcp"
)

// Upstream is a single aggregated MCP server behind the Gateway.
type Upstream struct {
	Namespace string
	Client    *mcp.Client

	mu      sync.RWMutex
	ready   bool
	server  mcp.Implementation
	lastErr error
}

func (u *Upstream) setReady(server mcp.Implementation) {
	u.mu.Lock()
	defer u.mu.Unlock()
	u.ready = true
	u.server = server
	u.lastErr = nil
}

func (u *Upstream) setErr(err error) {
	u.mu.Lock()
	defer u.mu.Unlock()
	u.ready = false
	u.lastErr = err
}

// Status is a snapshot of an upstream's connection state.
type Status struct {
	Namespace string `json:"namespace"`
	Ready     bool   `json:"ready"`
	Server    string `json:"server,omitempty"`
	Error     string `json:"error,omitempty"`
}

// Router aggregates upstreams and routes tool calls, backed by the Registry.
type Router struct {
	info     mcp.Implementation
	log      *slog.Logger
	registry *Registry

	mu  sync.RWMutex
	ups map[string]*Upstream // by namespace
}

// New creates a Router. info identifies Nexus to upstreams.
func New(info mcp.Implementation, log *slog.Logger) *Router {
	if log == nil {
		log = slog.Default()
	}
	return &Router{
		info:     info,
		log:      log,
		registry: NewRegistry(),
		ups:      make(map[string]*Upstream),
	}
}

// Registry exposes the tool catalog.
func (r *Router) Registry() *Registry { return r.registry }

// Add registers an upstream client under a namespace. It does not connect.
func (r *Router) Add(namespace string, client *mcp.Client) *Upstream {
	u := &Upstream{Namespace: namespace, Client: client}
	r.mu.Lock()
	r.ups[namespace] = u
	r.mu.Unlock()
	return u
}

// ConnectAll initializes every upstream and loads its tools, in parallel.
// Individual failures are logged and leave that upstream not-ready rather than
// failing the whole gateway (the data plane must stay up — ARCHITECTURE §5).
func (r *Router) ConnectAll(ctx context.Context) {
	r.mu.RLock()
	ups := make([]*Upstream, 0, len(r.ups))
	for _, u := range r.ups {
		ups = append(ups, u)
	}
	r.mu.RUnlock()

	var wg sync.WaitGroup
	for _, u := range ups {
		wg.Add(1)
		go func(u *Upstream) {
			defer wg.Done()
			if err := r.connect(ctx, u); err != nil {
				u.setErr(err)
				r.log.Warn("upstream connect failed", "namespace", u.Namespace, "error", err)
			}
		}(u)
	}
	wg.Wait()
}

func (r *Router) connect(ctx context.Context, u *Upstream) error {
	if err := u.Client.Initialize(ctx); err != nil {
		return err
	}
	u.setReady(u.Client.ServerInfo())
	// Refresh tools list on upstream-initiated change notifications.
	u.Client.SetNotificationHandler(func(m *mcp.Message) {
		if m.Method == mcp.MethodToolListChged {
			r.log.Info("upstream tools changed", "namespace", u.Namespace)
			// Best-effort refresh in the background.
			go func() {
				if err := r.RefreshTools(context.Background(), u); err != nil {
					r.log.Warn("tool refresh failed", "namespace", u.Namespace, "error", err)
				}
			}()
		}
	})
	return r.RefreshTools(ctx, u)
}

// RefreshTools reloads and re-indexes a single upstream's tools.
func (r *Router) RefreshTools(ctx context.Context, u *Upstream) error {
	defs, err := u.Client.ListTools(ctx)
	if err != nil {
		return err
	}
	r.registry.Replace(u.Namespace, defs)
	r.log.Info("indexed upstream tools", "namespace", u.Namespace, "count", len(defs))
	return nil
}

// ListTools returns the aggregated, namespaced tool definitions for agents.
func (r *Router) ListTools() []mcp.ToolDefinition {
	return r.registry.Definitions()
}

// CallTool routes a namespaced tool call to its owning upstream and returns the
// raw MCP result (passed through unmodified) or a JSON-RPC error.
func (r *Router) CallTool(ctx context.Context, qualified string, args json.RawMessage) (json.RawMessage, *mcp.RPCError, error) {
	namespace, name, ok := r.registry.Lookup(qualified)
	if !ok {
		return nil, &mcp.RPCError{Code: mcp.CodeMethodNotFound, Message: fmt.Sprintf("unknown tool %q", qualified)}, nil
	}
	r.mu.RLock()
	u := r.ups[namespace]
	r.mu.RUnlock()
	if u == nil {
		return nil, &mcp.RPCError{Code: mcp.CodeInternalError, Message: fmt.Sprintf("no upstream for namespace %q", namespace)}, nil
	}
	u.mu.RLock()
	ready := u.ready
	u.mu.RUnlock()
	if !ready {
		return nil, &mcp.RPCError{Code: mcp.CodeInternalError, Message: fmt.Sprintf("upstream %q is not ready", namespace)}, nil
	}
	return u.Client.CallTool(ctx, name, args)
}

// Statuses returns a snapshot of every upstream's connection state.
func (r *Router) Statuses() []Status {
	r.mu.RLock()
	defer r.mu.RUnlock()
	out := make([]Status, 0, len(r.ups))
	for _, u := range r.ups {
		u.mu.RLock()
		s := Status{Namespace: u.Namespace, Ready: u.ready, Server: u.server.Name}
		if u.lastErr != nil {
			s.Error = u.lastErr.Error()
		}
		u.mu.RUnlock()
		out = append(out, s)
	}
	return out
}

// Close shuts down every upstream client.
func (r *Router) Close() {
	r.mu.Lock()
	defer r.mu.Unlock()
	for _, u := range r.ups {
		_ = u.Client.Close()
	}
}