telemetry: record sub_action for tool executions; decorator extracts 'action'; add tests for keyword/positional extraction

main
David Sarno 2025-09-10 08:52:38 -07:00
parent 2fd74f5dab
commit 89714d022c
3 changed files with 120 additions and 4 deletions

View File

@ -356,14 +356,29 @@ 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

View File

@ -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

View File

@ -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