Coverage for agentos/log/formatter.py: 42%

1"""AgentOS logging — structured JSON formatter with trace context.""" 

2 

3from __future__ import annotations 

4 

5import importlib 

6import json 

7 

8_stdlib_logging = importlib.import_module("logging") 

9import os 

10import sys 

11import time 

12import uuid 

13from typing import IO, Optional 

14 

15 

16# ── Trace context ───────────────────────────────────────────────────────────── 

17 

18 

19class TraceContext: 

20 """Carries trace_id and span_id through a request lifecycle.""" 

21 

22 def __init__(self, trace_id: Optional[str] = None, span_id: Optional[str] = None): 

23 self.trace_id = trace_id or uuid.uuid4().hex[:16] 

24 self.span_id = span_id or uuid.uuid4().hex[:8] 

25 

26 

27# ── JSON Formatter ──────────────────────────────────────────────────────────── 

28 

29 

30class JSONFormatter(_stdlib_logging.Formatter): 

31 """Emits log records as JSON with trace context fields.""" 

32 

33 def __init__(self, fmt=None, datefmt=None, style="%", trace_ctx: Optional[TraceContext] = None): 

34 super().__init__(fmt, datefmt, style) 

35 self.trace_ctx = trace_ctx or TraceContext() 

36 

37 def format(self, record: _stdlib_logging.LogRecord) -> str: 

38 log_entry = { 

39 "timestamp": self.formatTime(record, self.datefmt or "%Y-%m-%dT%H:%M:%S.%fZ"), 

40 "level": record.levelname, 

41 "logger": record.name, 

42 "message": record.getMessage(), 

43 "pid": os.getpid(), 

44 "trace_id": self.trace_ctx.trace_id, 

45 "span_id": self.trace_ctx.span_id, 

46 } 

47 if record.exc_info and record.exc_info[0]: 

48 log_entry["exc_info"] = self.formatException(record.exc_info) 

49 extras = getattr(record, "_structured_extra", None) 

50 if extras and isinstance(extras, dict): 

51 log_entry.update(extras) 

52 return json.dumps(log_entry, default=str, ensure_ascii=False) 

53 

54 

55class _ExtraAdapter(_stdlib_logging.LoggerAdapter): 

56 """Logging adapter that merges extra dict into the JSON output.""" 

57 

58 def process(self, msg, kwargs): 

59 extra = kwargs.get("extra", {}) 

60 extra["_structured_extra"] = kwargs.pop("structured_extra", {}) 

61 kwargs["extra"] = extra 

62 return msg, kwargs 

63 

64 

65# ── Audit log ───────────────────────────────────────────────────────────────── 

66 

67 

68def audit_log(logger: _stdlib_logging.Logger, action: str, user_id: str, result: str, details: Optional[dict] = None): 

69 """Emit a structured audit log entry.""" 

70 extra = { 

71 "category": "AUDIT", 

72 "action": action, 

73 "user_id": user_id, 

74 "result": result, 

75 "details": details or {}, 

76 } 

77 logger.info(f"AUDIT {action} by {user_id}: {result}", extra={"structured_extra": extra}) 

78 

79 

80# ── Convenience helpers ────────────────────────────────────────────────────── 

81 

82 

83def setup_structured_logging( 

84 name: str, 

85 level: int = _stdlib_logging.INFO, 

86 stream: Optional[IO] = None, 

87 trace_ctx: Optional[TraceContext] = None, 

88) -> _stdlib_logging.Logger: 

89 """Create a logger with JSONFormatter attached. 

90 

91 Args: 

92 name: Logger name. 

93 level: Logging level (default INFO). 

94 stream: Output stream (default stderr). 

95 trace_ctx: Optional TraceContext for correlation. 

96 

97 Returns: 

98 Configured logger instance. 

99 """ 

100 logger = _stdlib_logging.getLogger(name) 

101 logger.setLevel(level) 

102 logger.propagate = False 

103 if not any(isinstance(h, _stdlib_logging.StreamHandler) and isinstance(h.formatter, JSONFormatter) for h in logger.handlers): 

104 handler = _stdlib_logging.StreamHandler(stream or sys.stderr) 

105 handler.setFormatter(JSONFormatter(trace_ctx=trace_ctx or TraceContext())) 

106 logger.addHandler(handler) 

107 return logger 

108 

109 

110def get_logger(name: str) -> _stdlib_logging.Logger: 

111 """Get or create a logger.""" 

112 return _stdlib_logging.getLogger(name)