otel_dispatch.py

#!/usr/bin/env python3
"""OTEL logging dispatcher for the copilot-custom-logging plugin.

One hook handler wired to all six Copilot CLI lifecycle events in
``hooks.json`` (SessionStart, SessionEnd, UserPromptSubmit, PreToolUse,
PostToolUse, Stop). On each event it:

  1. Reads the event payload (from ``$COPILOT_OTEL_INPUT`` or stdin).
  2. Loads the TOML config.
  3. Runs the shell commands configured for that event, exposing payload
     fields to them as ``COPILOT_HOOK_*`` environment variables.
  4. POSTs each command's stdout to ``{endpoint}/v1/logs`` as one OTLP log
     record -- an attribute ``copilot.<name>`` (default) or a ``body`` entry.

When ``otel.session_trace`` is enabled (default), it also emits one span per
session: SessionStart records the start time; SessionEnd POSTs the completed
span to ``{endpoint}/v1/traces``. Trace/span ids derive from ``session_id`` so
every log in a session attaches to its span, and a SessionEnd command with
``as = "span"`` adds its output to the span's attributes.

The process always exits 0 (fail-open) so telemetry never blocks the user's
tools; with no config it exits immediately without touching the network. All
logic lives here -- ``hooks.json`` embeds only a tiny bootstrap
(``hook_bootstrap.sh`` / ``.ps1``) that locates and runs this file.
"""

from __future__ import annotations

import datetime
import hashlib
import json
import os
import secrets
import shutil
import subprocess
import sys
import tempfile
import urllib.request

CONFIG_FILENAME = "copilot-otel-logging.toml"
SCOPE_NAME = "copilot-custom-logging"
SCOPE_VERSION = "1.0.0"
LOG_PREFIX = "[otel-logging]"

DEFAULT_SERVICE_NAME = "github-copilot-cli"
DEFAULT_TIMEOUT_SEC = 5

IS_WINDOWS = os.name == "nt"


def log_error(message: str) -> None:
    """Write a diagnostic line to stderr. Never raises."""
    try:
        print(f"{LOG_PREFIX} {message}", file=sys.stderr)
    except Exception:
        pass


# --------------------------------------------------------------------------
# Input / config loading
# --------------------------------------------------------------------------
def load_toml(path: str) -> dict:
    """Parse a TOML file with the stdlib (3.11+) or the ``tomli`` backport."""
    try:
        import tomllib as toml_reader  # Python 3.11+
    except ModuleNotFoundError:  # pragma: no cover - depends on runtime
        try:
            import tomli as toml_reader  # type: ignore
        except ModuleNotFoundError:
            raise RuntimeError(
                "TOML support requires Python 3.11+ (tomllib) or the 'tomli' package"
            )
    with open(path, "rb") as fh:
        return toml_reader.load(fh)


def find_config_path() -> str | None:
    """Return the first existing config path, or None.

    Order: ``$COPILOT_OTEL_CONFIG`` -> ``$COPILOT_HOME/<file>`` ->
    ``~/.copilot/<file>``.
    """
    explicit = os.environ.get("COPILOT_OTEL_CONFIG")
    if explicit:
        return explicit if os.path.isfile(explicit) else None

    candidates = []
    copilot_home = os.environ.get("COPILOT_HOME")
    if copilot_home:
        candidates.append(os.path.join(copilot_home, CONFIG_FILENAME))
    candidates.append(os.path.join(os.path.expanduser("~"), ".copilot", CONFIG_FILENAME))

    for path in candidates:
        if os.path.isfile(path):
            return path
    return None


def read_payload() -> dict:
    """Read the hook payload from ``$COPILOT_OTEL_INPUT`` or stdin."""
    raw = os.environ.get("COPILOT_OTEL_INPUT")
    if raw is None:
        try:
            raw = sys.stdin.read()
        except Exception:
            raw = ""
    raw = (raw or "").strip()
    if not raw:
        return {}
    try:
        return json.loads(raw)
    except Exception:
        return {}


# --------------------------------------------------------------------------
# Payload helpers
# --------------------------------------------------------------------------
def payload_get(payload: dict, key: str, default: str = "") -> str:
    """Return ``payload[key]`` coerced to a string (JSON for dict/list)."""
    value = payload.get(key)
    if value is None:
        return default
    if isinstance(value, (dict, list)):
        return json.dumps(value)
    return str(value)


def event_name(payload: dict) -> str:
    """Resolve the fired event's name (e.g. ``UserPromptSubmit``)."""
    return payload_get(payload, "hook_event_name")


def build_command_env(payload: dict, raw_payload: str) -> dict:
    """Base process env plus the exposed ``COPILOT_HOOK_*`` variables."""
    env = dict(os.environ)
    tool_result = payload.get("tool_result")
    if isinstance(tool_result, dict):
        tool_result_text = tool_result.get("text_result_for_llm", "")
    else:
        tool_result_text = ""

    env.update(
        {
            "COPILOT_HOOK_EVENT": event_name(payload),
            "COPILOT_HOOK_SESSION_ID": payload_get(payload, "session_id"),
            "COPILOT_HOOK_CWD": payload_get(payload, "cwd"),
            "COPILOT_HOOK_TIMESTAMP": payload_get(payload, "timestamp"),
            "COPILOT_HOOK_PROMPT": payload_get(payload, "prompt"),
            "COPILOT_HOOK_TOOL_NAME": payload_get(payload, "tool_name"),
            "COPILOT_HOOK_TOOL_RESULT": str(tool_result_text or ""),
            "COPILOT_HOOK_STOP_REASON": payload_get(payload, "stop_reason"),
            "COPILOT_HOOK_PAYLOAD": raw_payload,
        }
    )
    return env


# --------------------------------------------------------------------------
# Command execution + OTLP assembly
# --------------------------------------------------------------------------
def powershell_executable() -> str:
    """Return the PowerShell executable to use (prefers PowerShell 7 `pwsh`)."""
    return shutil.which("pwsh") or shutil.which("powershell") or "powershell"


def resolve_command(command: dict) -> tuple:
    """Pick the command string and shell for the current platform.

    Returns ``(cmd, use_powershell)``. On Windows, ``run_powershell`` is
    preferred and falls back to ``run``; on POSIX only ``run`` is used.
    ``(None, False)`` means skip the entry.
    """
    run = command.get("run")
    run_powershell = command.get("run_powershell")
    if IS_WINDOWS and run_powershell:
        return run_powershell, True
    return (run, False) if run else (None, False)


def run_commands(commands: list, env: dict, timeout_sec: int) -> list:
    """Run each command via the shell, capturing stdout.

    On Windows a command's ``run_powershell`` is preferred and executed with
    PowerShell; otherwise ``run`` is executed with the default shell (bash/sh
    on macOS/Linux, cmd on Windows). Returns a list of
    ``{"name", "output", "error"}`` dicts (one per command).
    """
    results = []
    for index, command in enumerate(commands):
        if not isinstance(command, dict):
            continue
        cmd, use_powershell = resolve_command(command)
        name = command.get("name") or f"command_{index}"
        if not cmd:
            continue
        target = str(command.get("as", "attribute")).strip().lower()
        if target not in ("body", "span"):
            target = "attribute"
        entry = {"name": name, "output": "", "error": "", "target": target}
        try:
            if use_powershell:
                completed = subprocess.run(
                    [
                        powershell_executable(),
                        "-NoProfile",
                        "-NonInteractive",
                        "-Command",
                        cmd,
                    ],
                    env=env,
                    capture_output=True,
                    text=True,
                    timeout=timeout_sec,
                )
            else:
                completed = subprocess.run(
                    cmd,
                    shell=True,
                    env=env,
                    capture_output=True,
                    text=True,
                    timeout=timeout_sec,
                )
            entry["output"] = (completed.stdout or "").strip()
            if completed.returncode != 0:
                entry["error"] = (
                    completed.stderr or ""
                ).strip() or f"exit code {completed.returncode}"
        except subprocess.TimeoutExpired:
            entry["error"] = f"timed out after {timeout_sec}s"
        except Exception as exc:  # noqa: BLE001 - surface any run failure as data
            entry["error"] = str(exc)
        results.append(entry)
    return results


def timestamp_to_unix_nano(payload: dict) -> int:
    """Event time (an ISO 8601 string) in unix nanoseconds; falls back to now()."""
    raw = payload.get("timestamp")
    if isinstance(raw, str) and raw.strip():
        text = raw.strip()
        if text.endswith("Z"):
            text = text[:-1] + "+00:00"
        try:
            parsed = datetime.datetime.fromisoformat(text)
            if parsed.tzinfo is None:
                parsed = parsed.replace(tzinfo=datetime.timezone.utc)
            return int(parsed.timestamp() * 1_000_000_000)
        except ValueError:
            pass
    return int(datetime.datetime.now(datetime.timezone.utc).timestamp() * 1_000_000_000)


def _attr(key: str, value: str) -> dict:
    return {"key": key, "value": {"stringValue": value}}


def _attr_key(name: str) -> str:
    """Turn a command name into an OTLP attribute key: ``copilot.<name>``."""
    slug = "_".join(str(name).split()).lower()
    return f"copilot.{slug}"


def _kvlist(fields: dict) -> dict:
    """Wrap a mapping as an OTLP kvlistValue (object) AnyValue."""
    return {
        "kvlistValue": {
            "values": [
                {"key": key, "value": {"stringValue": str(value)}}
                for key, value in fields.items()
            ]
        }
    }


def build_body(results: list) -> dict:
    """Build the log record body as an OTLP arrayValue.

    Only commands routed to the body (``as = "body"``) are included. Each
    element is an object with ``name`` and ``output``; an ``error`` field is
    added only when the command failed.
    """
    values = []
    for result in results:
        if result.get("target") != "body":
            continue
        fields = {"name": result["name"], "output": result["output"]}
        if result["error"]:
            fields["error"] = result["error"]
        values.append(_kvlist(fields))
    return {"arrayValue": {"values": values}}


def _routed_attributes(results: list, targets: tuple) -> list:
    """Build ``copilot.<name>`` attributes for results routed to ``targets``.

    A failed command additionally gets ``copilot.<name>.error``.
    """
    attributes = []
    for result in results:
        if result.get("target") not in targets:
            continue
        key = _attr_key(result["name"])
        attributes.append(_attr(key, result["output"]))
        if result["error"]:
            attributes.append(_attr(f"{key}.error", result["error"]))
    return attributes


def command_attributes(results: list) -> list:
    """Log-record attributes: commands routed to an attribute or the span."""
    return _routed_attributes(results, ("attribute", "span"))


def span_attributes(results: list) -> list:
    """Session-span attributes: commands with ``as = "span"`` (SessionEnd only)."""
    return _routed_attributes(results, ("span",))


def _session_hex(session_id: str) -> str:
    """Return the hex digits of a session id (drops UUID dashes etc.)."""
    return "".join(c for c in session_id.lower() if c in "0123456789abcdef")


def trace_id_for_session(session_id: str) -> str:
    """OTLP trace id for a session: its UUID (dashes removed), shared by all of
    the session's logs and its span. A non-UUID id falls back to a stable hash.
    """
    if not session_id:
        return secrets.token_hex(16)
    hex_id = _session_hex(session_id)
    if len(hex_id) >= 32:
        return hex_id[:32]
    return hashlib.sha256(session_id.encode("utf-8")).hexdigest()[:32]


def new_span_id() -> str:
    """Return a random 8-byte (16 hex char) OTLP span id for this event."""
    return secrets.token_hex(8)


def session_span_id(session_id: str) -> str:
    """OTLP span id for the session's root span: the session UUID's first 8
    bytes, shared by the span and every log record. Non-UUID falls back to a hash.
    """
    if not session_id:
        return secrets.token_hex(8)
    hex_id = _session_hex(session_id)
    if len(hex_id) >= 32:
        return hex_id[:16]
    return hashlib.sha256(("span:" + session_id).encode("utf-8")).hexdigest()[:16]


def session_marker_path(session_id: str) -> str:
    """Temp-file path where a session's start time is stashed between events."""
    safe = "".join(
        c if (c.isalnum() or c in "-_") else "_" for c in (session_id or "default")
    )
    return os.path.join(tempfile.gettempdir(), "copilot-otel-session-" + safe)


def write_session_start(session_id: str, start_unix_nano: int) -> None:
    """Record the session start time so SessionEnd can build a span."""
    try:
        with open(session_marker_path(session_id), "w") as fh:
            fh.write(str(start_unix_nano))
    except Exception as exc:  # noqa: BLE001
        log_error(f"failed to record session start: {exc}")


def read_session_start(session_id: str) -> int | None:
    """Read and remove a session's stored start time; None if not present."""
    path = session_marker_path(session_id)
    try:
        with open(path) as fh:
            value = int(fh.read().strip())
    except FileNotFoundError:
        return None
    except Exception:  # noqa: BLE001
        value = None
    try:
        os.remove(path)
    except OSError:
        pass
    return value


def build_trace_payload(
    session_id: str,
    cwd: str,
    start_unix_nano: int,
    end_unix_nano: int,
    service_name: str,
    end_reason: str = "",
    span_name: str = "copilot.session",
    extra_attributes: list | None = None,
) -> dict:
    """Assemble a single-span OTLP/HTTP-JSON trace for the whole session."""
    attributes = [_attr("copilot.session_id", session_id)]
    if cwd:
        attributes.append(_attr("copilot.cwd", cwd))
    if end_reason:
        attributes.append(_attr("copilot.session.end_reason", end_reason))
    if extra_attributes:
        attributes.extend(extra_attributes)
    return {
        "resourceSpans": [
            {
                "resource": {"attributes": [_attr("service.name", service_name)]},
                "scopeSpans": [
                    {
                        "scope": {"name": SCOPE_NAME, "version": SCOPE_VERSION},
                        "spans": [
                            {
                                "traceId": trace_id_for_session(session_id),
                                "spanId": session_span_id(session_id),
                                "name": span_name,
                                "kind": 1,  # SPAN_KIND_INTERNAL
                                "startTimeUnixNano": str(start_unix_nano),
                                "endTimeUnixNano": str(end_unix_nano),
                                "attributes": attributes,
                            }
                        ],
                    }
                ],
            }
        ]
    }


def build_otlp_payload(
    payload: dict,
    results: list,
    service_name: str,
    span_id: str | None = None,
    span_name: str = "",
) -> dict:
    """Assemble a single OTLP/HTTP-JSON log record for the event."""
    body = build_body(results)

    session_id = payload_get(payload, "session_id")
    attributes = [
        _attr("copilot.event", event_name(payload)),
        _attr("copilot.session_id", session_id),
        _attr("copilot.cwd", payload_get(payload, "cwd")),
    ]
    if span_name:
        attributes.append(_attr("copilot.session.name", span_name))
    tool_name = payload_get(payload, "tool_name")
    if tool_name:
        attributes.append(_attr("copilot.tool_name", tool_name))
    attributes.extend(command_attributes(results))

    time_unix_nano = str(timestamp_to_unix_nano(payload))
    return {
        "resourceLogs": [
            {
                "resource": {
                    "attributes": [_attr("service.name", service_name)]
                },
                "scopeLogs": [
                    {
                        "scope": {"name": SCOPE_NAME, "version": SCOPE_VERSION},
                        "logRecords": [
                            {
                                "timeUnixNano": time_unix_nano,
                                "observedTimeUnixNano": time_unix_nano,
                                "severityNumber": 9,
                                "severityText": "INFO",
                                "body": body,
                                "attributes": attributes,
                                "traceId": trace_id_for_session(session_id),
                                "spanId": span_id or new_span_id(),
                            }
                        ],
                    }
                ],
            }
        ]
    }


def expand_headers(headers: dict) -> dict:
    """Expand ``${ENV}`` / ``$ENV`` references in header values."""
    expanded = {}
    for key, value in headers.items():
        expanded[str(key)] = os.path.expandvars(str(value))
    return expanded


def signal_endpoint(endpoint: str, signal: str) -> str:
    """Return ``{endpoint}/v1/{signal}`` (e.g. ``.../v1/logs``)."""
    return endpoint.rstrip("/") + "/v1/" + signal


def post_otlp(
    endpoint: str, headers: dict, body: dict, timeout_sec: int, signal: str = "logs"
) -> None:
    """POST an OTLP JSON body to ``{endpoint}/v1/{signal}``. Raises on error."""
    data = json.dumps(body).encode("utf-8")
    request_headers = {"Content-Type": "application/json"}
    request_headers.update(expand_headers(headers))
    request = urllib.request.Request(
        signal_endpoint(endpoint, signal),
        data=data,
        headers=request_headers,
        method="POST",
    )
    with urllib.request.urlopen(request, timeout=timeout_sec) as response:
        status = getattr(response, "status", response.getcode())
        if not (200 <= status < 300):
            raise RuntimeError(f"HTTP {status}")


# --------------------------------------------------------------------------
# Entry point
# --------------------------------------------------------------------------
def main() -> int:
    payload = read_payload()
    raw_payload = json.dumps(payload)
    name = event_name(payload)
    if not name:
        return 0

    config_path = find_config_path()
    if not config_path:
        return 0

    try:
        config = load_toml(config_path)
    except Exception as exc:  # noqa: BLE001
        log_error(f"failed to load config {config_path}: {exc}")
        return 0

    otel = config.get("otel") or {}
    endpoint = otel.get("endpoint")
    if not endpoint:
        log_error("no otel.endpoint configured; skipping")
        return 0
    service_name = otel.get("service_name") or DEFAULT_SERVICE_NAME
    try:
        timeout_sec = int(otel.get("timeout_sec", DEFAULT_TIMEOUT_SEC))
    except (TypeError, ValueError):
        timeout_sec = DEFAULT_TIMEOUT_SEC
    headers = otel.get("headers") or {}
    session_trace = bool(otel.get("session_trace", True))
    span_name = otel.get("session_span_name") or "copilot.session"
    session_id = payload_get(payload, "session_id")

    # Run this event's configured commands (if any) up front, so their output
    # can feed both the log record and -- for SessionEnd, via ``as = "span"`` --
    # the session span's attributes.
    hooks = config.get("hooks") or {}
    commands = hooks.get(name)
    if isinstance(commands, dict):
        commands = [commands]
    env = build_command_env(payload, raw_payload)
    results = run_commands(commands, env, timeout_sec) if commands else []

    # Session span lifecycle: record start on SessionStart, emit the completed
    # span to /v1/traces on SessionEnd. Runs regardless of configured commands.
    if session_trace and session_id:
        event_nano = timestamp_to_unix_nano(payload)
        if name == "SessionStart":
            write_session_start(session_id, event_nano)
        elif name == "SessionEnd":
            start_nano = read_session_start(session_id)
            if start_nano is None:
                start_nano = event_nano
            trace_body = build_trace_payload(
                session_id,
                payload_get(payload, "cwd"),
                start_nano,
                event_nano,
                service_name,
                payload_get(payload, "reason"),
                span_name,
                span_attributes(results),
            )
            try:
                post_otlp(endpoint, headers, trace_body, timeout_sec, signal="traces")
            except Exception as exc:  # noqa: BLE001
                log_error(f"failed to export session trace to {endpoint}: {exc}")

    # Logs: export one record for this event's commands' output.
    if not results:
        return 0

    span_id = session_span_id(session_id) if (session_trace and session_id) else None
    span_name_attr = span_name if (session_trace and session_id) else ""
    otlp_body = build_otlp_payload(
        payload, results, service_name, span_id, span_name_attr
    )
    try:
        post_otlp(endpoint, headers, otlp_body, timeout_sec)
    except Exception as exc:  # noqa: BLE001
        log_error(f"failed to export logs to {endpoint}: {exc}")

    return 0


if __name__ == "__main__":
    try:
        sys.exit(main())
    except Exception as exc:  # noqa: BLE001 - never fail closed
        log_error(f"unexpected error: {exc}")
        sys.exit(0)