87 lines
3.6 KiB
Python
87 lines
3.6 KiB
Python
"""
|
|
Telemetry decorator for Unity MCP tools
|
|
"""
|
|
|
|
import functools
|
|
import time
|
|
import inspect
|
|
import logging
|
|
from typing import Callable, Any
|
|
from telemetry import record_tool_usage, record_milestone, MilestoneType
|
|
|
|
_log = logging.getLogger("unity-mcp-telemetry")
|
|
_decorator_log_count = 0
|
|
|
|
def telemetry_tool(tool_name: str):
|
|
"""Decorator to add telemetry tracking to MCP tools"""
|
|
def decorator(func: Callable) -> Callable:
|
|
@functools.wraps(func)
|
|
def _sync_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:
|
|
_log.info(f"telemetry_decorator sync: tool={tool_name}")
|
|
_decorator_log_count += 1
|
|
result = func(*args, **kwargs)
|
|
success = True
|
|
if tool_name == "manage_script" and kwargs.get("action") == "create":
|
|
record_milestone(MilestoneType.FIRST_SCRIPT_CREATION)
|
|
elif tool_name.startswith("manage_scene"):
|
|
record_milestone(MilestoneType.FIRST_SCENE_MODIFICATION)
|
|
record_milestone(MilestoneType.FIRST_TOOL_USAGE)
|
|
return result
|
|
except Exception as e:
|
|
error = str(e)
|
|
raise
|
|
finally:
|
|
duration_ms = (time.time() - start_time) * 1000
|
|
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:
|
|
_log.info(f"telemetry_decorator async: tool={tool_name}")
|
|
_decorator_log_count += 1
|
|
result = await func(*args, **kwargs)
|
|
success = True
|
|
if tool_name == "manage_script" and kwargs.get("action") == "create":
|
|
record_milestone(MilestoneType.FIRST_SCRIPT_CREATION)
|
|
elif tool_name.startswith("manage_scene"):
|
|
record_milestone(MilestoneType.FIRST_SCENE_MODIFICATION)
|
|
record_milestone(MilestoneType.FIRST_TOOL_USAGE)
|
|
return result
|
|
except Exception as e:
|
|
error = str(e)
|
|
raise
|
|
finally:
|
|
duration_ms = (time.time() - start_time) * 1000
|
|
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 |