From c3aacc151cfb99418cdab64cc96d7d9feec9e8a4 Mon Sep 17 00:00:00 2001 From: Jason Andrews Date: Mon, 13 Jul 2026 15:28:19 -0500 Subject: [PATCH] Add logging for arm mcp server testing --- mcp-local/server.py | 8 ++- mcp-local/tests/test_invocation_logger.py | 51 +++++++++++++++++ mcp-local/utils/invocation_logger.py | 70 ++++++++++++++++++----- 3 files changed, 111 insertions(+), 18 deletions(-) create mode 100644 mcp-local/tests/test_invocation_logger.py diff --git a/mcp-local/server.py b/mcp-local/server.py index b08e6d5..54d2671 100644 --- a/mcp-local/server.py +++ b/mcp-local/server.py @@ -28,7 +28,7 @@ from utils.migrate_ease_utils import run_migrate_ease_scan from utils.skopeo_tool import skopeo_help, skopeo_inspect from utils.llvm_mca_tool import mca_help, llvm_mca_analyze -from utils.invocation_logger import log_invocation_reason +from utils.invocation_logger import log_invocation_reason, log_tool_result from utils.error_handling import format_tool_error # Initialize the MCP server @@ -52,7 +52,7 @@ ) def knowledge_base_search(query: str, invocation_reason: Optional[str] = None) -> List[Dict[str, Any]]: # Log invocation reason if provided - log_invocation_reason( + entry_id = log_invocation_reason( tool="knowledge_base_search", reason=invocation_reason, args={"query": query}, @@ -67,7 +67,9 @@ def knowledge_base_search(query: str, invocation_reason: Optional[str] = None) - List of dictionaries with metadata including url and text snippets. """ try: - return arm_kb_search.search(query, SEARCH_RESOURCES) + results = arm_kb_search.search(query, SEARCH_RESOURCES) + log_tool_result(entry_id, "knowledge_base_search", results) + return results except Exception as e: return format_tool_error( tool="knowledge_base_search", diff --git a/mcp-local/tests/test_invocation_logger.py b/mcp-local/tests/test_invocation_logger.py new file mode 100644 index 0000000..76c66bd --- /dev/null +++ b/mcp-local/tests/test_invocation_logger.py @@ -0,0 +1,51 @@ +import json + +from utils import invocation_logger + + +def test_logs_paired_call_and_result(tmp_path, monkeypatch): + traffic_path = tmp_path / "mcp-traffic.jsonl" + monkeypatch.setattr(invocation_logger, "WORKSPACE_DIR", str(tmp_path)) + monkeypatch.setenv(invocation_logger.MCP_TRAFFIC_LOG_ENV, str(traffic_path)) + + entry_id = invocation_logger.log_invocation_reason( + tool="knowledge_base_search", + reason="Need current Arm documentation", + args={"query": "SME overview"}, + ) + result = [{"title": "SME guide", "score": 0.9}] + invocation_logger.log_tool_result(entry_id, "knowledge_base_search", result) + + entries = [json.loads(line) for line in traffic_path.read_text().splitlines()] + assert entries == [ + { + "id": entry_id, + "timestamp": entries[0]["timestamp"], + "tool": "knowledge_base_search", + "args": {"query": "SME overview"}, + "invocation_reason": "Need current Arm documentation", + }, + { + "id": entry_id, + "type": "result", + "tool": "knowledge_base_search", + "result": result, + }, + ] + + +def test_logs_call_without_invocation_reason(tmp_path, monkeypatch): + traffic_path = tmp_path / "mcp-traffic.jsonl" + monkeypatch.setattr(invocation_logger, "WORKSPACE_DIR", str(tmp_path)) + monkeypatch.setenv(invocation_logger.MCP_TRAFFIC_LOG_ENV, str(traffic_path)) + + entry_id = invocation_logger.log_invocation_reason( + tool="knowledge_base_search", + reason=None, + args={"query": "SVE2"}, + ) + + entry = json.loads(traffic_path.read_text()) + assert entry["id"] == entry_id + assert entry["invocation_reason"] is None + assert not (tmp_path / invocation_logger.LOG_FILE_NAME).exists() diff --git a/mcp-local/utils/invocation_logger.py b/mcp-local/utils/invocation_logger.py index 64d2e6c..3bfa4cf 100644 --- a/mcp-local/utils/invocation_logger.py +++ b/mcp-local/utils/invocation_logger.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import json import os import uuid from datetime import datetime, timezone @@ -23,38 +24,77 @@ LOG_FILE_NAME = "invocation_reasons.yaml" +MCP_TRAFFIC_LOG_ENV = "MCP_LOG_FILE" +MCP_TRAFFIC_LOG_DEFAULT = "/workspace/mcp-traffic.jsonl" def _now_iso() -> str: return datetime.now(timezone.utc).isoformat() -def log_invocation_reason(tool: str, reason: Optional[str], args: Optional[Dict[str, Any]] = None) -> None: +def log_invocation_reason( + tool: str, + reason: Optional[str], + args: Optional[Dict[str, Any]] = None, +) -> str: """ Append a YAML document with the tool invocation reason and metadata to /workspace/invocation_reasons.yaml. + Also append a JSONL call entry to MCP_LOG_FILE. - Each call writes a separate YAML document with fields: id, timestamp, tool, args, reason. + Returns the entry ID so the caller can pair the tool result with this invocation. Errors are swallowed to avoid impacting tool execution. """ - if not reason: - return + entry_id = str(uuid.uuid4()) + timestamp = _now_iso() - entry = { - "id": str(uuid.uuid4()), - "timestamp": _now_iso(), + if reason: + entry = { + "id": entry_id, + "timestamp": timestamp, + "tool": tool, + "args": args or {}, + "reason": str(reason), + } + + log_path = os.path.join(WORKSPACE_DIR, LOG_FILE_NAME) + + try: + os.makedirs(WORKSPACE_DIR, exist_ok=True) + with open(log_path, "a", encoding="utf-8") as f: + yaml.safe_dump(entry, f, explicit_start=True, sort_keys=False, allow_unicode=True) + except Exception: + pass + + traffic_entry = { + "id": entry_id, + "timestamp": timestamp, "tool": tool, "args": args or {}, - "reason": str(reason), + "invocation_reason": reason, } + traffic_path = os.environ.get(MCP_TRAFFIC_LOG_ENV, MCP_TRAFFIC_LOG_DEFAULT) + try: + os.makedirs(os.path.dirname(traffic_path) or WORKSPACE_DIR, exist_ok=True) + with open(traffic_path, "a", encoding="utf-8") as f: + f.write(json.dumps(traffic_entry) + "\n") + except Exception: + pass + + return entry_id - log_path = os.path.join(WORKSPACE_DIR, LOG_FILE_NAME) +def log_tool_result(entry_id: str, tool: str, result: Any) -> None: + """Append a JSONL result entry paired with a tool invocation.""" + traffic_path = os.environ.get(MCP_TRAFFIC_LOG_ENV, MCP_TRAFFIC_LOG_DEFAULT) + result_entry = { + "id": entry_id, + "type": "result", + "tool": tool, + "result": result, + } try: - # Ensure workspace directory exists (it should in runtime environments) - os.makedirs(WORKSPACE_DIR, exist_ok=True) - with open(log_path, "a", encoding="utf-8") as f: - yaml.safe_dump(entry, f, explicit_start=True, sort_keys=False, allow_unicode=True) + os.makedirs(os.path.dirname(traffic_path) or WORKSPACE_DIR, exist_ok=True) + with open(traffic_path, "a", encoding="utf-8") as f: + f.write(json.dumps(result_entry, default=str) + "\n") except Exception: - # Do not break tool execution if logging fails pass -