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

2

3 from __future__ import annotations

4

5 import importlib

6 import json

7

8_stdlib_logging = importlib.import_module("logging")

9 import os

10 import sys

11 import time

12 import uuid

13 from typing import IO, Optional

14

15

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

17

18

19 class 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

30 class 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

55 class _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

68 def 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

83 def 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

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

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

112 return _stdlib_logging.getLogger(name)