diff --git a/admin/server/mcp.go b/admin/server/mcp.go new file mode 100644 index 00000000000..55a382ae540 --- /dev/null +++ b/admin/server/mcp.go @@ -0,0 +1,428 @@ +package server + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "strings" + "time" + + "github.com/google/jsonschema-go/jsonschema" + "github.com/google/uuid" + "github.com/modelcontextprotocol/go-sdk/mcp" + "github.com/rilldata/rill/admin/database" + "github.com/rilldata/rill/admin/server/auth" + "github.com/rilldata/rill/runtime/ai" + "github.com/rilldata/rill/runtime/pkg/httputil" + "github.com/rilldata/rill/runtime/pkg/middleware" + "github.com/rilldata/rill/runtime/pkg/observability" +) + +// mcpInstructions are the instructions advertised to clients of the admin service's MCP server. +// They are a combination of the admin service's own instructions and the instructions from the runtime's AI tools that we proxy to. +const mcpInstructions = ` +# Rill Cloud MCP Server +This server provides access to several Rill projects. + +1. **List projects:** Use "list_projects" to discover the projects you have access to. +2. **Target a project:** Pass the "project" argument on every other tool call. Do this before the workflow below. + +` + ai.MCPInstructions + +const ( + // mcpProjectArg is the tool argument the MCP server uses to route a tool call to a project. + mcpProjectArg = "project" + mcpListProjectsLimit = 100 + + // mcpToolCallTimeout is deliberately higher than the runtime's own tool timeout, + // so the runtime's timeout applies first and returns a more useful message. + mcpToolCallTimeout = 6 * time.Minute +) + +// mcpForwardedTools are the runtime tools exposed on the MCP server. +// The runtime exposes more tools than these, but the rest require runtime.EditRepo, +// which is only granted for editable deployments, and a project's primary deployment is never editable. +var mcpForwardedTools = []string{ + ai.ListMetricsViewsName, + ai.GetMetricsViewName, + ai.QueryMetricsViewSummaryName, + ai.QueryMetricsViewName, +} + +// mcpSessionNamespace is a UUIDv5 namespace for the per-project session IDs derived in callRuntimeTool. +// The value is a randomly generated UUID with no meaning: UUIDv5 requires a namespace and any fixed value will do. +// It must not be changed, since that would orphan the sessions derived using the previous value. +var mcpSessionNamespace = uuid.MustParse("6f4a1c94-3d1e-4f6b-9a02-1c0f6b7d5e8a") + +// mcpHandler creates a handler for the MCP server. +// Unlike the per-project runtime proxy, it serves a single endpoint for all the projects a caller has access to: +// it answers the MCP protocol methods itself and forwards tool calls to the runtime of the project named in the tool arguments. +func (s *Server) mcpHandler() (http.Handler, error) { + adminTools := s.mcpAdminTools() + forwardedTools, err := mcpForwardedToolSpecs() + if err != nil { + return nil, err + } + for _, at := range adminTools { + for _, ft := range forwardedTools { + if at.spec.Name == ft.Name { + return nil, fmt.Errorf("mcp: tool %q is both served by the admin service and forwarded to the runtime", ft.Name) + } + } + } + + return mcp.NewStreamableHTTPHandler(func(r *http.Request) *mcp.Server { + srv := mcp.NewServer( + &mcp.Implementation{ + Name: "rill", + Title: "Rill Cloud MCP Server", + Version: s.admin.Version.String(), + }, + &mcp.ServerOptions{ + Instructions: mcpInstructions, + HasTools: true, + // Issue a session ID so clients pass it back on subsequent requests, enabling us to keep one AI session per project. + GetSessionID: uuid.NewString, + }, + ) + + srv.AddReceivingMiddleware(observability.MCPMiddleware()) + + // Separate timeout middleware is required since the SDK strips the request's deadline and cancellation before calling the tool handler. + srv.AddReceivingMiddleware(middleware.TimeoutMCPMiddleware(func(method, tool string) time.Duration { + return mcpToolCallTimeout + })) + + client := &mcpClient{ + sessionID: r.Header.Get("Mcp-Session-Id"), + userAgent: r.UserAgent(), + protocolVersion: r.Header.Get("Mcp-Protocol-Version"), + } + + for _, t := range adminTools { + srv.AddTool(t.spec, t.handler) + } + for _, t := range forwardedTools { + srv.AddTool(t, s.mcpForwardToolCall(t.Name, client)) + } + + return srv + }, &mcp.StreamableHTTPOptions{ + Stateless: true, // Required since the server serves multiple users. + }), nil +} + +// mcpRequireAuth rejects unauthenticated requests to the MCP server. +// It wraps the MCP handler instead of being applied inside it, so that clients receive the challenge on the initialization request and can start an OAuth flow. +func (s *Server) mcpRequireAuth(next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + // Unlike the per-project runtime proxy, we can't accept a runtime JWT here since it's scoped to a single deployment. + if auth.GetClaims(r.Context()).OwnerType() == auth.OwnerTypeAnon { + w.Header().Set("WWW-Authenticate", fmt.Sprintf("Bearer resource_metadata=%q", s.admin.URLs.OAuthProtectedResourceMetadata(r))) + httputil.WriteError(w, httputil.Errorf(http.StatusUnauthorized, "not authenticated")) + return + } + next.ServeHTTP(w, r) + }) +} + +// mcpAdminTools returns the tools served by the admin service itself. +// These are tools for concepts the runtime is not aware of, such as organizations and projects. +func (s *Server) mcpAdminTools() []mcpAdminTool { + return []mcpAdminTool{ + { + spec: &mcp.Tool{ + Name: "list_projects", + Title: "List Projects", + Description: `List the Rill projects you have access to. Pass a returned value as the "project" argument of other tools. The list may be incomplete; if you already know a project, you can pass it directly.`, + Annotations: &mcp.ToolAnnotations{ReadOnlyHint: true}, + InputSchema: &jsonschema.Schema{Type: "object"}, + }, + handler: s.mcpListProjects, + }, + } +} + +// mcpAdminTool is a tool served by the admin service itself instead of being forwarded to a project's runtime. +type mcpAdminTool struct { + spec *mcp.Tool + handler mcp.ToolHandler +} + +// mcpForwardedToolSpecs returns the specs of the tools forwarded to a project's runtime, using the tool definitions in runtime/ai as the source of truth. +// It injects the "project" argument into the input schema of each tool, which we intercept to route calls to the correct runtime. +func mcpForwardedToolSpecs() ([]*mcp.Tool, error) { + specs := ai.MCPToolSpecs() + + res := make([]*mcp.Tool, 0, len(mcpForwardedTools)) + for _, name := range mcpForwardedTools { + spec, ok := specs[name] + if !ok { + return nil, fmt.Errorf("mcp: no spec for tool %q", name) + } + + // Inject the "project" argument into the input schema. + schema, ok := spec.InputSchema.(*jsonschema.Schema) + if !ok { + return nil, fmt.Errorf("mcp: tool %q does not have a JSON schema", name) + } + if schema.Properties == nil { + schema.Properties = make(map[string]*jsonschema.Schema) + } + schema.Properties[mcpProjectArg] = &jsonschema.Schema{ + Type: "string", + Description: `The project to target, as returned by "list_projects".`, + } + schema.Required = append(schema.Required, mcpProjectArg) + + res = append(res, spec) + } + + return res, nil +} + +// mcpClient describes a proxy client making a request from the admin service to a runtime. +type mcpClient struct { + sessionID string + userAgent string + protocolVersion string +} + +// mcpForwardToolCall returns a handler that forwards a tool call to the runtime of the project named in its arguments. +func (s *Server) mcpForwardToolCall(name string, client *mcpClient) mcp.ToolHandler { + return func(ctx context.Context, req *mcp.CallToolRequest) (*mcp.CallToolResult, error) { + // Note we return tool errors instead of errors for anything the caller can act on, + // since an error is returned to the client as a protocol error it can't recover from. + org, project, args, err := takeMCPProjectArg(req.Params.Arguments) + if err != nil { + return mcpToolErrorf(`%s. Call "list_projects" to discover available projects.`, err), nil + } + + proj, depl, err := s.resolveDeploymentForOrgAndProject(ctx, org, project, "") + if err != nil { + return mcpToolErrorf("failed to find project %q: %s", org+"/"+project, err), nil + } + + jwt, err := s.issueEphemeralRuntimeToken(ctx, proj, depl, runtimeProxyAccessTokenTTL) + if err != nil { + if errors.Is(err, errNoProdAccess) { + return mcpToolErrorf("you do not have access to project %q", org+"/"+project), nil + } + return nil, err + } + + s.admin.Used.Deployment(depl.ID) + + return s.callRuntimeTool(ctx, depl, jwt, client, name, args) + } +} + +// callRuntimeTool calls a tool on a deployment's MCP server. +// It sends a plain JSON-RPC request instead of using the MCP SDK's client, which would add an initialization +// handshake to every tool call and would not let us pass through the session ID and user agent. +func (s *Server) callRuntimeTool(ctx context.Context, depl *database.Deployment, jwt string, client *mcpClient, name string, args json.RawMessage) (*mcp.CallToolResult, error) { + body, err := json.Marshal(map[string]any{ + "jsonrpc": "2.0", + "id": 1, + "method": "tools/call", + "params": mcp.CallToolParams{Name: name, Arguments: args}, + }) + if err != nil { + return nil, err + } + + url := runtimeHTTPHost(depl.RuntimeHost) + "/v1/instances/" + depl.RuntimeInstanceID + "/mcp" + req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(body)) + if err != nil { + return nil, err + } + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "application/json, text/event-stream") // The runtime requires both, even though it responds with JSON + req.Header.Set("Authorization", "Bearer "+jwt) + if client.protocolVersion != "" { + req.Header.Set("Mcp-Protocol-Version", client.protocolVersion) + } + if client.userAgent != "" { + req.Header.Set("User-Agent", client.userAgent) + } + if client.sessionID != "" { + // Derive a session ID per project instead of passing the client's through. + // A runtime may serve several projects from one catalog, where AI session IDs must be unique. + sessionID := uuid.NewSHA1(mcpSessionNamespace, []byte(depl.RuntimeInstanceID+"/"+client.sessionID)) + req.Header.Set("Mcp-Session-Id", sessionID.String()) + } + + // Uses http.DefaultClient to get caching/pooling of TCP connections. + res, err := http.DefaultClient.Do(req) + if err != nil { + return nil, err + } + defer res.Body.Close() + + data, err := io.ReadAll(res.Body) + if err != nil { + return nil, err + } + if res.StatusCode != http.StatusOK { + return nil, fmt.Errorf("runtime returned status %d: %s", res.StatusCode, strings.TrimSpace(string(data))) + } + + // The runtime responds with a single JSON-RPC message, or an array of them if it also emitted notifications. + msgs := []json.RawMessage{data} + if bytes.HasPrefix(bytes.TrimSpace(data), []byte("[")) { + if err := json.Unmarshal(data, &msgs); err != nil { + return nil, err + } + } + for _, msg := range msgs { + var rpc struct { + Result *mcp.CallToolResult `json:"result"` + Error *struct { + Message string `json:"message"` + } `json:"error"` + } + if err := json.Unmarshal(msg, &rpc); err != nil { + return nil, err + } + if rpc.Error != nil { + return nil, errors.New(rpc.Error.Message) + } + if rpc.Result != nil { + return rpc.Result, nil + } + } + return nil, errors.New("runtime did not return a result") +} + +// mcpListProjects lists the projects the caller has access to. +// It is a discovery convenience, not an access check: it may return fewer projects than the caller can access +// (for example, it does not cover public projects), and each tool call checks access for the project it targets. +func (s *Server) mcpListProjects(ctx context.Context, req *mcp.CallToolRequest) (*mcp.CallToolResult, error) { + claims := auth.GetClaims(ctx) + + var projs []*database.Project + switch claims.OwnerType() { + case auth.OwnerTypeUser: + var err error + projs, err = s.admin.DB.FindProjectsForUser(ctx, claims.OwnerID()) + if err != nil { + return nil, err + } + case auth.OwnerTypeService: + svc, err := s.admin.DB.FindService(ctx, claims.OwnerID()) + if err != nil { + return nil, err + } + // A service belongs to exactly one organization. + // If it can manage the organization's projects, it has access to all of them, not just those it's a member of. + if claims.OrganizationPermissions(ctx, svc.OrgID).ManageProjects { + projs, err = s.admin.DB.FindProjectsForOrganization(ctx, svc.OrgID, "", mcpListProjectsLimit) + if err != nil { + return nil, err + } + break + } + members, err := s.admin.DB.FindProjectMemberServicesForService(ctx, svc.ID) + if err != nil { + return nil, err + } + for _, m := range members { + proj, err := s.admin.DB.FindProject(ctx, m.ProjectID) + if err != nil { + return nil, err + } + projs = append(projs, proj) + } + case auth.OwnerTypeMagicAuthToken: + tkn, ok := claims.AuthTokenModel().(*database.MagicAuthToken) + if !ok { + return nil, errors.New("unexpected auth token model") + } + proj, err := s.admin.DB.FindProject(ctx, tkn.ProjectID) + if err != nil { + return nil, err + } + projs = append(projs, proj) + default: + return nil, fmt.Errorf("listing projects is not supported for %s tokens", claims.OwnerType()) + } + + if len(projs) > mcpListProjectsLimit { + projs = projs[:mcpListProjectsLimit] + } + + // Build the "/" slugs used for the "project" tool argument. + res := struct { + Projects []string `json:"projects"` + Note string `json:"note,omitempty"` + }{Projects: make([]string, 0, len(projs))} + orgNames := make(map[string]string) + for _, proj := range projs { + name, ok := orgNames[proj.OrganizationID] + if !ok { + org, err := s.admin.DB.FindOrganization(ctx, proj.OrganizationID) + if err != nil { + return nil, err + } + name = org.Name + orgNames[proj.OrganizationID] = name + } + res.Projects = append(res.Projects, name+"/"+proj.Name) + } + if len(res.Projects) == mcpListProjectsLimit { + res.Note = fmt.Sprintf("Only the first %d projects are shown.", mcpListProjectsLimit) + } + + data, err := json.Marshal(res) + if err != nil { + return nil, err + } + return &mcp.CallToolResult{ + Content: []mcp.Content{&mcp.TextContent{Text: string(data)}}, + StructuredContent: res, + }, nil +} + +// takeMCPProjectArg removes the project argument from raw tool arguments and returns the organization and project it names. +// It must be removed because the runtime's tools do not accept it. +func takeMCPProjectArg(args json.RawMessage) (org, project string, rest json.RawMessage, err error) { + var m map[string]json.RawMessage + if len(args) > 0 { + if err := json.Unmarshal(args, &m); err != nil { + return "", "", nil, fmt.Errorf("invalid arguments: %w", err) + } + } + + raw, ok := m[mcpProjectArg] + if !ok { + return "", "", nil, fmt.Errorf("missing the %q argument", mcpProjectArg) + } + var slug string + if err := json.Unmarshal(raw, &slug); err != nil { + return "", "", nil, fmt.Errorf("invalid %q argument: %w", mcpProjectArg, err) + } + org, project, ok = strings.Cut(slug, "/") + if !ok || org == "" || project == "" { + return "", "", nil, fmt.Errorf("invalid %q argument %q: expected the format \"/\"", mcpProjectArg, slug) + } + delete(m, mcpProjectArg) + + rest, err = json.Marshal(m) + if err != nil { + return "", "", nil, err + } + return org, project, rest, nil +} + +// mcpToolErrorf returns a tool call result representing an error. +// Unlike an error returned from a tool handler, it is presented to the model as a result it can act on. +func mcpToolErrorf(format string, args ...any) *mcp.CallToolResult { + return &mcp.CallToolResult{ + IsError: true, + Content: []mcp.Content{&mcp.TextContent{Text: fmt.Sprintf(format, args...)}}, + } +} diff --git a/admin/server/mcp_test.go b/admin/server/mcp_test.go new file mode 100644 index 00000000000..c164a489f62 --- /dev/null +++ b/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) + } +} diff --git a/admin/server/runtime_proxy.go b/admin/server/runtime_proxy.go index a1625346917..41a95528a87 100644 --- a/admin/server/runtime_proxy.go +++ b/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" +} diff --git a/admin/server/server.go b/admin/server/server.go index 244c6ea2898..143231187f6 100644 --- a/admin/server/server.go +++ b/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) diff --git a/runtime/ai/ai.go b/runtime/ai/ai.go index 85b6a4dcdb1..d7ac889d2a8 100644 --- a/runtime/ai/ai.go +++ b/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: "", diff --git a/runtime/ai/mcp.go b/runtime/ai/mcp.go index d057540551b..992c4202eb1 100644 --- a/runtime/ai/mcp.go +++ b/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 diff --git a/runtime/ai/mcp_test.go b/runtime/ai/mcp_test.go index d16a5121862..cb77d107f23 100644 --- a/runtime/ai/mcp_test.go +++ b/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. diff --git a/runtime/server/mcp.go b/runtime/server/mcp.go index bb982aefb58..3a3400f551f 100644 --- a/runtime/server/mcp.go +++ b/runtime/server/mcp.go @@ -32,13 +32,20 @@ func (s *Server) mcpHandler() http.Handler { // Get session ID (will be empty if it's the first request) sessionID := r.Header.Get("Mcp-Session-Id") + // It's just preliminary: for direct connections, the MCP server updates it with the actual user agent after the initialization handshake. + // Requests proxied by the admin service's unified MCP server don't carry that handshake, so the forwarded user agent is the only signal of the real client. + userAgent := r.UserAgent() + if userAgent == "" { + userAgent = "mcp/unknown" + } + // Create session sess, err := runner.Session(r.Context(), &ai.SessionOptions{ InstanceID: instanceID, SessionID: sessionID, CreateIfNotExists: true, Claims: auth.GetClaims(r.Context(), instanceID), - UserAgent: "mcp/unknown", // It's just preliminary: the MCP server updates it with the actual user agent after the initialization handshake. + UserAgent: userAgent, }) if err != nil { s.logger.Error("failed to create AI session for MCP", zap.String("instance_id", instanceID), zap.String("session_id", sessionID), zap.Error(err)) @@ -57,6 +64,9 @@ func (s *Server) mcpHandler() http.Handler { return srv }, &mcp.StreamableHTTPOptions{ Stateless: true, + // Respond with application/json instead of text/event-stream. + // The spec allows either, and it means the admin service's unified MCP server can forward tool calls with a plain HTTP request. + JSONResponse: true, }) }