Add logging for arm mcp server testing

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",

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()

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
-