diff --git a/UnityMcpBridge/UnityMcpServer~/src/telemetry.py b/UnityMcpBridge/UnityMcpServer~/src/telemetry.py index 5e2031a..413a7f9 100644 --- a/UnityMcpBridge/UnityMcpServer~/src/telemetry.py +++ b/UnityMcpBridge/UnityMcpServer~/src/telemetry.py @@ -356,13 +356,28 @@ def record_milestone(milestone: MilestoneType, data: Optional[Dict[str, Any]] = """Convenience function to record a milestone""" return get_telemetry().record_milestone(milestone, data) -def record_tool_usage(tool_name: str, success: bool, duration_ms: float, error: Optional[str] = None): - """Record tool usage telemetry""" +def record_tool_usage(tool_name: str, success: bool, duration_ms: float, error: Optional[str] = None, sub_action: Optional[str] = None): + """Record tool usage telemetry + + Args: + tool_name: Name of the tool invoked (e.g., 'manage_scene'). + success: Whether the tool completed successfully. + duration_ms: Execution duration in milliseconds. + error: Optional error message (truncated if present). + sub_action: Optional sub-action/operation within the tool (e.g., 'get_hierarchy'). + """ data = { "tool_name": tool_name, "success": success, "duration_ms": round(duration_ms, 2) } + + if sub_action is not None: + try: + data["sub_action"] = str(sub_action) + except Exception: + # Ensure telemetry is never disruptive + data["sub_action"] = "unknown" if error: data["error"] = str(error)[:200] # Limit error message length diff --git a/UnityMcpBridge/UnityMcpServer~/src/telemetry_decorator.py b/UnityMcpBridge/UnityMcpServer~/src/telemetry_decorator.py index 875c66a..3ea27d9 100644 --- a/UnityMcpBridge/UnityMcpServer~/src/telemetry_decorator.py +++ b/UnityMcpBridge/UnityMcpServer~/src/telemetry_decorator.py @@ -20,6 +20,15 @@ def telemetry_tool(tool_name: str): start_time = time.time() success = False error = None + # Extract sub-action (e.g., 'get_hierarchy') from bound args when available + sub_action = None + try: + sig = inspect.signature(func) + bound = sig.bind_partial(*args, **kwargs) + bound.apply_defaults() + sub_action = bound.arguments.get("action") + except Exception: + sub_action = None try: global _decorator_log_count if _decorator_log_count < 10: @@ -38,13 +47,22 @@ def telemetry_tool(tool_name: str): raise finally: duration_ms = (time.time() - start_time) * 1000 - record_tool_usage(tool_name, success, duration_ms, error) + record_tool_usage(tool_name, success, duration_ms, error, sub_action=sub_action) @functools.wraps(func) async def _async_wrapper(*args, **kwargs) -> Any: start_time = time.time() success = False error = None + # Extract sub-action (e.g., 'get_hierarchy') from bound args when available + sub_action = None + try: + sig = inspect.signature(func) + bound = sig.bind_partial(*args, **kwargs) + bound.apply_defaults() + sub_action = bound.arguments.get("action") + except Exception: + sub_action = None try: global _decorator_log_count if _decorator_log_count < 10: @@ -63,7 +81,7 @@ def telemetry_tool(tool_name: str): raise finally: duration_ms = (time.time() - start_time) * 1000 - record_tool_usage(tool_name, success, duration_ms, error) + record_tool_usage(tool_name, success, duration_ms, error, sub_action=sub_action) return _async_wrapper if inspect.iscoroutinefunction(func) else _sync_wrapper return decorator \ No newline at end of file diff --git a/tests/test_telemetry_subaction.py b/tests/test_telemetry_subaction.py new file mode 100644 index 0000000..c1c597e --- /dev/null +++ b/tests/test_telemetry_subaction.py @@ -0,0 +1,83 @@ +import importlib + + +def _get_decorator_module(): + # Import the telemetry_decorator module from the Unity MCP server src + mod = importlib.import_module("UnityMcpBridge.UnityMcpServer~.src.telemetry_decorator") + return mod + + +def test_subaction_extracted_from_keyword(monkeypatch): + td = _get_decorator_module() + + captured = {} + + def fake_record_tool_usage(tool_name, success, duration_ms, error, sub_action=None): + captured["tool_name"] = tool_name + captured["success"] = success + captured["error"] = error + captured["sub_action"] = sub_action + + # Silence milestones/logging in test + monkeypatch.setattr(td, "record_tool_usage", fake_record_tool_usage) + monkeypatch.setattr(td, "record_milestone", lambda *a, **k: None) + monkeypatch.setattr(td, "_decorator_log_count", 999) + + def dummy_tool(ctx, action: str, name: str = ""): + return {"success": True, "name": name} + + wrapped = td.telemetry_tool("manage_scene")(dummy_tool) + + resp = wrapped(None, action="get_hierarchy", name="Sample") + assert resp["success"] is True + assert captured["tool_name"] == "manage_scene" + assert captured["success"] is True + assert captured["error"] is None + assert captured["sub_action"] == "get_hierarchy" + + +def test_subaction_extracted_from_positionals(monkeypatch): + td = _get_decorator_module() + + captured = {} + + def fake_record_tool_usage(tool_name, success, duration_ms, error, sub_action=None): + captured["tool_name"] = tool_name + captured["sub_action"] = sub_action + + monkeypatch.setattr(td, "record_tool_usage", fake_record_tool_usage) + monkeypatch.setattr(td, "record_milestone", lambda *a, **k: None) + monkeypatch.setattr(td, "_decorator_log_count", 999) + + def dummy_tool(ctx, action: str, name: str = ""): + return True + + wrapped = td.telemetry_tool("manage_scene")(dummy_tool) + + _ = wrapped(None, "save", "MyScene") + assert captured["tool_name"] == "manage_scene" + assert captured["sub_action"] == "save" + + +def test_subaction_none_when_not_present(monkeypatch): + td = _get_decorator_module() + + captured = {} + + def fake_record_tool_usage(tool_name, success, duration_ms, error, sub_action=None): + captured["tool_name"] = tool_name + captured["sub_action"] = sub_action + + monkeypatch.setattr(td, "record_tool_usage", fake_record_tool_usage) + monkeypatch.setattr(td, "record_milestone", lambda *a, **k: None) + monkeypatch.setattr(td, "_decorator_log_count", 999) + + def dummy_tool_without_action(ctx, name: str): + return 123 + + wrapped = td.telemetry_tool("apply_text_edits")(dummy_tool_without_action) + _ = wrapped(None, name="X") + assert captured["tool_name"] == "apply_text_edits" + assert captured["sub_action"] is None + +