Annotation-based permission checking for tools

Author: labkey-adam

api/src/org/labkey/api/mcp/McpService.java

@@ -6,7 +6,6 @@
 import org.jetbrains.annotations.NotNull;
 import org.jspecify.annotations.NonNull;
 import org.labkey.api.data.Container;
-import org.labkey.api.module.McpProvider;
 import org.labkey.api.security.User;
 import org.labkey.api.services.ServiceRegistry;
 import org.labkey.api.util.HtmlString;
@@ -52,9 +51,9 @@ default User getUser(ToolContext toolContext)
         }
 
         // Every MCP resource should call this on every invocation
-        default void incrementResourceReadCount(String resource)
+        default void incrementResourceRequestCount(String resource)
         {
-            get().incrementResourceCount(resource);
+            get().incrementResourceRequestCount(resource);
         }
     }
 
@@ -70,27 +69,18 @@ static void setInstance(McpService service)
 
     boolean isReady();
 
-
-    default void register(McpImpl obj)
+    default void register(McpImpl mcp)
     {
-        ToolCallback[] tools = ToolCallbacks.from(obj);
-        if (null != tools && tools.length > 0)
-            registerTools(Arrays.asList(tools));
+        ToolCallback[] tools = ToolCallbacks.from(mcp);
+        if (tools.length > 0)
+            registerTools(Arrays.asList(tools), mcp);
 
-        var resources = new SyncMcpResourceProvider(List.of(obj)).getResourceSpecifications();
+        var resources = new SyncMcpResourceProvider(List.of(mcp)).getResourceSpecifications();
         if (null != resources && !resources.isEmpty())
             registerResources(resources);
     }
 
-
-    default void register(McpProvider mcp)
-    {
-        registerTools(mcp.getMcpTools());
-        registerPrompts(mcp.getMcpPrompts());
-        registerResources(mcp.getMcpResources());
-    }
-
-    void registerTools(@NotNull List<ToolCallback> tools);
+    void registerTools(@NotNull List<ToolCallback> tools, McpImpl mcp);
 
     void registerPrompts(@NotNull List<McpServerFeatures.SyncPromptSpecification> prompts);
 
@@ -106,7 +96,7 @@ default ChatClient getChat(HttpSession session, String agentName, Supplier<Strin
 
     void saveSessionContainer(ToolContext context, Container container);
 
-    void incrementResourceCount(String resource);
+    void incrementResourceRequestCount(String resource);
 
     ChatClient getChat(HttpSession session, String agentName, Supplier<String> systemPromptSupplier, boolean createIfNotExists);
 

api/src/org/labkey/api/security/RequiresNoPermission.java

@@ -21,12 +21,11 @@
 import java.lang.annotation.Target;
 
 /**
- * Indicates that an action does not require any kind of authentication or permission to invoke. Use with extreme
- * caution. Typically, actions marked with this annotation will handle their own permission checks in their own code path.
- * User: adam
- * Date: Dec 22, 2009
-*/
-public @Retention(java.lang.annotation.RetentionPolicy.RUNTIME) @Target(ElementType.TYPE)
+ * Indicates that an action class or an MCP tool method does not require any kind of authentication or permission to
+ * invoke. Use with extreme caution. Typically, actions marked with this annotation will handle their own permission
+ * checks in their own code path.
+ */
+public @Retention(java.lang.annotation.RetentionPolicy.RUNTIME) @Target({ElementType.TYPE, ElementType.METHOD})
 @interface RequiresNoPermission
 {
 }

api/src/org/labkey/api/security/RequiresPermission.java

@@ -22,10 +22,10 @@
 import java.lang.annotation.Target;
 
 /**
- * Specifies the required permission for an action. It does not imply that the user needs to be logged in or otherwise
- * authenticated.
+ * Specifies the required permission for an action class or an MCP tool method. It does not imply that the user needs
+ * to be logged in or otherwise authenticated.
  */
-@Retention(java.lang.annotation.RetentionPolicy.RUNTIME) @Target(ElementType.TYPE)
+@Retention(java.lang.annotation.RetentionPolicy.RUNTIME) @Target({ElementType.TYPE, ElementType.METHOD})
 public @interface RequiresPermission
 {
     Class<? extends Permission> value();

core/src/org/labkey/core/CoreMcp.java

@@ -6,13 +6,16 @@
 import org.labkey.api.data.Container;
 import org.labkey.api.data.ContainerManager;
 import org.labkey.api.mcp.McpService;
+import org.labkey.api.security.RequiresNoPermission;
+import org.labkey.api.security.RequiresPermission;
 import org.labkey.api.security.User;
 import org.labkey.api.security.permissions.ReadPermission;
 import org.labkey.api.settings.AppProps;
 import org.labkey.api.settings.LookAndFeelProperties;
 import org.labkey.api.study.Study;
 import org.labkey.api.study.StudyService;
 import org.labkey.api.util.HtmlString;
+import org.labkey.api.view.UnauthorizedException;
 import org.springframework.ai.chat.model.ToolContext;
 import org.springframework.ai.tool.annotation.Tool;
 import org.springframework.ai.tool.annotation.ToolParam;
@@ -26,6 +29,7 @@ public class CoreMcp implements McpService.McpImpl
 {
     @Tool(description = "This tool provides useful context information about the current user (name, userid), webserver " +
         "(name, url, description), and current container/folder (name, path, url, description) once the container is set via setContainer.")
+    @RequiresPermission(ReadPermission.class)
     String whereAmIWhoAmITalkingTo(ToolContext context)
     {
         var cu = getContext(context);
@@ -70,6 +74,7 @@ String whereAmIWhoAmITalkingTo(ToolContext context)
     }
 
     @Tool(description = "List the hierarchical path for every container in the server where the user has read permissions.")
+    @RequiresNoPermission
     String listContainers(ToolContext toolContext)
     {
         return ContainerManager.getAllChildren(ContainerManager.getRoot(), getUser(toolContext), ReadPermission.class)
@@ -80,8 +85,10 @@ String listContainers(ToolContext toolContext)
     }
 
     @Tool(description = "Every tool in this MCP requires a container path, e.g. MyProject/MyFolder. A container is also called a folder or project. " +
-        "Please prompt the user for a container path and use this tool to save the path for this session. Don't suggest a leading slash on the path " +
-        "because typing a slash in some LLM clients triggers custom shortcuts.")
+        "Please prompt the user for a container path and use this tool to save the path for this MCP session. The user can also change the container " +
+        "during the session using this tool. The user must have read permissions in the container, in other words, the path must be on the list that " +
+        "the listContainers tool returns. Don't suggest a leading slash on the path because typing a slash in some LLM clients triggers custom shortcuts.")
+    @RequiresNoPermission // Because we don't have a container yet, but tool will check for read permission before setting the container
     String setContainer(ToolContext context, @ToolParam(description = "Container path, e.g. MyProject/MyFolder") String containerPath)
     {
         final String message;
@@ -100,6 +107,10 @@ String setContainer(ToolContext context, @ToolParam(description = "Container pat
             }
             else
             {
+                // Must have read permission to set a container
+                if (!container.hasPermission(getUser(context), ReadPermission.class))
+                    throw new UnauthorizedException();
+
                 McpService.get().saveSessionContainer(context, container);
                 message = "Container has been set to " + container.getPath();
             }

query/src/org/labkey/query/controllers/QueryMcp.java

@@ -19,6 +19,8 @@
 import org.labkey.api.query.SchemaKey;
 import org.labkey.api.query.SimpleSchemaTreeVisitor;
 import org.labkey.api.query.UserSchema;
+import org.labkey.api.security.RequiresPermission;
+import org.labkey.api.security.permissions.ReadPermission;
 import org.labkey.api.view.NotFoundException;
 import org.labkey.api.writer.ContainerUser;
 import org.labkey.query.sql.SqlParser;
@@ -44,7 +46,7 @@ public class QueryMcp implements McpService.McpImpl
             description = "Provide documentation for LabKey SQL specific syntax")
     public McpSchema.ReadResourceResult getLabKeySQLDocumentation() throws IOException
     {
-        incrementResourceReadCount("LabKey SQL");
+        incrementResourceRequestCount("LabKey SQL");
         String markdown = IOUtils.resourceToString("org/labkey/query/controllers/LabKeySql.md", null, QueryController.class.getClassLoader());
         return new McpSchema.ReadResourceResult(List.of(
             new McpSchema.TextResourceContents(
@@ -56,20 +58,23 @@ public McpSchema.ReadResourceResult getLabKeySQLDocumentation() throws IOExcepti
     }
 
     @Tool(description = "Provide column metadata for a sql table. This tool will also return SQL source for saved queries.")
+    @RequiresPermission(ReadPermission.class)
     String listColumns(ToolContext toolContext, @ToolParam(description = "Fully qualified table name as it would appear in SQL e.g. \"schema\".\"table\"") String fullQuotedTableName)
     {
         var json = _listColumns(fullQuotedTableName, toolContext);
         return json.toString();
     }
 
     @Tool(description = "Provide list of tables within the provided schema.")
+    @RequiresPermission(ReadPermission.class)
     String listTables(ToolContext toolContext, @ToolParam(description = "Fully qualified schema name as it would appear in SQL e.g. \"schema\"") String quotedSchemaName)
     {
         var json = _listTables(quotedSchemaName, getContext(toolContext));
         return json.toString();
     }
 
     @Tool(description = "Provide list of database schemas")
+    @RequiresPermission(ReadPermission.class)
     String listSchemas(ToolContext toolContext)
     {
         ContainerUser cu = getContext(toolContext);
@@ -88,6 +93,7 @@ String listSchemas(ToolContext toolContext)
 
 
     @Tool(description = "Provide the SQL source for a saved query.")
+    @RequiresPermission(ReadPermission.class)
     String getSourceForSavedQuery(ToolContext toolContext, @ToolParam(description = "Fully qualified query name as it would appear in SQL e.g. \"schema\".\"table or query\"") String fullQuotedTableName)
     {
         var json = _listTables(fullQuotedTableName, getContext(toolContext));
@@ -309,7 +315,6 @@ static String normalizeIdentifier(String compoundIdentifier)
         return new SqlParser().parseIdentifier(compoundIdentifier).toSQLString(true).toLowerCase();
     }
 
-
     /** JSON schema example provided by GEMINI, using triple tick-marks to delimit the machine-readable structured data
      *
      * Here is the database schema in JSON format:
@@ -338,4 +343,3 @@ static String normalizeIdentifier(String compoundIdentifier)
      * }```
      */
 }
-