Serve a unified MCP server for all projects on the admin service

admin/server/mcp_test.go

@@ -0,0 +1,43 @@
+package server
+
+import (
+	"encoding/json"
+	"testing"
+
+	"github.com/google/jsonschema-go/jsonschema"
+	"github.com/rilldata/rill/runtime/ai"
+	"github.com/stretchr/testify/require"
+)
+
+func TestMCPForwardedToolSpecs(t *testing.T) {
+	tools, err := mcpForwardedToolSpecs()
+	require.NoError(t, err)
+	require.Len(t, tools, len(mcpForwardedTools))
+
+	for _, tool := range tools {
+		schema, ok := tool.InputSchema.(*jsonschema.Schema)
+		require.True(t, ok, "tool %q", tool.Name)
+		require.Contains(t, schema.Properties, mcpProjectArg, "tool %q", tool.Name)
+		require.Contains(t, schema.Required, mcpProjectArg, "tool %q", tool.Name)
+
+		// query_metrics_view has a hand-written schema, which must be extended and not replaced.
+		if tool.Name == ai.QueryMetricsViewName {
+			require.Contains(t, schema.Properties, "metrics_view")
+			require.NotEmpty(t, schema.Defs)
+		}
+	}
+}
+
+func TestTakeMCPProjectArg(t *testing.T) {
+	org, project, rest, err := takeMCPProjectArg(json.RawMessage(`{"project":"org/proj","metrics_view":"mv"}`))
+	require.NoError(t, err)
+	require.Equal(t, "org", org)
+	require.Equal(t, "proj", project)
+	require.JSONEq(t, `{"metrics_view":"mv"}`, string(rest))
+
+	// The project argument is required, and must name both an organization and a project.
+	for _, args := range []string{``, `{}`, `{"project":""}`, `{"project":"org"}`, `{"project":"org/"}`, `{"project":"/proj"}`, `{"project":123}`, `not json`} {
+		_, _, _, err := takeMCPProjectArg(json.RawMessage(args))
+		require.Error(t, err, "args %q", args)
+	}
+}

admin/server/runtime_proxy.go

@@ -2,6 +2,8 @@ package server
 
 import (
 	"bufio"
+	"context"
+	"errors"
 	"fmt"
 	"io"
 	"net/http"
@@ -32,32 +34,11 @@ func (s *Server) runtimeProxyForOrgAndProject(w http.ResponseWriter, r *http.Req
 	proxyPath := r.PathValue("path")
 	proxyRawQuery := r.URL.RawQuery
 
-	// Find the project we're proxying to
-	proj, err := s.admin.DB.FindProjectByName(r.Context(), org, project)
-	if err != nil {
-		return httputil.Error(http.StatusBadRequest, err)
-	}
-
-	// Find the deployment to proxy to.
+	// Find the project and deployment we're proxying to.
 	// If a branch was specified, use the deployment for that branch; otherwise use the project's primary deployment.
-	var depl *database.Deployment
-	if branch == "" {
-		if proj.PrimaryDeploymentID == nil {
-			return httputil.Errorf(http.StatusBadRequest, "no prod deployment for project")
-		}
-		depl, err = s.admin.DB.FindDeployment(r.Context(), *proj.PrimaryDeploymentID)
-		if err != nil {
-			return httputil.Error(http.StatusBadRequest, err)
-		}
-	} else {
-		depls, err := s.admin.DB.FindDeploymentsForProject(r.Context(), proj.ID, "", branch)
-		if err != nil {
-			return httputil.Error(http.StatusBadRequest, err)
-		}
-		if len(depls) == 0 {
-			return httputil.Errorf(http.StatusBadRequest, "no deployment for branch %q", branch)
-		}
-		depl = depls[0] // At most one deployment per branch is allowed
+	proj, depl, err := s.resolveDeploymentForOrgAndProject(r.Context(), org, project, branch)
+	if err != nil {
+		return err
 	}
 
 	// Prepare a JWT to use for the proxied request.
@@ -78,27 +59,15 @@ func (s *Server) runtimeProxyForOrgAndProject(w http.ResponseWriter, r *http.Req
 	}
 	// If a direct JWT was not provided, issue a new ephemeral runtime JWT for the proxied request.
 	if jwt == "" {
-		permissions := claims.ProjectPermissions(r.Context(), proj.OrganizationID, depl.ProjectID)
-		if proj.Public {
-			permissions.ReadProject = true
-			permissions.ReadProd = true
-		}
-		if !permissions.ReadProd {
+		jwt, err = s.issueEphemeralRuntimeToken(r.Context(), proj, depl, runtimeProxyAccessTokenTTL)
+		if errors.Is(err, errNoProdAccess) {
 			if claims.OwnerType() == auth.OwnerTypeAnon {
 				// This means no token was provided, so return instructions for how to initiate an OAuth flow.
 				// This is currently used by MCP clients that authenticate with OAuth.
 				w.Header().Set("WWW-Authenticate", fmt.Sprintf("Bearer resource_metadata=%q", s.admin.URLs.OAuthProtectedResourceMetadata(r)))
 			}
-			return httputil.Errorf(http.StatusUnauthorized, "does not have permission to access the production deployment")
+			return httputil.Error(http.StatusUnauthorized, err)
 		}
-
-		jwt, err = s.issueRuntimeToken(r.Context(), &issueRuntimeTokenOptions{
-			project:            proj,
-			deployment:         depl,
-			projectPermissions: permissions,
-			forOwner:           true,
-			ttl:                runtimeProxyAccessTokenTTL,
-		})
 		if err != nil {
 			return httputil.Error(http.StatusInternalServerError, err)
 		}
@@ -107,20 +76,8 @@ func (s *Server) runtimeProxyForOrgAndProject(w http.ResponseWriter, r *http.Req
 	// Track usage of the deployment
 	s.admin.Used.Deployment(depl.ID)
 
-	// Determine runtime host.
-	// NOTE: In production, the runtime host serves both the HTTP and gRPC servers.
-	// But in development, the two are presently on different ports, and depl.RuntimeHost is that of the gRPC server.
-	// Until we get both servers on the same port in development, this hack rewrites the runtime host to the HTTP server.
-	runtimeHost := depl.RuntimeHost
-	if strings.HasPrefix(runtimeHost, "http://localhost:") {
-		runtimeHost = os.Getenv("RILL_RUNTIME_AUTH_AUDIENCE_URL")
-		if runtimeHost == "" {
-			runtimeHost = "http://localhost:8081"
-		}
-	}
-
 	// Create the URL to proxy to by prepending `/v1/instances/{instanceID}` to the proxy path.
-	proxyURL, err := url.Parse(runtimeHost)
+	proxyURL, err := url.Parse(runtimeHTTPHost(depl.RuntimeHost))
 	if err != nil {
 		return httputil.Error(http.StatusInternalServerError, err)
 	}
@@ -201,3 +158,73 @@ func (s *Server) runtimeProxyForOrgAndProject(w http.ResponseWriter, r *http.Req
 
 	return nil
 }
+
+// errNoProdAccess is returned when the caller does not have access to a project's production deployment.
+// It is a sentinel because callers own the response and decide whether to emit an OAuth challenge.
+var errNoProdAccess = errors.New("does not have permission to access the production deployment")
+
+// resolveDeploymentForOrgAndProject resolves an org and project name to the deployment to serve requests from.
+// If branch is empty, it resolves the project's primary deployment.
+func (s *Server) resolveDeploymentForOrgAndProject(ctx context.Context, org, project, branch string) (*database.Project, *database.Deployment, error) {
+	proj, err := s.admin.DB.FindProjectByName(ctx, org, project)
+	if err != nil {
+		return nil, nil, httputil.Error(http.StatusBadRequest, err)
+	}
+
+	if branch == "" {
+		if proj.PrimaryDeploymentID == nil {
+			return nil, nil, httputil.Errorf(http.StatusBadRequest, "no prod deployment for project")
+		}
+		depl, err := s.admin.DB.FindDeployment(ctx, *proj.PrimaryDeploymentID)
+		if err != nil {
+			return nil, nil, httputil.Error(http.StatusBadRequest, err)
+		}
+		return proj, depl, nil
+	}
+
+	depls, err := s.admin.DB.FindDeploymentsForProject(ctx, proj.ID, "", branch)
+	if err != nil {
+		return nil, nil, httputil.Error(http.StatusBadRequest, err)
+	}
+	if len(depls) == 0 {
+		return nil, nil, httputil.Errorf(http.StatusBadRequest, "no deployment for branch %q", branch)
+	}
+	return proj, depls[0], nil // At most one deployment per branch is allowed
+}
+
+// issueEphemeralRuntimeToken checks that the caller can read a project's production deployment,
+// then issues a runtime JWT for it similar to the one that could be obtained by calling GetProject.
+// It returns errNoProdAccess if the caller does not have access.
+func (s *Server) issueEphemeralRuntimeToken(ctx context.Context, proj *database.Project, depl *database.Deployment, ttl time.Duration) (string, error) {
+	claims := auth.GetClaims(ctx)
+	permissions := claims.ProjectPermissions(ctx, proj.OrganizationID, depl.ProjectID)
+	if proj.Public {
+		permissions.ReadProject = true
+		permissions.ReadProd = true
+	}
+	if !permissions.ReadProd {
+		return "", errNoProdAccess
+	}
+
+	return s.issueRuntimeToken(ctx, &issueRuntimeTokenOptions{
+		project:            proj,
+		deployment:         depl,
+		projectPermissions: permissions,
+		forOwner:           true,
+		ttl:                ttl,
+	})
+}
+
+// runtimeHTTPHost returns the host to send HTTP requests to for a deployment.
+// NOTE: In production, the runtime host serves both the HTTP and gRPC servers.
+// But in development, the two are presently on different ports, and depl.RuntimeHost is that of the gRPC server.
+// Until we get both servers on the same port in development, this hack rewrites the runtime host to the HTTP server.
+func runtimeHTTPHost(runtimeHost string) string {
+	if !strings.HasPrefix(runtimeHost, "http://localhost:") {
+		return runtimeHost
+	}
+	if host := os.Getenv("RILL_RUNTIME_AUTH_AUDIENCE_URL"); host != "" {
+		return host
+	}
+	return "http://localhost:8081"
+}

admin/server/server.go

@@ -220,6 +220,20 @@ func (s *Server) HTTPHandler(ctx context.Context) (http.Handler, error) {
 	observability.MuxHandle(mux, "/v1/organizations/{org}/projects/{project}/runtime/{path...}", proxyHandler)        // Backwards compatibility
 	observability.MuxHandle(mux, "/v1/orgs/{org}/projects/{project}/branch/{branch}/runtime/{path...}", proxyHandler) // Branch-specific deployment
 
+	// Add MCP server.
+	// Note this is the admin service's MCP server, not a runtime MCP server (although this includes an ability to proxy certain tool calls to a project's runtime).
+	mcpHandler, err := s.mcpHandler()
+	if err != nil {
+		return nil, err
+	}
+	mcpHandler = observability.Middleware(
+		"mcp",
+		s.logger,
+		runtimeProxyCORSMiddleware(s.authenticator.HTTPMiddlewareLenient(s.mcpRequireAuth(mcpHandler))),
+	)
+	observability.MuxHandle(mux, "/v1/mcp", mcpHandler)
+	observability.MuxHandle(mux, "/v1/mcp/{$}", mcpHandler) // Avoids falling through to the gRPC transcoder on a trailing slash
+
 	// Add backwards compatibility alias for iframe endpoint
 	observability.MuxHandle(mux, "/v1/organizations/{org}/projects/{project}/iframe", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		r.URL.Path = strings.Replace(r.URL.Path, "/v1/organizations/", "/v1/orgs/", 1)

runtime/ai/ai.go

@@ -81,8 +81,10 @@ func NewRunner(rt *runtime.Runtime, activity *activity.Client) *Runner {
 
 // SessionOptions provides options for initializing a new session.
 type SessionOptions struct {
-	InstanceID        string
-	SessionID         string
+	InstanceID string
+	SessionID  string
+	// CreateIfNotExists creates the session if it does not exist.
+	// If SessionID is set, the session is created with that ID.
 	CreateIfNotExists bool
 	Claims            *runtime.SecurityClaims
 	UserAgent         string
@@ -109,9 +111,15 @@ func (r *Runner) Session(ctx context.Context, opts *SessionOptions) (res *Sessio
 	if opts.SessionID != "" {
 		session, err = catalog.FindAISession(ctx, opts.SessionID)
 		if err != nil {
-			return nil, fmt.Errorf("failed to find session %q: %w", opts.SessionID, err)
+			// If CreateIfNotExists is set, an unknown session ID is created below instead of erroring.
+			// This lets a caller that doesn't know Rill's session IDs address a stable session,
+			// which the unified MCP server in the admin service relies on to keep one session per project.
+			if !errors.Is(err, drivers.ErrNotFound) || !opts.CreateIfNotExists {
+				return nil, fmt.Errorf("failed to find session %q: %w", opts.SessionID, err)
+			}
 		}
-
+	}
+	if session != nil {
 		// Check access: you can access anonymous sessions, your own sessions, and shared sessions.
 		// For shared sessions, if you are not the owner, you can only see messages up to the SharedUntilMessageID (inclusive).
 		// For sessions without an owner (unauthenticated users using a public project), we don't check access and rely on security by obscurity (generally a decent trade-off, but specifically introduced to get citation links over MCP working for unauthenticated demos).
@@ -147,9 +155,13 @@ func (r *Runner) Session(ctx context.Context, opts *SessionOptions) (res *Sessio
 			}
 		}
 	}
-	if opts.SessionID == "" {
+	if session == nil {
+		id := opts.SessionID
+		if id == "" {
+			id = uuid.NewString()
+		}
 		session = &drivers.AISession{
-			ID:         uuid.NewString(),
+			ID:         id,
 			InstanceID: opts.InstanceID,
 			OwnerID:    opts.Claims.UserID,
 			Title:      "",

runtime/ai/mcp.go

@@ -10,7 +10,9 @@ import (
 	"go.uber.org/zap"
 )
 
-const mcpInstructions = `
+// MCPInstructions are the instructions advertised by the MCP server.
+// It is exported so the unified MCP server in the admin service can extend it instead of restating it.
+const MCPInstructions = `
 # Rill MCP Server
 This server exposes APIs for querying **metrics views**, which represent Rill's metrics layer.
 
@@ -31,6 +33,19 @@ If you have edit access, the server also exposes tools for inspecting and editin
 - **Write a file:** Use "write_file" to create, update or delete a file. If the file declares a Rill resource, it returns the resource's status and any errors encountered after reconciliation.
 `
 
+// MCPToolSpecs returns the specs of all registered tools, keyed by name.
+// The specs are freshly built and owned by the caller, so they may be mutated.
+// It is used by the unified MCP server in the admin service, which advertises the tools without being able to run them.
+func MCPToolSpecs() map[string]*mcp.Tool {
+	// Safe to pass a nil runtime: Spec() does not use it (asserted by TestMCPToolSpecs).
+	runner := NewRunner(nil, nil)
+	specs := make(map[string]*mcp.Tool, len(runner.Tools))
+	for name, t := range runner.Tools {
+		specs[name] = t.Spec
+	}
+	return specs
+}
+
 // MCPServer returns a new MCP server scoped to the current session.
 // Since it is scoped to the session, a new MCP server should be created for each client connection.
 // Using a separate MCP server for each client enables tailoring the server's instructions and available tools to the end user's claims.
@@ -43,7 +58,7 @@ func (s *Session) MCPServer(ctx context.Context) *mcp.Server {
 			Version: s.runner.Runtime.Version().String(),
 		},
 		&mcp.ServerOptions{
-			Instructions: mcpInstructions,
+			Instructions: MCPInstructions,
 			InitializedHandler: func(ctx context.Context, r *mcp.InitializedRequest) {
 				// Save user agent in the session
 				clientInfo := r.Session.InitializeParams().ClientInfo

runtime/ai/mcp_test.go

@@ -11,6 +11,58 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
+// TestMCPToolSpecs asserts that tool specs can be built without a runtime,
+// which is the invariant ai.MCPToolSpecs relies on to serve the admin service's unified MCP server.
+func TestMCPToolSpecs(t *testing.T) {
+	specs := ai.MCPToolSpecs()
+	require.NotEmpty(t, specs)
+	for name, spec := range specs {
+		require.Equal(t, name, spec.Name)
+		require.NotEmpty(t, spec.Description, "tool %q", name)
+		require.NotNil(t, spec.InputSchema, "tool %q", name)
+	}
+	require.Contains(t, specs, ai.QueryMetricsViewName)
+}
+
+// TestSessionCreateIfNotExists asserts that a caller can address a session by an ID of its own choosing,
+// which the admin service's unified MCP server relies on to keep one session per project.
+func TestSessionCreateIfNotExists(t *testing.T) {
+	rt, instanceID := testruntime.NewInstanceWithOptions(t, testruntime.InstanceOptions{})
+	claims := &runtime.SecurityClaims{UserID: uuid.NewString(), Permissions: []runtime.Permission{runtime.UseAI}}
+	r := ai.NewRunner(rt, activity.NewNoopClient())
+	sessionID := uuid.NewString()
+
+	// The session does not exist yet, so it is created with the given ID.
+	s, err := r.Session(t.Context(), &ai.SessionOptions{
+		InstanceID:        instanceID,
+		SessionID:         sessionID,
+		CreateIfNotExists: true,
+		Claims:            claims,
+		UserAgent:         "mcp-client",
+	})
+	require.NoError(t, err)
+	require.Equal(t, sessionID, s.ID())
+	require.NoError(t, s.Flush(t.Context()))
+
+	// A second call loads the same session instead of creating another one.
+	s2, err := r.Session(t.Context(), &ai.SessionOptions{
+		InstanceID:        instanceID,
+		SessionID:         sessionID,
+		CreateIfNotExists: true,
+		Claims:            claims,
+	})
+	require.NoError(t, err)
+	require.Equal(t, sessionID, s2.ID())
+
+	// Without CreateIfNotExists, an unknown session ID is still an error.
+	_, err = r.Session(t.Context(), &ai.SessionOptions{
+		InstanceID: instanceID,
+		SessionID:  uuid.NewString(),
+		Claims:     claims,
+	})
+	require.Error(t, err)
+}
+
 // TestDeveloperToolsMCPAccess verifies that the leaf developer tools are exposed to external
 // MCP clients (non-rill user agents) for callers with EditRepo, while the developer agents remain
 // restricted to first-party Rill clients.