diff --git a/Server/src/main.py b/Server/src/main.py index ba8ab8c82..04b7ccd1a 100644 --- a/Server/src/main.py +++ b/Server/src/main.py @@ -15,10 +15,11 @@ resolve_project_id_for_unity_instance, ) from core.config import config -from starlette.routing import WebSocketRoute +from starlette.routing import Route, WebSocketRoute from starlette.responses import JSONResponse import argparse import asyncio +import uvicorn # Fix to IPV4 Connection Issue #853 # Will disable features in ProactorEventLoop including subprocess pipes and named pipes @@ -304,8 +305,10 @@ def _build_instructions(project_scoped_tools: bool) -> str: Targeting Unity instances: - Use the resource mcpforunity://instances to list active Unity sessions (Name@hash). -- When multiple instances are connected, call set_active_instance with the exact Name@hash before using tools/resources to pin routing for the whole session. The server will error if multiple are connected and no active instance is set. -- Alternatively, pass unity_instance as a parameter on any individual tool call to route just that call (e.g. unity_instance="MyGame@abc123", unity_instance="abc" for a hash prefix, or unity_instance="6401" for a port number in stdio mode). This does not change the session default. +- The preferred flow is to pass unity_instance on any tool call or resource read to route that request explicitly (e.g. unity_instance="MyGame@abc123", unity_instance="abc" for a hash prefix, or unity_instance="6401" for a port number in stdio mode). +- HTTP clients can bind a whole MCP session to one Unity editor by connecting to /mcp/instance/{{Name@hash}} instead of the unbound /mcp endpoint. +- set_active_instance remains available as a compatibility path for simple single-agent workflows, but it should not be the default multi-agent HTTP pattern. +- When multiple instances are connected and no explicit target or compatibility default is set, Unity-backed calls will return a selection error instead of guessing. Important Workflows: @@ -366,6 +369,62 @@ def _normalize_instance_token(instance_token: str | None) -> tuple[str | None, s return None, instance_token +def _build_bound_mcp_path(base_path: str) -> str: + normalized = "/" + base_path.strip("/") + if normalized == "//": + normalized = "/" + return f"{normalized.rstrip('/')}/instance/{{instance:path}}" + + +def _attach_bound_mcp_route(app) -> str | None: + """Expose an HTTP MCP endpoint that binds a session to one Unity instance.""" + base_path = getattr(app.state, "path", None) + if not isinstance(base_path, str) or not base_path: + return None + + bound_path = _build_bound_mcp_path(base_path) + existing_paths = { + getattr(route, "path", None) + for route in app.router.routes + } + if bound_path in existing_paths: + return bound_path + + base_route = next( + ( + route for route in app.router.routes + if isinstance(route, Route) and route.path == base_path + ), + None, + ) + if base_route is None: + return None + + app.router.routes.append( + Route( + bound_path, + endpoint=base_route.endpoint, + methods=base_route.methods, + name=f"{base_route.name or 'mcp'}_bound_instance", + include_in_schema=False, + ) + ) + return bound_path + + +def create_http_app( + mcp: FastMCP, + *, + transport: str = "http", + path: str | None = None, +): + app = mcp.http_app(path=path, transport=transport) + bound_path = _attach_bound_mcp_route(app) + if bound_path: + logger.info("Registered instance-bound MCP HTTP endpoint at %s", bound_path) + return app + + def create_mcp_server(project_scoped_tools: bool) -> FastMCP: mcp = FastMCP( name="mcp-for-unity-server", @@ -891,16 +950,21 @@ def main(): # Determine transport mode if config.transport_mode == 'http': - # Use HTTP transport for FastMCP - transport = 'http' - # Use the parsed host and port from URL/args http_url = os.environ.get("UNITY_MCP_HTTP_URL", args.http_url) parsed_url = urlparse(http_url) host = args.http_host or os.environ.get( "UNITY_MCP_HTTP_HOST") or parsed_url.hostname or "127.0.0.1" port = args.http_port or _env_port or parsed_url.port or 8080 logger.info(f"Starting FastMCP with HTTP transport on {host}:{port}") - mcp.run(transport=transport, host=host, port=port) + app = create_http_app(mcp, transport="http") + uvicorn.run( + app, + host=host, + port=port, + timeout_graceful_shutdown=0, + lifespan="on", + ws="websockets-sansio", + ) else: # Use stdio transport for traditional MCP logger.info("Starting FastMCP with stdio transport") diff --git a/Server/src/services/custom_tool_service.py b/Server/src/services/custom_tool_service.py index 7407d9a9d..77a02330e 100644 --- a/Server/src/services/custom_tool_service.py +++ b/Server/src/services/custom_tool_service.py @@ -374,7 +374,7 @@ async def _handler(ctx: Context, **kwargs) -> MCPResponse: if not unity_instance: return MCPResponse( success=False, - message="No active Unity instance. Call set_active_instance with Name@hash from mcpforunity://instances.", + message="No active Unity instance. Pass unity_instance explicitly or call set_active_instance with Name@hash from mcpforunity://instances.", ) project_id = resolve_project_id_for_unity_instance(unity_instance) diff --git a/Server/src/services/registry/resource_registry.py b/Server/src/services/registry/resource_registry.py index 5c8e42607..f13787003 100644 --- a/Server/src/services/registry/resource_registry.py +++ b/Server/src/services/registry/resource_registry.py @@ -3,6 +3,11 @@ """ from typing import Callable, Any +from .unity_targeting import ( + add_optional_unity_instance_parameter, + append_unity_instance_query_template, +) + # Global registry to collect decorated resources _resource_registry: list[dict[str, Any]] = [] @@ -11,6 +16,7 @@ def mcp_for_unity_resource( uri: str, name: str | None = None, description: str | None = None, + unity_target: bool = True, **kwargs ) -> Callable: """ @@ -30,15 +36,21 @@ async def my_custom_resource(ctx: Context, ...): """ def decorator(func: Callable) -> Callable: resource_name = name if name is not None else func.__name__ + registered_func = func + registered_uri = uri + if unity_target: + registered_func = add_optional_unity_instance_parameter(func) + registered_uri = append_unity_instance_query_template(uri) _resource_registry.append({ - 'func': func, - 'uri': uri, + 'func': registered_func, + 'uri': registered_uri, 'name': resource_name, 'description': description, + 'unity_target': unity_target, 'kwargs': kwargs }) - return func + return registered_func return decorator diff --git a/Server/src/services/registry/tool_registry.py b/Server/src/services/registry/tool_registry.py index cfe147cde..6f6374e25 100644 --- a/Server/src/services/registry/tool_registry.py +++ b/Server/src/services/registry/tool_registry.py @@ -11,6 +11,8 @@ """ from typing import Callable, Any +from .unity_targeting import add_optional_unity_instance_parameter + # Global registry to collect decorated tools _tool_registry: list[dict[str, Any]] = [] @@ -94,8 +96,12 @@ def decorator(func: Callable) -> Callable: "Expected None or a non-empty string." ) + registered_func = func + if normalized_unity_target is not None: + registered_func = add_optional_unity_instance_parameter(func) + _tool_registry.append({ - 'func': func, + 'func': registered_func, 'name': tool_name, 'description': description, 'unity_target': normalized_unity_target, @@ -103,7 +109,7 @@ def decorator(func: Callable) -> Callable: 'kwargs': tool_kwargs, }) - return func + return registered_func return decorator diff --git a/Server/src/services/registry/unity_targeting.py b/Server/src/services/registry/unity_targeting.py new file mode 100644 index 000000000..31043b361 --- /dev/null +++ b/Server/src/services/registry/unity_targeting.py @@ -0,0 +1,70 @@ +from __future__ import annotations + +import functools +import inspect +from typing import Annotated, Any, Callable + +UnityInstanceParameter = Annotated[ + str | None, + "Target Unity instance (Name@hash, hash prefix, or port number in stdio mode).", +] + + +def add_optional_unity_instance_parameter( + func: Callable[..., Any], +) -> Callable[..., Any]: + """Expose an optional unity_instance kwarg without forwarding it downstream.""" + signature = inspect.signature(func) + if "unity_instance" in signature.parameters: + return func + + parameters = list(signature.parameters.values()) + unity_parameter = inspect.Parameter( + "unity_instance", + inspect.Parameter.KEYWORD_ONLY, + default=None, + annotation=UnityInstanceParameter, + ) + + insert_at = len(parameters) + for index, parameter in enumerate(parameters): + if parameter.kind == inspect.Parameter.VAR_KEYWORD: + insert_at = index + break + parameters.insert(insert_at, unity_parameter) + + wrapped_signature = signature.replace(parameters=parameters) + + if inspect.iscoroutinefunction(func): + @functools.wraps(func) + async def wrapper(*args, **kwargs): + kwargs.pop("unity_instance", None) + return await func(*args, **kwargs) + else: + @functools.wraps(func) + def wrapper(*args, **kwargs): + kwargs.pop("unity_instance", None) + return func(*args, **kwargs) + + wrapper.__signature__ = wrapped_signature + wrapper.__annotations__ = { + **getattr(func, "__annotations__", {}), + "unity_instance": UnityInstanceParameter, + } + return wrapper + + +def append_unity_instance_query_template(uri: str) -> str: + """Append unity_instance to a resource URI template if it is not present.""" + if "unity_instance" in uri: + return uri + + if "{?" not in uri: + return f"{uri}{{?unity_instance}}" + + prefix, _, suffix = uri.partition("{?") + query_suffix = suffix[:-1] if suffix.endswith("}") else suffix + query_names = [name for name in query_suffix.split(",") if name] + if "unity_instance" not in query_names: + query_names.append("unity_instance") + return f"{prefix}{{?{','.join(query_names)}}}" diff --git a/Server/src/services/resources/custom_tools.py b/Server/src/services/resources/custom_tools.py index aacda462f..54f024dd9 100644 --- a/Server/src/services/resources/custom_tools.py +++ b/Server/src/services/resources/custom_tools.py @@ -32,7 +32,7 @@ async def get_custom_tools(ctx: Context) -> CustomToolsResourceResponse | MCPRes if not unity_instance: return MCPResponse( success=False, - message="No active Unity instance. Call set_active_instance with Name@hash from mcpforunity://instances.", + message="No active Unity instance. Pass unity_instance explicitly or call set_active_instance with Name@hash from mcpforunity://instances.", ) project_id = resolve_project_id_for_unity_instance(unity_instance) diff --git a/Server/src/services/resources/gameobject.py b/Server/src/services/resources/gameobject.py index d4e1a2537..8f51b2367 100644 --- a/Server/src/services/resources/gameobject.py +++ b/Server/src/services/resources/gameobject.py @@ -42,6 +42,7 @@ def _validate_instance_id(instance_id: str) -> tuple[int | None, MCPResponse | N @mcp_for_unity_resource( uri="mcpforunity://scene/gameobject-api", name="gameobject_api", + unity_target=False, description="Documentation for GameObject resources. Use find_gameobjects tool to get instance IDs, then access resources below.\n\nURI: mcpforunity://scene/gameobject-api" ) async def get_gameobject_api_docs(_ctx: Context) -> MCPResponse: diff --git a/Server/src/services/resources/prefab.py b/Server/src/services/resources/prefab.py index 8e35bb82c..714cb4910 100644 --- a/Server/src/services/resources/prefab.py +++ b/Server/src/services/resources/prefab.py @@ -43,6 +43,7 @@ def _decode_prefab_path(encoded_path: str) -> str: @mcp_for_unity_resource( uri="mcpforunity://prefab-api", name="prefab_api", + unity_target=False, description="Documentation for Prefab resources. Use manage_asset action=search filterType=Prefab to find prefabs, then access resources below.\n\nURI: mcpforunity://prefab-api" ) async def get_prefab_api_docs(_ctx: Context) -> MCPResponse: diff --git a/Server/src/services/resources/tool_groups.py b/Server/src/services/resources/tool_groups.py index 669973d60..cb4d7c54b 100644 --- a/Server/src/services/resources/tool_groups.py +++ b/Server/src/services/resources/tool_groups.py @@ -18,6 +18,7 @@ @mcp_for_unity_resource( uri="mcpforunity://tool-groups", name="tool_groups", + unity_target=False, description=( "Available tool groups and their tools. " "Use manage_tools to activate/deactivate groups per session.\n\n" diff --git a/Server/src/services/resources/unity_instances.py b/Server/src/services/resources/unity_instances.py index 42ef4af5f..9bfb5f5bb 100644 --- a/Server/src/services/resources/unity_instances.py +++ b/Server/src/services/resources/unity_instances.py @@ -13,6 +13,7 @@ @mcp_for_unity_resource( uri="mcpforunity://instances", name="unity_instances", + unity_target=False, description="Lists all running Unity Editor instances with their details.\n\nURI: mcpforunity://instances" ) async def unity_instances(ctx: Context) -> dict[str, Any]: diff --git a/Server/src/services/tools/execute_custom_tool.py b/Server/src/services/tools/execute_custom_tool.py index 60c1fff89..c9db9c532 100644 --- a/Server/src/services/tools/execute_custom_tool.py +++ b/Server/src/services/tools/execute_custom_tool.py @@ -26,7 +26,7 @@ async def execute_custom_tool(ctx: Context, tool_name: str, parameters: dict | N if not unity_instance: return MCPResponse( success=False, - message="No active Unity instance. Call set_active_instance with Name@hash from mcpforunity://instances.", + message="No active Unity instance. Pass unity_instance explicitly or call set_active_instance with Name@hash from mcpforunity://instances.", ) project_id = resolve_project_id_for_unity_instance(unity_instance) diff --git a/Server/src/services/tools/set_active_instance.py b/Server/src/services/tools/set_active_instance.py index 6b90b351d..bfbb28c03 100644 --- a/Server/src/services/tools/set_active_instance.py +++ b/Server/src/services/tools/set_active_instance.py @@ -23,6 +23,17 @@ async def set_active_instance( ctx: Context, instance: Annotated[str, "Target instance (Name@hash, hash prefix, or port number in stdio mode)"] ) -> dict[str, Any]: + get_state = getattr(ctx, "get_state", None) + bound_instance = await get_state("bound_unity_instance") if callable(get_state) else None + if bound_instance: + return { + "success": False, + "error": ( + "set_active_instance is not available on an instance-bound MCP endpoint. " + "Use the bound endpoint target or switch back to the unbound /mcp endpoint." + ), + } + transport = (config.transport_mode or "stdio").lower() # Port number shorthand (stdio only) — resolve to Name@hash via pool discovery diff --git a/Server/src/transport/legacy/unity_connection.py b/Server/src/transport/legacy/unity_connection.py index c62e20fb1..d599aa09c 100644 --- a/Server/src/transport/legacy/unity_connection.py +++ b/Server/src/transport/legacy/unity_connection.py @@ -549,16 +549,27 @@ def _resolve_instance_id(self, instance_identifier: str | None, instances: list[ instance_identifier = self._default_instance_id logger.debug(f"Using default instance: {instance_identifier}") else: - # Use the most recently active instance - # Instances with no heartbeat (None) should be sorted last (use 0 as sentinel) - sorted_instances = sorted( - instances, - key=lambda inst: inst.last_heartbeat.timestamp() if inst.last_heartbeat else 0.0, - reverse=True, + if len(instances) == 1: + logger.info( + "No instance specified, auto-selecting sole instance: %s", + instances[0].id, + ) + return instances[0] + + suggestions = [ + { + "id": inst.id, + "path": inst.path, + "port": inst.port, + "suggest": f"Use unity_instance='{inst.id}'", + } + for inst in instances + ] + raise ConnectionError( + "Multiple Unity Editor instances are running. " + "Pass unity_instance explicitly or call set_active_instance first. " + f"Available instances: {suggestions}" ) - logger.info( - f"No instance specified, using most recent: {sorted_instances[0].id}") - return sorted_instances[0] identifier = instance_identifier.strip() diff --git a/Server/src/transport/plugin_hub.py b/Server/src/transport/plugin_hub.py index 1a71f30a7..a57bd724f 100644 --- a/Server/src/transport/plugin_hub.py +++ b/Server/src/transport/plugin_hub.py @@ -76,11 +76,11 @@ class InstanceSelectionRequiredError(RuntimeError): _SELECTION_REQUIRED = ( "Unity instance selection is required. " - "Call set_active_instance with Name@hash from mcpforunity://instances." + "Pass unity_instance explicitly or call set_active_instance with Name@hash from mcpforunity://instances." ) _MULTIPLE_INSTANCES = ( "Multiple Unity instances are connected. " - "Call set_active_instance with Name@hash from mcpforunity://instances." + "Pass unity_instance explicitly or call set_active_instance with Name@hash from mcpforunity://instances." ) def __init__(self, message: str | None = None): diff --git a/Server/src/transport/unity_instance_middleware.py b/Server/src/transport/unity_instance_middleware.py index 9c367ef46..4e3913f3f 100644 --- a/Server/src/transport/unity_instance_middleware.py +++ b/Server/src/transport/unity_instance_middleware.py @@ -7,6 +7,7 @@ from threading import RLock import logging import time +from urllib.parse import parse_qs, urlparse from fastmcp.server.middleware import Middleware, MiddlewareContext @@ -47,6 +48,19 @@ def set_unity_instance_middleware(middleware: 'UnityInstanceMiddleware') -> None _unity_instance_middleware = middleware +def _get_http_request_for_binding(): + """Resolve the current HTTP request using whichever FastMCP layout is installed.""" + try: + from fastmcp.dependencies import get_http_request + except Exception: + try: + from fastmcp.server.dependencies import get_http_request + except Exception as exc: # pragma: no cover - import-layout compatibility + raise RuntimeError("FastMCP HTTP request dependency is unavailable") from exc + + return get_http_request() + + class UnityInstanceMiddleware(Middleware): """ Middleware that manages per-session Unity instance selection. @@ -333,10 +347,8 @@ async def _resolve_user_id(self) -> str | None: from transport.unity_transport import _resolve_user_id_from_request return await _resolve_user_id_from_request() - async def _inject_unity_instance(self, context: MiddlewareContext) -> None: - """Inject active Unity instance and user_id into context if available.""" - ctx = context.fastmcp_context - + async def _inject_user_id(self, ctx) -> str | None: + """Inject request-scoped user_id when the HTTP transport requires it.""" # Resolve user_id from the HTTP request's API key header user_id = await self._resolve_user_id() if config.http_remote_hosted and user_id is None: @@ -345,25 +357,86 @@ async def _inject_unity_instance(self, context: MiddlewareContext) -> None: ) if user_id: await ctx.set_state("user_id", user_id) + return user_id + + @staticmethod + def _normalize_requested_instance(raw_value: object) -> str | None: + if raw_value is None: + return None + value = str(raw_value).strip() + return value or None + + def _extract_requested_instance(self, context: MiddlewareContext) -> str | None: + """Read an inline unity_instance override from tool args or resource query.""" + message = getattr(context, "message", None) + + message_arguments = getattr(message, "arguments", None) + if isinstance(message_arguments, dict) and "unity_instance" in message_arguments: + return self._normalize_requested_instance(message_arguments.pop("unity_instance")) + + message_uri = getattr(message, "uri", None) + if isinstance(message_uri, str) and message_uri: + query_values = parse_qs(urlparse(message_uri).query).get("unity_instance") + if query_values: + return self._normalize_requested_instance(query_values[-1]) + + return None + + def _extract_bound_instance(self) -> str | None: + """Read the fixed Unity target from an instance-bound HTTP endpoint.""" + try: + request = _get_http_request_for_binding() + except RuntimeError: + return None + + path_params = getattr(request, "path_params", None) + if isinstance(path_params, dict): + return self._normalize_requested_instance(path_params.get("instance")) + return None + + async def _inject_bound_unity_instance(self, ctx) -> str | None: + """Resolve and store a bound HTTP endpoint target for the current request.""" + bound_token = self._extract_bound_instance() + if not bound_token: + return None + + bound_instance = await self._resolve_instance_value(bound_token, ctx) + await ctx.set_state("bound_unity_instance", bound_instance) + return bound_instance + + async def _inject_unity_instance( + self, + context: MiddlewareContext, + *, + include_legacy_default: bool = True, + allow_autoselect: bool = True, + ) -> None: + """Inject active Unity instance and user_id into context if available.""" + ctx = context.fastmcp_context + + user_id = await self._inject_user_id(ctx) + bound_instance = await self._inject_bound_unity_instance(ctx) # Per-call routing: check if this tool call explicitly specifies unity_instance. # context.message.arguments is a mutable dict on CallToolRequestParams; resource - # reads use ReadResourceRequestParams which has no .arguments, so this is a no-op for them. - # We pop the key here so Pydantic's type_adapter.validate_python() never sees it. + # reads use ReadResourceRequestParams, so resource overrides come from the URI query. active_instance: str | None = None - msg_args = getattr(getattr(context, "message", None), "arguments", None) - if isinstance(msg_args, dict) and "unity_instance" in msg_args: - raw = msg_args.pop("unity_instance") - if raw is not None: - raw_str = str(raw).strip() - if raw_str: - # Raises ValueError with a user-friendly message on invalid input. - active_instance = await self._resolve_instance_value(raw_str, ctx) - logger.debug("Per-call unity_instance resolved to: %s", active_instance) + requested_instance = self._extract_requested_instance(context) + if bound_instance and requested_instance: + raise ValueError( + "Bound MCP endpoints do not accept unity_instance overrides. " + "Use the bound endpoint target or switch back to the unbound /mcp endpoint." + ) + if requested_instance: + # Raises ValueError with a user-friendly message on invalid input. + active_instance = await self._resolve_instance_value(requested_instance, ctx) + logger.debug("Per-call unity_instance resolved to: %s", active_instance) if not active_instance: + active_instance = bound_instance + if not active_instance and include_legacy_default: active_instance = await self.get_active_instance(ctx) - if not active_instance: + if not active_instance and allow_autoselect: active_instance = await self._maybe_autoselect_instance(ctx) if active_instance: # If using HTTP transport (PluginHub configured), validate session @@ -422,15 +495,19 @@ async def on_read_resource(self, context: MiddlewareContext, call_next): async def on_list_tools(self, context: MiddlewareContext, call_next): """Filter MCP tool listing to the Unity-enabled set when session data is available.""" try: - await self._inject_unity_instance(context) + await self._inject_user_id(context.fastmcp_context) + bound_instance = await self._inject_bound_unity_instance( + context.fastmcp_context, + ) except Exception as exc: # Re-raise authentication errors so callers get a proper auth failure if isinstance(exc, RuntimeError) and "authentication" in str(exc).lower(): raise _diag.warning( - "on_list_tools: _inject_unity_instance failed (%s: %s), continuing without instance", + "on_list_tools: user_id injection failed (%s: %s), continuing without instance", type(exc).__name__, exc, ) + bound_instance = None tools = await call_next(context) @@ -445,7 +522,10 @@ async def on_list_tools(self, context: MiddlewareContext, call_next): return tools self._refresh_tool_visibility_metadata_from_registry() - enabled_tool_names = await self._resolve_enabled_tool_names_for_context(context) + enabled_tool_names = await self._resolve_enabled_tool_names_for_context( + context, + active_instance=bound_instance, + ) if enabled_tool_names is None: _diag.debug("on_list_tools: no Unity session data, returning %d tools from FastMCP as-is", len(tools)) return tools @@ -470,10 +550,10 @@ def _should_filter_tool_listing(self) -> bool: async def _resolve_enabled_tool_names_for_context( self, context: MiddlewareContext, + active_instance: str | None, ) -> set[str] | None: ctx = context.fastmcp_context user_id = (await ctx.get_state("user_id")) if config.http_remote_hosted else None - active_instance = await ctx.get_state("unity_instance") project_hashes = self._resolve_candidate_project_hashes(active_instance) try: sessions_data = await PluginHub.get_sessions(user_id=user_id) diff --git a/Server/src/transport/unity_transport.py b/Server/src/transport/unity_transport.py index f3b1d3d3f..5cb9e193e 100644 --- a/Server/src/transport/unity_transport.py +++ b/Server/src/transport/unity_transport.py @@ -26,7 +26,10 @@ async def _resolve_user_id_from_request() -> str | None: if not ApiKeyService.is_initialized(): return None try: - from fastmcp.server.dependencies import get_http_headers + try: + from fastmcp.dependencies import get_http_headers + except ImportError: # pragma: no cover - compatibility with older FastMCP layouts + from fastmcp.server.dependencies import get_http_headers headers = get_http_headers(include_all=True) api_key = headers.get(API_KEY_HEADER.lower()) if not api_key: diff --git a/Server/tests/integration/test_inline_unity_instance.py b/Server/tests/integration/test_inline_unity_instance.py index 280ad47f9..130fbb244 100644 --- a/Server/tests/integration/test_inline_unity_instance.py +++ b/Server/tests/integration/test_inline_unity_instance.py @@ -355,6 +355,80 @@ async def test_resource_read_unaffected(monkeypatch): assert await ctx.get_state("unity_instance") == "Proj@abc123" +@pytest.mark.asyncio +async def test_resource_query_unity_instance_overrides_session(monkeypatch): + """Resource URI query targeting should override the legacy session selection for that read only.""" + instances = [ + SimpleNamespace(id="ProjA@aaa111", hash="aaa111"), + SimpleNamespace(id="ProjB@bbb222", hash="bbb222"), + ] + mw = _make_middleware(monkeypatch, pool_instances=instances) + + ctx = DummyContext() + ctx.client_id = "client-1" + await mw.set_active_instance(ctx, "ProjA@aaa111") + + resource_ctx = SimpleNamespace( + fastmcp_context=ctx, + message=SimpleNamespace( + uri="mcpforunity://scene/gameobject/-1?unity_instance=bbb222", + ), + ) + + await mw._inject_unity_instance(resource_ctx) + + assert await ctx.get_state("unity_instance") == "ProjB@bbb222" + assert await mw.get_active_instance(ctx) == "ProjA@aaa111" + + +@pytest.mark.asyncio +async def test_bound_endpoint_routes_to_bound_instance(monkeypatch): + """A bound endpoint should force the request onto its path-selected Unity instance.""" + instances = [ + SimpleNamespace(id="ProjA@aaa111", hash="aaa111"), + SimpleNamespace(id="ProjB@bbb222", hash="bbb222"), + ] + mw = _make_middleware(monkeypatch, pool_instances=instances) + + monkeypatch.setattr( + "transport.unity_instance_middleware._get_http_request_for_binding", + lambda: SimpleNamespace(path_params={"instance": "bbb222"}), + ) + + ctx = DummyContext() + ctx.client_id = "client-1" + await mw.set_active_instance(ctx, "ProjA@aaa111") + + tool_ctx = DummyMiddlewareContext(ctx, arguments={}) + await mw._inject_unity_instance(tool_ctx) + + assert await ctx.get_state("bound_unity_instance") == "ProjB@bbb222" + assert await ctx.get_state("unity_instance") == "ProjB@bbb222" + assert await mw.get_active_instance(ctx) == "ProjA@aaa111" + + +@pytest.mark.asyncio +async def test_bound_endpoint_rejects_inline_unity_instance_override(monkeypatch): + """Bound endpoints should reject per-call unity_instance overrides.""" + instances = [ + SimpleNamespace(id="ProjA@aaa111", hash="aaa111"), + SimpleNamespace(id="ProjB@bbb222", hash="bbb222"), + ] + mw = _make_middleware(monkeypatch, pool_instances=instances) + + monkeypatch.setattr( + "transport.unity_instance_middleware._get_http_request_for_binding", + lambda: SimpleNamespace(path_params={"instance": "aaa111"}), + ) + + ctx = DummyContext() + ctx.client_id = "client-1" + tool_ctx = DummyMiddlewareContext(ctx, arguments={"unity_instance": "bbb222"}) + + with pytest.raises(ValueError, match="Bound MCP endpoints do not accept unity_instance overrides"): + await mw._inject_unity_instance(tool_ctx) + + # --------------------------------------------------------------------------- # set_active_instance tool: port number support # --------------------------------------------------------------------------- @@ -407,6 +481,24 @@ async def test_set_active_instance_port_http_errors(monkeypatch): assert "not supported in HTTP transport mode" in result["error"] +@pytest.mark.asyncio +async def test_set_active_instance_rejects_bound_endpoint(monkeypatch): + """set_active_instance should not mutate routing on a bound endpoint.""" + monkeypatch.setattr(config, "transport_mode", "http") + monkeypatch.setattr(config, "http_remote_hosted", False) + + from services.tools.set_active_instance import set_active_instance + + ctx = DummyContext() + ctx.client_id = "client-1" + await ctx.set_state("bound_unity_instance", "Proj@abc123") + + result = await set_active_instance(ctx, instance="Proj@abc123") + + assert result["success"] is False + assert "instance-bound MCP endpoint" in result["error"] + + # --------------------------------------------------------------------------- # batch_execute rejects inner unity_instance # --------------------------------------------------------------------------- diff --git a/Server/tests/test_main_http_routes.py b/Server/tests/test_main_http_routes.py new file mode 100644 index 000000000..29ae17cfc --- /dev/null +++ b/Server/tests/test_main_http_routes.py @@ -0,0 +1,27 @@ +from types import SimpleNamespace +from starlette.routing import Route + +from main import _attach_bound_mcp_route, _build_bound_mcp_path + + +def test_build_bound_mcp_path_appends_instance_segment(): + assert _build_bound_mcp_path("/mcp") == "/mcp/instance/{instance:path}" + assert _build_bound_mcp_path("mcp") == "/mcp/instance/{instance:path}" + assert _build_bound_mcp_path("/nested/mcp/") == "/nested/mcp/instance/{instance:path}" + + +def test_attach_bound_mcp_route_registers_alias(): + async def _endpoint(_request): + return None + + base_route = Route("/mcp", endpoint=_endpoint, methods=["GET", "POST", "DELETE"]) + app = SimpleNamespace( + state=SimpleNamespace(path="/mcp"), + router=SimpleNamespace(routes=[base_route]), + ) + + bound_path = _attach_bound_mcp_route(app) + route_paths = [getattr(route, "path", None) for route in app.router.routes] + + assert bound_path == "/mcp/instance/{instance:path}" + assert route_paths == ["/mcp", "/mcp/instance/{instance:path}"] diff --git a/Server/tests/test_resource_registry_metadata.py b/Server/tests/test_resource_registry_metadata.py new file mode 100644 index 000000000..034e70c32 --- /dev/null +++ b/Server/tests/test_resource_registry_metadata.py @@ -0,0 +1,65 @@ +import inspect + +import pytest + +from services.registry import get_registered_resources, mcp_for_unity_resource +import services.registry.resource_registry as resource_registry_module + + +@pytest.fixture(autouse=True) +def restore_resource_registry_state(): + original_registry = list(resource_registry_module._resource_registry) + try: + yield + finally: + resource_registry_module._resource_registry[:] = original_registry + + +@pytest.mark.asyncio +async def test_resource_registry_defaults_to_explicit_unity_targeting(): + @mcp_for_unity_resource(uri="mcpforunity://scene/example", name="example_resource") + async def _example_resource(ctx): + return ctx + + registered_resources = get_registered_resources() + resource_info = next(item for item in registered_resources if item["name"] == "example_resource") + + assert resource_info["unity_target"] is True + assert resource_info["uri"] == "mcpforunity://scene/example{?unity_instance}" + + signature = inspect.signature(_example_resource) + assert "unity_instance" in signature.parameters + assert signature.parameters["unity_instance"].default is None + + result = await _example_resource("ctx", unity_instance="Project@abc123") + assert result == "ctx" + + +def test_resource_registry_supports_server_only_resources(): + @mcp_for_unity_resource( + uri="mcpforunity://instances", + name="instances_resource", + unity_target=False, + ) + def _instances_resource(): + return None + + registered_resources = get_registered_resources() + resource_info = next(item for item in registered_resources if item["name"] == "instances_resource") + + assert resource_info["unity_target"] is False + assert resource_info["uri"] == "mcpforunity://instances" + assert "unity_instance" not in inspect.signature(_instances_resource).parameters + + +def test_resource_registry_merges_existing_query_templates(): + @mcp_for_unity_resource( + uri="mcpforunity://scene/example{?cursor,page_size}", + name="templated_resource", + ) + def _templated_resource(): + return None + + registered_resources = get_registered_resources() + resource_info = next(item for item in registered_resources if item["name"] == "templated_resource") + assert resource_info["uri"] == "mcpforunity://scene/example{?cursor,page_size,unity_instance}" diff --git a/Server/tests/test_tool_registry_metadata.py b/Server/tests/test_tool_registry_metadata.py index 2d7329da1..c9f2afa05 100644 --- a/Server/tests/test_tool_registry_metadata.py +++ b/Server/tests/test_tool_registry_metadata.py @@ -1,3 +1,5 @@ +import inspect + import pytest from services.registry import get_registered_tools, mcp_for_unity_tool @@ -23,6 +25,20 @@ def _default_target_tool(): assert tool_info["unity_target"] == "_default_target_tool" +@pytest.mark.asyncio +async def test_tool_registry_exposes_explicit_unity_instance_parameter(): + @mcp_for_unity_tool() + async def _unity_targeted_tool(ctx, action: str): + return {"ctx": ctx, "action": action} + + signature = inspect.signature(_unity_targeted_tool) + assert "unity_instance" in signature.parameters + assert signature.parameters["unity_instance"].default is None + + result = await _unity_targeted_tool("ctx", "play", unity_instance="Project@abc123") + assert result == {"ctx": "ctx", "action": "play"} + + def test_tool_registry_supports_server_only_and_alias_targets(): @mcp_for_unity_tool(unity_target=None) def _server_only_tool(): @@ -38,6 +54,7 @@ def _manage_script_alias_tool(): assert server_only["unity_target"] is None assert alias_tool["unity_target"] == "manage_script" + assert "unity_instance" not in inspect.signature(_server_only_tool).parameters def test_tool_registry_does_not_leak_unity_target_into_tool_kwargs(): diff --git a/Server/tests/test_transport_characterization.py b/Server/tests/test_transport_characterization.py index 45e35e6a8..9d166cfd5 100644 --- a/Server/tests/test_transport_characterization.py +++ b/Server/tests/test_transport_characterization.py @@ -23,6 +23,7 @@ from transport.unity_instance_middleware import UnityInstanceMiddleware, get_unity_instance_middleware, set_unity_instance_middleware from transport.plugin_registry import PluginRegistry, PluginSession from transport.plugin_hub import PluginHub, NoUnitySessionError, InstanceSelectionRequiredError, PluginDisconnectedError +from transport.legacy.unity_connection import UnityConnectionPool from transport.models import ( RegisterMessage, RegisterToolsMessage, @@ -313,7 +314,7 @@ async def test_list_tools_filters_disabled_unity_tools_and_aliases(self, mock_co async def call_next(_ctx): return available_tools - with patch.object(middleware, "_inject_unity_instance", new=AsyncMock()): + with patch.object(middleware, "_inject_user_id", new=AsyncMock(return_value="user-123")): with patch("transport.unity_instance_middleware.PluginHub.is_configured", return_value=True): with patch("transport.unity_instance_middleware.get_registered_tools", return_value=_tool_registry_for_visibility_tests()): with patch("transport.unity_instance_middleware.PluginHub.get_sessions", new_callable=AsyncMock) as mock_get_sessions: @@ -367,7 +368,7 @@ async def test_list_tools_skips_filter_when_no_tools_registered_yet(self, mock_c async def call_next(_ctx): return original_tools - with patch.object(middleware, "_inject_unity_instance", new=AsyncMock()): + with patch.object(middleware, "_inject_user_id", new=AsyncMock(return_value="user-123")): with patch("transport.unity_instance_middleware.PluginHub.is_configured", return_value=True): with patch("transport.unity_instance_middleware.get_registered_tools", return_value=_tool_registry_for_visibility_tests()): with patch("transport.unity_instance_middleware.PluginHub.get_sessions", new_callable=AsyncMock) as mock_get_sessions: @@ -424,7 +425,7 @@ async def call_next(_ctx): disabled_tool = SimpleNamespace(name="_marker_tool_indicates_registration_sent") registered_tools = [disabled_tool] - with patch.object(middleware, "_inject_unity_instance", new=AsyncMock()): + with patch.object(middleware, "_inject_user_id", new=AsyncMock(return_value="user-123")): with patch("transport.unity_instance_middleware.PluginHub.is_configured", return_value=True): with patch("transport.unity_instance_middleware.get_registered_tools", return_value=_tool_registry_for_visibility_tests()): with patch("transport.unity_instance_middleware.PluginHub.get_sessions", new_callable=AsyncMock) as mock_get_sessions: @@ -474,7 +475,7 @@ async def test_list_tools_skips_filter_when_enabled_set_lookup_fails(self, mock_ async def call_next(_ctx): return original_tools - with patch.object(middleware, "_inject_unity_instance", new=AsyncMock()): + with patch.object(middleware, "_inject_user_id", new=AsyncMock(return_value="user-123")): with patch("transport.unity_instance_middleware.PluginHub.is_configured", return_value=True): with patch("transport.unity_instance_middleware.get_registered_tools", return_value=_tool_registry_for_visibility_tests()): with patch("transport.unity_instance_middleware.PluginHub.get_sessions", new_callable=AsyncMock) as mock_get_sessions: @@ -513,7 +514,7 @@ async def test_list_tools_uses_user_scoped_tool_lookup_in_hosted_mode(self, mock async def call_next(_ctx): return [SimpleNamespace(name="manage_scene")] - with patch.object(middleware, "_inject_unity_instance", new=AsyncMock()): + with patch.object(middleware, "_inject_user_id", new=AsyncMock(return_value="user-123")): with patch("transport.unity_instance_middleware.PluginHub.is_configured", return_value=True): with patch("transport.unity_instance_middleware.get_registered_tools", return_value=_tool_registry_for_visibility_tests()): with patch("transport.unity_instance_middleware.PluginHub.get_sessions", new_callable=AsyncMock) as mock_get_sessions: @@ -537,6 +538,10 @@ async def call_next(_ctx): @pytest.mark.asyncio async def test_list_tools_skips_filter_when_active_instance_hash_is_stale(self, mock_context, monkeypatch): + """ + Current behavior: tools/list ignores the legacy session-selected instance + and instead filters against the actual visible HTTP sessions. + """ middleware = UnityInstanceMiddleware() middleware_ctx = Mock() middleware_ctx.fastmcp_context = mock_context @@ -568,11 +573,14 @@ async def call_next(_ctx): ) } ) + mock_get_tools.return_value = [ + SimpleNamespace(name="manage_scene"), + ] filtered = await middleware.on_list_tools(middleware_ctx, call_next) - assert [tool.name for tool in filtered] == [tool.name for tool in original_tools] - mock_get_tools.assert_not_called() + assert [tool.name for tool in filtered] == ["manage_scene", "set_active_instance"] + mock_get_tools.assert_awaited_once_with("abc123", user_id=None) @pytest.mark.asyncio async def test_list_tools_hides_alias_when_target_tool_is_disabled(self, mock_context, monkeypatch): @@ -803,6 +811,21 @@ async def test_autoselect_handles_plugin_hub_connection_error(self, mock_context assert instance is None +class TestLegacyUnityConnectionPoolSelection: + """Test stdio instance resolution when no explicit target is provided.""" + + def test_resolve_instance_id_requires_selection_when_multiple_instances_exist(self): + pool = UnityConnectionPool() + pool._default_instance_id = None + instances = [ + SimpleNamespace(id="ProjectA@aaa111", name="ProjectA", hash="aaa111", path="/tmp/a", port=6401), + SimpleNamespace(id="ProjectB@bbb222", name="ProjectB", hash="bbb222", path="/tmp/b", port=6402), + ] + + with pytest.raises(ConnectionError, match="Multiple Unity Editor instances are running"): + pool._resolve_instance_id(None, instances) + + # ============================================================================ # PLUGIN REGISTRY TESTS # ============================================================================