Nexum-Expanded/runtime/src/nexum_runtime/cli.py

"""Command line entry point for the Nexum runtime."""

from __future__ import annotations

import argparse
import base64
import hmac
import json
import os
from dataclasses import replace
from pathlib import Path
from typing import Any

from .bundle import release_artifact_sha256, validate_bundle
from .capabilities import capability_catalog
from .executor import (
    TOOL_CALL_END,
    TOOL_CALL_START,
    execute_tool_calls,
    execute_tool_text,
    list_tools,
    parse_tool_calls,
    tool_schemas,
)
from .harness import NNFXHarnessRequest, NNFXHarnessSession
from .identity import candidate_identity_sha256, runtime_source_sha256
from .runner import NexumLocalRunner
from .server import serve
from .status import SELF_CORRECTION_STATUS, SELF_IMPROVEMENT_STATUS, runtime_health
from .templates import TEMPLATES, render_template, template_names
from .tooling.artifacts import ArtifactStore
from .tooling.events import EventLog
from .tooling.language_packs import (
    analyze_language_file,
    language_pack_catalog,
    verify_language_packs,
)
from .tooling.security import ApprovalStore
from .tooling.tasks import TaskStore


def _print_json(payload: object) -> None:
    print(json.dumps(payload, indent=2, sort_keys=True))


def _cmd_doctor(args: argparse.Namespace) -> int:
    model_dir = str(args.model or os.environ.get("NEXUM_MODEL_DIR", "")).strip()
    if not model_dir:
        payload = runtime_health()
        payload["error"] = "model directory is required"
        _print_json(payload)
        return 1
    report = validate_bundle(model_dir, deep=True)
    payload = runtime_health(bundle_ok=report.ok)
    payload["bundle"] = report.to_dict()
    _print_json(payload)
    return 0 if report.ok else 1


def _cmd_bundle_validate(args: argparse.Namespace) -> int:
    report = validate_bundle(args.model, deep=args.deep)
    _print_json(report.to_dict())
    return 0 if report.ok else 1


def _cmd_tools_list(_: argparse.Namespace) -> int:
    _print_json({"tools": list_tools()})
    return 0


def _cmd_tools_execute(args: argparse.Namespace) -> int:
    metadata_supplied = bool(
        args.call_id or args.idempotency_key or args.approval_id
    )
    if metadata_supplied:
        try:
            parsed = parse_tool_calls(args.tool_text)
        except ValueError as exc:
            _print_json({"ok": False, "error": f"invalid tool call: {exc}"})
            return 1
        if len(parsed) != 1:
            _print_json(
                {
                    "ok": False,
                    "error": "execution metadata requires exactly one tool call",
                }
            )
            return 1
        call = replace(
            parsed[0],
            call_id=args.call_id,
            idempotency_key=args.idempotency_key,
            approval_id=args.approval_id,
        )
        executed = execute_tool_calls(
            (call,),
            cwd=args.workspace,
            timeout_s=args.timeout_s,
            session_id=args.session_id,
        )
    else:
        executed = tuple(
            execute_tool_text(
                args.tool_text,
                cwd=args.workspace,
                timeout_s=args.timeout_s,
                session_id=args.session_id,
            )
        )
    results = [
        item.to_dict()
        for item in executed
    ]
    ok = bool(results) and all(item["ok"] for item in results)
    _print_json({"ok": ok, "results": results})
    return 0 if ok else 1


def _cmd_tools_exec_smoke(args: argparse.Namespace) -> int:
    workspace = Path(args.workspace).expanduser().resolve()
    workspace.mkdir(parents=True, exist_ok=True)
    call = f"{TOOL_CALL_START}[Bash(command='printf nexum-tool-ok')]{TOOL_CALL_END}"
    results = [
        item.to_dict()
        for item in execute_tool_text(
            call, cwd=str(workspace), timeout_s=args.timeout_s
        )
    ]
    ok = (
        len(results) == 1
        and bool(results[0]["ok"])
        and "nexum-tool-ok" in str(results[0]["output"])
    )
    _print_json({"ok": ok, "results": results})
    return 0 if ok else 1


def _cmd_status_self_correction(_: argparse.Namespace) -> int:
    _print_json(SELF_CORRECTION_STATUS)
    return 0


def _cmd_status_self_improvement(_: argparse.Namespace) -> int:
    _print_json(SELF_IMPROVEMENT_STATUS)
    return 0


def _configure_learning_state(args: argparse.Namespace) -> None:
    if str(args.state_dir).strip():
        os.environ["NEXUM_STATE_DIR"] = str(
            Path(args.state_dir).expanduser().resolve()
        )
    os.environ["NEXUM_STATE_NAMESPACE"] = str(args.state_namespace)


def _validated_candidate_identity(model_dir: str) -> str:
    normalized_model_dir = str(model_dir).strip()
    if not normalized_model_dir:
        raise ValueError("model directory is required")
    report = validate_bundle(normalized_model_dir, deep=True)
    if not report.ok:
        raise RuntimeError("Nexum model bundle validation failed")
    release_sha256 = release_artifact_sha256(
        normalized_model_dir,
        validated_report=report,
    )
    return candidate_identity_sha256(
        release_sha256,
        runtime_source_sha256(),
    )


def _cmd_learning_status(args: argparse.Namespace) -> int:
    _configure_learning_state(args)
    candidate_sha256 = _validated_candidate_identity(args.model)
    runner = NexumLocalRunner(args.model, device=args.device)
    _print_json(
        {
            **runner.learning_status(),
            "candidate_identity_sha256": candidate_sha256,
            "state_namespace": runner.state_namespace,
        }
    )
    return 0


def _cmd_learning_rollback(args: argparse.Namespace) -> int:
    _configure_learning_state(args)
    candidate_sha256 = _validated_candidate_identity(args.model)
    if not hmac.compare_digest(args.expected_candidate, candidate_sha256):
        raise ValueError("candidate identity changed before rollback")
    runner = NexumLocalRunner(args.model, device=args.device)
    _print_json(
        {
            **runner.rollback_learning(
                target_generation=args.target_generation,
                expected_generation=args.expected_generation,
            ),
            "candidate_identity_sha256": candidate_sha256,
            "state_namespace": runner.state_namespace,
        }
    )
    return 0


def _cmd_capabilities(_: argparse.Namespace) -> int:
    _print_json(capability_catalog())
    return 0


def _cmd_language_packs(args: argparse.Namespace) -> int:
    if bool(args.verify):
        verification = verify_language_packs()
        _print_json(verification)
        return 0 if bool(verification["ok"]) else 1
    _print_json(language_pack_catalog(args.query))
    return 0


def _cmd_language_inspect(args: argparse.Namespace) -> int:
    workspace = Path(args.workspace).expanduser().resolve()
    target = (workspace / args.path).resolve()
    try:
        relative = target.relative_to(workspace)
    except ValueError as exc:
        raise ValueError("language inspection path leaves workspace") from exc
    if relative.parts and relative.parts[0] == ".nexum":
        raise ValueError("workspace control paths are not source inputs")
    if not target.is_file():
        raise ValueError("language inspection target is not a file")
    _print_json(analyze_language_file(target, language=args.language))
    return 0


def _cmd_approvals_list(args: argparse.Namespace) -> int:
    records = [
        record.__dict__
        for record in ApprovalStore(args.workspace).list(
            session_id=args.session_id
        )
    ]
    _print_json({"approvals": records})
    return 0


def _cmd_approvals_decide(args: argparse.Namespace) -> int:
    record = ApprovalStore(args.workspace).decide(
        args.approval_id,
        approved=args.decision == "approve",
        session_id=args.session_id,
    )
    EventLog(args.workspace).append(
        "approval_decision",
        session_id=args.session_id,
        status=record.status,
        detail={"approval_id": record.approval_id, "tool": record.tool},
    )
    _print_json({"ok": True, "approval": record.__dict__})
    return 0


def _cmd_tasks_list(args: argparse.Namespace) -> int:
    records = [
        record.to_dict()
        for record in TaskStore(args.workspace).list(session_id=args.session_id)
    ]
    _print_json({"tasks": records})
    return 0


def _cmd_tasks_status(args: argparse.Namespace) -> int:
    _print_json(
        {
            "task": TaskStore(args.workspace)
            .status(args.task_id, session_id=args.session_id)
            .to_dict()
        }
    )
    return 0


def _cmd_tasks_cancel(args: argparse.Namespace) -> int:
    task = TaskStore(args.workspace).cancel(
        args.task_id, session_id=args.session_id
    )
    _print_json({"ok": task.status == "cancelled", "task": task.to_dict()})
    return 0 if task.status == "cancelled" else 1


def _cmd_artifacts_list(args: argparse.Namespace) -> int:
    records = [
        record.to_dict()
        for record in ArtifactStore(args.workspace).list(
            session_id=args.session_id
        )
    ]
    _print_json({"artifacts": records})
    return 0


def _cmd_artifacts_read(args: argparse.Namespace) -> int:
    record, data = ArtifactStore(args.workspace).read(
        args.artifact_id,
        offset=args.offset,
        length=args.length,
        session_id=args.session_id,
    )
    _print_json(
        {
            "artifact": record.to_dict(),
            "offset": args.offset,
            "data_base64": base64.b64encode(data).decode("ascii"),
        }
    )
    return 0


def _cmd_events_list(args: argparse.Namespace) -> int:
    events = [
        event.to_dict()
        for event in EventLog(args.workspace).read(
            session_id=args.session_id,
            after_sequence=args.after,
        )
    ]
    _print_json({"events": events})
    return 0


def _cmd_templates_list(_: argparse.Namespace) -> int:
    _print_json({"templates": template_names()})
    return 0


def _cmd_templates_show(args: argparse.Namespace) -> int:
    _print_json({"name": args.name, "template": TEMPLATES[args.name]})
    return 0


def _cmd_templates_write(args: argparse.Namespace) -> int:
    out = Path(args.out).expanduser()
    out.mkdir(parents=True, exist_ok=True)
    for name in template_names():
        (out / f"{name}.txt").write_text(render_template(name), encoding="utf-8")
    _print_json({"ok": True, "out": str(out), "templates": template_names()})
    return 0


def _cmd_serve(args: argparse.Namespace) -> int:
    serve(
        host=args.host,
        port=args.port,
        model_dir=args.model,
        workspace=args.workspace,
        device=args.device,
        api_key=os.environ.get("NEXUM_API_KEY", ""),
        enable_tools=bool(args.enable_tools),
        state_dir=args.state_dir,
        state_namespace=args.state_namespace,
    )
    return 0


def _json_file(path: str) -> Any:
    return json.loads(Path(path).expanduser().read_text(encoding="utf-8"))


def _completion_payload(args: argparse.Namespace) -> dict[str, Any]:
    payload: dict[str, Any] = {
        "model": "Nexum",
        "temperature": float(args.temperature),
        "session_id": str(args.session_id or ""),
    }
    if args.messages:
        messages = _json_file(args.messages)
        if not isinstance(messages, list):
            raise ValueError("messages file must contain a JSON list")
        payload["messages"] = messages
    else:
        payload["prompt"] = str(args.prompt or "")
    if args.tools:
        payload["tools"] = tool_schemas()
    if args.max_completion_tokens is not None:
        payload["max_completion_tokens"] = int(args.max_completion_tokens)
    return payload


def _cmd_chat(args: argparse.Namespace) -> int:
    response = NexumLocalRunner(args.model, device=args.device).chat(
        _completion_payload(args)
    )
    _print_json(response)
    return 0


def _cmd_agent(args: argparse.Namespace) -> int:
    runner = NexumLocalRunner(args.model, device=args.device)
    payload = _completion_payload(args)
    payload["workspace"] = args.workspace
    payload["timeout_s"] = float(args.timeout_s)
    request = NNFXHarnessRequest.from_json(payload)
    response = NNFXHarnessSession(
        complete=runner.chat,
        sign_observation=runner.sign_observation,
    ).run(request)
    _print_json(response.to_dict())
    return 0 if response.ok else 1


def _add_completion_arguments(
    parser: argparse.ArgumentParser, *, include_workspace: bool
) -> None:
    parser.add_argument(
        "--model", default=os.environ.get("NEXUM_MODEL_DIR", ""), required=False
    )
    parser.add_argument("--device", default=os.environ.get("NEXUM_DEVICE", "cuda:0"))
    source = parser.add_mutually_exclusive_group(required=True)
    source.add_argument("--prompt")
    source.add_argument("--messages", help="JSON file containing chat messages")
    parser.add_argument(
        "--tools", action="store_true", help="include the built-in tool schemas"
    )
    parser.add_argument("--temperature", type=float, default=0.0)
    parser.add_argument("--max-completion-tokens", type=int)
    parser.add_argument("--session-id", default="")
    if include_workspace:
        parser.add_argument(
            "--workspace", default=os.environ.get("NEXUM_WORKSPACE", ".")
        )
        parser.add_argument("--timeout-s", type=float, default=0.0)


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(prog="nexum", description="Nexum private runtime")
    sub = parser.add_subparsers(dest="command", required=True)

    doctor = sub.add_parser("doctor", help="show runtime health")
    doctor.add_argument("--model", default=os.environ.get("NEXUM_MODEL_DIR", ""))
    doctor.set_defaults(fn=_cmd_doctor)

    bundle = sub.add_parser("bundle", help="model folder operations")
    bundle_sub = bundle.add_subparsers(dest="bundle_command", required=True)
    bundle_validate = bundle_sub.add_parser(
        "validate", help="validate a flat model folder"
    )
    bundle_validate.add_argument("--model", required=True)
    bundle_validate.add_argument("--deep", action="store_true")
    bundle_validate.set_defaults(fn=_cmd_bundle_validate)

    tools = sub.add_parser("tools", help="tool operations")
    tools_sub = tools.add_subparsers(dest="tools_command", required=True)
    tools_list = tools_sub.add_parser("list", help="list tools")
    tools_list.set_defaults(fn=_cmd_tools_list)
    tools_execute = tools_sub.add_parser("execute", help="execute supplied tool text")
    tools_execute.add_argument("tool_text")
    tools_execute.add_argument("--workspace", default=".")
    tools_execute.add_argument("--timeout-s", type=float, default=0.0)
    tools_execute.add_argument("--session-id", default="")
    tools_execute.add_argument(
        "--call-id", default="", help="stable identifier for one exact tool call"
    )
    tools_execute.add_argument(
        "--idempotency-key",
        default="",
        help="caller-selected exactly-once key for one tool call",
    )
    tools_execute.add_argument(
        "--approval-id",
        default="",
        help="approved exact-action receipt used by one tool-call retry",
    )
    tools_execute.set_defaults(fn=_cmd_tools_execute)
    tools_smoke = tools_sub.add_parser(
        "exec-smoke", help="run a small terminal tool smoke"
    )
    tools_smoke.add_argument("--workspace", default=".")
    tools_smoke.add_argument("--timeout-s", type=float, default=0.0)
    tools_smoke.set_defaults(fn=_cmd_tools_exec_smoke)

    status = sub.add_parser("status", help="runtime feature status")
    status_sub = status.add_subparsers(dest="status_command", required=True)
    sc = status_sub.add_parser("self-correction", help="show self-correction status")
    sc.set_defaults(fn=_cmd_status_self_correction)
    si = status_sub.add_parser("self-improvement", help="show self-improvement status")
    si.set_defaults(fn=_cmd_status_self_improvement)

    learning = sub.add_parser("learning", help="inspect or restore shared learning")
    learning_sub = learning.add_subparsers(dest="learning_command", required=True)
    for learning_parser in (
        learning_sub.add_parser("status", help="show validated learning history"),
        learning_sub.add_parser("rollback", help="restore a retained generation"),
    ):
        learning_parser.add_argument(
            "--model",
            default=os.environ.get("NEXUM_MODEL_DIR", ""),
            required=False,
        )
        learning_parser.add_argument(
            "--device",
            default=os.environ.get("NEXUM_DEVICE", "cuda:0"),
        )
        learning_parser.add_argument(
            "--state-dir",
            default=os.environ.get("NEXUM_STATE_DIR", ""),
        )
        learning_parser.add_argument(
            "--state-namespace",
            default=os.environ.get("NEXUM_STATE_NAMESPACE", "default"),
        )
    learning_status = learning_sub.choices["status"]
    learning_status.set_defaults(fn=_cmd_learning_status)
    learning_rollback = learning_sub.choices["rollback"]
    learning_rollback.add_argument("--target-generation", type=int, required=True)
    learning_rollback.add_argument("--expected-generation", type=int, required=True)
    learning_rollback.add_argument("--expected-candidate", required=True)
    learning_rollback.set_defaults(fn=_cmd_learning_rollback)

    capabilities_parser = sub.add_parser(
        "capabilities", help="show runtime capability discovery"
    )
    capabilities_parser.set_defaults(fn=_cmd_capabilities)

    language_packs_parser = sub.add_parser(
        "language-packs",
        help="show source-intelligence packs and measured native toolchains",
    )
    language_packs_parser.add_argument(
        "--query",
        default="",
        help="filter by language name, alias, identifier, or extension",
    )
    language_packs_parser.add_argument(
        "--verify",
        action="store_true",
        help="exercise every declared source-intelligence backend",
    )
    language_packs_parser.set_defaults(fn=_cmd_language_packs)

    language_inspect_parser = sub.add_parser(
        "language-inspect",
        help="inspect one workspace source file with its language pack",
    )
    language_inspect_parser.add_argument("path")
    language_inspect_parser.add_argument("--language", default="")
    language_inspect_parser.add_argument("--workspace", default=".")
    language_inspect_parser.set_defaults(fn=_cmd_language_inspect)

    approvals = sub.add_parser("approvals", help="manage exact action approvals")
    approvals_sub = approvals.add_subparsers(dest="approvals_command", required=True)
    approvals_list = approvals_sub.add_parser("list", help="list approval records")
    approvals_list.add_argument("--workspace", default=".")
    approvals_list.add_argument("--session-id", default="")
    approvals_list.set_defaults(fn=_cmd_approvals_list)
    approvals_decide = approvals_sub.add_parser(
        "decide", help="approve or deny one exact action"
    )
    approvals_decide.add_argument("approval_id")
    approvals_decide.add_argument("decision", choices=("approve", "deny"))
    approvals_decide.add_argument("--session-id", required=True)
    approvals_decide.add_argument("--workspace", default=".")
    approvals_decide.set_defaults(fn=_cmd_approvals_decide)

    tasks = sub.add_parser("tasks", help="inspect durable task lifecycle")
    tasks_sub = tasks.add_subparsers(dest="tasks_command", required=True)
    tasks_list = tasks_sub.add_parser("list", help="list durable tasks")
    tasks_list.add_argument("--workspace", default=".")
    tasks_list.add_argument("--session-id", default="")
    tasks_list.set_defaults(fn=_cmd_tasks_list)
    tasks_status = tasks_sub.add_parser("status", help="show one durable task")
    tasks_status.add_argument("task_id")
    tasks_status.add_argument("--workspace", default=".")
    tasks_status.add_argument("--session-id", default="")
    tasks_status.set_defaults(fn=_cmd_tasks_status)
    tasks_cancel = tasks_sub.add_parser("cancel", help="cancel one running task")
    tasks_cancel.add_argument("task_id")
    tasks_cancel.add_argument("--workspace", default=".")
    tasks_cancel.add_argument("--session-id", default="")
    tasks_cancel.set_defaults(fn=_cmd_tasks_cancel)

    artifacts = sub.add_parser("artifacts", help="inspect complete tool artifacts")
    artifacts_sub = artifacts.add_subparsers(dest="artifacts_command", required=True)
    artifacts_list = artifacts_sub.add_parser("list", help="list artifacts")
    artifacts_list.add_argument("--workspace", default=".")
    artifacts_list.add_argument("--session-id", default="")
    artifacts_list.set_defaults(fn=_cmd_artifacts_list)
    artifacts_read = artifacts_sub.add_parser("read", help="read an artifact range")
    artifacts_read.add_argument("artifact_id")
    artifacts_read.add_argument("--workspace", default=".")
    artifacts_read.add_argument("--session-id", default="")
    artifacts_read.add_argument("--offset", type=int, default=0)
    artifacts_read.add_argument("--length", type=int)
    artifacts_read.set_defaults(fn=_cmd_artifacts_read)

    events = sub.add_parser("events", help="inspect ordered runtime receipts")
    events.add_argument("--workspace", default=".")
    events.add_argument("--session-id", default="")
    events.add_argument("--after", type=int, default=0)
    events.set_defaults(fn=_cmd_events_list)

    templates = sub.add_parser("templates", help="template operations")
    templates_sub = templates.add_subparsers(dest="templates_command", required=True)
    templates_list = templates_sub.add_parser("list", help="list templates")
    templates_list.set_defaults(fn=_cmd_templates_list)
    templates_show = templates_sub.add_parser("show", help="show a template")
    templates_show.add_argument("name", choices=template_names())
    templates_show.set_defaults(fn=_cmd_templates_show)
    templates_write = templates_sub.add_parser(
        "write", help="write templates to a folder"
    )
    templates_write.add_argument("--out", required=True)
    templates_write.set_defaults(fn=_cmd_templates_write)

    chat = sub.add_parser("chat", help="run one Nexum completion")
    _add_completion_arguments(chat, include_workspace=False)
    chat.set_defaults(fn=_cmd_chat)

    agent = sub.add_parser("agent", help="run the model-driven NNF X tool loop")
    _add_completion_arguments(agent, include_workspace=True)
    agent.set_defaults(fn=_cmd_agent)

    server = sub.add_parser("serve", help="run HTTP server")
    server.add_argument("--host", default="127.0.0.1")
    server.add_argument("--port", type=int, default=8080)
    server.add_argument("--model", default=os.environ.get("NEXUM_MODEL_DIR", ""))
    server.add_argument("--workspace", default=os.environ.get("NEXUM_WORKSPACE", "."))
    server.add_argument("--device", default=os.environ.get("NEXUM_DEVICE", "cuda:0"))
    server.add_argument("--state-dir", default=os.environ.get("NEXUM_STATE_DIR", ""))
    server.add_argument(
        "--state-namespace",
        default=os.environ.get("NEXUM_STATE_NAMESPACE", "default"),
    )
    server.add_argument(
        "--enable-tools",
        action="store_true",
        default=os.environ.get("NEXUM_ENABLE_TOOLS", "").lower()
        in {"1", "true", "yes"},
        help="allow authenticated HTTP tool execution and agent endpoints",
    )
    server.set_defaults(fn=_cmd_serve)

    return parser


def main(argv: list[str] | None = None) -> int:
    args = build_parser().parse_args(argv)
    return int(args.fn(args))


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