agent_deply/pam_deploy_graph/cli.py

"""PAM 部署 Agent 的命令行入口。"""

from __future__ import annotations

import argparse
import json
from dataclasses import asdict

from .agent import PamDeployAgent
from .checkpoint_store import load_agent_state, redact_mapping
from .interactive import run_interactive_chat
from .langgraph_runtime import LangGraphDeploymentRuntime, LangGraphRunResult
from .llm import build_llm_client
from .mcp_factory import build_mcp_runner_from_config
from .params_loader import load_params_file


def add_llm_args(parser: argparse.ArgumentParser) -> None:
    """为子命令追加真实 LLM 配置参数。"""
    parser.add_argument("--llm-base-url")
    parser.add_argument("--llm-api-key")
    parser.add_argument("--llm-model")
    parser.add_argument("--llm-action-analysis-prompt-file")


def add_mcp_args(parser: argparse.ArgumentParser) -> None:
    """为需要执行 MCP action 的子命令追加 MCP 配置参数。"""
    parser.add_argument("--mcp-config", help="MCP client JSON 配置文件路径")


def add_action_analysis_arg(parser: argparse.ArgumentParser) -> None:
    """为执行类子命令追加 action 后诊断开关。"""
    parser.add_argument("--analyze-actions", action="store_true", help="每个 action 后追加 LLM/规则诊断建议")


def require_confirm(args: argparse.Namespace) -> None:
    """真实执行前强制要求命令行显式传入 --confirm。"""
    if not getattr(args, "confirm", False):
        raise SystemExit("Refusing to execute actions without --confirm.")


def print_pause_payload(agent: PamDeployAgent, state) -> None:
    """输出 checkpoint 和待确认信息，便于用户续跑或确认。"""
    if state.pending_confirmation:
        print(json.dumps({"confirmation": agent.build_confirmation_request(state)}, ensure_ascii=False, indent=2))
    if state.checkpoint_path:
        print(json.dumps({"checkpoint": state.checkpoint_path}, ensure_ascii=False, indent=2))


def run_graph_once(agent: PamDeployAgent, state, *, flow: str = "deploy") -> LangGraphRunResult:
    """用 LangGraph runtime 执行一次状态，返回图执行结果。"""
    runtime = LangGraphDeploymentRuntime(agent=agent, flow=flow)  # type: ignore[arg-type]
    return runtime.start(state)


def print_graph_result(agent: PamDeployAgent, result: LangGraphRunResult) -> None:
    """输出 LangGraph 执行结果、报告和暂停信息。"""
    state = result.state
    if result.report:
        print(result.report)
    elif state is not None:
        print(agent.render_report(state))
    if result.interrupted and result.confirmation:
        print(json.dumps({"confirmation": result.confirmation}, ensure_ascii=False, indent=2))
    if state is not None:
        print_pause_payload(agent, state)


def main() -> None:
    """解析 CLI 参数并分发到对应命令。"""
    parser = argparse.ArgumentParser(prog="pam-deploy-agent")
    sub = parser.add_subparsers(dest="command", required=True)

    preview = sub.add_parser("preview")
    preview.add_argument("--config", required=True)
    preview.add_argument("--strategy", default="hybrid_node_mcp", choices=["hybrid_node_mcp", "script_only", "fake"])

    analyze = sub.add_parser("analyze")
    analyze.add_argument("--text", required=True)
    analyze.add_argument("--config")
    add_llm_args(analyze)

    chat = sub.add_parser("chat")
    chat.add_argument("--config", required=True)
    chat.add_argument("--strategy", default="fake", choices=["hybrid_node_mcp", "script_only", "fake"])
    chat.add_argument("--target-ip", action="append", default=[])
    chat.add_argument("--checkpoint")
    add_llm_args(chat)
    add_mcp_args(chat)
    add_action_analysis_arg(chat)

    run = sub.add_parser("run-global")
    run.add_argument("--config", required=True)
    run.add_argument("--strategy", default="fake", choices=["hybrid_node_mcp", "script_only", "fake"])
    run.add_argument("--checkpoint")
    run.add_argument("--confirm", action="store_true")
    add_llm_args(run)
    add_mcp_args(run)
    add_action_analysis_arg(run)

    deploy = sub.add_parser("run-deploy")
    deploy.add_argument("--config", required=True)
    deploy.add_argument("--strategy", default="fake", choices=["hybrid_node_mcp", "script_only", "fake"])
    deploy.add_argument("--target-ip", action="append", default=[])
    deploy.add_argument("--checkpoint")
    deploy.add_argument("--confirm", action="store_true")
    add_llm_args(deploy)
    add_mcp_args(deploy)
    add_action_analysis_arg(deploy)

    resume = sub.add_parser("resume")
    resume.add_argument("--checkpoint", required=True)
    resume.add_argument("--confirm", action="store_true")
    add_llm_args(resume)
    add_mcp_args(resume)
    add_action_analysis_arg(resume)

    confirm = sub.add_parser("confirm")
    confirm.add_argument("--checkpoint", required=True)
    confirm.add_argument("--decision", required=True, choices=["approve", "reject"])
    confirm.add_argument("--note", default="")
    confirm.add_argument("--confirm", action="store_true")
    add_llm_args(confirm)
    add_mcp_args(confirm)
    add_action_analysis_arg(confirm)

    args = parser.parse_args()
    params = load_params_file(args.config) if getattr(args, "config", None) else {}
    llm_client = None
    if args.command != "preview":
        llm_client = build_llm_client(
            base_url=getattr(args, "llm_base_url", None),
            api_key=getattr(args, "llm_api_key", None),
            model=getattr(args, "llm_model", None),
            action_analysis_prompt_path=getattr(args, "llm_action_analysis_prompt_file", None),
        )
    mcp_runner = None
    if getattr(args, "mcp_config", None):
        mcp_runner = build_mcp_runner_from_config(args.mcp_config)
    agent = PamDeployAgent(
        llm_client=llm_client,
        mcp_runner=mcp_runner,
        action_analysis_enabled=bool(getattr(args, "analyze_actions", False)),
    )

    if args.command == "analyze":
        result = agent.analyze_request(args.text, params)
        payload = redact_mapping({key: asdict(value) for key, value in result.items()})
        print(json.dumps(payload, ensure_ascii=False, indent=2))
        return

    if args.command == "chat":
        run_interactive_chat(
            agent=agent,
            params=params,
            strategy=args.strategy,
            checkpoint_path=args.checkpoint,
            target_ips=args.target_ip,
        )
        return

    if args.command == "preview":
        print(agent.preview(params, args.strategy))
        return

    require_confirm(args)
    if args.command == "run-global":
        state = agent.create_state(
            params=params,
            execution_strategy=args.strategy,
            checkpoint_path=args.checkpoint,
        )
        result = run_graph_once(agent, state, flow="global")
        if result.state is not None:
            print(json.dumps({"events": result.state.events}, ensure_ascii=False, indent=2))
            print_pause_payload(agent, result.state)
        return

    if args.command == "resume":
        state = load_agent_state(args.checkpoint)
        state.checkpoint_path = state.checkpoint_path or args.checkpoint
        if state.paused:
            state = agent.resume_state(state)
        result = run_graph_once(agent, state, flow="deploy")
        print_graph_result(agent, result)
        return

    if args.command == "confirm":
        state = load_agent_state(args.checkpoint)
        state.checkpoint_path = state.checkpoint_path or args.checkpoint
        runtime = LangGraphDeploymentRuntime(agent=agent, flow="deploy")
        first = runtime.start(state)
        if first.interrupted:
            result = runtime.resume(approved=args.decision == "approve", note=args.note)
            print_graph_result(agent, result)
            return
        print_graph_result(agent, first)
        return

    state = agent.create_state(
        params=params,
        execution_strategy=args.strategy,
        checkpoint_path=args.checkpoint,
        target_ips=args.target_ip,
    )
    result = run_graph_once(agent, state, flow="deploy")
    print_graph_result(agent, result)


if __name__ == "__main__":
    main()