Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 5 additions & 1 deletion Makefile
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
15 changes: 15 additions & 0 deletions examples/mcp_server.py
Original file line number Diff line number Diff line change
@@ -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", "<your-api-key>"]
```

This implementation is a work in progress and will likely change dramatically in the near future.
"""
4 changes: 4 additions & 0 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down
225 changes: 225 additions & 0 deletions stackone_ai/server.py
Original file line number Diff line number Diff line change
@@ -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()
16 changes: 16 additions & 0 deletions stackone_ai/toolset.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
62 changes: 62 additions & 0 deletions uv.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.