JWeBot/packages/jwebot-extension-host/src/jwebot_extension_host/__main__.py

"""Single-request JSON-line entrypoint for the JWeBot extension host."""

from __future__ import annotations

import argparse
import json
import logging
import os
import sys
import time
from collections.abc import Mapping
from pathlib import Path

from jwebot_kernel import RuntimeTelemetry

from .runtime import (
    CapabilityUnavailable,
    ExtensionRuntime,
    HostError,
    process_request,
)


MAX_REQUEST_BYTES = 256 * 1024
MAX_SERVICE_CALLS = 8


def _scrub_environment() -> None:
    retained = {
        key: value
        for key, value in os.environ.items()
        if key in {"LANG", "LC_ALL", "TZ"}
    }
    os.environ.clear()
    os.environ.update(retained)


def _read_message(stream: object) -> dict[str, object]:
    raw = stream.readline(MAX_REQUEST_BYTES + 1)
    if not raw or len(raw) > MAX_REQUEST_BYTES:
        raise HostError("request payload is empty or too large")
    try:
        value = json.loads(raw)
    except json.JSONDecodeError as error:
        raise HostError("request payload is invalid") from error
    if not isinstance(value, dict):
        raise HostError("request payload must be an object")
    return value


def _write_message(
    value: Mapping[str, object],
    stream: object,
) -> None:
    payload = (
        json.dumps(
            value,
            ensure_ascii=False,
            sort_keys=True,
            separators=(",", ":"),
        )
        + "\n"
    )
    encoded = payload.encode("utf-8")
    if len(encoded) > MAX_REQUEST_BYTES:
        raise HostError("response payload is too large")
    stream.write(encoded)
    stream.flush()


class _ProtocolServiceBroker:
    def __init__(
        self,
        input_stream: object,
        output_stream: object,
        telemetry: RuntimeTelemetry,
    ) -> None:
        self.calls = 0
        self.input_stream = input_stream
        self.output_stream = output_stream
        self.telemetry = telemetry

    def __call__(
        self,
        profile: str,
        operation: str,
        payload: Mapping[str, object],
    ) -> Mapping[str, object]:
        started = time.monotonic()
        self.calls += 1
        if self.calls > MAX_SERVICE_CALLS:
            raise CapabilityUnavailable("service call limit exceeded")
        call_id = str(self.calls)
        _write_message(
            {
                "status": "service_request",
                "call_id": call_id,
                "profile": profile,
                "operation": operation,
                "payload": dict(payload),
            },
            self.output_stream,
        )
        response = _read_message(self.input_stream)
        if response.get("call_id") != call_id:
            raise CapabilityUnavailable("service response id mismatch")
        if response.get("op") == "service_error":
            self.telemetry.elapsed(
                "host.service.wait",
                started,
                status="error",
                plugin_id=profile,
                operation=operation,
                error_type=str(response.get("error", "ServiceError")),
            )
            raise CapabilityUnavailable("kernel service call failed")
        if response.get("op") != "service_result":
            raise CapabilityUnavailable("kernel service response is invalid")
        result = response.get("result")
        if not isinstance(result, Mapping):
            raise CapabilityUnavailable("kernel service result is invalid")
        self.telemetry.elapsed(
            "host.service.wait",
            started,
            plugin_id=profile,
            operation=operation,
        )
        return dict(result)


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--release-dir", type=Path, required=True)
    parser.add_argument("--allowed-signers", type=Path)
    parser.add_argument("--publisher")
    parser.add_argument("--allow-unsigned-local", action="store_true")
    args = parser.parse_args()
    if args.allow_unsigned_local:
        if args.allowed_signers is not None or args.publisher is not None:
            parser.error(
                "--allow-unsigned-local cannot include publisher trust options"
            )
    elif args.allowed_signers is None:
        parser.error("signed host requires --allowed-signers")
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s %(levelname)s %(message)s",
    )
    _scrub_environment()
    telemetry = RuntimeTelemetry(component="extension-host")
    process_started = time.monotonic()
    input_stream = sys.stdin.buffer
    output_stream = sys.stdout.buffer
    try:
        service_broker = _ProtocolServiceBroker(
            input_stream,
            output_stream,
            telemetry,
        )
        load_started = time.monotonic()
        runtime = ExtensionRuntime.load(
            args.release_dir,
            allowed_signers=args.allowed_signers,
            expected_publisher=args.publisher,
            allow_unsigned=args.allow_unsigned_local,
            service_broker=service_broker,
        )
        telemetry.elapsed(
            "host.runtime.load",
            load_started,
            plugin_id=runtime.manifest.plugin_id,
            process="extension-host",
        )
        request = _read_message(input_stream)
        event_value = request.get("event")
        event_id = (
            event_value.get("event_id", "")
            if isinstance(event_value, Mapping)
            else ""
        )
        telemetry.emit(
            "host.request.received",
            event_id=event_id,
            plugin_id=runtime.manifest.plugin_id,
            operation=str(request.get("op", "")),
            process="extension-host",
        )
        response = process_request(runtime, request)
    except HostError as error:
        logging.warning("extension host request failed error_type=%s", type(error).__name__)
        telemetry.elapsed(
            "host.process.complete",
            process_started,
            status="error",
            process="extension-host",
            error_type=type(error).__name__,
        )
        response = {"status": "error", "error": type(error).__name__}
    except Exception as error:
        logging.error("extension host failed error_type=%s", type(error).__name__)
        telemetry.elapsed(
            "host.process.complete",
            process_started,
            status="error",
            process="extension-host",
            error_type=type(error).__name__,
        )
        response = {"status": "error", "error": "HostUnavailable"}
    else:
        telemetry.elapsed(
            "host.process.complete",
            process_started,
            plugin_id=str(response.get("plugin_id", "")),
            operation=str(request.get("op", "")),
            process="extension-host",
        )
    _write_message(response, output_stream)
    return 0 if response.get("status") == "ok" else 1


if __name__ == "__main__":
    raise SystemExit(main())