From 3d74151afa09d5e5c1df0db67afc9517d85ce4b4 Mon Sep 17 00:00:00 2001 From: abhijitt-code Date: Mon, 10 Aug 2026 16:51:28 +0530 Subject: [PATCH] Remove FastMCP and own the low-level MCP server (Phase 0). Build tools/resources/prompts on mcp.server.lowlevel.Server like the TS SDK, pin mcp to <2, and keep Studio tool forms working with $ref input schemas. Co-authored-by: Cursor --- .gitignore | 62 ++ nitrostack.log | 65 -- nitrostack/core/app.py | 1036 ++++++++++++++++++-------------- nitrostack/core/mcp_server.py | 66 ++ nitrostack/static/app.js | 29 +- nitrostack/testing/__init__.py | 116 ++-- pyproject.toml | 2 +- requirements.txt | 2 +- tests/test_initial_tool.py | 4 +- tests/test_tasks.py | 13 +- tests/test_widget_metadata.py | 10 +- 11 files changed, 800 insertions(+), 605 deletions(-) create mode 100644 .gitignore delete mode 100644 nitrostack.log create mode 100644 nitrostack/core/mcp_server.py diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..d6c4633 --- /dev/null +++ b/.gitignore @@ -0,0 +1,62 @@ +# Virtual environments +.venv/ +venv/ +env/ +.env/ +.ENV/ + +# Secrets / local env +.env +.env.* +!.env.example + +# Python +__pycache__/ +*.py[cod] +*$py.class +*.so +.Python +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +*.egg-info/ +.installed.cfg +*.egg +MANIFEST +.mypy_cache/ +.ruff_cache/ +.pytest_cache/ +.coverage +htmlcov/ +.tox/ +.nox/ +.cache/ + +# Logs +*.log +nitrostack.log + +# IDE / OS +.idea/ +.vscode/ +*.swp +*.swo +.DS_Store +Thumbs.db + +# Local scaffold / scratch projects (not part of the SDK) +adv-proj/ +new-mcp-server/ + +# Node widgets (if installed locally in examples/templates) +node_modules/ +.next/ diff --git a/nitrostack.log b/nitrostack.log deleted file mode 100644 index fbc40c6..0000000 --- a/nitrostack.log +++ /dev/null @@ -1,65 +0,0 @@ -2026-06-10 13:22:50,079 [INFO] (nitrostack): Executing calculation: 12.5 add 3.5 -2026-06-10 13:23:13,039 [INFO] (nitrostack): Executing calculation: 12.5 add 3.5 -2026-06-10 13:23:13,044 [INFO] (nitrostack): Executing calculation: 10.0 divide 0.0 -2026-06-10 13:23:13,046 [INFO] (nitrostack): Converting temperature: 100.0 from celsius to fahrenheit -2026-06-10 13:23:39,823 [INFO] (nitrostack): Executing calculation: 12.5 add 3.5 -2026-06-10 13:23:39,827 [INFO] (nitrostack): Executing calculation: 10.0 divide 0.0 -2026-06-10 13:23:39,828 [INFO] (nitrostack): Converting temperature: 100.0 from celsius to fahrenheit -2026-06-10 13:24:09,252 [INFO] (nitrostack): Executing calculation: 12.5 add 3.5 -2026-06-10 13:24:09,260 [INFO] (nitrostack): Executing calculation: 10.0 divide 0.0 -2026-06-10 13:24:09,261 [INFO] (nitrostack): Converting temperature: 100.0 from celsius to fahrenheit -2026-06-10 13:24:35,992 [INFO] (nitrostack): Executing calculation: 12.5 add 3.5 -2026-06-10 13:24:35,997 [INFO] (nitrostack): Executing calculation: 10.0 divide 0.0 -2026-06-10 13:24:35,999 [INFO] (nitrostack): Converting temperature: 100.0 from celsius to fahrenheit -2026-06-10 13:25:00,299 [INFO] (nitrostack): Executing calculation: 12.5 add 3.5 -2026-06-10 13:25:00,303 [INFO] (nitrostack): Executing calculation: 10.0 divide 0.0 -2026-06-10 13:25:00,304 [INFO] (nitrostack): Converting temperature: 100.0 from celsius to fahrenheit -2026-06-10 13:25:32,026 [INFO] (nitrostack): Executing calculation: 12.5 add 3.5 -2026-06-10 13:25:32,032 [INFO] (nitrostack): Executing calculation: 10.0 divide 0.0 -2026-06-10 13:25:32,033 [INFO] (nitrostack): Converting temperature: 100.0 from celsius to fahrenheit -2026-06-10 13:35:53,144 [INFO] (nitrostack): Executing calculation: 12.5 add 3.5 -2026-06-10 13:35:53,151 [INFO] (nitrostack): Executing calculation: 10.0 divide 0.0 -2026-06-10 13:35:53,153 [INFO] (nitrostack): Converting temperature: 100.0 from celsius to fahrenheit -2026-06-10 13:36:34,272 [INFO] (nitrostack): Executing calculation: 12.5 add 3.5 -2026-06-10 13:36:34,283 [INFO] (nitrostack): Executing calculation: 10.0 divide 0.0 -2026-06-10 13:36:34,285 [INFO] (nitrostack): Converting temperature: 100.0 from celsius to fahrenheit -2026-06-10 14:33:53,465 [INFO] (nitrostack): Executing calculation: 10.0 divide 15.0 -2026-06-10 14:34:00,907 [INFO] (nitrostack): Executing calculation: 10.0 add 15.0 -2026-06-10 14:34:03,731 [INFO] (nitrostack): Executing calculation: 10.0 add 15.0 -2026-06-10 14:37:22,534 [INFO] (nitrostack): Executing calculation: 12.5 add 3.5 -2026-06-10 14:37:22,538 [INFO] (nitrostack): Executing calculation: 10.0 divide 0.0 -2026-06-10 14:37:22,540 [INFO] (nitrostack): Converting temperature: 100.0 from celsius to fahrenheit -2026-06-10 14:41:33,833 [INFO] (nitrostack): Executing calculation: 10.0 add 20.0 -2026-06-10 14:48:06,730 [INFO] (nitrostack): Executing calculation: 10.0 add 20.0 -2026-06-10 14:52:34,751 [INFO] (nitrostack): Converting temperature: 35.0 from celsius to fahrenheit -2026-06-10 14:52:39,316 [INFO] (nitrostack): Converting temperature: 35.0 from celsius to celsius -2026-06-10 14:52:42,765 [INFO] (nitrostack): Converting temperature: 35.0 from celsius to kelvin -2026-06-10 15:07:19,747 [INFO] (nitrostack): Executing calculation: 12.5 add 3.5 -2026-06-10 15:07:19,752 [INFO] (nitrostack): Executing calculation: 10.0 divide 0.0 -2026-06-10 15:07:19,753 [INFO] (nitrostack): Converting temperature: 100.0 from celsius to fahrenheit -2026-06-10 15:20:07,388 [INFO] (nitrostack): Executing calculation: 12.5 add 3.5 -2026-06-10 15:20:07,393 [INFO] (nitrostack): Executing calculation: 10.0 divide 0.0 -2026-06-10 15:20:07,395 [INFO] (nitrostack): Converting temperature: 100.0 from celsius to fahrenheit -2026-06-10 15:39:46,729 [INFO] (nitrostack): Executing calculation: 12.5 add 3.5 -2026-06-10 15:39:46,738 [INFO] (nitrostack): Executing calculation: 10.0 divide 0.0 -2026-06-10 15:39:46,740 [INFO] (nitrostack): Converting temperature: 100.0 from celsius to fahrenheit -2026-06-10 15:45:21,942 [INFO] (nitrostack): Executing calculation: 12.5 add 3.5 -2026-06-10 15:45:21,947 [INFO] (nitrostack): Executing calculation: 10.0 divide 0.0 -2026-06-10 15:45:21,948 [INFO] (nitrostack): Converting temperature: 100.0 from celsius to fahrenheit -2026-06-12 11:55:23,249 [INFO] (nitrostack): Executing calculation: 12.5 add 3.5 -2026-06-12 11:55:23,259 [INFO] (nitrostack): Executing calculation: 10.0 divide 0.0 -2026-06-12 11:55:23,262 [INFO] (nitrostack): Converting temperature: 100.0 from celsius to fahrenheit -2026-06-12 13:05:52,734 [INFO] (nitrostack): Executing calculation: 12.5 add 3.5 -2026-06-12 13:05:52,756 [INFO] (nitrostack): Executing calculation: 10.0 divide 0.0 -2026-06-12 13:05:52,760 [INFO] (nitrostack): Converting temperature: 100.0 from celsius to fahrenheit -2026-06-12 13:32:04,654 [INFO] (nitrostack): Executing calculation: 10.0 add 20.0 -2026-06-17 12:46:21,443 [INFO] (nitrostack): Executing calculation: 12.5 add 3.5 -2026-06-17 12:46:21,452 [INFO] (nitrostack): Executing calculation: 10.0 divide 0.0 -2026-06-17 12:46:21,454 [INFO] (nitrostack): Converting temperature: 100.0 from celsius to fahrenheit -2026-06-18 12:25:18,921 [INFO] (nitrostack): Executing calculation: 10.0 multiply 20.0 -2026-06-18 12:25:24,581 [INFO] (nitrostack): Converting temperature: 35.0 from celsius to celsius -2026-06-18 12:25:28,765 [INFO] (nitrostack): Converting temperature: 35.0 from celsius to fahrenheit -2026-06-18 12:26:45,327 [INFO] (nitrostack): Converting temperature: 35.0 from celsius to kelvin -2026-06-18 13:27:32,721 [INFO] (nitrostack): Executing calculation: 12.5 add 3.5 -2026-06-18 13:27:32,729 [INFO] (nitrostack): Executing calculation: 10.0 divide 0.0 -2026-06-18 13:27:32,731 [INFO] (nitrostack): Converting temperature: 100.0 from celsius to fahrenheit diff --git a/nitrostack/core/app.py b/nitrostack/core/app.py index b694ab7..a91c238 100644 --- a/nitrostack/core/app.py +++ b/nitrostack/core/app.py @@ -1,37 +1,30 @@ import os import sys +import re +import json import uuid +import asyncio import inspect +import datetime from dataclasses import dataclass, field -from typing import Any, Callable, Dict, List, Set, Type, Optional -from mcp.server.fastmcp import FastMCP -from mcp.types import PromptMessage as McpPromptMessage, TextContent +from typing import Any, Callable, Dict, List, Optional, Pattern, Set, Tuple, Type + +import mcp.types as types +from mcp.server.lowlevel.server import request_ctx +from mcp.server.lowlevel.helper_types import ReadResourceContents +from mcp.server.stdio import stdio_server from pydantic import BaseModel, create_model -from nitrostack.core.context import ExecutionContext +from nitrostack.core.context import ExecutionContext, TaskContext from nitrostack.core.decorators import ToolConfig, ResourceConfig, PromptConfig from nitrostack.core.di import DIContainer +from nitrostack.core.errors import PromptNotFoundError, ResourceNotFoundError +from nitrostack.core.mcp_server import NitroStackMcpServer from nitrostack.core.pipeline import run_pipeline from nitrostack.core.additional_decorators import HealthCheckRegistry +from nitrostack.core.task import TaskRegistry from nitrostack.events.event_emitter import EventEmitter -# Monkeypatch FastMCP.call_tool to support returning CreateTaskResult without conversion -_original_fastmcp_call_tool = FastMCP.call_tool - -async def _custom_fastmcp_call_tool(self, name: str, arguments: dict[str, Any]): - t = self._tool_manager.get_tool(name) - if not t: - return await _original_fastmcp_call_tool(self, name, arguments) - context = self.get_context() - result = await self._tool_manager.call_tool( - name, arguments, context=context, convert_result=False - ) - import mcp.types as types - if isinstance(result, types.CreateTaskResult): - return result - return t.fn_metadata.convert_result(result) - -FastMCP.call_tool = _custom_fastmcp_call_tool @dataclass class ServerConfig: @@ -39,6 +32,7 @@ class ServerConfig: version: str = "1.0.0" transport_type: Optional[str] = None + def mcp_app(module: Type, server: ServerConfig): """ Decorator to declare the main application class. @@ -75,15 +69,39 @@ def get_pydantic_model(schema: Any) -> Type[BaseModel]: t = list elif prop_type == "object": t = dict - + default = ... if name in required else None fields[name] = (t, default) return create_model("DynamicInputModel", **fields) - + # Return a default empty model if invalid or empty return create_model("EmptyInputModel") +@dataclass +class _ToolEntry: + config: ToolConfig + input_model: Type[BaseModel] + instance: Any + method: Callable + + +@dataclass +class _ResourceEntry: + config: ResourceConfig + instance: Any + method: Callable + param_names: List[str] = field(default_factory=list) + pattern: Optional[Pattern] = None + + +@dataclass +class _PromptEntry: + config: PromptConfig + instance: Any + method: Callable + + class McpApplication: def __init__(self, app_class: Type): self.app_class = app_class @@ -95,13 +113,26 @@ def __init__(self, app_class: Type): self.server_config = ServerConfig(name=app_class._mcp_module_config.name or "mcp-server") else: raise ValueError("Invalid application class. Must be decorated with @mcp_app or @module.") - self.mcp_server: Optional[FastMCP] = None + + self.mcp_server: Optional[NitroStackMcpServer] = None + + # nitrostack owns these registries directly (no FastMCP-managed tool/resource + # manager in between) so that any number of low-level `Server` instances can be + # wired against the same registered tools/resources/prompts (see + # `create_configured_mcp_server`). + self._tools: Dict[str, _ToolEntry] = {} + self._resources: Dict[str, _ResourceEntry] = {} + self._resource_templates: List[_ResourceEntry] = [] + self._prompts: Dict[str, _PromptEntry] = {} + self._initial_tools: List[Tuple[Any, Callable, ToolConfig]] = [] + self._bootstrap() def _bootstrap(self) -> None: - # 1. Initialize FastMCP - self.mcp_server = FastMCP( - name=self.server_config.name + # 1. Construct the low-level server directly (no FastMCP) + self.mcp_server = NitroStackMcpServer( + name=self.server_config.name, + version=self.server_config.version, ) # 2. Resolve Module Tree recursively and register services @@ -122,7 +153,6 @@ def _bootstrap(self) -> None: container.resolve(controller) # 3. Discover decorated methods on all instances in the container - initial_tools = [] for token, instance in list(container._instances.items()): # Scan members of this instance for name, member in inspect.getmembers(instance): @@ -130,8 +160,6 @@ def _bootstrap(self) -> None: if hasattr(member, "_mcp_tool_config"): tool_config: ToolConfig = getattr(member, "_mcp_tool_config") self._register_tool(instance, member, tool_config) - if tool_config.is_initial: - initial_tools.append((instance, member, tool_config)) # Discover Resources if hasattr(member, "_mcp_resource_config"): @@ -155,33 +183,444 @@ def _bootstrap(self) -> None: # 4. Register Built-in Health check Resource if any checks exist if HealthCheckRegistry.get_checks(): - @self.mcp_server.resource("health://status", name="Health Status", description="System health status check") - def health_status_resource() -> str: - import json - results = HealthCheckRegistry.run_all() - return json.dumps(results) - - # 5. Register Task Support hook & endpoints - low_level_server = self.mcp_server._mcp_server - original_get_caps = low_level_server.get_capabilities - - def custom_get_capabilities(notification_options, experimental_capabilities): - caps = original_get_caps(notification_options, experimental_capabilities) - import mcp.types as types - caps.tasks = types.ServerTasksCapability( - list=types.TasksListCapability(), - cancel=types.TasksCancelCapability(), - requests=types.ServerTasksRequestsCapability() + self._register_health_resource() + + # 5. Wire every protocol handler (tools/resources/prompts/tasks/notifications) + # onto the owned low-level server. This is factored out so additional server + # instances (e.g. one per HTTP session) can be configured identically. + self._setup_handlers(self.mcp_server) + + def _resolve_modules(self, module_class: Type, resolved_modules: Set[Type]) -> None: + if module_class in resolved_modules: + return + resolved_modules.add(module_class) + + mod_config = getattr(module_class, "_mcp_module_config", None) + if mod_config: + for imp in mod_config.imports: + self._resolve_modules(imp, resolved_modules) + + # ------------------------------------------------------------------ + # Registration (pure bookkeeping — no protocol/server calls here) + # ------------------------------------------------------------------ + + def _register_tool(self, instance: Any, method: Callable, tool_config: ToolConfig) -> None: + input_model = get_pydantic_model(tool_config.input_schema) + entry = _ToolEntry(config=tool_config, input_model=input_model, instance=instance, method=method) + self._tools[tool_config.name] = entry + if tool_config.is_initial: + self._initial_tools.append((instance, method, tool_config)) + + def _register_resource(self, instance: Any, method: Callable, resource_config: ResourceConfig) -> None: + param_names = re.findall(r"\{([^}]+)\}", resource_config.uri) + entry = _ResourceEntry(config=resource_config, instance=instance, method=method, param_names=param_names) + + if param_names: + # Build a matching regex from the URI template, e.g. "a://b/{id}" -> + # "^a://b/(?P[^/]+)$", preserving the existing template-matching semantics. + regex_str = re.escape(resource_config.uri) + for pname in param_names: + regex_str = regex_str.replace(re.escape("{" + pname + "}"), f"(?P<{pname}>[^/]+)") + entry.pattern = re.compile(f"^{regex_str}$") + self._resource_templates.append(entry) + else: + self._resources[resource_config.uri] = entry + + def _register_prompt(self, instance: Any, method: Callable, prompt_config: PromptConfig) -> None: + self._prompts[prompt_config.name] = _PromptEntry(config=prompt_config, instance=instance, method=method) + + def _register_health_resource(self) -> None: + config = ResourceConfig( + uri="health://status", + name="Health Status", + description="System health status check", + ) + + async def health_status_resource(context: ExecutionContext) -> str: + results = HealthCheckRegistry.run_all() + return json.dumps(results) + + self._resources[config.uri] = _ResourceEntry(config=config, instance=None, method=health_status_resource) + + # ------------------------------------------------------------------ + # Protocol handler wiring (owned low-level `mcp.server.lowlevel.Server`) + # ------------------------------------------------------------------ + + def _setup_handlers(self, server: NitroStackMcpServer) -> None: + @server.list_tools() + async def _list_tools() -> List[types.Tool]: + return [self._build_tool_definition(entry) for entry in self._tools.values()] + + @server.call_tool() + async def _call_tool(name: str, arguments: Optional[Dict[str, Any]]): + return await self._call_tool(name, arguments or {}) + + @server.list_resources() + async def _list_resources() -> List[types.Resource]: + return [self._build_resource_definition(entry) for entry in self._resources.values()] + + @server.list_resource_templates() + async def _list_resource_templates() -> List[types.ResourceTemplate]: + return [self._build_resource_template_definition(entry) for entry in self._resource_templates] + + @server.read_resource() + async def _read_resource(uri: Any): + return await self._read_resource(str(uri)) + + @server.subscribe_resource() + async def _subscribe_resource(uri: Any) -> None: + if str(uri) not in self._resources: + raise ResourceNotFoundError(str(uri)) + + @server.unsubscribe_resource() + async def _unsubscribe_resource(uri: Any) -> None: + return None + + @server.list_prompts() + async def _list_prompts() -> List[types.Prompt]: + return [self._build_prompt_definition(entry) for entry in self._prompts.values()] + + @server.get_prompt() + async def _get_prompt(name: str, arguments: Optional[Dict[str, str]]) -> types.GetPromptResult: + return await self._get_prompt(name, arguments or {}) + + self._register_task_handlers(server) + self._register_initialized_handler(server) + + def create_configured_mcp_server(self) -> NitroStackMcpServer: + """ + Build a fresh, fully-wired low-level server instance sharing this + application's tool/resource/prompt registries. A single shared server is + sufficient for the Python `mcp` SDK's HTTP session managers (unlike the TS + SDK's `Server`, `.run()` here does not bind a transport for the instance's + lifetime), but this factory is kept available for isolated/stateless use + cases (e.g. Phase 3 extensions) that want their own server instance. + """ + server = NitroStackMcpServer(name=self.server_config.name, version=self.server_config.version) + self._setup_handlers(server) + return server + + # ------------------------------------------------------------------ + # Definition builders (registry entry -> wire-level `mcp.types` objects) + # ------------------------------------------------------------------ + + def _wrap_tool_input_schema(self, input_schema: Dict[str, Any]) -> Dict[str, Any]: + """ + Wrap a Pydantic model JSON schema under `{"input": ...}` the way FastMCP + (and NitroStudio) expect. + + Studio's form generator resolves `properties.input.$ref` against `$defs` + to render individual fields. Inlining the model under `properties.input` + makes Studio treat it as a single JSON-object blob — so we always emit the + `$ref` + `$defs` shape (matching FastMCP's previous wire contract). + """ + nested_defs = dict(input_schema.get("$defs") or input_schema.get("definitions") or {}) + model_schema = {k: v for k, v in input_schema.items() if k not in ("$defs", "definitions")} + model_name = model_schema.get("title") or "Input" + # Avoid colliding with an existing nested def of the same name + if model_name in nested_defs: + model_name = f"{model_name}Input" + + return { + "type": "object", + "title": f"{model_name}Arguments", + "properties": {"input": {"$ref": f"#/$defs/{model_name}"}}, + "required": ["input"], + "$defs": {**nested_defs, model_name: model_schema}, + } + + def _build_tool_definition(self, entry: _ToolEntry) -> types.Tool: + cfg = entry.config + wrapped_schema = self._wrap_tool_input_schema(entry.input_model.model_json_schema()) + + meta: Dict[str, Any] = { + "is_initial": cfg.is_initial, + "visibility": cfg.visibility, + "task_support": cfg.task_support, + **(cfg.metadata or {}), + } + widget_route = getattr(entry.method, "_mcp_widget", None) + if widget_route: + meta["ui/template"] = widget_route + meta["ui"] = {"resourceUri": widget_route} + meta["openai/outputTemplate"] = widget_route + if cfg.invocation: + meta["openai/toolInvocation/invoking"] = cfg.invocation.invoking + meta["openai/toolInvocation/invoked"] = cfg.invocation.invoked + if cfg.examples: + meta["examples"] = { + "input": cfg.examples.input, + "output": cfg.examples.output, + "description": cfg.examples.description, + } + + app_mode = os.environ.get("NITROSTACK_APP_MODE", "mcp") + if app_mode == "openai": + meta["openai/type"] = "function" + meta["openai/function"] = { + "name": cfg.name, + "description": cfg.description, + "parameters": wrapped_schema, + } + elif app_mode == "mcpapps": + meta["_meta"] = {"ui": {"title": cfg.title or cfg.name, "description": cfg.description}} + + annotations = types.ToolAnnotations( + readOnlyHint=cfg.annotations.read_only_hint, + destructiveHint=cfg.annotations.destructive_hint, + idempotentHint=cfg.annotations.idempotent_hint, + openWorldHint=cfg.annotations.open_world_hint, + ) + + execution = None + if cfg.task_support and cfg.task_support != "forbidden": + execution = types.ToolExecution(taskSupport=cfg.task_support) + + return types.Tool( + name=cfg.name, + title=cfg.title, + description=cfg.description, + inputSchema=wrapped_schema, + annotations=annotations, + execution=execution, + **{"_meta": meta}, + ) + + def _build_resource_definition(self, entry: _ResourceEntry) -> types.Resource: + cfg = entry.config + return types.Resource( + uri=cfg.uri, + name=cfg.name, + title=cfg.title, + description=cfg.description, + mimeType=cfg.mime_type, + size=cfg.size, + ) + + def _build_resource_template_definition(self, entry: _ResourceEntry) -> types.ResourceTemplate: + cfg = entry.config + return types.ResourceTemplate( + uriTemplate=cfg.uri, + name=cfg.name, + title=cfg.title, + description=cfg.description, + mimeType=cfg.mime_type, + ) + + def _build_prompt_definition(self, entry: _PromptEntry) -> types.Prompt: + cfg = entry.config + arguments = [ + types.PromptArgument(name=arg.name, description=arg.description, required=arg.required) + for arg in cfg.arguments + ] + return types.Prompt(name=cfg.name, description=cfg.description, arguments=arguments or None) + + # ------------------------------------------------------------------ + # Request dispatch + # ------------------------------------------------------------------ + + def _pipeline_stages(self, method: Callable) -> Tuple[List[Type], List[Type], List[Type], List[Type], List[Type]]: + return ( + getattr(method, "_mcp_guards", []), + getattr(method, "_mcp_middleware", []), + getattr(method, "_mcp_interceptors", []), + getattr(method, "_mcp_pipes", []), + getattr(method, "_mcp_filters", []), + ) + + @staticmethod + def _to_call_tool_result(result: Any) -> types.CallToolResult: + if isinstance(result, types.CallToolResult): + return result + if isinstance(result, BaseModel): + result = result.model_dump() + if isinstance(result, dict): + if "content" in result and "isError" in result: + return types.CallToolResult(**result) + return types.CallToolResult( + content=[types.TextContent(type="text", text=json.dumps(result, indent=2, default=str))], + structuredContent=result, + isError=False, + ) + return types.CallToolResult( + content=[types.TextContent(type="text", text=str(result))], + isError=False, + ) + + async def _call_tool(self, name: str, arguments: Dict[str, Any]): + entry = self._tools.get(name) + if entry is None: + return types.CallToolResult( + content=[types.TextContent(type="text", text=f"Tool '{name}' not found")], + isError=True, + ) + + cfg = entry.config + input_instance = entry.input_model.model_validate(arguments.get("input", {})) + guards, middleware, interceptors, pipes, filters = self._pipeline_stages(entry.method) + + # Detect task-augmented invocation via the request context's public + # `experimental.task_metadata` field (populated by the low-level server + # from `req.params.task`) rather than reaching into private state. + task_metadata = None + rc = request_ctx.get(None) + if rc is not None and getattr(rc, "experimental", None) is not None: + task_metadata = rc.experimental.task_metadata + + is_task = (task_metadata is not None) or (cfg.task_support == "required") + if cfg.task_support == "forbidden": + is_task = False + + if is_task: + task_id = f"task_{uuid.uuid4().hex[:12]}" + ttl = task_metadata.ttl if task_metadata else 300 + TaskRegistry.create_task(task_id, ttl=ttl) + + async def background_execution(): + task_ctx = ExecutionContext( + request_id=str(uuid.uuid4()), + tool_name=cfg.name, + metadata={"input": input_instance}, + ) + task_ctx.task = TaskContext(task_id) + try: + result = await run_pipeline( + handler=entry.method, + handler_instance=entry.instance, + args=(input_instance, task_ctx), + kwargs={}, + context=task_ctx, + guards=guards, + middleware=middleware, + interceptors=interceptors, + pipes=pipes, + filters=filters, + param_name="input", + param_type=entry.input_model, + ) + TaskRegistry.complete_task(task_id, self._to_call_tool_result(result)) + except Exception as e: + TaskRegistry.fail_task(task_id, e) + + asyncio.create_task(background_execution()) + now = datetime.datetime.now(datetime.timezone.utc) + return types.CreateTaskResult( + task=types.Task( + taskId=task_id, + status="working", + statusMessage="Task started", + createdAt=now, + lastUpdatedAt=now, + ttl=ttl, + pollInterval=5, + ) ) - return caps - low_level_server.get_capabilities = custom_get_capabilities + ctx = ExecutionContext(request_id=str(uuid.uuid4()), tool_name=cfg.name, metadata={"input": input_instance}) + result = await run_pipeline( + handler=entry.method, + handler_instance=entry.instance, + args=(input_instance, ctx), + kwargs={}, + context=ctx, + guards=guards, + middleware=middleware, + interceptors=interceptors, + pipes=pipes, + filters=filters, + param_name="input", + param_type=entry.input_model, + ) + return self._to_call_tool_result(result) + + async def _read_resource(self, uri: str) -> List[ReadResourceContents]: + entry = self._resources.get(uri) + path_kwargs: Dict[str, str] = {} + + if entry is None: + for template_entry in self._resource_templates: + match = template_entry.pattern.match(uri) if template_entry.pattern else None + if match: + entry = template_entry + path_kwargs = match.groupdict() + break + + if entry is None: + raise ResourceNotFoundError(uri) + + cfg = entry.config + ctx = ExecutionContext(request_id=str(uuid.uuid4()), metadata=dict(path_kwargs)) + guards, middleware, interceptors, pipes, filters = self._pipeline_stages(entry.method) + + result = await run_pipeline( + handler=entry.method, + handler_instance=entry.instance, + args=(), + kwargs={**path_kwargs, "context": ctx}, + context=ctx, + guards=guards, + middleware=middleware, + interceptors=interceptors, + pipes=pipes, + filters=filters, + ) - # Class-level monkeypatch handles call_tool mapping + # Support Discriminated Union formats for return (Section 1.2) + if isinstance(result, dict) and "type" in result and "data" in result: + result = result["data"] + + if isinstance(result, BaseModel): + result = result.model_dump() + + if isinstance(result, bytes): + return [ReadResourceContents(content=result, mime_type=cfg.mime_type or "application/octet-stream")] + + if isinstance(result, str): + return [ReadResourceContents(content=result, mime_type=cfg.mime_type or "text/plain")] + + text = json.dumps(result, indent=2, default=str) + return [ReadResourceContents(content=text, mime_type=cfg.mime_type or "application/json")] + + async def _get_prompt(self, name: str, arguments: Dict[str, str]) -> types.GetPromptResult: + entry = self._prompts.get(name) + if entry is None: + raise PromptNotFoundError(name) + + cfg = entry.config + args_dict = dict(arguments or {}) + ctx = ExecutionContext(request_id=str(uuid.uuid4()), metadata=args_dict) + guards, middleware, interceptors, pipes, filters = self._pipeline_stages(entry.method) + + raw_messages = await run_pipeline( + handler=entry.method, + handler_instance=entry.instance, + args=(args_dict, ctx), + kwargs={}, + context=ctx, + guards=guards, + middleware=middleware, + interceptors=interceptors, + pipes=pipes, + filters=filters, + ) + + if not isinstance(raw_messages, (list, tuple)): + raw_messages = [raw_messages] + + messages: List[types.PromptMessage] = [] + for msg in raw_messages: + role = msg.role if hasattr(msg, "role") else msg.get("role") + content = msg.content if hasattr(msg, "content") else msg.get("content") + messages.append(types.PromptMessage(role=role, content=types.TextContent(type="text", text=content))) - import mcp.types as types - from nitrostack.core.task import TaskRegistry + return types.GetPromptResult(description=cfg.description, messages=messages) + # ------------------------------------------------------------------ + # Task subsystem — registered directly on the low-level server's public + # `request_handlers`/`notification_handlers` dicts (no FastMCP reach-through). + # ------------------------------------------------------------------ + + def _register_task_handlers(self, server: NitroStackMcpServer) -> None: async def handle_list_tasks(req): tasks_list = [ types.Task( @@ -191,7 +630,7 @@ async def handle_list_tasks(req): createdAt=t.created_at, lastUpdatedAt=t.last_updated_at, ttl=t.ttl, - pollInterval=t.poll_interval + pollInterval=t.poll_interval, ) for t in TaskRegistry.list_tasks() ] @@ -202,10 +641,7 @@ async def handle_get_task(req): t = TaskRegistry.get_task(task_id) if not t: raise types.McpError( - types.ErrorData( - code=types.INVALID_PARAMS, - message=f"Task {task_id} not found" - ) + types.ErrorData(code=types.INVALID_PARAMS, message=f"Task {task_id} not found") ) return types.GetTaskResult( taskId=t.task_id, @@ -214,7 +650,7 @@ async def handle_get_task(req): createdAt=t.created_at, lastUpdatedAt=t.last_updated_at, ttl=t.ttl, - pollInterval=t.poll_interval + pollInterval=t.poll_interval, ) async def handle_cancel_task(req): @@ -222,10 +658,7 @@ async def handle_cancel_task(req): t = TaskRegistry.get_task(task_id) if not t: raise types.McpError( - types.ErrorData( - code=types.INVALID_PARAMS, - message=f"Task {task_id} not found" - ) + types.ErrorData(code=types.INVALID_PARAMS, message=f"Task {task_id} not found") ) TaskRegistry.cancel_task(task_id) t = TaskRegistry.get_task(task_id) @@ -236,7 +669,7 @@ async def handle_cancel_task(req): createdAt=t.created_at, lastUpdatedAt=t.last_updated_at, ttl=t.ttl, - pollInterval=t.poll_interval + pollInterval=t.poll_interval, ) async def handle_get_task_payload(req): @@ -244,10 +677,7 @@ async def handle_get_task_payload(req): t = TaskRegistry.get_task(task_id) if not t: raise types.McpError( - types.ErrorData( - code=types.INVALID_PARAMS, - message=f"Task {task_id} not found" - ) + types.ErrorData(code=types.INVALID_PARAMS, message=f"Task {task_id} not found") ) await t.done_event.wait() if t.status == "completed": @@ -255,71 +685,40 @@ async def handle_get_task_payload(req): elif t.status == "cancelled": return types.CallToolResult( content=[types.TextContent(type="text", text="Task was cancelled.")], - isError=True + isError=True, ) else: return types.CallToolResult( content=[types.TextContent(type="text", text=str(t.error or t.status_message))], - isError=True + isError=True, ) - low_level_server.request_handlers[types.ListTasksRequest] = handle_list_tasks - low_level_server.request_handlers[types.GetTaskRequest] = handle_get_task - low_level_server.request_handlers[types.CancelTaskRequest] = handle_cancel_task - low_level_server.request_handlers[types.GetTaskPayloadRequest] = handle_get_task_payload - - # 6. Register App Mode formatters for tool listings - original_list_tools_handler = low_level_server.request_handlers.get(types.ListToolsRequest) - if original_list_tools_handler: - async def custom_list_tools_handler(req): - res = await original_list_tools_handler(req) - app_mode = os.environ.get("NITROSTACK_APP_MODE", "mcp") - if app_mode in ("openai", "mcpapps"): - for tool in res.tools: - if tool.meta is None: - tool.meta = {} - if app_mode == "openai": - tool.meta["openai/type"] = "function" - tool.meta["openai/function"] = { - "name": tool.name, - "description": tool.description, - "parameters": tool.inputSchema - } - elif app_mode == "mcpapps": - tool.meta["_meta"] = { - "ui": { - "title": tool.title or tool.name, - "description": tool.description - } - } - return res - low_level_server.request_handlers[types.ListToolsRequest] = custom_list_tools_handler - - # 7. Register a notification handler for initialized to auto-call @initial_tool annotated tools + server.request_handlers[types.ListTasksRequest] = handle_list_tasks + server.request_handlers[types.GetTaskRequest] = handle_get_task + server.request_handlers[types.CancelTaskRequest] = handle_cancel_task + server.request_handlers[types.GetTaskPayloadRequest] = handle_get_task_payload + server.has_task_support = True + + def _register_initialized_handler(self, server: NitroStackMcpServer) -> None: async def handle_initialized(notification: types.InitializedNotification): - for inst, memb, config in initial_tools: + for instance, method, config in self._initial_tools: try: - input_model = get_pydantic_model(config.input_schema) + entry = self._tools[config.name] try: - input_inst = input_model() + input_inst = entry.input_model() except Exception: input_inst = None - + ctx = ExecutionContext( request_id=f"initial-tool-{uuid.uuid4().hex[:8]}", tool_name=config.name, - metadata={} + metadata={}, ) - - guards = getattr(memb, "_mcp_guards", []) - middleware = getattr(memb, "_mcp_middleware", []) - interceptors = getattr(memb, "_mcp_interceptors", []) - pipes = getattr(memb, "_mcp_pipes", []) - filters = getattr(memb, "_mcp_filters", []) + guards, middleware, interceptors, pipes, filters = self._pipeline_stages(method) await run_pipeline( - handler=memb, - handler_instance=inst, + handler=method, + handler_instance=instance, args=(input_inst, ctx), kwargs={}, context=ctx, @@ -329,356 +728,64 @@ async def handle_initialized(notification: types.InitializedNotification): pipes=pipes, filters=filters, param_name="input", - param_type=input_model + param_type=entry.input_model, ) except Exception as e: - import sys - print(f"Error executing initial tool '{config.name}': {e}", file=sys.stderr) - - low_level_server.notification_handlers[types.InitializedNotification] = handle_initialized - - def _resolve_modules(self, module_class: Type, resolved_modules: Set[Type]) -> None: - if module_class in resolved_modules: - return - resolved_modules.add(module_class) - - mod_config = getattr(module_class, "_mcp_module_config", None) - if mod_config: - for imp in mod_config.imports: - self._resolve_modules(imp, resolved_modules) - - def _register_tool(self, instance: Any, method: Callable, tool_config: ToolConfig) -> None: - input_model = get_pydantic_model(tool_config.input_schema) - - # Define the wrapper function that executes pipeline stages - async def tool_wrapper(input: input_model) -> Any: - # Check if task execution is requested or required - task_metadata = None - try: - from mcp.server.lowlevel.server import request_ctx - req = request_ctx.get().request - task_metadata = getattr(req.params, "task", None) if req and hasattr(req, "params") else None - except Exception: - pass - - is_task = (task_metadata is not None) or (tool_config.task_support == "required") - if tool_config.task_support == "forbidden": - is_task = False - - if is_task: - task_id = f"task_{uuid.uuid4().hex[:12]}" - ttl = task_metadata.ttl if task_metadata else 300 - from nitrostack.core.task import TaskRegistry - TaskRegistry.create_task(task_id, ttl=ttl) - - async def background_execution(): - task_ctx = ExecutionContext( - request_id=str(uuid.uuid4()), - tool_name=tool_config.name, - metadata={"input": input} - ) - from nitrostack.core.context import TaskContext - task_ctx.task = TaskContext(task_id) - try: - guards = getattr(method, "_mcp_guards", []) - middleware = getattr(method, "_mcp_middleware", []) - interceptors = getattr(method, "_mcp_interceptors", []) - pipes = getattr(method, "_mcp_pipes", []) - filters = getattr(method, "_mcp_filters", []) - - result = await run_pipeline( - handler=method, - handler_instance=instance, - args=(input, task_ctx), - kwargs={}, - context=task_ctx, - guards=guards, - middleware=middleware, - interceptors=interceptors, - pipes=pipes, - filters=filters, - param_name="input", - param_type=input_model - ) - if isinstance(result, BaseModel): - result_dump = result.model_dump() - else: - result_dump = result - - import mcp.types as types - import json - if isinstance(result_dump, types.CallToolResult): - final_result = result_dump - elif isinstance(result_dump, dict): - if "content" in result_dump and "isError" in result_dump: - final_result = types.CallToolResult(**result_dump) - else: - final_result = types.CallToolResult( - content=[types.TextContent(type="text", text=json.dumps(result_dump, indent=2))], - structuredContent=result_dump, - isError=False - ) - else: - final_result = types.CallToolResult( - content=[types.TextContent(type="text", text=str(result_dump))], - isError=False - ) - TaskRegistry.complete_task(task_id, final_result) - except Exception as e: - TaskRegistry.fail_task(task_id, e) - - import asyncio - asyncio.create_task(background_execution()) - import datetime - import mcp.types as types - now = datetime.datetime.now(datetime.timezone.utc) - return types.CreateTaskResult( - task=types.Task( - taskId=task_id, - status="working", - statusMessage="Task started", - createdAt=now, - lastUpdatedAt=now, - ttl=ttl, - pollInterval=5 - ) - ) + sys.stderr.write(f"Error executing initial tool '{config.name}': {e}\n") + sys.stderr.flush() - # Initialize ExecutionContext - ctx = ExecutionContext( - request_id=str(uuid.uuid4()), - tool_name=tool_config.name, - metadata={"input": input} - ) + server.notification_handlers[types.InitializedNotification] = handle_initialized - # Retrieve pipeline decorators - guards = getattr(method, "_mcp_guards", []) - middleware = getattr(method, "_mcp_middleware", []) - interceptors = getattr(method, "_mcp_interceptors", []) - pipes = getattr(method, "_mcp_pipes", []) - filters = getattr(method, "_mcp_filters", []) - - # Run through pipeline runner - result = await run_pipeline( - handler=method, - handler_instance=instance, - args=(input, ctx), - kwargs={}, - context=ctx, - guards=guards, - middleware=middleware, - interceptors=interceptors, - pipes=pipes, - filters=filters, - param_name="input", - param_type=input_model - ) - - if isinstance(result, BaseModel): - return result.model_dump() - return result - - # Register metadata - meta = { - "is_initial": tool_config.is_initial, - "visibility": tool_config.visibility, - "task_support": tool_config.task_support, - **(tool_config.metadata or {}) - } - widget_route = getattr(method, "_mcp_widget", None) - if widget_route: - meta["ui/template"] = widget_route - meta["ui"] = {"resourceUri": widget_route} - meta["openai/outputTemplate"] = widget_route - - if tool_config.invocation: - meta["openai/toolInvocation/invoking"] = tool_config.invocation.invoking - meta["openai/toolInvocation/invoked"] = tool_config.invocation.invoked - if tool_config.examples: - meta["examples"] = { - "input": tool_config.examples.input, - "output": tool_config.examples.output, - "description": tool_config.examples.description - } - - # Add to FastMCP - self.mcp_server.add_tool( - tool_wrapper, - name=tool_config.name, - title=tool_config.title, - description=tool_config.description, - meta=meta - ) - - def _register_resource(self, instance: Any, method: Callable, resource_config: ResourceConfig) -> None: - import re - # Find all template variables in braces, e.g. {result_id} - param_names = re.findall(r"\{([^}]+)\}", resource_config.uri) - sig_str = ", ".join(param_names) - - exec_locals = {} - exec_globals = { - "run_pipeline": run_pipeline, - "method": method, - "instance": instance, - "ExecutionContext": ExecutionContext, - "uuid": uuid, - } - - # Build dynamic wrapper signature to expose correct parameters to clients/studios - code = f""" -async def resource_wrapper({sig_str}): - import uuid - args_dict = {{ - {", ".join(f"'{name}': {name}" for name in param_names)} - }} - ctx = ExecutionContext( - request_id=str(uuid.uuid4()), - metadata=args_dict - ) - - guards = getattr(method, "_mcp_guards", []) - middleware = getattr(method, "_mcp_middleware", []) - interceptors = getattr(method, "_mcp_interceptors", []) - pipes = getattr(method, "_mcp_pipes", []) - filters = getattr(method, "_mcp_filters", []) - - result = await run_pipeline( - handler=method, - handler_instance=instance, - args=(), - kwargs={{**args_dict, "context": ctx}}, - context=ctx, - guards=guards, - middleware=middleware, - interceptors=interceptors, - pipes=pipes, - filters=filters - ) - - # Support Discriminated Union formats for return (Section 1.2) - if isinstance(result, dict) and "type" in result and "data" in result: - return result["data"] - return result -""" - exec(code, exec_globals, exec_locals) - resource_wrapper = exec_locals["resource_wrapper"] - - # Register resource decorator on FastMCP - resource_decorator = self.mcp_server.resource( - resource_config.uri, - name=resource_config.name, - title=resource_config.title, - description=resource_config.description, - mime_type=resource_config.mime_type - ) - resource_decorator(resource_wrapper) - - def _register_prompt(self, instance: Any, method: Callable, prompt_config: PromptConfig) -> None: - arg_names = [arg.name for arg in prompt_config.arguments] - sig_parts = [] - for arg in prompt_config.arguments: - sig_part = f"{arg.name}: str" - if not arg.required: - sig_part += " = None" - sig_parts.append(sig_part) - sig_str = ", ".join(sig_parts) - - exec_locals = {} - exec_globals = { - "run_pipeline": run_pipeline, - "method": method, - "instance": instance, - "ExecutionContext": ExecutionContext, - "uuid": uuid, - "McpPromptMessage": McpPromptMessage, - "TextContent": TextContent, - } - - # Build dynamic wrapper signature to expose correct parameters to clients/studios - code = f""" -async def prompt_wrapper({sig_str}): - args_dict = {{ - {", ".join(f"'{name}': {name}" for name in arg_names)} - }} - ctx = ExecutionContext( - request_id=str(uuid.uuid4()), - metadata=args_dict, - ) - - guards = getattr(method, "_mcp_guards", []) - middleware = getattr(method, "_mcp_middleware", []) - interceptors = getattr(method, "_mcp_interceptors", []) - pipes = getattr(method, "_mcp_pipes", []) - filters = getattr(method, "_mcp_filters", []) - - raw_messages = await run_pipeline( - handler=method, - handler_instance=instance, - args=(args_dict, ctx), - kwargs={{}}, - context=ctx, - guards=guards, - middleware=middleware, - interceptors=interceptors, - pipes=pipes, - filters=filters - ) - - mcp_messages = [] - from collections.abc import Iterable - if not isinstance(raw_messages, Iterable) or isinstance(raw_messages, (dict, str, bytes)): - raw_list = [raw_messages] - else: - raw_list = list(raw_messages) - - for msg in raw_list: - role = msg.role if hasattr(msg, "role") else msg.get("role") - content = msg.content if hasattr(msg, "content") else msg.get("content") - mcp_messages.append({{ - "role": role, - "content": {{ - "type": "text", - "text": content - }} - }}) - return mcp_messages -""" - exec(code, exec_globals, exec_locals) - prompt_wrapper = exec_locals["prompt_wrapper"] - - # Register prompt decorator on FastMCP - prompt_decorator = self.mcp_server.prompt( - name=prompt_config.name, - description=prompt_config.description - ) - prompt_decorator(prompt_wrapper) + # ------------------------------------------------------------------ + # Transports + # ------------------------------------------------------------------ def get_combined_app(self) -> Any: + """ + Minimal Starlette app wiring the owned low-level server directly to both the + modern Streamable HTTP transport (`/mcp`) and the legacy SSE transport + (`/sse` + `/mcp/messages`). Full session lifecycle management, auth + middleware, and dual-mode orchestration are Phase 3's job — this is the + foundation it extends. + """ + import contextlib from starlette.applications import Starlette - from starlette.routing import Route - combined_app = Starlette() - - # Add custom Starlette routes registered via custom_route - try: - http_app = self.mcp_server.streamable_http_app() - for route in http_app.routes: - combined_app.routes.append(route) - except Exception: - pass + from starlette.routing import Mount, Route + from starlette.responses import Response + from mcp.server.streamable_http_manager import StreamableHTTPSessionManager + from mcp.server.sse import SseServerTransport + + session_manager = StreamableHTTPSessionManager(app=self.mcp_server, stateless=False) + sse_transport = SseServerTransport("/mcp/messages") + + async def handle_streamable_http(scope, receive, send) -> None: + await session_manager.handle_request(scope, receive, send) + + async def handle_sse(request): + async with sse_transport.connect_sse(request.scope, request.receive, request._send) as streams: + await self.mcp_server.run(streams[0], streams[1], self.mcp_server.create_initialization_options()) + return Response() + + @contextlib.asynccontextmanager + async def lifespan(app): + async with session_manager.run(): + yield + + return Starlette( + routes=[ + # `handle_streamable_http` and `handle_post_message` are raw ASGI apps + # (scope, receive, send), so they must be `Mount`ed rather than used as + # `Route` endpoints (which expect a `Request -> Response` function). + Mount("/mcp", app=handle_streamable_http), + Route("/sse", endpoint=handle_sse, methods=["GET"]), + Mount("/mcp/messages", app=sse_transport.handle_post_message), + ], + lifespan=lifespan, + ) - try: - sse_app = self.mcp_server.sse_app() - for route in sse_app.routes: - if not any(r.path == route.path for r in combined_app.routes): - combined_app.routes.append(route) - if route.path == "/messages": - if not any(r.path == "/mcp/messages" for r in combined_app.routes): - combined_app.routes.append(Route("/mcp/messages", endpoint=route.endpoint, methods=getattr(route, "methods", None))) - except Exception: - pass - - return combined_app + async def _run_stdio(self) -> None: + async with stdio_server() as (read_stream, write_stream): + await self.mcp_server.run(read_stream, write_stream, self.mcp_server.create_initialization_options()) async def start(self) -> None: """Starts the MCP application based on transport configurations.""" @@ -694,7 +801,6 @@ async def start(self) -> None: node_env = os.environ.get("NODE_ENV", "development") port = int(os.environ.get("PORT") or os.environ.get("MCP_SERVER_PORT") or 8000) - # Stdio mode and Dual mode if transport == "http": # Run combined HTTP Server import uvicorn @@ -705,11 +811,12 @@ async def start(self) -> None: elif transport == "dual" or (node_env == "production" and not transport): # dual mode: stdio + HTTP import threading + def run_http(): - import asyncio + import asyncio as _asyncio import uvicorn - loop = asyncio.new_event_loop() - asyncio.set_event_loop(loop) + loop = _asyncio.new_event_loop() + _asyncio.set_event_loop(loop) app = self.get_combined_app() config = uvicorn.Config(app, host="0.0.0.0", port=port, log_level="warning") server = uvicorn.Server(config) @@ -717,17 +824,16 @@ def run_http(): http_thread = threading.Thread(target=run_http, daemon=True) http_thread.start() - + # Run stdio in the main thread from nitrostack.transports.stdio import safe_stdio_transport with safe_stdio_transport(): - await self.mcp_server.run_stdio_async() + await self._run_stdio() else: # Default Stdio from nitrostack.transports.stdio import safe_stdio_transport with safe_stdio_transport(): - await self.mcp_server.run_stdio_async() - + await self._run_stdio() class McpApplicationFactory: diff --git a/nitrostack/core/mcp_server.py b/nitrostack/core/mcp_server.py new file mode 100644 index 0000000..183406f --- /dev/null +++ b/nitrostack/core/mcp_server.py @@ -0,0 +1,66 @@ +""" +Low-level MCP server ownership for nitrostack. + +Mirrors how the TypeScript SDK builds directly on `@modelcontextprotocol/sdk`'s +low-level `Server` class (see `NitroStackServer` in the TS core) instead of a +higher-level convenience wrapper: nitrostack owns registration/dispatch and +only declares the protocol-level capabilities it actually implements. +""" +from typing import Any, Dict, Optional + +import mcp.types as types +from mcp.server.lowlevel import Server as LowLevelServer +from mcp.server.lowlevel.server import NotificationOptions + + +class NitroStackMcpServer(LowLevelServer): + """ + Thin subclass of the official low-level MCP `Server`. + + The base `Server.get_capabilities()` only advertises a capability when a + handler for the corresponding request type has been registered, and it + always reports `resources.subscribe=False`. This subclass declares the + same static capability set the TypeScript SDK's `mcpServerOptions` uses: + `listChanged=True` for tools/resources/prompts, `resources.subscribe=True` + (nitrostack always registers subscribe/unsubscribe handlers), and the + `tasks` capability whenever nitrostack's task subsystem handlers are + registered on this server. + """ + + def __init__(self, name: str, version: Optional[str] = None): + super().__init__(name=name, version=version) + self.has_task_support: bool = False + + def get_capabilities( + self, + notification_options: NotificationOptions, + experimental_capabilities: Dict[str, Dict[str, Any]], + ) -> types.ServerCapabilities: + caps = super().get_capabilities(notification_options, experimental_capabilities) + + if caps.resources is not None: + caps.resources.subscribe = True + + if self.has_task_support: + caps.tasks = types.ServerTasksCapability( + list=types.TasksListCapability(), + cancel=types.TasksCancelCapability(), + requests=types.ServerTasksRequestsCapability( + tools=types.TasksToolsCapability(call=types.TasksCallCapability()) + ), + ) + + return caps + + def create_initialization_options(self, **kwargs: Any): + # Match TS's static `listChanged: true` for tools/resources/prompts by + # default, unless the caller explicitly supplies notification_options. + kwargs.setdefault( + "notification_options", + NotificationOptions( + prompts_changed=True, + resources_changed=True, + tools_changed=True, + ), + ) + return super().create_initialization_options(**kwargs) diff --git a/nitrostack/static/app.js b/nitrostack/static/app.js index da510da..eb329b4 100644 --- a/nitrostack/static/app.js +++ b/nitrostack/static/app.js @@ -412,17 +412,30 @@ function generateFormFromSchema(schema) { // We resolve the Pydantic properties schema if (!schema) return; - const defs = schema.$defs || {}; + const defs = schema.$defs || schema.definitions || {}; let properties = {}; let required = []; - // Pydantic wraps fields under properties.input referencing definition model - if (schema.properties && schema.properties.input && schema.properties.input.$ref) { - const refPath = schema.properties.input.$ref.split('/').pop(); - const model = defs[refPath] || {}; - properties = model.properties || {}; - required = model.required || []; - } else if (schema.properties) { + // NitroStack tools wrap the real Pydantic model under `properties.input`. + // That may be a `$ref` into `$defs` (preferred) or an inlined object schema. + // Either way, expand the nested model so Studio shows real fields (a, b, ...) + // instead of a single "input" JSON blob. + if (schema.properties && schema.properties.input) { + const inputProp = schema.properties.input; + let model = null; + if (inputProp.$ref) { + const refPath = inputProp.$ref.split('/').pop(); + model = defs[refPath] || {}; + } else if (inputProp.properties) { + model = inputProp; + } + if (model) { + properties = model.properties || {}; + required = model.required || []; + } + } + + if (!Object.keys(properties).length && schema.properties) { properties = schema.properties || {}; required = schema.required || []; } diff --git a/nitrostack/testing/__init__.py b/nitrostack/testing/__init__.py index 9b3486b..7899d4b 100644 --- a/nitrostack/testing/__init__.py +++ b/nitrostack/testing/__init__.py @@ -1,10 +1,19 @@ -from typing import Type, Any +import json +from typing import Type, Any, Dict, Optional + +import mcp.types as types + from nitrostack.core.app import mcp_app, McpApplicationFactory, ServerConfig, McpApplication + class NitroTestingModule: """ In-process test harness for testing NitroStack applications without spinning up real transports or subprocesses (Section 15). + + Dispatches directly through the owned low-level `mcp.server.lowlevel.Server`'s + registered `request_handlers` (the same dict the real stdio/HTTP transports use), + rather than any FastMCP-specific convenience method. """ @classmethod async def create(cls, app_module: Type) -> "NitroTestingModule": @@ -19,68 +28,65 @@ class TestApp: def __init__(self, app: McpApplication): self.app = app - async def call_tool(self, name: str, arguments: dict) -> Any: + @staticmethod + def _extract_text(text_val: Optional[str]) -> Any: + if text_val is None: + return None + try: + return json.loads(text_val) + except Exception: + return text_val + + async def call_tool(self, name: str, arguments: Dict[str, Any]) -> Any: """Calls a tool by name in the test harness, returning the raw or deserialized result.""" if not self.app.mcp_server: raise RuntimeError("Application has not been bootstrapped.") - res = await self.app.mcp_server.call_tool(name, arguments) - - # FastMCP returns a list of ContentBlocks, or a tuple of (list[ContentBlock], meta) - content_list = res - if isinstance(res, tuple) and len(res) > 0: - content_list = res[0] - - if isinstance(content_list, list) and len(content_list) > 0: - block = content_list[0] - if hasattr(block, "text"): - text_val = block.text - import json - try: - return json.loads(text_val) - except Exception: - return text_val - return res + + handler = self.app.mcp_server.request_handlers[types.CallToolRequest] + request = types.CallToolRequest( + method="tools/call", + params=types.CallToolRequestParams(name=name, arguments=arguments), + ) + response = await handler(request) + result = response.root + + content_list = getattr(result, "content", None) or [] + if content_list and hasattr(content_list[0], "text"): + extracted = self._extract_text(content_list[0].text) + if extracted is not None: + return extracted + return result async def read_resource(self, uri: str) -> Any: """Reads a resource by URI in the test harness, returning raw text or deserialized JSON.""" if not self.app.mcp_server: raise RuntimeError("Application has not been bootstrapped.") - res = await self.app.mcp_server.read_resource(uri) - - # FastMCP returns a list of ReadResourceContents or a result object - content_list = [] - res_obj = res[0] if isinstance(res, tuple) and len(res) > 0 else res - - if isinstance(res_obj, list): - content_list = res_obj - elif hasattr(res_obj, "contents"): - content_list = res_obj.contents - - if content_list and len(content_list) > 0: - content_block = content_list[0] - text_val = None - if hasattr(content_block, "content"): - text_val = content_block.content - elif hasattr(content_block, "text"): - text_val = content_block.text - - if text_val is not None: - import json - try: - return json.loads(text_val) - except Exception: - return text_val - return res - - async def get_prompt(self, name: str, arguments: dict) -> Any: + + handler = self.app.mcp_server.request_handlers[types.ReadResourceRequest] + request = types.ReadResourceRequest( + method="resources/read", + params=types.ReadResourceRequestParams(uri=uri), + ) + response = await handler(request) + result = response.root + + content_list = result.contents or [] + if content_list and hasattr(content_list[0], "text"): + extracted = self._extract_text(content_list[0].text) + if extracted is not None: + return extracted + return result + + async def get_prompt(self, name: str, arguments: Dict[str, str]) -> Any: """Retrieves a prompt by name in the test harness, returning prompt messages.""" if not self.app.mcp_server: raise RuntimeError("Application has not been bootstrapped.") - res = await self.app.mcp_server.get_prompt(name, arguments) - - res_obj = res[0] if isinstance(res, tuple) and len(res) > 0 else res - - # FastMCP get_prompt returns a GetPromptResult which has messages list - if hasattr(res_obj, "messages"): - return res_obj.messages - return res + + handler = self.app.mcp_server.request_handlers[types.GetPromptRequest] + request = types.GetPromptRequest( + method="prompts/get", + params=types.GetPromptRequestParams(name=name, arguments=arguments), + ) + response = await handler(request) + result = response.root + return result.messages diff --git a/pyproject.toml b/pyproject.toml index 55d134e..9029cf2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -9,7 +9,7 @@ description = "NitroStack Python SDK — A Python-idiomatic port of the NitroSta readme = "README.md" requires-python = ">=3.10" dependencies = [ - "mcp>=1.0.0", + "mcp>=1.0.0,<2.0.0", "pydantic>=2.0.0", "starlette>=0.30.0", "uvicorn>=0.20.0", diff --git a/requirements.txt b/requirements.txt index aecb838..e7d3739 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,4 @@ -mcp>=1.0.0 +mcp>=1.0.0,<2.0.0 pydantic>=2.0.0 starlette>=0.30.0 uvicorn>=0.20.0 diff --git a/tests/test_initial_tool.py b/tests/test_initial_tool.py index d105d98..e9d494d 100644 --- a/tests/test_initial_tool.py +++ b/tests/test_initial_tool.py @@ -46,8 +46,8 @@ async def _test_initial_tool_hook(): # Check that called is False initially assert ExecutionState.called is False - # Retrieve LowLevelServer - server = harness.app.mcp_server._mcp_server + # Retrieve the owned low-level Server directly (no FastMCP wrapper in between) + server = harness.app.mcp_server # Get the InitializedNotification handler handler = server.notification_handlers[types.InitializedNotification] diff --git a/tests/test_tasks.py b/tests/test_tasks.py index 25ade6e..c5abc94 100644 --- a/tests/test_tasks.py +++ b/tests/test_tasks.py @@ -10,6 +10,7 @@ from nitrostack import module, injectable, tool, ExecutionContext, NitroTestingModule import mcp.types as types from mcp.server.lowlevel.server import request_ctx, RequestContext +from mcp.server.experimental.request_context import Experimental class AsyncTaskInput(BaseModel): duration: float @@ -56,7 +57,7 @@ async def _test_all_tasks(): ) ) - handler = harness.app.mcp_server._mcp_server.request_handlers[types.CallToolRequest] + handler = harness.app.mcp_server.request_handlers[types.CallToolRequest] # Manually set request_ctx to simulate LowLevelServer._handle_request token = request_ctx.set(RequestContext( @@ -64,6 +65,7 @@ async def _test_all_tasks(): meta=None, session=None, lifespan_context=None, + experimental=Experimental(task_metadata=req.params.task), request=req )) try: @@ -82,7 +84,7 @@ async def _test_all_tasks(): method="tasks/list", params=types.PaginatedRequestParams() ) - list_handler = harness.app.mcp_server._mcp_server.request_handlers[types.ListTasksRequest] + list_handler = harness.app.mcp_server.request_handlers[types.ListTasksRequest] list_res = await list_handler(list_req) print("Tasks list count:", len(list_res.tasks)) assert len(list_res.tasks) >= 1 @@ -94,7 +96,7 @@ async def _test_all_tasks(): method="tasks/get", params=types.GetTaskRequestParams(taskId=task_id) ) - get_handler = harness.app.mcp_server._mcp_server.request_handlers[types.GetTaskRequest] + get_handler = harness.app.mcp_server.request_handlers[types.GetTaskRequest] get_res = await get_handler(get_req) print("Task status:", get_res.status) assert get_res.status == "working" @@ -105,7 +107,7 @@ async def _test_all_tasks(): method="tasks/result", params=types.GetTaskPayloadRequestParams(taskId=task_id) ) - result_handler = harness.app.mcp_server._mcp_server.request_handlers[types.GetTaskPayloadRequest] + result_handler = harness.app.mcp_server.request_handlers[types.GetTaskPayloadRequest] result_res = await result_handler(result_req) print("Result response content:", result_res.content) assert result_res.isError is False @@ -132,6 +134,7 @@ async def _test_all_tasks(): meta=None, session=None, lifespan_context=None, + experimental=Experimental(task_metadata=req_cancel.params.task), request=req_cancel )) try: @@ -155,7 +158,7 @@ async def _test_all_tasks(): method="tasks/cancel", params=types.CancelTaskRequestParams(taskId=task_id_cancel) ) - cancel_handler = harness.app.mcp_server._mcp_server.request_handlers[types.CancelTaskRequest] + cancel_handler = harness.app.mcp_server.request_handlers[types.CancelTaskRequest] cancel_res = await cancel_handler(cancel_req) print("Cancellation response status:", cancel_res.status) assert cancel_res.status == "cancelled" diff --git a/tests/test_widget_metadata.py b/tests/test_widget_metadata.py index 1860c60..302186d 100644 --- a/tests/test_widget_metadata.py +++ b/tests/test_widget_metadata.py @@ -8,6 +8,7 @@ from nitrostack import injectable, tool, widget, module, ExecutionContext from nitrostack.testing import NitroTestingModule from pydantic import BaseModel +import mcp.types as types class DummyInput(BaseModel): pass @@ -36,9 +37,12 @@ async def main(): # 1. Initialize test harness harness = await NitroTestingModule.create(WidgetTestModule) - # 2. Extract tools from FastMCP server - tools = await harness.app.mcp_server.list_tools() - + # 2. Extract tools by invoking the registered `tools/list` handler directly + # (the owned low-level Server has no FastMCP-style `list_tools()` convenience method) + list_tools_handler = harness.app.mcp_server.request_handlers[types.ListToolsRequest] + list_result = await list_tools_handler(None) + tools = list_result.root.tools + # Find our tool target_tool = None for t in tools: