Source code for taimoe.platform.proxy.headers

"""Build Taimoe + W3C Trace Context headers for L4 proxy mode.

Produces two layers of headers in one call:

- **Taimoe-native** ``X-Taimoe-*`` headers for audit / governance / routing
  consumed by the Taimoe Gateway and audit logs.
- **W3C Trace Context** (``traceparent`` / ``tracestate``) so the same trace
  flows transparently into any downstream service or APM that speaks OTel.

IDs are W3C-format (32-hex trace, 16-hex span) — no Taimoe-only prefixes,
so OTel collectors accept them without translation.
"""

from __future__ import annotations

from typing import Final

from taimoe.platform._ids import (
    new_span_id,
    new_trace_id,
    normalize_span_id,
    normalize_trace_id,
)

# --- Header name constants --------------------------------------------------
# Single source of truth — backend audit middleware and tests should import
# these instead of hard-coding the strings.

HEADER_TRACE_ID: Final = "X-Taimoe-Trace-Id"
HEADER_SPAN_ID: Final = "X-Taimoe-Span-Id"
HEADER_PARENT_SPAN_ID: Final = "X-Taimoe-Parent-Span-Id"
HEADER_RUNTIME_ID: Final = "X-Taimoe-Runtime-Id"
HEADER_AGENT_ID: Final = "X-Taimoe-Agent-Id"
HEADER_PARENT_AGENT_ID: Final = "X-Taimoe-Parent-Agent-Id"
HEADER_STEP: Final = "X-Taimoe-Step"
HEADER_SESSION_ID: Final = "X-Taimoe-Session-Id"
HEADER_USER_ID: Final = "X-Taimoe-User-Id"

HEADER_TRACEPARENT: Final = "traceparent"
HEADER_TRACESTATE: Final = "tracestate"

# W3C traceparent format: "{version}-{trace_id}-{span_id}-{flags}"
_TRACEPARENT_VERSION: Final = "00"
_TRACEPARENT_SAMPLED_FLAG: Final = "01"
_TRACEPARENT_NOT_SAMPLED_FLAG: Final = "00"
_TRACESTATE_VENDOR: Final = "taimoe"

_MAX_ID_FIELD_LEN: Final = 256  # generous cap for runtime_id / agent_id / etc.


def _validate_id_field(name: str, value: str) -> str:
    """Reject empty, CRLF-bearing, or oversized HTTP header values."""
    if not isinstance(value, str) or not value or not value.strip():
        raise ValueError(f"{name} must be a non-empty string.")
    if "\r" in value or "\n" in value:
        raise ValueError(f"{name} must not contain CR or LF characters.")
    if len(value) > _MAX_ID_FIELD_LEN:
        raise ValueError(f"{name} exceeds the {_MAX_ID_FIELD_LEN}-char limit.")
    return value




def build_traceparent(
    trace_id: str, span_id: str, *, sampled: bool = True
) -> str:
    """Build a W3C ``traceparent`` value from a (trace_id, span_id) pair.

    Args:
        trace_id: 32 lowercase hex chars.
        span_id: 16 lowercase hex chars.
        sampled: Whether the trace should be exported by downstream APMs.

    Raises:
        ValueError: if either ID is not in W3C format.
    """
    flags = _TRACEPARENT_SAMPLED_FLAG if sampled else _TRACEPARENT_NOT_SAMPLED_FLAG
    return f"{_TRACEPARENT_VERSION}-{trace_id}-{span_id}-{flags}"





def build_tracestate(*, runtime_id: str | None = None) -> str:
    """Build a minimal ``tracestate`` carrying our vendor segment.

    For now we only stamp the runtime; we can add more vendor keys later
    without breaking parsers (downstream services preserve unknown vendors).
    """
    if not runtime_id:
        return f"{_TRACESTATE_VENDOR}="
    return f"{_TRACESTATE_VENDOR}=runtime:{runtime_id}"





def build_taimoe_headers(
    *,
    runtime_id: str,
    agent_id: str,
    trace_id: str | None = None,
    span_id: str | None = None,
    parent_span_id: str | None = None,
    parent_agent_id: str | None = None,
    step: str | None = None,
    session_id: str | None = None,
    user_id: str | None = None,
    sampled: bool = True,
    include_traceparent: bool = True,
) -> dict[str, str]:
    """Compose Taimoe-native + W3C trace headers for an outbound HTTP call.

    Args:
        runtime_id: ID of the runtime making the call. Required.
        agent_id: ID of the agent making the call. Required.
        trace_id: 32-hex trace ID. Generated if absent.
        span_id: 16-hex span ID for this call. Generated if absent.
        parent_span_id: 16-hex span ID of the caller, when known.
        parent_agent_id: Agent ID of the caller, when known.
        step: Lifecycle step label (e.g. ``"pre_call"``).
        session_id: End-user session correlation ID.
        user_id: End-user identity (already privacy-filtered upstream).
        sampled: Whether to set the W3C sampled flag.
        include_traceparent: When False, only emit ``X-Taimoe-*`` headers.
            Useful when the caller already manages W3C headers itself.

    Returns:
        A flat ``dict[str, str]`` suitable for ``httpx``/``requests``.
    """
    runtime_id = _validate_id_field("runtime_id", runtime_id)
    agent_id = _validate_id_field("agent_id", agent_id)
    resolved_trace_id = normalize_trace_id(trace_id) if trace_id else new_trace_id()
    resolved_span_id = normalize_span_id(span_id) if span_id else new_span_id()

    headers: dict[str, str] = {
        HEADER_TRACE_ID: resolved_trace_id,
        HEADER_SPAN_ID: resolved_span_id,
        HEADER_RUNTIME_ID: runtime_id,
        HEADER_AGENT_ID: agent_id,
    }

    if parent_span_id is not None:
        headers[HEADER_PARENT_SPAN_ID] = normalize_span_id(parent_span_id)
    for header_name, raw_value in (
        (HEADER_PARENT_AGENT_ID, parent_agent_id),
        (HEADER_STEP, step),
        (HEADER_SESSION_ID, session_id),
        (HEADER_USER_ID, user_id),
    ):
        if raw_value is None:
            continue
        headers[header_name] = _validate_id_field(header_name, raw_value)

    if include_traceparent:
        headers[HEADER_TRACEPARENT] = build_traceparent(
            resolved_trace_id, resolved_span_id, sampled=sampled
        )
        headers[HEADER_TRACESTATE] = build_tracestate(runtime_id=runtime_id)

    return headers