6krrt/logs.py

"""Structured logging for the dispatcher.

The service had no logger. It had four ``print(..., file=sys.stderr)`` calls,
all of them failure paths, so a request that WORKED said nothing at all — not
the category, the tier, the candidates, the model chosen, the cost, or the
latency. Watching a live agent session showed only uvicorn's access line, which
does not even name the model that served it.

Three things this module provides, in order of how much they matter:

**A trace id per request.** One request touches classification, filtering,
ranking, dispatch, verification and possibly a retry, and until now nothing
tied those together. The id is carried in a ContextVar so helpers read it
without every signature growing a parameter, and it is printed on every line.

**logfmt, not prose.** ``route id=r7f3a91 cat=coding_general pick=...`` reads
fine in ``journalctl`` and greps without a JSON parser. Prose reads better once
and aggregates never.

**Real journald priorities.** systemd strips a ``<N>`` prefix off a log line and
records the message at that priority (``SyslogLevelPrefix`` defaults to true),
which is what makes ``journalctl -p warning`` mean something. Without it every
line lands at PRIORITY 6 and severity cannot be filtered at all.

The prefix is emitted ONLY when ``$JOURNAL_STREAM`` is set, which systemd
exports when it owns our stderr and a terminal does not — so running uvicorn in
a foreground shell prints clean lines rather than a literal ``<6>`` on each one.
"""

from __future__ import annotations

import contextvars
import logging
import os
import secrets
import sys
from typing import Any, Optional, TextIO

LOGGER_NAME = "llm_router"
LEVEL_ENV = "LLM_ROUTER_LOG_LEVEL"

log = logging.getLogger(LOGGER_NAME)

LEVELS = {
    "debug": logging.DEBUG,
    "info": logging.INFO,
    "warning": logging.WARNING,
    "error": logging.ERROR,
}

# syslog priorities, which is what journald speaks. 5 (notice) and 1 (alert)
# are unused: nothing here is between info and warning, or above error.
PRIORITIES = {
    logging.DEBUG: 7,
    logging.INFO: 6,
    logging.WARNING: 4,
    logging.ERROR: 3,
    logging.CRITICAL: 2,
}

NO_TRACE = "-"

_trace: contextvars.ContextVar[str] = contextvars.ContextVar("trace_id", default=NO_TRACE)


# --- trace ids ------------------------------------------------------------

def new_trace() -> str:
    """Start a new request trace and return its id."""
    trace_id = "r" + secrets.token_hex(3)
    _trace.set(trace_id)
    return trace_id


def current_trace() -> str:
    return _trace.get()


def set_trace(trace_id: str) -> None:
    """Adopt an existing trace id in this context.

    Needed on the streaming path. A StreamingResponse's generator is iterated
    from a different context than the endpoint that built it, so the ContextVar
    set in the endpoint is NOT visible inside the generator — it must capture
    the id as a local and re-set it here, or every streamed request logs no id
    at all. That is the path all agent traffic takes.
    """
    _trace.set(trace_id)


# --- rendering ------------------------------------------------------------

def render_value(value: Any) -> str:
    """One logfmt value: unambiguous, single-line, cheap to read."""
    if value is None:
        return "-"
    if isinstance(value, bool):
        return "1" if value else "0"
    if isinstance(value, float):
        # Six significant figures. Enough for microdollars and kWh in
        # scientific notation, without 0.0007830000000000001.
        return f"{value:.6g}"
    text = str(value)
    if text == "":
        return '""'
    if any(c in text for c in ' "\n\t='):
        text = text.replace("\\", "\\\\").replace('"', '\\"')
        text = text.replace("\n", " ").replace("\t", " ")
        return f'"{text}"'
    return text


def render(name: str, fields: dict[str, Any], trace_id: Optional[str] = None) -> str:
    parts = [name, f"id={trace_id or current_trace()}"]
    parts += [f"{key}={render_value(value)}" for key, value in fields.items()]
    return " ".join(parts)


def _emit(
    level: int, name: str, fields: dict[str, Any], trace_id: Optional[str] = None
) -> None:
    """Emit one logfmt line, if this level is enabled.

    The level check comes first so DEBUG formatting costs nothing when DEBUG is
    off — this runs on every request.
    """
    if log.isEnabledFor(level):
        log.log(level, "%s", render(name, fields, trace_id))


def event(level: int, name: str, **fields: Any) -> None:
    _emit(level, name, fields)


def debug(name: str, **fields: Any) -> None:
    _emit(logging.DEBUG, name, fields)


def info(name: str, **fields: Any) -> None:
    _emit(logging.INFO, name, fields)


def warning(name: str, **fields: Any) -> None:
    _emit(logging.WARNING, name, fields)


def error(name: str, **fields: Any) -> None:
    _emit(logging.ERROR, name, fields)


class Bound:
    """An emitter pinned to one trace id, for code the ContextVar cannot reach.

    A StreamingResponse's generator is resumed by starlette through a
    threadpool, and every ``next()`` runs in a FRESH COPY of the caller's
    context — so a trace id set on one resumption is gone by the next, and the
    ``finally`` block that logs the dispatch line sees nothing. Measured: the
    streamed dispatch line logged ``id=-`` until this existed.

    Carrying the id explicitly is the only thing that survives that, and
    streaming is the path all agent traffic takes.
    """

    __slots__ = ("trace_id",)

    def __init__(self, trace_id: str):
        self.trace_id = trace_id

    def debug(self, name: str, **fields: Any) -> None:
        _emit(logging.DEBUG, name, fields, self.trace_id)

    def info(self, name: str, **fields: Any) -> None:
        _emit(logging.INFO, name, fields, self.trace_id)

    def warning(self, name: str, **fields: Any) -> None:
        _emit(logging.WARNING, name, fields, self.trace_id)

    def error(self, name: str, **fields: Any) -> None:
        _emit(logging.ERROR, name, fields, self.trace_id)


def bind(trace_id: Optional[str] = None) -> Bound:
    """Pin an emitter to a trace id, defaulting to the current one."""
    return Bound(trace_id or current_trace())


def enabled_for_debug() -> bool:
    """For callers that must do real work to produce debug fields."""
    return log.isEnabledFor(logging.DEBUG)


# --- configuration --------------------------------------------------------

class PriorityFormatter(logging.Formatter):
    """Plain message, optionally prefixed with a journald priority."""

    def __init__(self, prefix: bool):
        super().__init__("%(message)s")
        self.prefix = prefix

    def format(self, record: logging.LogRecord) -> str:
        message = super().format(record)
        if not self.prefix:
            return message
        return f"<{PRIORITIES.get(record.levelno, 6)}>{message}"


def under_journald(stream: Optional[TextIO] = None) -> bool:
    """Whether this stream really is the journal.

    systemd sets $JOURNAL_STREAM to the "device:inode" of the stream it
    connected, and the variable is INHERITED by every child — including one
    whose stderr has since been redirected to a file or a pipe. Presence alone
    is therefore a false positive, and it fired immediately: a foreground
    uvicorn started from a systemd-managed session, with its output redirected
    to a file, wrote a literal <7> on every line.

    Comparing the value against the actual fd is what systemd documents, and it
    is the only check that answers the real question.
    """
    expected = os.environ.get("JOURNAL_STREAM")
    if not expected:
        return False
    device, _, inode = expected.partition(":")
    try:
        stat = os.fstat((stream if stream is not None else sys.stderr).fileno())
        return stat.st_dev == int(device) and stat.st_ino == int(inode)
    except (AttributeError, OSError, ValueError):
        # No fileno at all (a StringIO under test), or an unparseable value.
        return False


def resolve_level(configured: str) -> int:
    """The env var wins, so a live service can be debugged without a repo edit.

    config.yaml is the documented home of the setting -- every knob belongs
    there -- but flipping it means editing a tracked file and leaving a stray
    diff. LLM_ROUTER_LOG_LEVEL goes in a systemd drop-in instead.
    """
    override = os.environ.get(LEVEL_ENV)
    if override:
        level = LEVELS.get(override.strip().lower())
        if level is not None:
            return level
        print(
            f"{LEVEL_ENV}={override!r} is not one of {sorted(LEVELS)}; "
            f"falling back to logging.level={configured!r}",
            file=sys.stderr,
        )
    return LEVELS.get((configured or "info").strip().lower(), logging.INFO)


def configure(
    level: str = "info",
    *,
    stream: Optional[TextIO] = None,
    journald: Optional[bool] = None,
) -> logging.Logger:
    """Install the one handler this logger gets. Idempotent.

    ``propagate`` is off: uvicorn configures its own loggers and a root handler
    installed by anything else would print every line a second time.
    """
    resolved = resolve_level(level)
    log.setLevel(resolved)
    log.propagate = False
    for existing in list(log.handlers):
        log.removeHandler(existing)

    target = stream if stream is not None else sys.stderr
    handler = logging.StreamHandler(target)
    handler.setLevel(resolved)
    handler.setFormatter(
        PriorityFormatter(under_journald(target) if journald is None else journald)
    )
    log.addHandler(handler)
    return log