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 

10import sys 

11import uuid 

12from typing import IO, Optional 

13 

14 

15# ── Trace context ───────────────────────────────────────────────────────────── 

16 

17 

18class TraceContext: 

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

20 

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

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

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

24 

25 

26# ── JSON Formatter ──────────────────────────────────────────────────────────── 

27 

28 

29class JSONFormatter(_stdlib_logging.Formatter): 

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

31 

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

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

34 self.trace_ctx = trace_ctx or TraceContext() 

35 

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

37 log_entry = { 

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

39 "level": record.levelname, 

40 "logger": record.name, 

41 "message": record.getMessage(), 

42 "pid": os.getpid(), 

43 "trace_id": self.trace_ctx.trace_id, 

44 "span_id": self.trace_ctx.span_id, 

45 } 

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

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

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

49 if extras and isinstance(extras, dict): 

50 log_entry.update(extras) 

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

52 

53 

54class _ExtraAdapter(_stdlib_logging.LoggerAdapter): 

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

56 

57 def process(self, msg, kwargs): 

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

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

60 kwargs["extra"] = extra 

61 return msg, kwargs 

62 

63 

64# ── Audit log ───────────────────────────────────────────────────────────────── 

65 

66 

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

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

69 extra = { 

70 "category": "AUDIT", 

71 "action": action, 

72 "user_id": user_id, 

73 "result": result, 

74 "details": details or {}, 

75 } 

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

77 

78 

79# ── Convenience helpers ────────────────────────────────────────────────────── 

80 

81 

82def setup_structured_logging( 

83 name: str, 

84 level: int = _stdlib_logging.INFO, 

85 stream: Optional[IO] = None, 

86 trace_ctx: Optional[TraceContext] = None, 

87) -> _stdlib_logging.Logger: 

88 """Create a logger with JSONFormatter attached. 

89 

90 Args: 

91 name: Logger name. 

92 level: Logging level (default INFO). 

93 stream: Output stream (default stderr). 

94 trace_ctx: Optional TraceContext for correlation. 

95 

96 Returns: 

97 Configured logger instance. 

98 """ 

99 logger = _stdlib_logging.getLogger(name) 

100 logger.setLevel(level) 

101 logger.propagate = False 

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

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

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

105 logger.addHandler(handler) 

106 return logger 

107 

108 

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

110 """Get or create a logger.""" 

111 return _stdlib_logging.getLogger(name)