diff --git a/Makefile b/Makefile index d355a01..3eb5ba1 100644 --- a/Makefile +++ b/Makefile @@ -12,7 +12,7 @@ test: uv run pytest test-tools: - uv run pytest stackone_ai + uv run pytest tests test-examples: uv run pytest examples @@ -27,3 +27,7 @@ docs-serve: docs-build: uv run scripts/build_docs.py uv run mkdocs build + +mcp-inspector: + uv sync --all-extras + npx @modelcontextprotocol/inspector stackmcp diff --git a/examples/mcp_server.py b/examples/mcp_server.py new file mode 100644 index 0000000..c564ecf --- /dev/null +++ b/examples/mcp_server.py @@ -0,0 +1,15 @@ +""" +This package can also be used as a Model Context Protocol (MCP) server. + +To add this server to and MCP client like Claude Code, use: + +```bash +# install the package +uv pip install stackone-ai + +# add the server to Claude Code +claude mcp add stackone uv stackmcp ["--api-key", ""] +``` + +This implementation is a work in progress and will likely change dramatically in the near future. +""" diff --git a/pyproject.toml b/pyproject.toml index e9aa2c5..ae5777f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -19,8 +19,12 @@ dependencies = [ "pydantic>=2.10.6", "requests>=2.32.3", "langchain-core>=0.1.0", + "mcp[cli]>=1.3.0", ] +[project.scripts] +stackmcp = "stackone_ai.server:cli" + [build-system] requires = ["hatchling"] build-backend = "hatchling.build" diff --git a/stackone_ai/server.py b/stackone_ai/server.py new file mode 100644 index 0000000..ef216ec --- /dev/null +++ b/stackone_ai/server.py @@ -0,0 +1,225 @@ +import argparse +import asyncio +import logging +import os +import sys +from typing import Any, TypeVar + +import mcp.types as types +from mcp.server import NotificationOptions, Server +from mcp.server.models import InitializationOptions +from mcp.server.stdio import stdio_server +from mcp.shared.exceptions import McpError +from mcp.types import EmbeddedResource, ErrorData, ImageContent, TextContent, Tool +from pydantic import ValidationError + +from stackone_ai import StackOneToolSet +from stackone_ai.models import StackOneAPIError, StackOneError + +# Set up logging +logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s") +logger = logging.getLogger("stackone.mcp") + +app: Server = Server("stackone-ai") +toolset: StackOneToolSet | None = None + +NO_ACCOUNT_ID_PREFIXES = [ + "stackone_", +] + +# Type variables for function annotations +T = TypeVar("T") +R = TypeVar("R") + + +def tool_needs_account_id(tool_name: str) -> bool: + for prefix in NO_ACCOUNT_ID_PREFIXES: + if tool_name.startswith(prefix): + return False + + # By default, assume all other tools need account_id + return True + + +@app.list_tools() # type: ignore[misc] +async def list_tools() -> list[Tool]: + """List all available StackOne tools as MCP tools.""" + if not toolset: + logger.error("Toolset not initialized") + raise McpError( + ErrorData( + code=types.INTERNAL_ERROR, + message="Toolset not initialized, please check your STACKONE_API_KEY.", + ) + ) + + try: + mcp_tools: list[Tool] = [] + tools = toolset.get_tools() + # Convert to a list if it's not already iterable in the expected way + tool_list = list(tools.tools) if hasattr(tools, "tools") else [] + + for tool in tool_list: + # Convert StackOne tool parameters to MCP schema + properties = {} + required = [] + + # Add account_id parameter only for tools that need it + if tool_needs_account_id(tool.name): + properties["account_id"] = { + "type": "string", + "description": "The StackOne account ID to use for this tool call", + } + + for name, details in tool.parameters.properties.items(): + if isinstance(details, dict): + prop = { + "type": details.get("type", "string"), + "description": details.get("description", ""), + } + if not details.get("nullable", False): + required.append(name) + properties[name] = prop + + schema = {"type": "object", "properties": properties} + if required: + schema["required"] = required + + mcp_tools.append(Tool(name=tool.name, description=tool.description, inputSchema=schema)) + + logger.info(f"Listed {len(mcp_tools)} tools") + return mcp_tools + except Exception as e: + logger.error(f"Error listing tools: {str(e)}", exc_info=True) + raise McpError( + ErrorData( + code=types.INTERNAL_ERROR, + message=f"Error listing tools: {str(e)}", + ) + ) from e + + +@app.call_tool() # type: ignore[misc] +async def call_tool( + name: str, arguments: dict[str, Any] +) -> list[TextContent | ImageContent | EmbeddedResource]: + """Execute a StackOne tool and return its result.""" + if not toolset: + logger.error("Toolset not initialized") + raise McpError( + ErrorData( + code=types.INTERNAL_ERROR, + message="Server configuration error: Toolset not initialized", + ) + ) + + try: + tool = toolset.get_tool(name) + if not tool: + logger.warning(f"Tool not found: {name}") + raise McpError( + ErrorData( + code=types.INVALID_PARAMS, + message=f"Tool not found: {name}", + ) + ) + + if "account_id" in arguments: + tool.set_account_id(arguments.pop("account_id")) + + if tool_needs_account_id(name) and tool.get_account_id() is None: + logger.warning(f"Tool {name} needs account_id but none provided") + raise McpError( + ErrorData( + code=types.INVALID_PARAMS, + message=f"Tool {name} needs account_id but none provided", + ) + ) + + result = tool.execute(arguments) + return [TextContent(type="text", text=str(result))] + + except ValidationError as e: + logger.warning(f"Invalid parameters for tool {name}: {str(e)}") + raise McpError( + ErrorData( + code=types.INVALID_PARAMS, + message=f"Invalid parameters for tool {name}: {str(e)}", + ) + ) from e + except StackOneAPIError as e: + logger.error(f"API error: {str(e)}") + raise McpError( + ErrorData( + code=types.INTERNAL_ERROR, + message=f"API error: {str(e)}", + ) + ) from e + except StackOneError as e: + logger.error(f"Error: {str(e)}") + raise McpError( + ErrorData( + code=types.INTERNAL_ERROR, + message=f"Error: {str(e)}", + ) + ) from e + except Exception as e: + logger.error(f"Unexpected error: {str(e)}", exc_info=True) + raise McpError( + ErrorData( + code=types.INTERNAL_ERROR, + message="An unexpected error occurred. Please try again later.", + ) + ) from e + + +async def main(api_key: str | None = None) -> None: + """Run the MCP server.""" + + if not api_key: + api_key = os.getenv("STACKONE_API_KEY") + if not api_key: + raise ValueError("STACKONE_API_KEY not found in environment variables") + + global toolset + toolset = StackOneToolSet(api_key=api_key) + logger.info("StackOne toolset initialized successfully") + + async with stdio_server() as (read_stream, write_stream): + await app.run( + read_stream, + write_stream, + InitializationOptions( + server_name="stackone-ai", + server_version="0.1.0", + capabilities=app.get_capabilities( + notification_options=NotificationOptions(), + experimental_capabilities={}, + ), + ), + ) + + +def cli() -> None: + """CLI entry point for the MCP server.""" + parser = argparse.ArgumentParser(description="StackOne AI MCP Server") + parser.add_argument("--api-key", help="StackOne API key (can also be set via STACKONE_API_KEY env var)") + parser.add_argument( + "--log-level", + default="INFO", + choices=["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"], + help="Set the logging level", + ) + args = parser.parse_args() + + logger.setLevel(args.log_level) + + try: + asyncio.run(main(args.api_key)) + except Exception as e: + logger.critical(f"Failed to start server: {str(e)}", exc_info=True) + sys.exit(1) + + +if __name__ == "__main__": + cli() diff --git a/stackone_ai/toolset.py b/stackone_ai/toolset.py index 6520070..4e2c372 100644 --- a/stackone_ai/toolset.py +++ b/stackone_ai/toolset.py @@ -106,6 +106,22 @@ def _matches_filter(self, tool_name: str, filter_pattern: str | list[str]) -> bo return matches_positive and not matches_negative + def get_tool(self, name: str, *, account_id: str | None = None) -> StackOneTool | None: + """Get a specific tool by name + + Args: + name: Name of the tool to retrieve + account_id: Optional account ID override. If not provided, uses the one from initialization + + Returns: + The tool if found, None otherwise + + Raises: + ToolsetLoadError: If there is an error loading the tools + """ + tools = self.get_tools(name, account_id=account_id) + return tools.get_tool(name) + def get_tools( self, filter_pattern: str | list[str] | None = None, *, account_id: str | None = None ) -> Tools: diff --git a/uv.lock b/uv.lock index df43c2d..96f999e 100644 --- a/uv.lock +++ b/uv.lock @@ -874,6 +874,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/56/95/9377bcb415797e44274b51d46e3249eba641711cf3348050f76ee7b15ffc/httpx-0.27.2-py3-none-any.whl", hash = "sha256:7bb2708e112d8fdd7829cd4243970f0c223274051cb35ee80c03301ee29a3df0", size = 76395 }, ] +[[package]] +name = "httpx-sse" +version = "0.4.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/4c/60/8f4281fa9bbf3c8034fd54c0e7412e66edbab6bc74c4996bd616f8d0406e/httpx-sse-0.4.0.tar.gz", hash = "sha256:1e81a3a3070ce322add1d3529ed42eb5f70817f45ed6ec915ab753f961139721", size = 12624 } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e1/9b/a181f281f65d776426002f330c31849b86b31fc9d848db62e16f03ff739f/httpx_sse-0.4.0-py3-none-any.whl", hash = "sha256:f329af6eae57eaa2bdfd962b42524764af68075ea87370a2de920af5341e318f", size = 7819 }, +] + [[package]] name = "huggingface-hub" version = "0.29.1" @@ -1324,6 +1333,31 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/8f/8e/9ad090d3553c280a8060fbf6e24dc1c0c29704ee7d1c372f0c174aa59285/matplotlib_inline-0.1.7-py3-none-any.whl", hash = "sha256:df192d39a4ff8f21b1895d72e6a13f5fcc5099f00fa84384e0ea28c2cc0653ca", size = 9899 }, ] +[[package]] +name = "mcp" +version = "1.3.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "anyio" }, + { name = "httpx" }, + { name = "httpx-sse" }, + { name = "pydantic" }, + { name = "pydantic-settings" }, + { name = "sse-starlette" }, + { name = "starlette" }, + { name = "uvicorn" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/6b/b6/81e5f2490290351fc97bf46c24ff935128cb7d34d68e3987b522f26f7ada/mcp-1.3.0.tar.gz", hash = "sha256:f409ae4482ce9d53e7ac03f3f7808bcab735bdfc0fba937453782efb43882d45", size = 150235 } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d0/d2/a9e87b506b2094f5aa9becc1af5178842701b27217fa43877353da2577e3/mcp-1.3.0-py3-none-any.whl", hash = "sha256:2829d67ce339a249f803f22eba5e90385eafcac45c94b00cab6cef7e8f217211", size = 70672 }, +] + +[package.optional-dependencies] +cli = [ + { name = "python-dotenv" }, + { name = "typer" }, +] + [[package]] name = "mdurl" version = "0.1.2" @@ -2292,6 +2326,19 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/51/b2/b2b50d5ecf21acf870190ae5d093602d95f66c9c31f9d5de6062eb329ad1/pydantic_core-2.27.2-cp313-cp313-win_arm64.whl", hash = "sha256:ac4dbfd1691affb8f48c2c13241a2e3b60ff23247cbcf981759c768b6633cf8b", size = 1885186 }, ] +[[package]] +name = "pydantic-settings" +version = "2.8.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "pydantic" }, + { name = "python-dotenv" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/ca/a2/ad2511ede77bb424f3939e5148a56d968cdc6b1462620d24b2a1f4ab65b4/pydantic_settings-2.8.0.tar.gz", hash = "sha256:88e2ca28f6e68ea102c99c3c401d6c9078e68a5df600e97b43891c34e089500a", size = 83347 } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c1/a9/3b9642025174bbe67e900785fb99c9bfe91ea584b0b7126ff99945c24a0e/pydantic_settings-2.8.0-py3-none-any.whl", hash = "sha256:c782c7dc3fb40e97b238e713c25d26f64314aece2e91abcff592fcac15f71820", size = 30746 }, +] + [[package]] name = "pygments" version = "2.19.1" @@ -2745,6 +2792,19 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/e9/44/75a9c9421471a6c4805dbf2356f7c181a29c1879239abab1ea2cc8f38b40/sniffio-1.3.1-py3-none-any.whl", hash = "sha256:2f6da418d1f1e0fddd844478f41680e794e6051915791a034ff65e5f100525a2", size = 10235 }, ] +[[package]] +name = "sse-starlette" +version = "2.2.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "anyio" }, + { name = "starlette" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/71/a4/80d2a11af59fe75b48230846989e93979c892d3a20016b42bb44edb9e398/sse_starlette-2.2.1.tar.gz", hash = "sha256:54470d5f19274aeed6b2d473430b08b4b379ea851d953b11d7f1c4a2c118b419", size = 17376 } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d9/e0/5b8bd393f27f4a62461c5cf2479c75a2cc2ffa330976f9f00f5f6e4f50eb/sse_starlette-2.2.1-py3-none-any.whl", hash = "sha256:6410a3d3ba0c89e7675d4c273a301d64649c03a5ef1ca101f10b47f895fd0e99", size = 10120 }, +] + [[package]] name = "stack-data" version = "0.6.3" @@ -2765,6 +2825,7 @@ version = "0.0.2" source = { editable = "." } dependencies = [ { name = "langchain-core" }, + { name = "mcp", extra = ["cli"] }, { name = "pydantic" }, { name = "requests" }, ] @@ -2800,6 +2861,7 @@ requires-dist = [ { name = "crewai", marker = "extra == 'examples'", specifier = ">=0.102.0" }, { name = "langchain-core", specifier = ">=0.1.0" }, { name = "langchain-openai", marker = "extra == 'examples'", specifier = ">=0.3.6" }, + { name = "mcp", extras = ["cli"], specifier = ">=1.3.0" }, { name = "mkdocs-terminal", marker = "extra == 'docs'", specifier = ">=4.7.0" }, { name = "openai", marker = "extra == 'examples'", specifier = ">=1.63.2" }, { name = "pydantic", specifier = ">=2.10.6" },