kimi_subagent_router

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""kimi-subagent-router — 按 subagent 类型路由模型别名与思考强度。

子命令:
  mcp                     以 MCP stdio 服务运行（插件 manifest 使用此模式）
  route --type NAME       返回某 subagent 类型应使用的模型别名与思考强度
  list                    列出全部路由映射
  doctor [--json]         检查 subagent-models.json 与 config.toml 一致性
  print-config            打印需并入 config.toml 的 TOML 片段（手动安装用）
  setup-config [--home H] 自动把片段并入 config.toml（标记块、幂等、先备份）
  remove-config [--home H] 移除已并入的片段
  init-json [--path P]    写出默认 subagent-models.json

所有 subagent -> 模型/思考强度映射一律从外部 subagent-models.json 读取；
该文件缺失时使用内置默认值。配置文件查找顺序:
  1. $KIMI_SUBAGENT_MODELS 指定的路径
  2. <home>/subagent-models.json（home 依次取 --home 参数、$KIMI_CODE_HOME、~/.kimi-code）
  3. <插件根目录>/subagent-models.json
  4. 内置默认值
"""

import argparse
import copy
import datetime
import json
import os
import re
import shutil
import sys

PLUGIN_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))

MARKER_BEGIN = "# >>> kimi-subagent-router >>>"
MARKER_END = "# <<< kimi-subagent-router <<<"

# ---------------------------------------------------------------------------
# 内置默认值：仅在 subagent-models.json 不存在时使用（用户文件始终优先）。
# model 一律指向"模型变体别名"；思考强度由该别名在 config.toml 里注册的
# default_effort 绑定（Agent 工具没有 effort 参数）。
# ---------------------------------------------------------------------------
DEFAULT_CONFIG = {
    "subagents": {
        "cheap-coder": {
            "model": "k2.8-high",
            "effort": "high",
            "subagent_type": "coder",
            "description": "Mechanical, pattern-fixed coding: bulk renames, formatting, boilerplate, trivial test stubs (K2.8 high)",
        },
        "cheap-explore": {
            "model": "k2.8-low",
            "effort": "low",
            "subagent_type": "explore",
            "description": "Read-only exploration: locate files, read structure, search symbols (K2.8 low)",
        },
        "k2.8-medium": {
            "model": "k2.8-high",
            "effort": "high",
            "subagent_type": "coder",
            "description": "Local code change scoped to one file (K2.8 high)",
        },
        "k2.8-debug": {
            "model": "k2.8-max",
            "effort": "max",
            "subagent_type": "coder",
            "description": "Simple feature work, routine tests, known-symptom debugging (K2.8 max; non-obvious bugs belong to k3-planner)",
        },
        "k3-planner": {
            "model": "k3-256k-max",
            "effort": "max",
            "subagent_type": "plan",
            "description": "Complex design, architecture, cross-module refactors, hard bugs, critical code (K3-256k max; preferred heavy tier)",
        },
        "k3-max": {
            "model": "k3-max",
            "effort": "max",
            "subagent_type": "plan",
            "description": "Hardest tasks that ALSO need more than 256k context (kimi-code/k3, 1M window, max)",
        },
        "doc-writer": {
            "model": "alibaba-cn/qwen3.8-flash",
            "effort": "xhigh",
            "subagent_type": "coder",
            "exclude_tools": ["shell", "background_task"],
            "description": "Documentation: README, API docs, prose and comment cleanup (third-party alibaba-cn/qwen3.8-flash, xhigh tier)",
        },
    },
    "default_model": None,
}

# 插件可自动注册的"内置管理模型变体"：变体别名 -> 底层模型 ID。
# 仅用于 setup-config 生成 [models] 条目；用户的映射仍以 JSON 配置为准。
VARIANT_BASE_MODEL = {
    "k2.8-low": "kimi-for-coding",
    "k2.8-high": "kimi-for-coding",
    "k2.8-max": "kimi-for-coding",
    "k3-256k-high": "k3-256k",
    "k3-256k-max": "k3-256k",
    "k3-max": "k3",
}

# 当 config.toml 中找不到底层模型条目时的字段模板（Kimi Code 内置托管模型）。
BUILTIN_MODEL_FALLBACKS = {
    "kimi-for-coding": {
        "provider": "managed:kimi-code",
        "model": "kimi-for-coding",
        "max_context_size": 1048576,
        "capabilities": ["thinking", "always_thinking", "image_in", "video_in", "tool_use"],
        "support_efforts": ["low", "high", "max"],
    },
    "k3-256k": {
        "provider": "managed:kimi-code",
        "model": "k3-256k",
        "max_context_size": 262144,
        "capabilities": ["thinking", "always_thinking", "image_in", "video_in", "tool_use"],
        "support_efforts": ["low", "high", "max"],
    },
    "k3": {
        "provider": "managed:kimi-code",
        "model": "k3",
        "max_context_size": 1048576,
        "capabilities": ["thinking", "always_thinking", "image_in", "video_in", "tool_use"],
        "support_efforts": ["low", "high", "max"],
    },
}

# [secondary_model] 池默认值的兜底：与临时环境基准一致（K3-256k + max）。
FALLBACK_POOL_DEFAULT = "k3-256k-max"

_BARE_KEY = re.compile(r"^[A-Za-z0-9_-]+$")


# ---------------------------------------------------------------------------
# 配置读取
# ---------------------------------------------------------------------------

def find_models_json(home_arg=None):
    env = os.environ.get("KIMI_SUBAGENT_MODELS")
    if env:
        return env, "$KIMI_SUBAGENT_MODELS"
    homes = []
    if home_arg:
        homes.append(home_arg)
    if os.environ.get("KIMI_CODE_HOME"):
        homes.append(os.environ["KIMI_CODE_HOME"])
    homes.append(os.path.join(os.path.expanduser("~"), ".kimi-code"))
    for home in homes:
        p = os.path.join(home, "subagent-models.json")
        if os.path.isfile(p):
            return p, os.path.join(os.path.basename(home) or home, "subagent-models.json")
    p = os.path.join(PLUGIN_ROOT, "subagent-models.json")
    if os.path.isfile(p):
        return p, os.path.join("<plugin>", "subagent-models.json")
    return None, "builtin defaults"


def load_routes(home_arg=None):
    """返回 (routes_dict, source_label, path_or_None)。文件缺失用内置默认值。"""
    path, source = find_models_json(home_arg)
    if path is None:
        return copy.deepcopy(DEFAULT_CONFIG), source, None
    try:
        with open(path, "r", encoding="utf-8") as f:
            data = json.load(f)
    except Exception as e:  # noqa: BLE001 - 报错信息要直达用户
        raise SystemExit(f"Failed to parse config file {path}: {e}")
    if not isinstance(data, dict) or not isinstance(data.get("subagents"), dict):
        raise SystemExit(f"Invalid config file {path}: expected an object containing a \"subagents\" object")
    return data, source, path


def resolve_home(home_arg):
    if home_arg:
        return home_arg
    return os.environ.get("KIMI_CODE_HOME") or os.path.join(
        os.path.expanduser("~"), ".kimi-code")


def config_toml_path(home_arg):
    return os.path.join(resolve_home(home_arg), "config.toml")


def read_config_toml(home_arg):
    """返回 (text, parsed_or_None)。parsed 在无法解析时为 None。"""
    path = config_toml_path(home_arg)
    if not os.path.isfile(path):
        return "", None
    with open(path, "r", encoding="utf-8") as f:
        text = f.read()
    parsed = None
    try:
        import tomllib
        with open(path, "rb") as f:
            parsed = tomllib.load(f)
    except Exception:  # noqa: BLE001 - 文本扫描兜底
        pass
    return text, parsed


def extract_models(parsed, text):
    """返回 {alias: entry_dict}，含 overrides 子表（如有）。"""
    if parsed is not None and isinstance(parsed.get("models"), dict):
        return parsed["models"]
    models = {}
    for m in re.finditer(r'^\[models\.("([^"]+)"|([A-Za-z0-9_-]+))\]\s*$', text, re.M):
        alias = m.group(2) or m.group(3)
        body = text[m.end():]
        nxt = re.search(r'^\[', body, re.M)
        sect = body[: nxt.start()] if nxt else body
        entry = {}
        for line in sect.splitlines():
            kv = re.match(r'^\s*([A-Za-z_]+)\s*=\s*(.+?)\s*$', line)
            if kv:
                entry[kv.group(1)] = kv.group(2).strip().strip('"')
        models[alias] = entry
    return models


def extract_secondary_model(parsed, text):
    if parsed is not None:
        sm = parsed.get("secondary_model")
        return sm if isinstance(sm, dict) else None
    m = re.search(r'^\[secondary_model\]\s*$', text, re.M)
    if not m:
        return None
    body = text[m.end():]
    nxt = re.search(r'^\[(?!models\.)', body, re.M)
    sect = body[: nxt.start()] if nxt else body
    out = {"models": {}}
    for line in sect.splitlines():
        kv = re.match(r'^\s*([A-Za-z_]+|"(?:[^"\\]|\\.)*")\s*=\s*(.+?)\s*$', line)
        if not kv:
            continue
        key, val = kv.group(1).strip('"'), kv.group(2).strip().strip('"')
        if key == "models":
            continue
        if key in ("force",):
            out[key] = val.lower() == "true"
        else:
            out[key] = val
    mm = re.search(r'^\[secondary_model\.models\]\s*$', text, re.M)
    if mm:
        body = text[mm.end():]
        nxt = re.search(r'^\[', body, re.M)
        sect = body[: nxt.start()] if nxt else body
        for line in sect.splitlines():
            kv = re.match(r'^\s*("(?:[^"\\]|\\.)*"|[A-Za-z0-9_-]+)\s*=\s*(.+?)\s*$', line)
            if kv:
                out["models"][kv.group(1).strip('"')] = kv.group(2).strip().strip('"')
    return out


# ---------------------------------------------------------------------------
# TOML 生成
# ---------------------------------------------------------------------------

def toml_str(s):
    return '"' + str(s).replace("\\", "\\\\").replace('"', "'").replace("\n", " ") + '"'


def toml_key(alias):
    return alias if _BARE_KEY.match(alias) else toml_str(alias)


def collect_aliases(routes):
    """JSON 配置中出现的全部非空 model 别名 -> {effort, users[]}。"""
    used = {}
    for name, r in routes.get("subagents", {}).items():
        alias = r.get("model")
        if not alias:
            continue
        slot = used.setdefault(alias, {"effort": r.get("effort"), "users": []})
        if r.get("effort"):
            slot["effort"] = r.get("effort")
        slot["users"].append(name)
    return used


def render_model_alias(alias, src, effort, warnings):
    out = [f"[models.{toml_key(alias)}]"]
    out.append(f"provider = {toml_str(src['provider'])}")
    out.append(f"model = {toml_str(src['model'])}")
    out.append(f"max_context_size = {int(src['max_context_size'])}")
    caps = src.get("capabilities") or []
    if caps:
        out.append("capabilities = [ " + ", ".join(toml_str(c) for c in caps) + " ]")
    se = src.get("support_efforts") or []
    if se:
        out.append("support_efforts = [ " + ", ".join(toml_str(e) for e in se) + " ]")
    if effort:
        if se and effort not in se:
            warnings.append(
                f"别名 {alias}: effort={effort} 不在 support_efforts {se} 中，已跳过 default_effort。")
        else:
            out += ["", f"[models.{toml_key(alias)}.overrides]",
                    f"default_effort = {toml_str(effort)}"]
    out.append("")
    return out


def _resolve_variant_src(alias, existing_models):
    """变体 [models.<alias>] 条目的字段来源解析（generate_block 与 adopt 折叠路径共用）。

    返回 (src, status)：src 为可用的字段 dict 时 status 是 "existing"（来自 config.toml
    已有条目）或 "fallback"（内置模板）；src 为 None 时 status 是 "not-variant" /
    "incomplete" / "missing" / "no-template"，调用方据此前缀决定警告与是否跳过。"""
    base_model_id = VARIANT_BASE_MODEL.get(alias)
    if not base_model_id:
        return None, "not-variant"
    src = None
    for _, fields in (existing_models or {}).items():
        if fields.get("model") == base_model_id:
            src = fields
            break
    if src is not None and not (src.get("provider") and src.get("max_context_size")):
        return BUILTIN_MODEL_FALLBACKS.get(base_model_id), "incomplete"
    if src is None:
        return BUILTIN_MODEL_FALLBACKS.get(base_model_id), "missing"
    return src, "existing"


def generate_block(routes, existing_models):
    """生成可并入 config.toml 的标记块。返回 (text, warnings)。"""
    warnings = []
    used = collect_aliases(routes)
    lines = [
        MARKER_BEGIN,
        "# Managed by the kimi-subagent-router plugin. To change routing, edit subagent-models.json",
        "# and re-run: python3 scripts/kimi_subagent_router.py setup-config",
        "",
    ]
    generated = set()
    for alias, info in used.items():
        if alias in existing_models:
            continue
        src, status = _resolve_variant_src(alias, existing_models)
        if status == "not-variant":
            warnings.append(
                f"别名 {alias} 未在 [models] 注册，且不是内置托管变体；无法自动生成。"
                "请先按 README「切换到第三方模型」手动注册该别名。")
            continue
        if status == "incomplete":
            warnings.append(
                f"config.toml 中 {VARIANT_BASE_MODEL[alias]} 的已有条目缺少 provider/max_context_size 字段，"
                f"别名 {alias} 已改用插件内置字段模板，请核对 max_context_size。")
        elif status == "missing":
            warnings.append(
                f"config.toml 中未找到底层模型 {VARIANT_BASE_MODEL[alias]} 的条目，别名 {alias} 已使用"
                "插件内置字段模板，请核对 max_context_size。")
        if src is None:
            warnings.append(f"别名 {alias} 缺少可用的字段模板，已跳过生成。")
            continue
        lines += render_model_alias(alias, src, info["effort"], warnings)
        generated.add(alias)

    registered = set(existing_models or {}) | generated
    pool_aliases = [a for a in used if a in registered]
    if not pool_aliases:
        warnings.append("没有可写入 [secondary_model] 池的模型别名。")
        return "\n".join(lines).rstrip() + "\n" + MARKER_END + "\n", warnings

    default_model = routes.get("default_model")
    if not default_model:
        default_model = FALLBACK_POOL_DEFAULT
        warnings.append(
            f"subagent-models.json 的 default_model 为空，池默认值回退为 {FALLBACK_POOL_DEFAULT}。")
    if default_model not in registered:
        fallback = pool_aliases[0]
        warnings.append(
            f"default_model={default_model} 未注册到 [models]，已回退为 {fallback}。")
        default_model = fallback

    lines += ["[secondary_model]", f"default_model = {toml_str(default_model)}", "",
              "[secondary_model.models]"]
    for alias in pool_aliases:
        info = used[alias]
        users = ", ".join(info["users"])
        hint = f"For subagents: {users}. Thinking tier = this alias's default_effort."
        lines.append(f"{toml_key(alias)} = {toml_str(hint)}")
    lines += ["", MARKER_END]
    return "\n".join(lines) + "\n", warnings


def find_marked_block(text):
    begin = text.find(MARKER_BEGIN)
    end = text.find(MARKER_END)
    if begin == -1 or end == -1 or end < begin:
        return None
    end += len(MARKER_END)
    while end < len(text) and text[end] == "\n":
        end += 1
    return begin, end


# ---------------------------------------------------------------------------
# 命令实现
# ---------------------------------------------------------------------------

def cmd_route(args):
    routes, source, path = load_routes()
    name = args.type
    entry = routes.get("subagents", {}).get(name)
    if entry is None:
        known = sorted(routes.get("subagents", {}).keys())
        print(json.dumps({
            "error": f"Unknown subagent type: {name}",
            "known_subagent_types": known,
            "config_source": source,
        }, ensure_ascii=False, indent=2))
        return 1
    model = entry.get("model")
    effort = entry.get("effort")
    dispatch_args = {
        "subagent_type": name,
        "prompt": "<self-contained task brief: context, goal, files in scope, acceptance criteria, expected output>",
        "description": "<3-5 word summary>",
    }
    notes = [
        "The Agent tool has no effort parameter: the thinking tier is bound to this alias's default_effort.",
    ]
    if model:
        dispatch_args["model"] = model
    else:
        notes.append("model is null: omit the model argument so [secondary_model].default_model takes over.")
        dm = routes.get("default_model")
        if dm:
            model = dm
            dispatch_args["model"] = dm
            notes.append(f"Fell back to default_model from subagent-models.json: {dm}")
    result = {
        "subagent_type": name,
        "model": model,
        "effort": effort,
        "subagent_type_inherits": entry.get("subagent_type"),
        "description": entry.get("description"),
        "exclude_tools": entry.get("exclude_tools"),
        "config_source": source,
        "dispatch": {"tool": "Agent", "arguments": dispatch_args},
        "notes": notes,
    }
    print(json.dumps(result, ensure_ascii=False, indent=2))
    return 0


def cmd_list(args):
    routes, source, path = load_routes()
    out = {
        "config_file": path,
        "config_source": source,
        "default_model": routes.get("default_model"),
        "subagents": routes.get("subagents", {}),
    }
    print(json.dumps(out, ensure_ascii=False, indent=2))
    return 0


def run_doctor(home_arg):
    routes, source, path = load_routes(home_arg)
    issues = []

    text, parsed = read_config_toml(home_arg)
    cfg_path = config_toml_path(home_arg)
    if parsed is None and text.strip():
        issues.append(("warn", f"{cfg_path} 无法解析为 TOML，已退化为文本扫描，检查结果可能不全。"))
    if not os.path.isfile(cfg_path):
        issues.append(("error", f"找不到 {cfg_path}；请先运行 setup-config 或手动添加配置。"))
        return issues, source, cfg_path

    models = extract_models(parsed, text)
    sm = extract_secondary_model(parsed, text)

    for name, r in routes.get("subagents", {}).items():
        alias = r.get("model")
        eff = r.get("effort")
        if not alias:
            issues.append(("info", f"{name}: 未配置 model，派生时将回退到 [secondary_model].default_model。"))
            continue
        if alias not in models:
            issues.append(("error",
                           f"{name}: 模型别名 {alias} 未注册到 [models]。运行 setup-config 或按 README 手动注册。"))
            continue
        entry = models[alias] or {}
        overrides = entry.get("overrides") or {}
        actual_effort = overrides.get("default_effort") or entry.get("default_effort")
        se = entry.get("support_efforts") or []
        if eff:
            if actual_effort != eff:
                issues.append(("warn",
                               f"{name}: 声明 effort={eff}，但别名 {alias} 的实际默认档位为 "
                               f"{actual_effort or '（未声明）'}；实际派发以别名档位为准。"
                               "请重新运行 setup-config 或修正 [models 条目的 overrides].default_effort。"))
            if se and eff not in se:
                issues.append(("error", f"{name}: effort={eff} 不在 {alias} 的 support_efforts {se} 中。"))
        if sm:
            pool = sm.get("models") or {}
            if sm.get("default_model") != alias and alias not in pool:
                issues.append(("error",
                               f"{name}: 别名 {alias} 不在 [secondary_model] 池中；Agent 工具的 model 参数"
                               "只接受池别名或 \"primary\"。运行 setup-config。"))

    if sm:
        if sm.get("default_effort"):
            issues.append(("warn",
                           "[secondary_model].default_effort 已设置：它会覆盖所有变体别名档位，"
                           "导致按类型路由的思考强度失效，建议删除该字段。"))
        if sm.get("force"):
            issues.append(("warn",
                           "[secondary_model].force=true 时 Agent 工具会拒绝 model 参数，路由失效；"
                           "建议改为 false。"))
        if not sm.get("default_model"):
            issues.append(("error", "[secondary_model] 缺少 default_model。"))
    else:
        issues.append(("error",
                       "未找到 [secondary_model] 配置：Agent 工具没有 model 参数，subagent 将继承主 Agent 模型。"
                       "运行 setup-config。"))
    return issues, source, cfg_path


def cmd_doctor(args):
    issues, source, cfg_path = run_doctor(args.home)
    levels = {"error": 0, "warn": 0, "info": 0}
    for level, _ in issues:
        levels[level] = levels.get(level, 0) + 1
    if args.json:
        print(json.dumps({
            "config_file": cfg_path,
            "config_source": source,
            "ok": levels["error"] == 0,
            "issues": [{"level": l, "message": m} for l, m in issues],
        }, ensure_ascii=False, indent=2))
    else:
        icon = {"error": "[错误]", "warn": "[警告]", "info": "[信息]"}
        print(f"配置文件: {cfg_path}")
        print(f"映射来源: {source}")
        if not issues:
            print("一切正常。")
        for level, msg in issues:
            print(f"{icon.get(level, level)} {msg}")
        print(f"\n汇总: {levels['error']} 错误 / {levels['warn']} 警告 / {levels['info']} 信息")
    return 1 if levels["error"] else 0


def cmd_print_config(args):
    routes, source, _ = load_routes(args.home)
    text, _ = read_config_toml(args.home)
    span = find_marked_block(text)
    if span:
        text = text[: span[0]] + text[span[1]:]
    try:
        import tomllib
        parsed = tomllib.loads(text) if text.strip() else None
    except Exception:  # noqa: BLE001
        parsed = None
    existing = extract_models(parsed, text)
    block, warnings = generate_block(routes, existing)
    for w in warnings:
        print(f"# 警告: {w}", file=sys.stderr)
    sys.stdout.write(block)
    return 0


def _backup(path):
    stamp = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
    bak = f"{path}.kimi-subagent-router.bak-{stamp}"
    shutil.copy2(path, bak)
    return bak


def _expected_variant_entry(alias, search_space, effort):
    """adopt 折叠路径：generate_block 在 search_space（existing 去掉已折叠候选后的样子）
    下会为 alias 生成的 [models.<alias>] 内容，以解析后的 dict 表示；无法生成时为 None。"""
    src, status = _resolve_variant_src(alias, search_space)
    if src is None or status == "not-variant":
        return None
    entry = {"provider": src["provider"], "model": src["model"],
             "max_context_size": int(src["max_context_size"])}
    caps = src.get("capabilities") or []
    if caps:
        entry["capabilities"] = list(caps)
    se = src.get("support_efforts") or []
    if se:
        entry["support_efforts"] = list(se)
    if effort and (not se or effort in se):
        entry["overrides"] = {"default_effort": effort}
    return entry


def _variant_table_span(lines, alias):
    """text 中 [models.<alias>] 表（含可选 [models.<alias>.overrides] 子表）的行区间。
    结束于下一个不以该表头前缀开头的 '[' 行；找不到表头时为 None。"""
    key = toml_key(alias)
    head = re.compile(r'^\[models\.' + re.escape(key) + r'(\.|\])')
    start = None
    for i, line in enumerate(lines):
        if head.match(line):
            start = i
            break
    if start is None:
        return None
    end = len(lines)
    for j in range(start + 1, len(lines)):
        if lines[j].startswith("[") and not head.match(lines[j]):
            end = j
            break
    return start, end


def _collapse_seam(prefix, suffix):
    """只折叠拼接缝处的多余空行（prefix 尾部 / suffix 头部），不触碰文件其余内容。
    返回归一化后的 (prefix, suffix)。"""
    return (re.sub(r"\n{3,}$", "\n\n", prefix),
            re.sub(r"^\n{2,}", "\n", suffix))


def strip_identical_variant_tables(text, existing, routes):
    """把与本插件生成内容完全一致的未标记 [models.<variant>] 表从 text 中移除，并从
    existing 中剔除，使其随后由 generate_block 在标记块内重新生成；用户改过的表保持
    不动，避免被冻结在标记块之外。返回 (new_text, new_existing, folded_aliases)。"""
    used = collect_aliases(routes)
    folded = {a for a in used if a in (existing or {}) and a in VARIANT_BASE_MODEL}
    if not folded:
        return text, existing, []
    # 固定点：期望内容须按 generate_block 实际运行时的来源解析（在"existing 去掉全部
    # 已折叠候选"后）计算；某候选不一致时逐出候选集并重算，直到稳定。
    while True:
        search_space = {k: v for k, v in (existing or {}).items() if k not in folded}
        mismatch = [a for a in sorted(folded)
                    if _expected_variant_entry(a, search_space, used[a]["effort"]) != existing.get(a)]
        if not mismatch:
            break
        folded -= set(mismatch)
    if not folded:
        return text, existing, []
    lines = text.splitlines(keepends=True)
    spans = []
    for alias in folded:
        span = _variant_table_span(lines, alias)
        if span is None:
            folded.discard(alias)  # parsed 有表但文本找不到表头（理论不会发生）：不折
            continue
        spans.append(span)
    if not folded:
        return text, existing, []
    for start, end in sorted(spans, reverse=True):
        prefix, suffix = _collapse_seam("".join(lines[:start]), "".join(lines[end:]))
        lines = (prefix + suffix).splitlines(keepends=True)
    new_existing = {k: v for k, v in (existing or {}).items() if k not in folded}
    return "".join(lines), new_existing, sorted(folded)


def try_adopt_secondary_model(text, parsed, existing, routes, block, warnings):
    """接管路径：外部工具（如 CLI 设置写入）重写 config.toml 时会剥掉标记注释、
    但保留段内容。若未标记的 [secondary_model] 段与本插件将生成的内容完全一致，
    则同时把与本插件生成内容一致的未标记 [models.<variant>] 表折叠回标记块内，
    移除旧段并在原位置接入完整标记块。返回 (new_text, action, warnings)，否则 None。"""
    try:
        import tomllib
        block_sm = tomllib.loads(block).get("secondary_model")
    except Exception:  # noqa: BLE001 - 生成块理论上必然可解析，解析失败则保持拒绝
        return None
    if block_sm is None:
        return None
    if extract_secondary_model(parsed, text) != block_sm:
        return None
    text, existing, folded = strip_identical_variant_tables(text, existing, routes)
    if folded:
        # 变体已移出 existing，重新生成块使其回到标记块内（[secondary_model] 池内容不变）。
        block, warnings = generate_block(routes, existing)
    lines = text.splitlines(keepends=True)
    start = None
    for i, line in enumerate(lines):
        if re.match(r'^\[secondary_model\]', line):
            start = i
            break
    if start is None:
        return None
    end = len(lines)
    for j in range(start + 1, len(lines)):
        if lines[j].startswith("[") and not lines[j].startswith("[secondary_model"):
            end = j
            break
    prefix, suffix = _collapse_seam("".join(lines[:start]), "".join(lines[end:]))
    new_text = prefix + block + suffix
    return new_text, "已接管被外部工具重写的 [secondary_model] 段（内容与本插件生成的一致）", warnings


def cmd_setup_config(args):
    routes, source, _ = load_routes(args.home)
    cfg_path = config_toml_path(args.home)
    home = os.path.dirname(cfg_path)
    os.makedirs(home, exist_ok=True)
    text = ""
    if os.path.isfile(cfg_path):
        with open(cfg_path, "r", encoding="utf-8") as f:
            text = f.read()

    # 剥掉旧标记块再解析：块内由本插件生成的变体不算用户已有条目，
    # 否则更新时会把它们误判为"已注册"而从新块中遗漏。
    span = find_marked_block(text)
    base_text = (text[: span[0]] + text[span[1]:]) if span else text
    try:
        import tomllib
        parsed = tomllib.loads(base_text) if base_text.strip() else None
    except Exception:  # noqa: BLE001
        parsed = None

    existing = extract_models(parsed, base_text)
    block, warnings = generate_block(routes, existing)

    if span:
        new_text = text[: span[0]] + block + text[span[1]:]
        action = "已更新已有标记块"
    else:
        conflict_models = []
        for alias in collect_aliases(routes):
            if alias in existing:
                continue
            if re.search(r'^\[models\.("?' + re.escape(alias) + r'"?)\]', base_text, re.M):
                conflict_models.append(alias)
        has_sm = re.search(r'^\[secondary_model(\.|\])', base_text, re.M)
        if conflict_models or has_sm:
            adopted = None
            if has_sm and parsed is not None and not conflict_models:
                adopted = try_adopt_secondary_model(base_text, parsed, existing, routes, block, warnings)
            if adopted is None:
                print("检测到现有配置冲突，为避免破坏 config.toml，未写入。请手动合并：", file=sys.stderr)
                print(f"  冲突: [secondary_model] 段已存在: {bool(has_sm)}; 冲突别名: {conflict_models}", file=sys.stderr)
                print("  可手动添加以下片段（放在已有 [secondary_model] 与 [models] 表中对应位置）:", file=sys.stderr)
                sys.stderr.write(block)
                return 2
            new_text, action, warnings = adopted
        else:
            header = "" if not text.strip() else ("\n" if not text.endswith("\n") else "\n")
            new_text = text + header + "\n" + block
            action = "已追加标记块"

    bak = _backup(cfg_path) if os.path.isfile(cfg_path) else None
    try:
        try:
            import tomllib
            tomllib.loads(new_text)
        except Exception as e:  # noqa: BLE001
            raise SystemExit(f"生成的 TOML 校验失败: {e}")
        with open(cfg_path, "w", encoding="utf-8") as f:
            f.write(new_text)
    except Exception:
        if bak:
            shutil.copy2(bak, cfg_path)
        raise

    print(f"{action}: {cfg_path}")
    if bak:
        print(f"备份: {bak}")
    for w in warnings:
        print(f"警告: {w}")
    print("下一步: 重启 Kimi Code 或执行 /reload 使配置生效。")
    return 0


def cmd_remove_config(args):
    cfg_path = config_toml_path(args.home)
    if not os.path.isfile(cfg_path):
        print(f"找不到 {cfg_path}，无需移除。")
        return 0
    with open(cfg_path, "r", encoding="utf-8") as f:
        text = f.read()
    span = find_marked_block(text)
    if not span:
        print("未找到 kimi-subagent-router 标记块，无需移除。")
        return 0
    bak = _backup(cfg_path)
    new_text = text[: span[0]] + text[span[1]:]
    new_text = re.sub(r"\n{3,}", "\n\n", new_text).rstrip() + "\n"
    if not new_text.strip():
        os.remove(cfg_path)
        print(f"config.toml 已为空，已删除；备份: {bak}")
        return 0
    with open(cfg_path, "w", encoding="utf-8") as f:
        f.write(new_text)
    print(f"已移除标记块: {cfg_path}")
    print(f"备份: {bak}")
    return 0


def cmd_init_json(args):
    target = args.path or os.path.join(PLUGIN_ROOT, "subagent-models.json")
    if os.path.exists(target) and not args.force:
        print(f"{target} 已存在；加 --force 覆盖。", file=sys.stderr)
        return 1
    with open(target, "w", encoding="utf-8") as f:
        json.dump(DEFAULT_CONFIG, f, ensure_ascii=False, indent=2)
        f.write("\n")
    print(f"已写出默认配置: {target}")
    return 0


# ---------------------------------------------------------------------------
# MCP stdio 服务（newline-delimited JSON-RPC 2.0）
# ---------------------------------------------------------------------------

MCP_TOOLS = [
    {
        "name": "route_subagent_model",
        "description": (
            "Resolve a subagent type to the model alias and thinking tier to dispatch it with, plus ready-to-use "
            "Agent tool arguments. Call this immediately before every Agent/AgentSwarm dispatch for all subagent "
            "types EXCEPT k3-planner: the k3-planner fast path dispatches directly with model \"k3-256k-max\" and "
            "skips this tool, resolving live only when the user said they edited the mapping or a dispatch failed "
            "with an unknown subagent type. For every other type the mapping lives in subagent-models.json and the "
            "user may have edited it, so never dispatch from a remembered alias. "
            "The Agent tool has no effort parameter — the thinking tier is bound to the returned alias's default_effort."),
        "inputSchema": {
            "type": "object",
            "properties": {
                "subagent_type": {
                    "type": "string",
                    "description": "Subagent name, e.g. cheap-coder / cheap-explore / k2.8-medium / k2.8-debug / k3-planner / k3-max / doc-writer, or any custom key in subagent-models.json",
                },
                "task": {
                    "type": "string",
                    "description": "Optional task text, recorded for debugging",
                },
            },
            "required": ["subagent_type"],
        },
    },
    {
        "name": "list_subagent_routes",
        "description": "List every subagent -> model alias / thinking tier mapping and which config file it came from. Use it when route_subagent_model reports an unknown subagent type.",
        "inputSchema": {"type": "object", "properties": {}},
    },
    {
        "name": "doctor",
        "description": "Check subagent-models.json against config.toml: alias registration, [secondary_model] pool membership, effort tiers. Use it when a dispatch fails or a route looks wrong.",
        "inputSchema": {"type": "object", "properties": {}},
    },
    {
        "name": "reload_routes",
        "description": "Re-read subagent-models.json and report how many routes are live. Every tool call already re-reads the file, so this is only an explicit confirmation.",
        "inputSchema": {"type": "object", "properties": {}},
    },
]


def _mcp_result(payload, is_error=False):
    result = {"content": [{"type": "text", "text": json.dumps(payload, ensure_ascii=False, indent=2)}]}
    if is_error:
        result["isError"] = True
    return result


def handle_mcp_call(name, arguments):
    args_obj = arguments or {}
    if name == "route_subagent_model":
        st = args_obj.get("subagent_type", "")
        routes, source, path = load_routes()
        entry = routes.get("subagents", {}).get(st)
        if entry is None:
            return _mcp_result({
                "error": f"Unknown subagent type: {st}",
                "hint": "Call list_subagent_routes to see the live names.",
                "known_subagent_types": sorted(routes.get("subagents", {}).keys()),
            }, is_error=True)
        model = entry.get("model")
        dispatch_args = {
            "subagent_type": st,
            "prompt": "<self-contained task brief: context, goal, files in scope, acceptance criteria, expected output>",
            "description": "<3-5 word summary>",
        }
        notes = ["The Agent tool has no effort parameter: the thinking tier is bound to this alias's default_effort."]
        if model:
            dispatch_args["model"] = model
        else:
            notes.append("model is null: omit the model argument so [secondary_model].default_model takes over.")
        return _mcp_result({
            "subagent_type": st,
            "model": model,
            "effort": entry.get("effort"),
            "description": entry.get("description"),
            "config_source": source,
            "dispatch": {"tool": "Agent", "arguments": dispatch_args},
            "notes": notes,
        })
    if name == "list_subagent_routes":
        routes, source, path = load_routes()
        return _mcp_result({
            "config_file": path,
            "config_source": source,
            "default_model": routes.get("default_model"),
            "subagents": routes.get("subagents", {}),
        })
    if name == "doctor":
        class _A:  # noqa: D401 - 复用 CLI 的 doctor
            home = None
            json = True
        issues, source, cfg_path = run_doctor(None)
        return _mcp_result({
            "config_file": cfg_path,
            "config_source": source,
            "ok": not any(l == "error" for l, _ in issues),
            "issues": [{"level": l, "message": m} for l, m in issues],
        })
    if name == "reload_routes":
        routes, source, path = load_routes()
        return _mcp_result({
            "reloaded": True,
            "config_file": path,
            "config_source": source,
            "subagent_count": len(routes.get("subagents", {})),
        })
    return _mcp_result({"error": f"Unknown tool: {name}"}, is_error=True)


def handle_mcp_message(msg):
    method = msg.get("method")
    msg_id = msg.get("id")
    if method == "initialize":
        return {
            "jsonrpc": "2.0", "id": msg_id,
            "result": {
                "protocolVersion": msg.get("params", {}).get("protocolVersion", "2025-06-18"),
                "capabilities": {"tools": {"listChanged": False}},
                "serverInfo": {"name": "kimi-subagent-router", "version": "1.0.0"},
            },
        }
    if method == "ping":
        return {"jsonrpc": "2.0", "id": msg_id, "result": {}}
    if method == "tools/list":
        return {"jsonrpc": "2.0", "id": msg_id, "result": {"tools": MCP_TOOLS}}
    if method == "tools/call":
        params = msg.get("params", {})
        try:
            result = handle_mcp_call(params.get("name", ""), params.get("arguments", {}))
        except SystemExit as e:
            result = _mcp_result({"error": str(e)}, is_error=True)
        except Exception as e:  # noqa: BLE001
            result = _mcp_result({"error": f"{type(e).__name__}: {e}"}, is_error=True)
        return {"jsonrpc": "2.0", "id": msg_id, "result": result}
    if msg_id is None:  # 通知（notifications/initialized 等）：忽略
        return None
    return {"jsonrpc": "2.0", "id": msg_id,
            "error": {"code": -32601, "message": f"Method not found: {method}"}}


def run_mcp():
    out = sys.stdout
    for line in sys.stdin:
        line = line.strip()
        if not line:
            continue
        try:
            msg = json.loads(line)
        except json.JSONDecodeError:
            resp = {"jsonrpc": "2.0", "id": None,
                    "error": {"code": -32700, "message": "Parse error"}}
            out.write(json.dumps(resp, ensure_ascii=False) + "\n")
            out.flush()
            continue
        resp = handle_mcp_message(msg)
        if resp is not None:
            out.write(json.dumps(resp, ensure_ascii=False) + "\n")
            out.flush()


# ---------------------------------------------------------------------------

def main(argv=None):
    parser = argparse.ArgumentParser(prog="kimi_subagent_router",
                                     description="按 subagent 类型路由模型别名与思考强度")
    sub = parser.add_subparsers(dest="cmd", required=True)

    sub.add_parser("mcp", help="以 MCP stdio 服务运行")

    p = sub.add_parser("route", help="返回某 subagent 类型的模型别名与思考强度")
    p.add_argument("--type", required=True, help="subagent 类型名")
    p.add_argument("--task", default="", help="可选任务描述")

    sub.add_parser("list", help="列出全部路由映射")

    p = sub.add_parser("doctor", help="检查配置一致性")
    p.add_argument("--home", default=None, help="KIMI_CODE_HOME 路径（默认自动解析）")
    p.add_argument("--json", action="store_true", help="以 JSON 输出")

    p = sub.add_parser("print-config", help="打印需并入 config.toml 的 TOML 片段")
    p.add_argument("--home", default=None)

    p = sub.add_parser("setup-config", help="自动把片段并入 config.toml（幂等、先备份）")
    p.add_argument("--home", default=None)

    p = sub.add_parser("remove-config", help="移除已并入的片段")
    p.add_argument("--home", default=None)

    p = sub.add_parser("init-json", help="写出默认 subagent-models.json")
    p.add_argument("--path", default=None)
    p.add_argument("--force", action="store_true")

    args = parser.parse_args(argv)
    dispatch = {
        "mcp": lambda a: (run_mcp(), 0)[1],
        "route": cmd_route,
        "list": cmd_list,
        "doctor": cmd_doctor,
        "print-config": cmd_print_config,
        "setup-config": cmd_setup_config,
        "remove-config": cmd_remove_config,
        "init-json": cmd_init_json,
    }
    return dispatch[args.cmd](args)


if __name__ == "__main__":
    sys.exit(main())