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

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 # noqa: E402 

10import sys # noqa: E402 

11import uuid # noqa: E402 

12from typing import IO # noqa: E402 

13 

14# ── Trace context ───────────────────────────────────────────────────────────── 

15 

16 

17class TraceContext: 

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

19 

20 def __init__(self, trace_id: str | None = None, span_id: str | None = None): 

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

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

23 

24 

25# ── JSON Formatter ──────────────────────────────────────────────────────────── 

26 

27 

28class JSONFormatter(_stdlib_logging.Formatter): 

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

30 

31 def __init__(self, fmt=None, datefmt=None, style="%", trace_ctx: TraceContext | None = None): 

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

33 self.trace_ctx = trace_ctx or TraceContext() 

34 

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

36 log_entry = { 

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

38 "level": record.levelname, 

39 "logger": record.name, 

40 "message": record.getMessage(), 

41 "pid": os.getpid(), 

42 "trace_id": self.trace_ctx.trace_id, 

43 "span_id": self.trace_ctx.span_id, 

44 } 

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

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

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

48 if extras and isinstance(extras, dict): 

49 log_entry.update(extras) 

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

51 

52 

53class _ExtraAdapter(_stdlib_logging.LoggerAdapter): 

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

55 

56 def process(self, msg, kwargs): 

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

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

59 kwargs["extra"] = extra 

60 return msg, kwargs 

61 

62 

63# ── Audit log ───────────────────────────────────────────────────────────────── 

64 

65 

66def audit_log( 

67 logger: _stdlib_logging.Logger, 

68 action: str, 

69 user_id: str, 

70 result: str, 

71 details: dict | None = None, 

72): 

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

74 extra = { 

75 "category": "AUDIT", 

76 "action": action, 

77 "user_id": user_id, 

78 "result": result, 

79 "details": details or {}, 

80 } 

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

82 

83 

84# ── Convenience helpers ────────────────────────────────────────────────────── 

85 

86 

87def setup_structured_logging( 

88 name: str, 

89 level: int = _stdlib_logging.INFO, 

90 stream: IO | None = None, 

91 trace_ctx: TraceContext | None = None, 

92) -> _stdlib_logging.Logger: 

93 """Create a logger with JSONFormatter attached. 

94 

95 Args: 

96 name: Logger name. 

97 level: Logging level (default INFO). 

98 stream: Output stream (default stderr). 

99 trace_ctx: Optional TraceContext for correlation. 

100 

101 Returns: 

102 Configured logger instance. 

103 """ 

104 logger = _stdlib_logging.getLogger(name) 

105 logger.setLevel(level) 

106 logger.propagate = False 

107 if not any( 

108 isinstance(h, _stdlib_logging.StreamHandler) and isinstance(h.formatter, JSONFormatter) 

109 for h in logger.handlers 

110 ): 

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

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

113 logger.addHandler(handler) 

114 return logger 

115 

116 

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

118 """Get or create a logger.""" 

119 return _stdlib_logging.getLogger(name)