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"""
|
"""Convenience function to record a milestone"""
|
||||||
return get_telemetry().record_milestone(milestone, data)
|
return get_telemetry().record_milestone(milestone, data)
|
||||||
|
|
||||||
def record_tool_usage(tool_name: str, success: bool, duration_ms: float, error: Optional[str] = None):
|
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"""
|
"""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 = {
|
data = {
|
||||||
"tool_name": tool_name,
|
"tool_name": tool_name,
|
||||||
"success": success,
|
"success": success,
|
||||||
"duration_ms": round(duration_ms, 2)
|
"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:
|
if error:
|
||||||
data["error"] = str(error)[:200] # Limit error message length
|
data["error"] = str(error)[:200] # Limit error message length
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -20,6 +20,15 @@ def telemetry_tool(tool_name: str):
|
||||||
start_time = time.time()
|
start_time = time.time()
|
||||||
success = False
|
success = False
|
||||||
error = None
|
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:
|
try:
|
||||||
global _decorator_log_count
|
global _decorator_log_count
|
||||||
if _decorator_log_count < 10:
|
if _decorator_log_count < 10:
|
||||||
|
|
@ -38,13 +47,22 @@ def telemetry_tool(tool_name: str):
|
||||||
raise
|
raise
|
||||||
finally:
|
finally:
|
||||||
duration_ms = (time.time() - start_time) * 1000
|
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)
|
@functools.wraps(func)
|
||||||
async def _async_wrapper(*args, **kwargs) -> Any:
|
async def _async_wrapper(*args, **kwargs) -> Any:
|
||||||
start_time = time.time()
|
start_time = time.time()
|
||||||
success = False
|
success = False
|
||||||
error = None
|
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:
|
try:
|
||||||
global _decorator_log_count
|
global _decorator_log_count
|
||||||
if _decorator_log_count < 10:
|
if _decorator_log_count < 10:
|
||||||
|
|
@ -63,7 +81,7 @@ def telemetry_tool(tool_name: str):
|
||||||
raise
|
raise
|
||||||
finally:
|
finally:
|
||||||
duration_ms = (time.time() - start_time) * 1000
|
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 _async_wrapper if inspect.iscoroutinefunction(func) else _sync_wrapper
|
||||||
return decorator
|
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