Skip to content

Commit f667582

Browse files
dsarnosakurachanclaude
authored
Feature/session based instance routing (#369)
* Add support for multiple Unity instances * fix port detection * add missing unity_instance parameter * add instance params for resources * Fix CodeRabbit review feedback - Fix partial framed response handling in port discovery Add _recv_exact() helper to ensure complete frame reading Prevents healthy Unity instances from being misidentified as offline - Remove unused default_conn variables in server.py (2 files) Fixes Ruff F841 lint error that would block CI/CD - Preserve sync/async nature of resources in wrapper Check if original function is coroutine before wrapping Prevents 'dict object is not awaitable' runtime errors - Fix reconnection to preserve instance_id Add instance_id tracking to UnityConnection dataclass Reconnection now targets the same Unity instance instead of any available one Prevents operations from being applied to wrong project - Add instance logging to manage_asset for debugging Helps troubleshoot multi-instance scenarios 🤖 Generated with Claude Code Co-Authored-By: Claude <[email protected]> * Fix CodeRabbit feedback: reconnection fallback and annotations safety Address 3 CodeRabbit review comments: 1. Critical: Guard reconnection fallback to prevent wrong instance routing - When instance_id is set but rediscovery fails, now raises ConnectionError - Added 'from e' to preserve exception chain for better debugging - Prevents silently connecting to different Unity instance - Ensures multi-instance routing integrity 2. Minor: Guard __annotations__ access in resource registration - Use getattr(func, '__annotations__', {}) instead of direct access - Prevents AttributeError for functions without type hints 3. Minor: Remove unused get_type_hints import - Clean up unused import in resources/__init__.py All changes applied to both Server/ and MCPForUnity/ directories. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <[email protected]> * Fix instance sorting and logging issues - Fix sorting logic for instances without heartbeat data: use epoch timestamp instead of current time to properly deprioritize instances with None last_heartbeat - Use logger.exception() instead of logger.error() in disconnect_all() to include stack traces for better debugging 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <[email protected]> * update uv.lock to prepare for merging into main * Restore Python 3.10 lockfiles and package list_unity_instances tool * Deduplicate Unity instance discovery by port * Scope status-file reload checks to the active instance * refactor: implement FastMCP middleware for session-based instance routing Replaces module-level session_state.py with UnityInstanceMiddleware class that follows FastMCP best practices. Middleware intercepts all tool calls via on_call_tool hook and injects active Unity instance into request state. Key changes: - Add UnityInstanceMiddleware class with on_call_tool hook - Tools now use ctx.get_state("unity_instance") instead of direct session_state calls - Remove unity_instance parameter from all tool schemas to prevent LLM hallucination - Convert list_unity_instances tool to unity_instances resource (read-only data) - Update error messages to reference unity://instances resource - Add set_state/get_state methods to DummyContext test helper - All 67 tests passing (55 passed, 5 skipped, 7 xpassed) Architecture benefits: - Centralized session management in middleware - Standard FastMCP patterns (middleware + request state) - Cleaner separation of concerns - Prevents AI hallucination of invalid instance IDs * fix: convert resource templates to static resources for discoverability Convert MCP resources from URI templates with query parameters to static resources to fix discoverability in MCP clients like Claude Code. Changes: - Remove {?force_refresh} from unity://instances - Remove {?unity_instance} from mcpforunity://menu-items - Remove {?unity_instance} from mcpforunity://tests - Keep {mode} path parameter in mcpforunity://tests/{mode} (legitimate) Root cause: Query parameters {?param} trigger ResourceTemplate registration, which are listed via resources/templates/list instead of resources/list. Claude Code's ListMcpResourcesTool only queries resources/list, making templates undiscoverable. Solution: Remove optional query parameters from URIs. Instance routing is handled by middleware/context, and force_refresh was cache control that doesn't belong in resource identity. Impact: Resources now discoverable via standard resources/list endpoint and work with all MCP clients including Claude Code and Cursor. Requires FastMCP >=2.13.0 for proper RFC 6570 query parameter support. * feat: improve material properties and sync Server resources Material Property Improvements (ManageAsset.cs): - Add GetMainColorPropertyName() helper that auto-detects shader color properties - Tries _BaseColor (URP), _Color (Standard), _MainColor, _Tint, _TintColor - Update both named and array color property handling to use auto-detection - Add warning messages when color properties don't exist on materials - Split HasProperty check from SetColor to enable error reporting This fixes the issue where simple color array format [r,g,b,a] defaulted to _Color property, causing silent failures with URP Lit shader which uses _BaseColor. Server Resource Sync: - Sync Server/resources with MCPForUnity/UnityMcpServer~/src/resources - Remove query parameters from resource URIs for discoverability - Use session-based instance routing via get_unity_instance_from_context() * fix: repair instance routing and simplify get_unity_instance_from_context PROBLEM: Instance routing was failing - scripts went to wrong Unity instances. Script1 (intended: ramble) -> went to UnityMCPTests ❌ Script2 (intended: UnityMCPTests) -> went to ramble ❌ ROOT CAUSE: Two incompatible approaches for accessing active instance: 1. Middleware: ctx.set_state() / ctx.get_state() - used by most tools 2. Legacy: ctx.request_context.meta - used by script tools Script tools were reading from wrong location, middleware had no effect. FIX: 1. Updated get_unity_instance_from_context() to read from ctx.get_state() 2. Removed legacy request_context.meta code path (98 lines removed) 3. Single source of truth: middleware state only TESTING: - Added comprehensive test suite (21 tests) covering all scenarios - Tests middleware state management, session isolation, race conditions - Tests reproduce exact 4-script failure scenario - All 88 tests pass (76 passed + 5 skipped + 7 xpassed) - Verified fix with live 4-script test: 100% success rate Files changed: - Server/tools/__init__.py: Simplified from 75 lines to 15 lines - MCPForUnity/UnityMcpServer~/src/tools/__init__.py: Same simplification - tests/test_instance_routing_comprehensive.py: New comprehensive test suite * refactor: standardize instance extraction and remove dead imports - Standardize all 18 tools to use get_unity_instance_from_context() helper instead of direct ctx.get_state() calls for consistency - Remove dead session_state imports from with_unity_instance decorator that would cause ModuleNotFoundError at runtime - Update README.md with concise instance routing documentation * fix: critical timezone and import bugs from code review - Remove incorrect port safety check that treated reclaimed ports as errors (GetPortWithFallback may legitimately return same port if it became available) - Fix timezone-aware vs naive datetime mixing in unity_connection.py sorting (use timestamp() for comparison to avoid TypeError) - Normalize all datetime comparisons in port_discovery.py to UTC (file_mtime and last_heartbeat now consistently timezone-aware) - Add missing send_with_unity_instance import in Server/tools/manage_script.py (was causing NameError at runtime on lines 108 and 488) All 88 tests pass (76 passed + 5 skipped + 7 xpassed) --------- Co-authored-by: Sakura <[email protected]> Co-authored-by: Claude <[email protected]>
1 parent 5e4b554 commit f667582

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

58 files changed

+3022
-412
lines changed

MCPForUnity/Editor/Helpers/PortManager.cs

Lines changed: 8 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -60,14 +60,17 @@ public static int GetPortWithFallback()
6060
if (IsDebugEnabled()) Debug.Log($"<b><color=#2EA3FF>MCP-FOR-UNITY</color></b>: Stored port {storedConfig.unity_port} became available after short wait");
6161
return storedConfig.unity_port;
6262
}
63-
// Prefer sticking to the same port; let the caller handle bind retries/fallbacks
64-
return storedConfig.unity_port;
63+
// Port is still busy after waiting - find a new available port instead
64+
if (IsDebugEnabled()) Debug.Log($"<b><color=#2EA3FF>MCP-FOR-UNITY</color></b>: Stored port {storedConfig.unity_port} is occupied by another instance, finding alternative...");
65+
int newPort = FindAvailablePort();
66+
SavePort(newPort);
67+
return newPort;
6568
}
6669

6770
// If no valid stored port, find a new one and save it
68-
int newPort = FindAvailablePort();
69-
SavePort(newPort);
70-
return newPort;
71+
int foundPort = FindAvailablePort();
72+
SavePort(foundPort);
73+
return foundPort;
7174
}
7275

7376
/// <summary>

MCPForUnity/Editor/MCPForUnityBridge.cs

Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -362,7 +362,24 @@ public static void Start()
362362
}
363363
catch (SocketException se) when (se.SocketErrorCode == SocketError.AddressAlreadyInUse && attempt >= maxImmediateRetries)
364364
{
365+
// Port is occupied by another instance, get a new available port
366+
int oldPort = currentUnityPort;
365367
currentUnityPort = PortManager.GetPortWithFallback();
368+
369+
// GetPortWithFallback() may return the same port if it became available during wait
370+
// or a different port if switching to an alternative
371+
if (IsDebugEnabled())
372+
{
373+
if (currentUnityPort == oldPort)
374+
{
375+
McpLog.Info($"Port {oldPort} became available, proceeding");
376+
}
377+
else
378+
{
379+
McpLog.Info($"Port {oldPort} occupied, switching to port {currentUnityPort}");
380+
}
381+
}
382+
366383
listener = new TcpListener(IPAddress.Loopback, currentUnityPort);
367384
listener.Server.SetSocketOption(
368385
SocketOptionLevel.Socket,
@@ -474,6 +491,22 @@ public static void Stop()
474491
try { AssemblyReloadEvents.afterAssemblyReload -= OnAfterAssemblyReload; } catch { }
475492
try { EditorApplication.quitting -= Stop; } catch { }
476493

494+
// Clean up status file when Unity stops
495+
try
496+
{
497+
string statusDir = Path.Combine(Environment.GetFolderPath(Environment.SpecialFolder.UserProfile), ".unity-mcp");
498+
string statusFile = Path.Combine(statusDir, $"unity-mcp-status-{ComputeProjectHash(Application.dataPath)}.json");
499+
if (File.Exists(statusFile))
500+
{
501+
File.Delete(statusFile);
502+
if (IsDebugEnabled()) McpLog.Info($"Deleted status file: {statusFile}");
503+
}
504+
}
505+
catch (Exception ex)
506+
{
507+
if (IsDebugEnabled()) McpLog.Warn($"Failed to delete status file: {ex.Message}");
508+
}
509+
477510
if (IsDebugEnabled()) McpLog.Info("MCPForUnityBridge stopped.");
478511
}
479512

@@ -1184,13 +1217,38 @@ private static void WriteHeartbeat(bool reloading, string reason = null)
11841217
}
11851218
Directory.CreateDirectory(dir);
11861219
string filePath = Path.Combine(dir, $"unity-mcp-status-{ComputeProjectHash(Application.dataPath)}.json");
1220+
1221+
// Extract project name from path
1222+
string projectName = "Unknown";
1223+
try
1224+
{
1225+
string projectPath = Application.dataPath;
1226+
if (!string.IsNullOrEmpty(projectPath))
1227+
{
1228+
// Remove trailing /Assets or \Assets
1229+
projectPath = projectPath.TrimEnd('/', '\\');
1230+
if (projectPath.EndsWith("Assets", StringComparison.OrdinalIgnoreCase))
1231+
{
1232+
projectPath = projectPath.Substring(0, projectPath.Length - 6).TrimEnd('/', '\\');
1233+
}
1234+
projectName = Path.GetFileName(projectPath);
1235+
if (string.IsNullOrEmpty(projectName))
1236+
{
1237+
projectName = "Unknown";
1238+
}
1239+
}
1240+
}
1241+
catch { }
1242+
11871243
var payload = new
11881244
{
11891245
unity_port = currentUnityPort,
11901246
reloading,
11911247
reason = reason ?? (reloading ? "reloading" : "ready"),
11921248
seq = heartbeatSeq,
11931249
project_path = Application.dataPath,
1250+
project_name = projectName,
1251+
unity_version = Application.unityVersion,
11941252
last_heartbeat = DateTime.UtcNow.ToString("O")
11951253
};
11961254
File.WriteAllText(filePath, JsonConvert.SerializeObject(payload), new System.Text.UTF8Encoding(false));

MCPForUnity/Editor/Tools/ManageAsset.cs

Lines changed: 50 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -911,7 +911,7 @@ private static bool ApplyMaterialProperties(Material mat, JObject properties)
911911
// Example: Set color property
912912
if (properties["color"] is JObject colorProps)
913913
{
914-
string propName = colorProps["name"]?.ToString() ?? "_Color"; // Default main color
914+
string propName = colorProps["name"]?.ToString() ?? GetMainColorPropertyName(mat); // Auto-detect if not specified
915915
if (colorProps["value"] is JArray colArr && colArr.Count >= 3)
916916
{
917917
try
@@ -922,10 +922,20 @@ private static bool ApplyMaterialProperties(Material mat, JObject properties)
922922
colArr[2].ToObject<float>(),
923923
colArr.Count > 3 ? colArr[3].ToObject<float>() : 1.0f
924924
);
925-
if (mat.HasProperty(propName) && mat.GetColor(propName) != newColor)
925+
if (mat.HasProperty(propName))
926926
{
927-
mat.SetColor(propName, newColor);
928-
modified = true;
927+
if (mat.GetColor(propName) != newColor)
928+
{
929+
mat.SetColor(propName, newColor);
930+
modified = true;
931+
}
932+
}
933+
else
934+
{
935+
Debug.LogWarning(
936+
$"Material '{mat.name}' with shader '{mat.shader.name}' does not have color property '{propName}'. " +
937+
$"Color not applied. Common color properties: _BaseColor (URP), _Color (Standard)"
938+
);
929939
}
930940
}
931941
catch (Exception ex)
@@ -938,7 +948,8 @@ private static bool ApplyMaterialProperties(Material mat, JObject properties)
938948
}
939949
else if (properties["color"] is JArray colorArr) //Use color now with examples set in manage_asset.py
940950
{
941-
string propName = "_Color";
951+
// Auto-detect the main color property for the shader
952+
string propName = GetMainColorPropertyName(mat);
942953
try
943954
{
944955
if (colorArr.Count >= 3)
@@ -949,10 +960,20 @@ private static bool ApplyMaterialProperties(Material mat, JObject properties)
949960
colorArr[2].ToObject<float>(),
950961
colorArr.Count > 3 ? colorArr[3].ToObject<float>() : 1.0f
951962
);
952-
if (mat.HasProperty(propName) && mat.GetColor(propName) != newColor)
963+
if (mat.HasProperty(propName))
953964
{
954-
mat.SetColor(propName, newColor);
955-
modified = true;
965+
if (mat.GetColor(propName) != newColor)
966+
{
967+
mat.SetColor(propName, newColor);
968+
modified = true;
969+
}
970+
}
971+
else
972+
{
973+
Debug.LogWarning(
974+
$"Material '{mat.name}' with shader '{mat.shader.name}' does not have color property '{propName}'. " +
975+
$"Color not applied. Common color properties: _BaseColor (URP), _Color (Standard)"
976+
);
956977
}
957978
}
958979
}
@@ -1140,6 +1161,27 @@ string ResolvePropertyName(string name)
11401161
return modified;
11411162
}
11421163

1164+
/// <summary>
1165+
/// Auto-detects the main color property name for a material's shader.
1166+
/// Tries common color property names in order: _BaseColor (URP), _Color (Standard), etc.
1167+
/// </summary>
1168+
private static string GetMainColorPropertyName(Material mat)
1169+
{
1170+
if (mat == null || mat.shader == null)
1171+
return "_Color";
1172+
1173+
// Try common color property names in order of likelihood
1174+
string[] commonColorProps = { "_BaseColor", "_Color", "_MainColor", "_Tint", "_TintColor" };
1175+
foreach (var prop in commonColorProps)
1176+
{
1177+
if (mat.HasProperty(prop))
1178+
return prop;
1179+
}
1180+
1181+
// Fallback to _Color if none found
1182+
return "_Color";
1183+
}
1184+
11431185
/// <summary>
11441186
/// Applies properties from JObject to a PhysicsMaterial.
11451187
/// </summary>

MCPForUnity/UnityMcpServer~/src/models.py

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
from typing import Any
2+
from datetime import datetime
23
from pydantic import BaseModel
34

45

@@ -7,3 +8,28 @@ class MCPResponse(BaseModel):
78
message: str | None = None
89
error: str | None = None
910
data: Any | None = None
11+
12+
13+
class UnityInstanceInfo(BaseModel):
14+
"""Information about a Unity Editor instance"""
15+
id: str # "ProjectName@hash" or fallback to hash
16+
name: str # Project name extracted from path
17+
path: str # Full project path (Assets folder)
18+
hash: str # 8-char hash of project path
19+
port: int # TCP port
20+
status: str # "running", "reloading", "offline"
21+
last_heartbeat: datetime | None = None
22+
unity_version: str | None = None
23+
24+
def to_dict(self) -> dict[str, Any]:
25+
"""Convert to dictionary for JSON serialization"""
26+
return {
27+
"id": self.id,
28+
"name": self.name,
29+
"path": self.path,
30+
"hash": self.hash,
31+
"port": self.port,
32+
"status": self.status,
33+
"last_heartbeat": self.last_heartbeat.isoformat() if self.last_heartbeat else None,
34+
"unity_version": self.unity_version
35+
}

0 commit comments

Comments
 (0)