"""Audit logging service for MCP tool invocations. [Task]: T058 [From]: specs/004-ai-chatbot/tasks.md This module provides audit logging for all MCP tool invocations to track usage patterns, detect abuse, and maintain compliance records. """ import logging import json from datetime import datetime from typing import Any, Optional from uuid import UUID from sqlmodel import Session from core.database import engine # Configure audit logger audit_logger = logging.getLogger("audit") audit_logger.setLevel(logging.INFO) # Audit log handler (separate from main logs) audit_handler = logging.FileHandler("logs/audit.log") audit_handler.setFormatter(logging.Formatter( '%(asctime)s | %(levelname)s | %(message)s' )) audit_logger.addHandler(audit_handler) def log_tool_invocation( tool_name: str, user_id: str | UUID, args: dict[str, Any], result: dict[str, Any], conversation_id: Optional[str | UUID] = None, execution_time_ms: Optional[float] = None, error: Optional[str] = None ) -> None: """Log an MCP tool invocation for audit purposes. [From]: specs/004-ai-chatbot/spec.md - NFR-018 Args: tool_name: Name of the tool that was invoked user_id: ID of the user who invoked the tool args: Arguments passed to the tool result: Result returned by the tool conversation_id: Optional conversation context execution_time_ms: Optional execution time in milliseconds error: Optional error message if invocation failed """ log_entry = { "timestamp": datetime.utcnow().isoformat(), "tool_name": tool_name, "user_id": str(user_id), "conversation_id": str(conversation_id) if conversation_id else None, "success": error is None, "error": error, "execution_time_ms": execution_time_ms, "args_summary": _summarize_args(tool_name, args), "result_summary": _summarize_result(result) } # Log to file audit_logger.info(json.dumps(log_entry)) # Also log to database for querying (if needed) _persist_audit_log(log_entry) def _summarize_args(tool_name: str, args: dict[str, Any]) -> dict[str, Any]: """Create a summary of tool arguments for logging. [From]: T058 - Add audit logging for all MCP tool invocations Args: tool_name: Name of the tool args: Full arguments dict Returns: Summarized arguments (sanitized for sensitive data) """ # Don't log full user content for privacy if "message" in args: return {"message_length": len(str(args.get("message", "")))} # For task operations, log relevant info if tool_name in ["add_task", "update_task", "complete_task", "delete_task"]: summary = {} if "task_id" in args: summary["task_id"] = str(args["task_id"]) if "title" in args: summary["title"] = args["title"][:50] # Truncate long titles if "completed" in args: summary["completed"] = args["completed"] if "priority" in args: summary["priority"] = args["priority"] return summary # For list_tasks, log filters if tool_name == "list_tasks": summary = {} if "status" in args: summary["status"] = args["status"] if "limit" in args: summary["limit"] = args["limit"] return summary # Default: return all args (tool-specific sanitization could be added) return args def _summarize_result(result: dict[str, Any]) -> dict[str, Any]: """Create a summary of tool result for logging. [From]: T058 - Add audit logging for all MCP tool invocations Args: result: Full result dict from tool Returns: Summarized result """ if not isinstance(result, dict): return {"result_type": type(result).__name__} summary = {} # Extract key fields if "success" in result: summary["success"] = result["success"] if "error" in result: summary["error"] = result["error"] if "task" in result: task = result["task"] summary["task_id"] = task.get("id") summary["task_title"] = task.get("title", "")[:50] if task.get("title") else None if "tasks" in result: tasks = result.get("tasks", []) summary["task_count"] = len(tasks) if isinstance(tasks, list) else 0 if "updated_count" in result: summary["updated_count"] = result["updated_count"] if "deleted_count" in result: summary["deleted_count"] = result["deleted_count"] if "message" in result: # Truncate long messages msg = result["message"] summary["message"] = msg[:100] + "..." if len(msg) > 100 else msg return summary def _persist_audit_log(log_entry: dict) -> None: """Persist audit log to database for querying. [From]: T058 - Add audit logging for all MCP tool invocations Args: log_entry: The audit log entry to persist """ # Note: This could be extended to write to an audit_logs table # For now, file-based logging is sufficient pass def get_user_activity_summary( user_id: str | UUID, limit: int = 100 ) -> list[dict[str, Any]]: """Get a summary of user activity from audit logs. [From]: T058 - Add audit logging for all MCP tool invocations Args: user_id: User ID to get activity for limit: Maximum number of entries to return Returns: List of audit log entries for the user """ # Read audit log file and filter by user_id try: with open("logs/audit.log", "r") as f: user_entries = [] for line in f: try: entry = json.loads(line.split(" | ", 2)[-1]) if entry.get("user_id") == str(user_id): user_entries.append(entry) if len(user_entries) >= limit: break except (json.JSONDecodeError, IndexError): continue return user_entries except FileNotFoundError: return [] # Decorator for automatic audit logging of MCP tools def audit_log(tool_name: Optional[str] = None): """Decorator to automatically log MCP tool invocations. [From]: T058 - Add audit logging for all MCP tool invocations Args: tool_name: Optional override for tool name (defaults to function name) Usage: @audit_log("add_task") async def add_task(user_id: str, title: str, ...): ... """ import functools import time def decorator(func): @functools.wraps(func) async def wrapper(*args, **kwargs): name = tool_name or func.__name__ start_time = time.time() # Extract user_id from args/kwargs user_id = kwargs.get("user_id") or (args[0] if args else None) try: result = await func(*args, **kwargs) execution_time = (time.time() - start_time) * 1000 log_tool_invocation( tool_name=name, user_id=user_id or "unknown", args=kwargs, result=result, execution_time_ms=execution_time ) return result except Exception as e: execution_time = (time.time() - start_time) * 1000 log_tool_invocation( tool_name=name, user_id=user_id or "unknown", args=kwargs, result={}, execution_time_ms=execution_time, error=str(e) ) raise return wrapper return decorator __all__ = [ "log_tool_invocation", "get_user_activity_summary", "audit_log", ]