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

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
67 changes: 64 additions & 3 deletions MCPForUnity/Editor/Helpers/ProjectIdentityUtility.cs
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@ namespace MCPForUnity.Editor.Helpers
internal static class ProjectIdentityUtility
{
private const string SessionPrefKey = EditorPrefKeys.WebSocketSessionId;
private static bool _legacyKeyCleared;
private static string _cachedProjectName = "Unknown";
private static string _cachedProjectHash = "default";
private static bool _cacheScheduled;
Expand Down Expand Up @@ -69,6 +70,7 @@ private static void UpdateIdentityCache()
/// </summary>
public static string GetProjectHash()
{
EnsureIdentityCache();
return _cachedProjectHash;
}

Expand Down Expand Up @@ -122,16 +124,21 @@ private static string ComputeProjectName(string dataPath)

/// <summary>
/// Retrieves a persistent session id for the plugin, creating one if absent.
/// The session id is unique per project (scoped by project hash).
/// </summary>
public static string GetOrCreateSessionId()
{
try
{
string sessionId = EditorPrefs.GetString(SessionPrefKey, string.Empty);
// Make the session ID project-specific by including the project hash in the key
string projectHash = GetProjectHash();
string projectSpecificKey = $"{SessionPrefKey}_{projectHash}";

string sessionId = EditorPrefs.GetString(projectSpecificKey, string.Empty);
if (string.IsNullOrEmpty(sessionId))
{
sessionId = Guid.NewGuid().ToString();
EditorPrefs.SetString(SessionPrefKey, sessionId);
EditorPrefs.SetString(projectSpecificKey, sessionId);
}
return sessionId;
}
Expand All @@ -149,15 +156,69 @@ public static void ResetSessionId()
{
try
{
if (EditorPrefs.HasKey(SessionPrefKey))
// Clear the project-specific session ID
string projectHash = GetProjectHash();
string projectSpecificKey = $"{SessionPrefKey}_{projectHash}";

if (EditorPrefs.HasKey(projectSpecificKey))
{
EditorPrefs.DeleteKey(projectSpecificKey);
}

if (!_legacyKeyCleared && EditorPrefs.HasKey(SessionPrefKey))
{
EditorPrefs.DeleteKey(SessionPrefKey);
_legacyKeyCleared = true;
}
}
catch
{
// Ignore
}
}

private static void EnsureIdentityCache()
{
// When Application.dataPath is unavailable (e.g., batch mode) we fall back to
// hashing the current working directory/Assets path so each project still
// derives a deterministic, per-project session id rather than sharing "default".
if (!string.IsNullOrEmpty(_cachedProjectHash) && _cachedProjectHash != "default")
{
return;
}

UpdateIdentityCache();

if (!string.IsNullOrEmpty(_cachedProjectHash) && _cachedProjectHash != "default")
{
return;
}

string fallback = TryComputeFallbackProjectHash();
if (!string.IsNullOrEmpty(fallback))
{
_cachedProjectHash = fallback;
}
}

private static string TryComputeFallbackProjectHash()
{
try
{
string workingDirectory = Directory.GetCurrentDirectory();
if (string.IsNullOrEmpty(workingDirectory))
{
return "default";
}

// Normalise trailing separators so hashes remain stable
workingDirectory = workingDirectory.TrimEnd(Path.DirectorySeparatorChar, Path.AltDirectorySeparatorChar);
return ComputeProjectHash(Path.Combine(workingDirectory, "Assets"));
}
catch
{
return "default";
}
}
}
}
13 changes: 11 additions & 2 deletions Server/plugin_hub.py
Original file line number Diff line number Diff line change
Expand Up @@ -222,10 +222,19 @@ async def _resolve_session_id(cls, unity_instance: Optional[str]) -> str:
retry_ms = float(getattr(config, "reload_retry_ms", 250))
sleep_seconds = max(0.05, retry_ms / 1000.0)

# Allow callers to provide either just the hash or Name@hash
target_hash: Optional[str] = None
if unity_instance:
if "@" in unity_instance:
_, _, suffix = unity_instance.rpartition("@")
target_hash = suffix or None
else:
target_hash = unity_instance

async def _try_once() -> Optional[str]:
# Prefer a specific Unity instance if one was requested
if unity_instance:
session_id = await cls._registry.get_session_id_by_hash(unity_instance)
if target_hash:
session_id = await cls._registry.get_session_id_by_hash(target_hash)
if session_id:
return session_id

Expand Down
99 changes: 77 additions & 22 deletions Server/resources/unity_instances.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,12 @@
from fastmcp import Context
from registry import mcp_for_unity_resource
from unity_connection import get_unity_connection_pool
from plugin_hub import PluginHub
from unity_transport import _is_http_transport as _core_is_http_transport


def _is_http_transport() -> bool:
return _core_is_http_transport()


@mcp_for_unity_resource(
Expand All @@ -20,42 +26,91 @@ async def unity_instances(ctx: Context) -> dict[str, Any]:
Returns information about each instance including:
- id: Unique identifier (ProjectName@hash)
- name: Project name
- path: Full project path
- path: Full project path (stdio only)
- hash: 8-character hash of project path
- port: TCP port number
- status: Current status (running, reloading, etc.)
- last_heartbeat: Last heartbeat timestamp
- port: TCP port number (stdio only)
- status: Current status (running, reloading, etc.) (stdio only)
- last_heartbeat: Last heartbeat timestamp (stdio only)
- unity_version: Unity version (if available)
- connected_at: Connection timestamp (HTTP only)

Returns:
Dictionary containing list of instances and metadata
"""
await ctx.info("Listing Unity instances")

try:
pool = get_unity_connection_pool()
instances = pool.discover_all_instances(force_refresh=False)
if _is_http_transport():
# HTTP/WebSocket transport: query PluginHub
sessions_data = await PluginHub.get_sessions()
sessions = sessions_data.get("sessions", {})

# Check for duplicate project names
name_counts = {}
for inst in instances:
name_counts[inst.name] = name_counts.get(inst.name, 0) + 1
instances = []
for session_id, session_info in sessions.items():
project = session_info.get("project") or session_info.get("project_name")
project_hash = session_info.get("hash")

duplicates = [name for name, count in name_counts.items() if count > 1]
if not project or not project_hash:
raise ValueError(
"PluginHub session missing required 'project' or 'hash' fields."
)

result = {
"success": True,
"instance_count": len(instances),
"instances": [inst.to_dict() for inst in instances],
}
instances.append({
"id": f"{project}@{project_hash}",
"name": project,
"hash": project_hash,
"unity_version": session_info.get("unity_version"),
"connected_at": session_info.get("connected_at"),
"session_id": session_id,
})

# Check for duplicate project names
name_counts = {}
for inst in instances:
name_counts[inst["name"]] = name_counts.get(inst["name"], 0) + 1

duplicates = [name for name, count in name_counts.items() if count > 1]

result = {
"success": True,
"transport": "http",
"instance_count": len(instances),
"instances": instances,
}

if duplicates:
result["warning"] = (
f"Multiple instances found with duplicate project names: {duplicates}. "
f"Use full format (e.g., 'ProjectName@hash') to specify which instance."
)

return result
else:
# Stdio/TCP transport: query connection pool
pool = get_unity_connection_pool()
instances = pool.discover_all_instances(force_refresh=False)

# Check for duplicate project names
name_counts = {}
for inst in instances:
name_counts[inst.name] = name_counts.get(inst.name, 0) + 1

duplicates = [name for name, count in name_counts.items() if count > 1]

result = {
"success": True,
"transport": "stdio",
"instance_count": len(instances),
"instances": [inst.to_dict() for inst in instances],
}

if duplicates:
result["warning"] = (
f"Multiple instances found with duplicate project names: {duplicates}. "
f"Use full format (e.g., 'ProjectName@hash') to specify which instance."
)
if duplicates:
result["warning"] = (
f"Multiple instances found with duplicate project names: {duplicates}. "
f"Use full format (e.g., 'ProjectName@hash') to specify which instance."
)

return result
return result

except Exception as e:
await ctx.error(f"Error listing Unity instances: {e}")
Expand Down
126 changes: 126 additions & 0 deletions Server/tests/integration/test_instance_routing_comprehensive.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@

from unity_instance_middleware import UnityInstanceMiddleware
from tools import get_unity_instance_from_context
from tools.set_active_instance import set_active_instance as set_active_instance_tool


class TestInstanceRoutingBasics:
Expand Down Expand Up @@ -168,6 +169,131 @@ def test_tool_category_respects_active_instance(self, tool_category, tool_names)
pass # Placeholder for category-level test


class TestInstanceRoutingHTTP:
"""Validate HTTP-specific routing helpers."""

@pytest.mark.asyncio
async def test_set_active_instance_http_transport(self, monkeypatch):
"""set_active_instance should enumerate PluginHub sessions under HTTP."""
middleware = UnityInstanceMiddleware()
ctx = Mock(spec=Context)
ctx.session_id = "http-session"
state_storage = {}
ctx.set_state = Mock(side_effect=lambda k, v: state_storage.__setitem__(k, v))
ctx.get_state = Mock(side_effect=lambda k: state_storage.get(k))

monkeypatch.setenv("UNITY_MCP_TRANSPORT", "http")
fake_sessions = {
"sessions": {
"sess-1": {
"project": "Ramble",
"hash": "8e29de57",
"unity_version": "6000.2.10f1",
"connected_at": "2025-11-21T03:30:03.682353+00:00",
}
}
}
monkeypatch.setattr(
"tools.set_active_instance.PluginHub.get_sessions",
AsyncMock(return_value=fake_sessions),
)
monkeypatch.setattr(
"tools.set_active_instance.get_unity_instance_middleware",
lambda: middleware,
)

result = await set_active_instance_tool(ctx, "Ramble@8e29de57")

assert result["success"] is True
assert middleware.get_active_instance(ctx) == "Ramble@8e29de57"

@pytest.mark.asyncio
async def test_set_active_instance_http_hash_only(self, monkeypatch):
"""Hash-only selection should resolve via PluginHub registry."""
middleware = UnityInstanceMiddleware()
ctx = Mock(spec=Context)
ctx.session_id = "http-session-2"
state_storage = {}
ctx.set_state = Mock(side_effect=lambda k, v: state_storage.__setitem__(k, v))
ctx.get_state = Mock(side_effect=lambda k: state_storage.get(k))

monkeypatch.setenv("UNITY_MCP_TRANSPORT", "http")
fake_sessions = {
"sessions": {
"sess-99": {
"project": "UnityMCPTests",
"hash": "cc8756d4",
"unity_version": "2021.3.45f2",
"connected_at": "2025-11-21T03:37:01.501022+00:00",
}
}
}
monkeypatch.setattr(
"tools.set_active_instance.PluginHub.get_sessions",
AsyncMock(return_value=fake_sessions),
)
monkeypatch.setattr(
"tools.set_active_instance.get_unity_instance_middleware",
lambda: middleware,
)

result = await set_active_instance_tool(ctx, "cc8756d4")

assert result["success"] is True
assert middleware.get_active_instance(ctx) == "UnityMCPTests@cc8756d4"

@pytest.mark.asyncio
async def test_set_active_instance_http_hash_missing(self, monkeypatch):
"""Unknown hashes should surface a clear error."""
middleware = UnityInstanceMiddleware()
ctx = Mock(spec=Context)
ctx.session_id = "http-session-3"

monkeypatch.setenv("UNITY_MCP_TRANSPORT", "http")
fake_sessions = {"sessions": {}}
monkeypatch.setattr(
"tools.set_active_instance.PluginHub.get_sessions",
AsyncMock(return_value=fake_sessions),
)
monkeypatch.setattr(
"tools.set_active_instance.get_unity_instance_middleware",
lambda: middleware,
)

result = await set_active_instance_tool(ctx, "deadbeef")

assert result["success"] is False
assert "No Unity instances" in result["error"]

@pytest.mark.asyncio
async def test_set_active_instance_http_hash_ambiguous(self, monkeypatch):
"""Ambiguous hash prefixes should mirror stdio error messaging."""
middleware = UnityInstanceMiddleware()
ctx = Mock(spec=Context)
ctx.session_id = "http-session-4"

monkeypatch.setenv("UNITY_MCP_TRANSPORT", "http")
fake_sessions = {
"sessions": {
"sess-a": {"project": "ProjA", "hash": "abc12345"},
"sess-b": {"project": "ProjB", "hash": "abc98765"},
}
}
monkeypatch.setattr(
"tools.set_active_instance.PluginHub.get_sessions",
AsyncMock(return_value=fake_sessions),
)
monkeypatch.setattr(
"tools.set_active_instance.get_unity_instance_middleware",
lambda: middleware,
)

result = await set_active_instance_tool(ctx, "abc")

assert result["success"] is False
assert "matches multiple instances" in result["error"]


class TestInstanceRoutingRaceConditions:
"""Test for race conditions and timing issues."""

Expand Down
Loading
Loading