Merge pull request #92 from useshortcut/amcd/per-tool-filtering

Allow per tool filtering.

Author: andreasmcdermott

README.md

@@ -129,67 +129,72 @@ Or you can edit the local JSON file directly:
 
 ### Stories
 
-- **get-story** - Get a single Shortcut story by ID
-- **search-stories** - Find Shortcut stories with filtering and search options
-- **get-story-branch-name** - Get the recommended branch name (based on workspace settings) for a specific story.
-- **create-story** - Create a new Shortcut story
-- **update-story** - Update an existing Shortcut story
-- **upload-file-to-story** - Upload a file and link it to a story
-- **assign-current-user-as-owner** - Assign the current user as the owner of a story
-- **unassign-current-user-as-owner** - Unassign the current user as the owner of a story
-- **create-story-comment** - Create a comment on a story
-- **add-task-to-story** - Add a task to a story
-- **update-task** - Update a task in a story
-- **add-relation-to-story** - Add a story relationship (relates to, blocks, duplicates, etc.)
-- **add-external-link-to-story** - Add an external link to a Shortcut story
-- **remove-external-link-from-story** - Remove an external link from a Shortcut story
-- **set-story-external-links** - Replace all external links on a story with a new set of links
-- **get-stories-by-external-link** - Find all stories that contain a specific external link
+- **stories-get-by-id** - Get a single Shortcut story by ID
+- **stories-search** - Find Shortcut stories with filtering and search options
+- **stories-get-branch-name** - Get the recommended branch name (based on workspace settings) for a specific story.
+- **stories-create** - Create a new Shortcut story
+- **stories-update** - Update an existing Shortcut story
+- **stories-upload-file** - Upload a file and link it to a story
+- **stories-assign-current-user** - Assign the current user as the owner of a story
+- **stories-unassign-current-user** - Unassign the current user as the owner of a story
+- **stories-create-comment** - Create a comment on a story
+- **stories-add-task** - Add a task to a story
+- **stories-update-task** - Update a task in a story
+- **stories-add-relation** - Add a story relationship (relates to, blocks, duplicates, etc.)
+- **stories-add-external-link** - Add an external link to a Shortcut story
+- **stories-remove-external-link** - Remove an external link from a Shortcut story
+- **stories-set-external-links** - Replace all external links on a story with a new set of links
+- **stories-get-by-external-link** - Find all stories that contain a specific external link
 
 ### Epics
 
-- **get-epic** - Get a Shortcut epic by ID
-- **search-epics** - Find Shortcut epics with filtering and search options
-- **create-epic** - Create a new Shortcut epic
+- **epics-get-by-id** - Get a Shortcut epic by ID
+- **epics-search** - Find Shortcut epics with filtering and search options
+- **epics-create** - Create a new Shortcut epic
 
 ### Iterations
 
-- **get-iteration-stories** - Get stories in a specific iteration by iteration ID
-- **get-iteration** - Get a Shortcut iteration by ID
-- **search-iterations** - Find Shortcut iterations with filtering and search options
-- **create-iteration** - Create a new Shortcut iteration with start/end dates
-- **get-active-iterations** - Get active iterations for the current user based on team memberships
-- **get-upcoming-iterations** - Get upcoming iterations for the current user based on team memberships
+- **iterations-get-stories** - Get stories in a specific iteration by iteration ID
+- **iterations-get-by-id** - Get a Shortcut iteration by ID
+- **iterations-search** - Find Shortcut iterations with filtering and search options
+- **iterations-create** - Create a new Shortcut iteration with start/end dates
+- **iterations-get-active** - Get active iterations for the current user based on team memberships
+- **iterations-get-upcoming** - Get upcoming iterations for the current user based on team memberships
 
 ### Objectives
 
-- **get-objective** - Get a Shortcut objective by ID
-- **search-objectives** - Find Shortcut objectives with filtering and search options
+- **objectives-get-by-id** - Get a Shortcut objective by ID
+- **objectives-search** - Find Shortcut objectives with filtering and search options
 
 ### Teams
 
-- **get-team** - Get a Shortcut team by ID
-- **list-teams** - List all Shortcut teams
+- **teams-get-by-id** - Get a Shortcut team by ID
+- **teams-list** - List all Shortcut teams
 
 ### Workflows
 
-- **get-default-workflow** - Get the default workflow for a specific team or the workspace default
-- **get-workflow** - Get a Shortcut workflow by ID
-- **list-workflows** - List all Shortcut workflows
+- **workflows-get-default** - Get the default workflow for a specific team or the workspace default
+- **workflows-get-by-id** - Get a Shortcut workflow by ID
+- **workflows-list** - List all Shortcut workflows
 
 ### Users
 
-- **get-current-user** - Get the current user information
-- **get-current-user-teams** - Get a list of teams where the current user is a member
-- **list-users** - Get all workspace users
+- **users-get-current** - Get the current user information
+- **users-get-current-teams** - Get a list of teams where the current user is a member
+- **users-list** - Get all workspace users
 
 ### Documents
 
-- **create-document** - Create a new document in Shortcut with HTML content
+- **documents-create** - Create a new document in Shortcut with HTML content
 
 ## Limit tools
 
-You can limit the tools available to the LLM by setting the `SHORTCUT_TOOLS` environment variable to a comma-separated list of entity types.
+You can limit the tools available to the LLM by setting the `SHORTCUT_TOOLS` environment variable to a comma-separated list.
+
+- Tools can be limited by entity type by just adding the entity, eg `stories` or `epics`.
+- Individual tools can also be limitied by their full name, eg `stories-get-by-id` or `epics-search`.
+
+By default, all tools are enabled.
 
 Example:
 
@@ -204,14 +209,14 @@ Example:
       ],
       "env": {
         "SHORTCUT_API_TOKEN": "<YOUR_SHORTCUT_API_TOKEN>",
-        "SHORTCUT_TOOLS": "stories,epics"
+        "SHORTCUT_TOOLS": "stories,epics,iterations-create"
       }
     }
   }
 }
 ```
 
-The following values are accepted:
+The following values are accepted in addition to the full tool names listed above under [Available Tools](#available-tools):
 
 - `users`
 - `stories`

src/mcp/CustomMcpServer.test.ts

@@ -0,0 +1,50 @@
+import { describe, expect, test } from "bun:test";
+import { CustomMcpServer } from "./CustomMcpServer";
+
+describe("CustomMcpServer", () => {
+	test("should allow read only tools when readonly is true", () => {
+		const server = new CustomMcpServer({ readonly: true, tools: null });
+		expect(
+			server.addToolWithReadAccess("test", "test", async () => {
+				return { content: [] };
+			}),
+		).not.toBeNull();
+	});
+
+	test("should not allow write tools when readonly is true", () => {
+		const server = new CustomMcpServer({ readonly: true, tools: null });
+		expect(
+			server.addToolWithWriteAccess("test", "test", async () => {
+				return { content: [] };
+			}),
+		).toBeNull();
+	});
+
+	test("should allow write tools when readonly is false", () => {
+		const server = new CustomMcpServer({ readonly: false, tools: null });
+		expect(
+			server.addToolWithWriteAccess("test", "test", async () => {
+				return { content: [] };
+			}),
+		).not.toBeNull();
+	});
+
+	test("should only allow tools in the tools list", () => {
+		const server = new CustomMcpServer({ readonly: false, tools: ["test"] });
+		expect(
+			server.addToolWithReadAccess("test", "test", async () => {
+				return { content: [] };
+			}),
+		).not.toBeNull();
+		expect(
+			server.addToolWithReadAccess("test-sub-tool", "test", async () => {
+				return { content: [] };
+			}),
+		).not.toBeNull();
+		expect(
+			server.addToolWithReadAccess("test2", "test", async () => {
+				return { content: [] };
+			}),
+		).toBeNull();
+	});
+});

src/mcp/CustomMcpServer.ts

@@ -0,0 +1,100 @@
+import {
+	McpServer,
+	type RegisteredTool,
+	type ToolCallback,
+} from "@modelcontextprotocol/sdk/server/mcp.js";
+import type { ToolAnnotations } from "@modelcontextprotocol/sdk/types.js";
+import type { ZodRawShape } from "zod";
+import { name, version } from "../../package.json";
+
+export class CustomMcpServer extends McpServer {
+	private readonly: boolean;
+	private tools: Set<string>;
+
+	constructor({ readonly, tools }: { readonly: boolean; tools: string[] | null | undefined }) {
+		super({ name, version });
+		this.readonly = readonly;
+		this.tools = new Set(tools || []);
+	}
+
+	shouldAddTool(name: string) {
+		console.log("Checking tool:", name, this.tools.size);
+		if (!this.tools.size) return true;
+		const [entityType] = name.split("-");
+		if (this.tools.has(entityType) || this.tools.has(name)) return true;
+		return false;
+	}
+
+	// Overloads for addToolWithWriteAccess to match all variants of the base tool() method
+	addToolWithWriteAccess(name: string, cb: ToolCallback): RegisteredTool | null;
+	addToolWithWriteAccess(
+		name: string,
+		description: string,
+		cb: ToolCallback,
+	): RegisteredTool | null;
+	addToolWithWriteAccess<Args extends ZodRawShape>(
+		name: string,
+		paramsSchemaOrAnnotations: Args | ToolAnnotations,
+		cb: ToolCallback<Args>,
+	): RegisteredTool | null;
+	addToolWithWriteAccess<Args extends ZodRawShape>(
+		name: string,
+		description: string,
+		paramsSchemaOrAnnotations: Args | ToolAnnotations,
+		cb: ToolCallback<Args>,
+	): RegisteredTool | null;
+	addToolWithWriteAccess<Args extends ZodRawShape>(
+		name: string,
+		paramsSchema: Args,
+		annotations: ToolAnnotations,
+		cb: ToolCallback<Args>,
+	): RegisteredTool | null;
+	addToolWithWriteAccess<Args extends ZodRawShape>(
+		name: string,
+		description: string,
+		paramsSchema: Args,
+		annotations: ToolAnnotations,
+		cb: ToolCallback<Args>,
+	): RegisteredTool | null;
+	addToolWithWriteAccess(...args: any[]): RegisteredTool | null {
+		if (this.readonly) return null;
+		if (!this.shouldAddTool(args[0])) return null;
+		return (super.tool as any)(...args);
+	}
+
+	// Overloads for addToolWithReadAccess to match all variants of the base tool() method
+	addToolWithReadAccess(name: string, cb: ToolCallback): RegisteredTool | null;
+	addToolWithReadAccess(name: string, description: string, cb: ToolCallback): RegisteredTool | null;
+	addToolWithReadAccess<Args extends ZodRawShape>(
+		name: string,
+		paramsSchemaOrAnnotations: Args | ToolAnnotations,
+		cb: ToolCallback<Args>,
+	): RegisteredTool | null;
+	addToolWithReadAccess<Args extends ZodRawShape>(
+		name: string,
+		description: string,
+		paramsSchemaOrAnnotations: Args | ToolAnnotations,
+		cb: ToolCallback<Args>,
+	): RegisteredTool | null;
+	addToolWithReadAccess<Args extends ZodRawShape>(
+		name: string,
+		paramsSchema: Args,
+		annotations: ToolAnnotations,
+		cb: ToolCallback<Args>,
+	): RegisteredTool | null;
+	addToolWithReadAccess<Args extends ZodRawShape>(
+		name: string,
+		description: string,
+		paramsSchema: Args,
+		annotations: ToolAnnotations,
+		cb: ToolCallback<Args>,
+	): RegisteredTool | null;
+	addToolWithReadAccess(...args: any[]): RegisteredTool | null {
+		if (!this.shouldAddTool(args[0])) return null;
+		return (super.tool as any)(...args);
+	}
+
+	tool(): never {
+		throw new Error("Call addToolWithReadAccess or addToolWithWriteAccess instead.");
+	}
+}

src/server.ts

@@ -1,9 +1,7 @@
-import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
 import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js";
 import { ShortcutClient } from "@shortcut/client";
 import { ShortcutClientWrapper } from "@/client/shortcut";
-import { name, version } from "../package.json";
-
+import { CustomMcpServer } from "./mcp/CustomMcpServer";
 import { DocumentTools } from "./tools/documents";
 import { EpicTools } from "./tools/epics";
 import { IterationTools } from "./tools/iterations";
@@ -15,9 +13,10 @@ import { WorkflowTools } from "./tools/workflows";
 
 let apiToken = process.env.SHORTCUT_API_TOKEN;
 let isReadonly = process.env.SHORTCUT_READONLY === "true";
-let enabledTools = process.env.SHORTCUT_TOOLS?.length
-	? process.env.SHORTCUT_TOOLS.split(",").map((tool) => tool.trim())
-	: null;
+let enabledTools = (process.env.SHORTCUT_TOOLS || "")
+	.split(",")
+	.map((tool) => tool.trim())
+	.filter(Boolean);
 
 // If a setting is provided as an argument, use it instead of the environment variable.
 if (process.argv.length >= 3) {
@@ -27,7 +26,11 @@ if (process.argv.length >= 3) {
 		.forEach(([name, value]) => {
 			if (name === "SHORTCUT_API_TOKEN") apiToken = value;
 			if (name === "SHORTCUT_READONLY") isReadonly = value === "true";
-			if (name === "SHORTCUT_TOOLS") enabledTools = value.split(",").map((tool) => tool.trim());
+			if (name === "SHORTCUT_TOOLS")
+				enabledTools = value
+					.split(",")
+					.map((tool) => tool.trim())
+					.filter(Boolean);
 		});
 }
 
@@ -36,23 +39,18 @@ if (!apiToken) {
 	process.exit(1);
 }
 
-const server = new McpServer({ name, version });
+const server = new CustomMcpServer({ readonly: isReadonly, tools: enabledTools });
 const client = new ShortcutClientWrapper(new ShortcutClient(apiToken));
 
-// Helper function to check if a tool should be enabled. All tools are enabled by default unless specified otherwise.
-const areToolsEnabled = (toolName: string) => {
-	return !enabledTools || enabledTools.includes(toolName);
-};
-
 // The order these are created impacts the order they are listed to the LLM. Most important tools should be at the top.
-if (areToolsEnabled("users")) UserTools.create(client, server, isReadonly);
-if (areToolsEnabled("stories")) StoryTools.create(client, server, isReadonly);
-if (areToolsEnabled("iterations")) IterationTools.create(client, server, isReadonly);
-if (areToolsEnabled("epics")) EpicTools.create(client, server, isReadonly);
-if (areToolsEnabled("objectives")) ObjectiveTools.create(client, server, isReadonly);
-if (areToolsEnabled("teams")) TeamTools.create(client, server, isReadonly);
-if (areToolsEnabled("workflows")) WorkflowTools.create(client, server, isReadonly);
-if (areToolsEnabled("documents")) DocumentTools.create(client, server, isReadonly);
+UserTools.create(client, server);
+StoryTools.create(client, server);
+IterationTools.create(client, server);
+EpicTools.create(client, server);
+ObjectiveTools.create(client, server);
+TeamTools.create(client, server);
+WorkflowTools.create(client, server);
+DocumentTools.create(client, server);
 
 async function startServer() {
 	try {

src/tools/base.ts

@@ -125,10 +125,7 @@ export type SimplifiedKind = "simple" | "list" | "full";
  * Base class for all tools.
  */
 export class BaseTools {
-	constructor(
-		protected client: ShortcutClientWrapper,
-		protected isReadonly = false,
-	) {}
+	constructor(protected client: ShortcutClientWrapper) {}
 
 	private renameEntityProps<T extends Record<string, unknown>>(entity: T) {
 		if (!entity || typeof entity !== "object") return entity;