telemetry: record sub_action for tool executions; decorator extracts 'action'; add tests for keyword/positional extraction
parent
2fd74f5dab
commit
89714d022c
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
||||
|
||||
Loading…
Reference in New Issue