diff --git a/CHANGELOG.md b/CHANGELOG.md index 1caf74fe0..38d24a1e1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,6 +11,9 @@ Only write entries that are worth mentioning to users. ## Unreleased +- Core: Support custom context compaction via plugins — plugins can declare a `compaction.entrypoint` in `plugin.json` to provide their own compaction implementation; use `loop_control.compaction_plugin` to select which plugin to use +- Core: Add `loop_control.compaction_model` config option to use a dedicated model for context compaction + ## 1.25.0 (2026-03-23) - Core: Add plugin system (Skills + Tools) — plugins extend Kimi Code CLI with custom tools packaged as `plugin.json`; tools are commands that run in isolated subprocesses and return their stdout to the agent; plugins support automatic credential injection via `inject` configuration diff --git a/docs/en/release-notes/changelog.md b/docs/en/release-notes/changelog.md index 53c38771e..2706e29e9 100644 --- a/docs/en/release-notes/changelog.md +++ b/docs/en/release-notes/changelog.md @@ -4,6 +4,9 @@ This page documents the changes in each Kimi Code CLI release. ## Unreleased +- Core: Support custom context compaction via plugins — plugins can declare a `compaction.entrypoint` in `plugin.json` to provide their own compaction implementation; use `loop_control.compaction_plugin` to select which plugin to use +- Core: Add `loop_control.compaction_model` config option to use a dedicated model for context compaction + ## 1.25.0 (2026-03-23) - Core: Add plugin system (Skills + Tools) — plugins extend Kimi Code CLI with custom tools packaged as `plugin.json`; tools are commands that run in isolated subprocesses and return their stdout to the agent; plugins support automatic credential injection via `inject` configuration diff --git a/docs/zh/release-notes/changelog.md b/docs/zh/release-notes/changelog.md index bf584aa36..09559763c 100644 --- a/docs/zh/release-notes/changelog.md +++ b/docs/zh/release-notes/changelog.md @@ -4,6 +4,9 @@ ## 未发布 +- Core:支持通过插件提供自定义上下文压缩实现——插件可在 `plugin.json` 中声明 `compaction.entrypoint` 提供自定义压缩器;通过 `loop_control.compaction_plugin` 配置项指定使用哪个插件的压缩器 +- Core:新增 `loop_control.compaction_model` 配置项,可为上下文压缩指定专用模型 + ## 1.25.0 (2026-03-23) - Core:新增插件系统(Skills + Tools)——插件通过 `plugin.json` 为 Kimi Code CLI 扩展自定义工具;工具是在独立子进程中运行的命令,其 stdout 返回给 Agent;插件支持通过 `inject` 配置自动注入凭证 diff --git a/examples/custom-kimi-soul/main.py b/examples/custom-kimi-soul/main.py index bfdef01e7..943ca94c9 100644 --- a/examples/custom-kimi-soul/main.py +++ b/examples/custom-kimi-soul/main.py @@ -35,6 +35,7 @@ async def create( config=config, oauth=OAuthManager(config), llm=llm, + compaction_llm=None, session=session, yolo=True, ) diff --git a/examples/kimi-psql/main.py b/examples/kimi-psql/main.py index 31f1994bc..38bbbf56f 100644 --- a/examples/kimi-psql/main.py +++ b/examples/kimi-psql/main.py @@ -276,6 +276,7 @@ async def create_psql_soul(llm: LLM | None, conninfo: str) -> KimiSoul: config=config, oauth=OAuthManager(config), llm=llm, + compaction_llm=None, session=session, yolo=True, # Auto-approve read-only SQL queries ) diff --git a/src/kimi_cli/app.py b/src/kimi_cli/app.py index b3c7853a5..473d7c536 100644 --- a/src/kimi_cli/app.py +++ b/src/kimi_cli/app.py @@ -16,7 +16,13 @@ from kimi_cli.auth.oauth import OAuthManager from kimi_cli.cli import InputFormat, OutputFormat from kimi_cli.config import Config, LLMModel, LLMProvider, load_config -from kimi_cli.llm import augment_provider_with_env_vars, create_llm, model_display_name +from kimi_cli.exception import ConfigError +from kimi_cli.llm import ( + augment_provider_credentials_with_env_vars, + augment_provider_with_env_vars, + create_llm, + model_display_name, +) from kimi_cli.session import Session from kimi_cli.share import get_share_dir from kimi_cli.soul import run_soul @@ -143,26 +149,28 @@ async def create( oauth = OAuthManager(config) - model: LLMModel | None = None - provider: LLMProvider | None = None - - # try to use config file - if not model_name and config.default_model: - # no --model specified && default model is set in config - model = config.models[config.default_model] - provider = config.providers[model.provider] - if model_name and model_name in config.models: - # --model specified && model is set in config - model = config.models[model_name] - provider = config.providers[model.provider] - - if not model: + selected_model = model_name or config.default_model + if selected_model and selected_model in config.models: + model = config.models[selected_model] + provider = config.providers.get(model.provider) + if provider is None: + logger.warning( + "Provider {provider!r} for model {model!r} missing; using placeholder", + provider=model.provider, + model=selected_model, + ) + model = LLMModel(provider="", model="", max_context_size=100_000) + provider = LLMProvider(type="kimi", base_url="", api_key=SecretStr("")) + else: + if selected_model: + logger.warning( + "Model {model!r} not found in config, using placeholder", + model=selected_model, + ) model = LLMModel(provider="", model="", max_context_size=100_000) provider = LLMProvider(type="kimi", base_url="", api_key=SecretStr("")) # try overwrite with environment variables - assert provider is not None - assert model is not None env_overrides = augment_provider_with_env_vars(provider, model) # determine thinking mode @@ -183,10 +191,52 @@ async def create( logger.info("Using LLM model: {model}", model=model) logger.info("Thinking mode: {thinking}", thinking=thinking) + compaction_llm = None + if config.loop_control.compaction_model is not None: + compaction_model_name = config.loop_control.compaction_model + compaction_model = config.models.get(compaction_model_name) + if compaction_model is None: + logger.warning( + "Compaction model {model!r} not found in config, skipping", + model=compaction_model_name, + ) + else: + if llm is not None and compaction_model.max_context_size < llm.max_context_size: + raise ConfigError( + "Compaction model " + f"{compaction_model_name!r} has max_context_size " + f"{compaction_model.max_context_size}, smaller than active model " + f"{selected_model!r} ({model.max_context_size})" + ) + compaction_provider = config.providers.get(compaction_model.provider) + if compaction_provider is None: + logger.warning( + "Compaction provider {provider!r} not found in config, skipping", + provider=compaction_model.provider, + ) + else: + compaction_provider = compaction_provider.model_copy(deep=True) + compaction_model = compaction_model.model_copy(deep=True) + augment_provider_credentials_with_env_vars(compaction_provider) + compaction_llm = create_llm( + compaction_provider, + compaction_model, + thinking=thinking, + session_id=session.id, + oauth=oauth, + ) + if compaction_llm is not None: + logger.info( + "Using compaction LLM model: {model}", + model=compaction_model, + ) + if startup_progress is not None: startup_progress("Scanning workspace...") - runtime = await Runtime.create(config, oauth, llm, session, yolo, skills_dir) + runtime = await Runtime.create( + config, oauth, llm, compaction_llm, session, yolo, skills_dir + ) runtime.notifications.recover() runtime.background_tasks.reconcile() _cleanup_stale_foreground_subagents(runtime) @@ -205,6 +255,20 @@ async def create( except Exception: logger.debug("Failed to refresh plugin configs, skipping") + if config.loop_control.compaction_plugin is not None: + from kimi_cli.plugin import PluginError + from kimi_cli.plugin.compaction import resolve_plugin_compactor + from kimi_cli.plugin.manager import get_plugins_dir + + try: + runtime.compaction = resolve_plugin_compactor( + get_plugins_dir(), config.loop_control.compaction_plugin + ) + except PluginError as exc: + raise ConfigError( + f"Invalid compaction plugin {config.loop_control.compaction_plugin!r}: {exc}" + ) from exc + if agent_file is None: agent_file = DEFAULT_AGENT_FILE if startup_progress is not None: diff --git a/src/kimi_cli/config.py b/src/kimi_cli/config.py index db1859734..52475eccc 100644 --- a/src/kimi_cli/config.py +++ b/src/kimi_cli/config.py @@ -12,6 +12,7 @@ SecretStr, ValidationError, field_serializer, + field_validator, model_validator, ) from tomlkit.exceptions import TOMLKitError @@ -78,6 +79,10 @@ class LoopControl(BaseModel): """Maximum number of retries in one step""" max_ralph_iterations: int = Field(default=0, ge=-1) """Extra iterations after the first turn in Ralph mode. Use -1 for unlimited.""" + compaction_model: str | None = Field(default=None) + """Optional model name to use for context compaction.""" + compaction_plugin: str | None = Field(default=None) + """Installed plugin name to use for context compaction.""" reserved_context_size: int = Field(default=50_000, ge=1000) """Reserved token count for LLM response generation. Auto-compaction triggers when either context_tokens + reserved_context_size >= max_context_size or @@ -87,6 +92,14 @@ class LoopControl(BaseModel): Auto-compaction triggers when context_tokens >= max_context_size * compaction_trigger_ratio or when context_tokens + reserved_context_size >= max_context_size.""" + @field_validator("compaction_model", "compaction_plugin", mode="before") + @classmethod + def normalize_optional_compaction_name(cls, value: object) -> object: + if isinstance(value, str): + value = value.strip() + return value or None + return value + class BackgroundConfig(BaseModel): """Background task runtime configuration.""" @@ -207,6 +220,13 @@ class Config(BaseModel): def validate_model(self) -> Self: if self.default_model and self.default_model not in self.models: raise ValueError(f"Default model {self.default_model} not found in models") + if ( + self.loop_control.compaction_model + and self.loop_control.compaction_model not in self.models + ): + raise ValueError( + f"Compaction model {self.loop_control.compaction_model} not found in models" + ) for model in self.models.values(): if model.provider not in self.providers: raise ValueError(f"Provider {model.provider} not found in providers") diff --git a/src/kimi_cli/llm.py b/src/kimi_cli/llm.py index fba2f10e9..8dfa9cfdf 100644 --- a/src/kimi_cli/llm.py +++ b/src/kimi_cli/llm.py @@ -94,6 +94,36 @@ def augment_provider_with_env_vars(provider: LLMProvider, model: LLMModel) -> di return applied +def augment_provider_credentials_with_env_vars(provider: LLMProvider) -> dict[str, str]: + """Override provider credentials/base URL from environment variables without changing model. + + This is used for secondary model selections, such as compaction, where the configured model + alias should remain stable even if the main chat model is overridden from the environment. + """ + + applied: dict[str, str] = {} + + match provider.type: + case "kimi": + if base_url := os.getenv("KIMI_BASE_URL"): + provider.base_url = base_url + applied["KIMI_BASE_URL"] = base_url + if api_key := os.getenv("KIMI_API_KEY"): + provider.api_key = SecretStr(api_key) + applied["KIMI_API_KEY"] = "******" + case "openai_legacy" | "openai_responses": + if base_url := os.getenv("OPENAI_BASE_URL"): + provider.base_url = base_url + applied["OPENAI_BASE_URL"] = base_url + if api_key := os.getenv("OPENAI_API_KEY"): + provider.api_key = SecretStr(api_key) + applied["OPENAI_API_KEY"] = "******" + case _: + pass + + return applied + + def _kimi_default_headers(provider: LLMProvider, oauth: OAuthManager | None) -> dict[str, str]: headers = {"User-Agent": USER_AGENT} if oauth: diff --git a/src/kimi_cli/plugin/__init__.py b/src/kimi_cli/plugin/__init__.py index 7359c3c43..5e9fdc6d3 100644 --- a/src/kimi_cli/plugin/__init__.py +++ b/src/kimi_cli/plugin/__init__.py @@ -4,7 +4,7 @@ from pathlib import Path from typing import Any -from pydantic import BaseModel, ConfigDict, Field +from pydantic import BaseModel, ConfigDict, Field, field_validator class PluginError(Exception): @@ -27,6 +27,23 @@ class PluginToolSpec(BaseModel): parameters: dict[str, object] = Field(default_factory=dict) +class PluginCompactionSpec(BaseModel): + """In-process compaction hook: a class importable from the plugin directory.""" + + model_config = ConfigDict(extra="forbid") + + entrypoint: str + """Dotted path ``module.Class`` resolved with the plugin directory on ``sys.path``.""" + + @field_validator("entrypoint") + @classmethod + def entrypoint_must_include_class(cls, value: str) -> str: + cleaned = value.strip() + if "." not in cleaned: + raise ValueError("compaction.entrypoint must look like 'module.ClassName'") + return cleaned + + class PluginSpec(BaseModel): """Parsed representation of a plugin.json file.""" @@ -38,6 +55,7 @@ class PluginSpec(BaseModel): config_file: str | None = None inject: dict[str, str] = Field(default_factory=dict) tools: list[PluginToolSpec] = Field(default_factory=list) # pyright: ignore[reportUnknownVariableType] + compaction: PluginCompactionSpec | None = None runtime: PluginRuntime | None = None diff --git a/src/kimi_cli/plugin/compaction.py b/src/kimi_cli/plugin/compaction.py new file mode 100644 index 000000000..50071426e --- /dev/null +++ b/src/kimi_cli/plugin/compaction.py @@ -0,0 +1,216 @@ +"""Load an explicitly selected compaction implementation from an installed plugin.""" + +from __future__ import annotations + +import importlib.util +import sys +from hashlib import sha1 +from pathlib import Path +from threading import Lock +from types import ModuleType +from typing import cast + +from kimi_cli.plugin import PLUGIN_JSON, PluginError, parse_plugin_json +from kimi_cli.soul.compaction import Compaction + +_IMPORT_LOCK = Lock() +_PLUGIN_PATH_REFS: dict[str, tuple[int, bool]] = {} + + +def _plugin_package_name(plugin_dir: Path) -> str: + digest = sha1(str(plugin_dir.resolve()).encode("utf-8")).hexdigest()[:12] + return f"_kimi_plugin_compaction_{digest}" + + +def _resolve_module_file(plugin_dir: Path, module_path: str) -> tuple[Path, bool]: + module_parts = module_path.split(".") + file_base = plugin_dir.joinpath(*module_parts) + module_file = file_base.with_suffix(".py") + if module_file.is_file(): + return module_file, False + + package_init = file_base / "__init__.py" + if package_init.is_file(): + return package_init, True + + raise PluginError( + f"Compaction module {module_path!r} not found in plugin directory {plugin_dir}" + ) + + +def _ensure_package_module(package_name: str, package_dir: Path) -> None: + module = sys.modules.get(package_name) + if module is None: + module = ModuleType(package_name) + module.__file__ = str(package_dir / "__init__.py") + module.__package__ = package_name + module.__path__ = [str(package_dir)] + sys.modules[package_name] = module + + +def _is_module_from_any_plugin_dir(module: ModuleType) -> bool: + module_file = getattr(module, "__file__", None) + if module_file is None: + return False + return any(parent.name == "plugins" for parent in Path(module_file).resolve().parents) + + +def _local_top_level_module_names(plugin_dir: Path) -> set[str]: + names: set[str] = set() + for child in plugin_dir.iterdir(): + if child.is_file() and child.suffix == ".py": + names.add(child.stem) + elif child.is_dir() and (child / "__init__.py").is_file(): + names.add(child.name) + return names + + +def _purge_conflicting_top_level_modules(plugin_dir: Path) -> None: + for name in _local_top_level_module_names(plugin_dir): + module = sys.modules.get(name) + if module is not None and _is_module_from_any_plugin_dir(module): + sys.modules.pop(name, None) + + +def _load_plugin_module(plugin_dir: Path, module_path: str) -> ModuleType: + module_file, is_package = _resolve_module_file(plugin_dir, module_path) + plugin_root = str(plugin_dir.resolve()) + package_name = _plugin_package_name(plugin_dir) + module_parts = module_path.split(".") + module_name = ".".join((package_name, *module_parts)) + + with _IMPORT_LOCK: + _ensure_package_module(package_name, plugin_dir) + _purge_conflicting_top_level_modules(plugin_dir) + parent_dir = plugin_dir + parent_package = package_name + for part in module_parts[:-1]: + parent_dir = parent_dir / part + parent_package = f"{parent_package}.{part}" + _ensure_package_module(parent_package, parent_dir) + + sys.modules.pop(module_name, None) + spec = importlib.util.spec_from_file_location( + module_name, + module_file, + submodule_search_locations=[str(module_file.parent)] if is_package else None, + ) + if spec is None or spec.loader is None: + raise PluginError(f"Failed to load compaction module {module_path!r}") + + module = importlib.util.module_from_spec(spec) + sys.modules[module_name] = module + _acquire_plugin_path_locked(plugin_root) + try: + spec.loader.exec_module(module) + except Exception as exc: + sys.modules.pop(module_name, None) + raise PluginError(f"Failed to import compaction module {module_path!r}: {exc}") from exc + finally: + _release_plugin_path_locked(plugin_root) + + return module + + +def _acquire_plugin_path_locked(plugin_root: str) -> None: + count, inserted = _PLUGIN_PATH_REFS.get(plugin_root, (0, plugin_root not in sys.path)) + if count == 0 and inserted: + sys.path.append(plugin_root) + _PLUGIN_PATH_REFS[plugin_root] = (count + 1, inserted) + + +def _acquire_plugin_path(plugin_dir: Path) -> str: + plugin_root = str(plugin_dir.resolve()) + with _IMPORT_LOCK: + _acquire_plugin_path_locked(plugin_root) + return plugin_root + + +def _release_plugin_path_locked(plugin_root: str) -> None: + count, inserted = _PLUGIN_PATH_REFS[plugin_root] + if count <= 1: + _PLUGIN_PATH_REFS.pop(plugin_root, None) + if inserted and plugin_root in sys.path: + sys.path.remove(plugin_root) + return + _PLUGIN_PATH_REFS[plugin_root] = (count - 1, inserted) + + +def _release_plugin_path(plugin_root: str) -> None: + with _IMPORT_LOCK: + _release_plugin_path_locked(plugin_root) + + +def _wrap_compactor_for_lazy_imports(instance: Compaction, plugin_dir: Path) -> Compaction: + original_compact = getattr(instance, "compact", None) + if original_compact is None or not callable(original_compact): + raise PluginError("Compaction object has no callable compact()") + + instance.__kimi_plugin_dir__ = str(plugin_dir.resolve()) + + async def compact(*args, **kwargs): + plugin_root = _acquire_plugin_path(plugin_dir) + try: + with _IMPORT_LOCK: + _purge_conflicting_top_level_modules(plugin_dir) + return await original_compact(*args, **kwargs) + finally: + _release_plugin_path(plugin_root) + + instance.compact = compact + return instance + + +def instantiate_plugin_compactor(plugin_dir: Path, entrypoint: str) -> Compaction: + module_path, _, class_name = entrypoint.rpartition(".") + if not module_path: + raise PluginError(f"Invalid compaction entrypoint: {entrypoint!r}") + + module = _load_plugin_module(plugin_dir, module_path) + + try: + cls = getattr(module, class_name) + except AttributeError as exc: + raise PluginError( + f"Compaction class {class_name!r} not found in module {module_path!r}" + ) from exc + + try: + instance = cls() + except Exception as exc: + raise PluginError( + "Failed to initialize compaction class " + f"{class_name!r} from module {module_path!r}: {exc}" + ) from exc + if not getattr(instance, "compact", None) or not callable(instance.compact): + raise PluginError(f"Compaction object from {entrypoint!r} has no callable compact()") + return cast(Compaction, _wrap_compactor_for_lazy_imports(instance, plugin_dir)) + + +def resolve_plugin_compactor(plugins_dir: Path, plugin_name: str | None) -> Compaction | None: + """Load one explicitly selected plugin compactor. + + Returns ``None`` when no plugin compactor is configured. + Raises ``PluginError`` when the selected plugin is missing, invalid, or does not declare + a compaction entrypoint. + """ + if plugin_name is None: + return None + if plugin_name == "": + raise PluginError("Plugin name cannot be empty") + if not plugins_dir.is_dir(): + raise PluginError(f"Plugins directory not found: {plugins_dir}") + + plugin_dir = (plugins_dir / plugin_name).resolve() + if not plugin_dir.is_relative_to(plugins_dir.resolve()): + raise PluginError(f"Invalid plugin name: {plugin_name}") + + plugin_json = plugin_dir / PLUGIN_JSON + if not plugin_dir.is_dir() or not plugin_json.is_file(): + raise PluginError(f"Plugin {plugin_name!r} not found in {plugins_dir}") + + spec = parse_plugin_json(plugin_json) + if spec.compaction is None: + raise PluginError(f"Plugin {plugin_name!r} does not declare compaction.entrypoint") + + return instantiate_plugin_compactor(plugin_dir, spec.compaction.entrypoint) diff --git a/src/kimi_cli/soul/agent.py b/src/kimi_cli/soul/agent.py index 115dbf54b..333ea0400 100644 --- a/src/kimi_cli/soul/agent.py +++ b/src/kimi_cli/soul/agent.py @@ -1,7 +1,7 @@ from __future__ import annotations import asyncio -from dataclasses import asdict, dataclass +from dataclasses import asdict, dataclass, field from datetime import datetime from pathlib import Path from typing import TYPE_CHECKING, Any, Literal @@ -28,6 +28,7 @@ resolve_skills_roots, ) from kimi_cli.soul.approval import Approval, ApprovalState +from kimi_cli.soul.compaction import Compaction, SimpleCompaction from kimi_cli.soul.denwarenji import DenwaRenji from kimi_cli.soul.toolset import KimiToolset from kimi_cli.subagents.models import AgentTypeDefinition, ToolPolicy @@ -90,6 +91,8 @@ class Runtime: background_tasks: BackgroundTaskManager skills: dict[str, Skill] additional_dirs: list[KaosPath] + compaction_llm: LLM | None = None + compaction: Compaction = field(default_factory=SimpleCompaction) subagent_store: SubagentStore | None = None approval_runtime: ApprovalRuntime | None = None root_wire_hub: RootWireHub | None = None @@ -113,6 +116,7 @@ async def create( config: Config, oauth: OAuthManager, llm: LLM | None, + compaction_llm: LLM | None, session: Session, yolo: bool, skills_dir: KaosPath | None = None, @@ -215,12 +219,32 @@ def _on_approval_change() -> None: ), skills=skills_by_name, additional_dirs=additional_dirs, + compaction_llm=compaction_llm, subagent_store=SubagentStore(session), approval_runtime=ApprovalRuntime(), root_wire_hub=RootWireHub(), role="root", ) + def _new_compaction(self) -> Compaction: + compaction = type(self.compaction)() + plugin_dir = getattr(self.compaction, "__kimi_plugin_dir__", None) + if plugin_dir is not None: + from kimi_cli.plugin.compaction import _wrap_compactor_for_lazy_imports + + return _wrap_compactor_for_lazy_imports(compaction, Path(plugin_dir)) + return compaction + + def resolve_compaction_llm(self, *, active_llm: LLM | None = None) -> LLM | None: + active_llm = self.llm if active_llm is None else active_llm + if active_llm is None: + return self.compaction_llm + if self.compaction_llm is None: + return active_llm + if self.compaction_llm.max_context_size < active_llm.max_context_size: + return active_llm + return self.compaction_llm + def copy_for_subagent( self, *, @@ -229,10 +253,11 @@ def copy_for_subagent( llm_override: LLM | None = None, ) -> Runtime: """Clone runtime for a subagent.""" + active_llm = llm_override if llm_override is not None else self.llm return Runtime( config=self.config, oauth=self.oauth, - llm=llm_override if llm_override is not None else self.llm, + llm=active_llm, session=self.session, builtin_args=self.builtin_args, denwa_renji=DenwaRenji(), # subagent must have its own DenwaRenji @@ -244,6 +269,8 @@ def copy_for_subagent( skills=self.skills, # Share the same list reference so /add-dir mutations propagate to all agents additional_dirs=self.additional_dirs, + compaction_llm=self.resolve_compaction_llm(active_llm=active_llm), + compaction=self._new_compaction(), subagent_store=self.subagent_store, approval_runtime=self.approval_runtime, root_wire_hub=self.root_wire_hub, diff --git a/src/kimi_cli/soul/kimisoul.py b/src/kimi_cli/soul/kimisoul.py index 18d80de17..0adde05c1 100644 --- a/src/kimi_cli/soul/kimisoul.py +++ b/src/kimi_cli/soul/kimisoul.py @@ -47,7 +47,6 @@ from kimi_cli.soul.agent import Agent, Runtime from kimi_cli.soul.compaction import ( CompactionResult, - SimpleCompaction, estimate_text_tokens, should_auto_compact, ) @@ -134,7 +133,7 @@ def __init__( self._approval = agent.runtime.approval self._context = context self._loop_control = agent.runtime.config.loop_control - self._compaction = SimpleCompaction() # TODO: maybe configurable and composable + self._compaction = agent.runtime.compaction for tool in agent.toolset.tools: if tool.name == SendDMail_NAME: @@ -504,7 +503,7 @@ def _build_slash_commands(self) -> list[SlashCommand[Any]]: command_name = f"{FLOW_COMMAND_PREFIX}{skill.name}" if command_name in seen_names: logger.warning( - "Skipping prompt flow slash command /{name}: name already registered", + "Skipping flow slash command /{name}: name already registered", name=command_name, ) continue @@ -789,13 +788,14 @@ async def compact_context(self, custom_instruction: str = "") -> None: ChatProviderError: When the chat provider returns an error. """ - chat_provider = self._runtime.llm.chat_provider if self._runtime.llm is not None else None + compaction_llm = self._runtime.resolve_compaction_llm() + chat_provider = compaction_llm.chat_provider if compaction_llm is not None else None async def _run_compaction_once() -> CompactionResult: - if self._runtime.llm is None: + if compaction_llm is None: raise LLMNotSet() return await self._compaction.compact( - self._context.history, self._runtime.llm, custom_instruction=custom_instruction + self._context.history, compaction_llm, custom_instruction=custom_instruction ) @tenacity.retry( diff --git a/tests/core/test_config.py b/tests/core/test_config.py index 545374169..a0d5a9dbc 100644 --- a/tests/core/test_config.py +++ b/tests/core/test_config.py @@ -31,6 +31,8 @@ def test_default_config_dump(): "max_steps_per_turn": 100, "max_retries_per_step": 3, "max_ralph_iterations": 0, + "compaction_model": None, + "compaction_plugin": None, "reserved_context_size": 50000, "compaction_trigger_ratio": 0.85, }, @@ -109,6 +111,25 @@ def test_load_config_reserved_context_size_too_low(): load_config_from_string('{"loop_control": {"reserved_context_size": 500}}') +def test_load_config_compaction_model(): + config = load_config_from_string( + '{"loop_control": {"compaction_model": "compact"}, "models": {"compact": {"provider": "p", "model": "compact-model", "max_context_size": 4096}}, "providers": {"p": {"type": "_echo", "base_url": "", "api_key": ""}}}' + ) + assert config.loop_control.compaction_model == "compact" + + +def test_load_config_compaction_model_strips_whitespace(): + config = load_config_from_string( + '{"loop_control": {"compaction_model": " compact "}, "models": {"compact": {"provider": "p", "model": "compact-model", "max_context_size": 4096}}, "providers": {"p": {"type": "_echo", "base_url": "", "api_key": ""}}}' + ) + assert config.loop_control.compaction_model == "compact" + + +def test_load_config_compaction_model_requires_known_model(): + with pytest.raises(ConfigError, match="Compaction model missing not found in models"): + load_config_from_string('{"loop_control": {"compaction_model": "missing"}}') + + def test_load_config_compaction_trigger_ratio(): config = load_config_from_string('{"loop_control": {"compaction_trigger_ratio": 0.8}}') assert config.loop_control.compaction_trigger_ratio == 0.8 @@ -119,6 +140,16 @@ def test_load_config_compaction_trigger_ratio_default(): assert config.loop_control.compaction_trigger_ratio == 0.85 +def test_load_config_compaction_plugin(): + config = load_config_from_string('{"loop_control": {"compaction_plugin": "alpha-plugin"}}') + assert config.loop_control.compaction_plugin == "alpha-plugin" + + +def test_load_config_compaction_plugin_empty_becomes_none(): + config = load_config_from_string('{"loop_control": {"compaction_plugin": " "}}') + assert config.loop_control.compaction_plugin is None + + def test_load_config_compaction_trigger_ratio_too_low(): with pytest.raises(ConfigError, match="compaction_trigger_ratio"): load_config_from_string('{"loop_control": {"compaction_trigger_ratio": 0.3}}') diff --git a/tests/core/test_plugin_compaction.py b/tests/core/test_plugin_compaction.py new file mode 100644 index 000000000..6608fabfc --- /dev/null +++ b/tests/core/test_plugin_compaction.py @@ -0,0 +1,366 @@ +from __future__ import annotations + +import json +from pathlib import Path + +import pytest +from kosong.chat_provider.echo import EchoChatProvider +from kosong.message import Message +from kosong.tooling.empty import EmptyToolset + +from kimi_cli.config import LLMModel +from kimi_cli.llm import LLM +from kimi_cli.plugin import PluginError, parse_plugin_json +from kimi_cli.plugin.compaction import resolve_plugin_compactor +from kimi_cli.soul import _current_wire +from kimi_cli.soul.agent import Agent, Runtime +from kimi_cli.soul.context import Context +from kimi_cli.soul.kimisoul import KimiSoul +from kimi_cli.wire import Wire +from kimi_cli.wire.types import CompactionBegin, CompactionEnd, TextPart + + +def _write_alpha_compactor(plugin_root: Path) -> None: + plugin_root.mkdir(parents=True, exist_ok=True) + (plugin_root / "alpha_compaction.py").write_text( + """ +from collections.abc import Sequence + +from kosong.message import Message + +from kimi_cli.llm import LLM +from kimi_cli.soul.compaction import CompactionResult +from kimi_cli.wire.types import TextPart + + +class AlphaCompaction: + PLUGIN_MARK = "alpha" + + def __init__(self) -> None: + self.calls = 0 + self.last_model = None + + async def compact( + self, messages: Sequence[Message], llm: LLM, *, custom_instruction: str = "" + ) -> CompactionResult: + self.calls += 1 + self.last_model = llm.model_config.model if llm.model_config is not None else None + return CompactionResult( + messages=[Message(role="user", content=[TextPart(text=f"plugin compacted via {self.last_model}")])], + usage=None, + ) +""".strip(), + encoding="utf-8", + ) + + +def _write_beta_compactor(plugin_root: Path) -> None: + plugin_root.mkdir(parents=True, exist_ok=True) + (plugin_root / "beta_compaction.py").write_text( + """ +from collections.abc import Sequence + +from kosong.message import Message + +from kimi_cli.llm import LLM +from kimi_cli.soul.compaction import CompactionResult + + +class BetaCompaction: + PLUGIN_MARK = "beta" + + async def compact( + self, messages: Sequence[Message], llm: LLM, *, custom_instruction: str = "" + ) -> CompactionResult: + return CompactionResult(messages=messages, usage=None) +""".strip(), + encoding="utf-8", + ) + + +def _write_plugin(plugin_root: Path, *, name: str, entrypoint: str) -> None: + plugin_root.mkdir(parents=True, exist_ok=True) + (plugin_root / "plugin.json").write_text( + json.dumps( + { + "name": name, + "version": "1.0.0", + "compaction": {"entrypoint": entrypoint}, + } + ), + encoding="utf-8", + ) + + +def _make_test_llm(model_name: str, *, max_context_size: int = 100_000) -> LLM: + return LLM( + chat_provider=EchoChatProvider(), + max_context_size=max_context_size, + capabilities=set(), + model_config=LLMModel( + provider="_echo", + model=model_name, + max_context_size=max_context_size, + ), + ) + + +def _make_soul(runtime: Runtime, tmp_path: Path) -> tuple[KimiSoul, Context]: + agent = Agent( + name="Plugin Compaction Agent", + system_prompt="System prompt.", + toolset=EmptyToolset(), + runtime=runtime, + ) + context = Context(file_backend=tmp_path / "history.jsonl") + return KimiSoul(agent, context=context), context + + +def test_parse_plugin_json_rejects_compaction_entrypoint_without_dot(tmp_path: Path) -> None: + path = tmp_path / "plugin.json" + path.write_text( + json.dumps( + { + "name": "bad", + "version": "1.0.0", + "compaction": {"entrypoint": "NoDot"}, + } + ), + encoding="utf-8", + ) + with pytest.raises(PluginError, match="Invalid plugin.json schema"): + parse_plugin_json(path) + + +def test_resolve_plugin_compactor_returns_none_when_unconfigured(tmp_path: Path) -> None: + assert resolve_plugin_compactor(tmp_path / "plugins", None) is None + + +def test_resolve_plugin_compactor_rejects_empty_plugin_name(tmp_path: Path) -> None: + (tmp_path / "plugins").mkdir() + with pytest.raises(PluginError, match="Plugin name cannot be empty"): + resolve_plugin_compactor(tmp_path / "plugins", "") + + +@pytest.mark.asyncio +async def test_resolve_plugin_compactor_loads_selected_plugin(tmp_path: Path) -> None: + plugins = tmp_path / "plugins" + pdir = plugins / "alpha-plugin" + _write_alpha_compactor(pdir) + _write_plugin( + pdir, + name="alpha-plugin", + entrypoint="alpha_compaction.AlphaCompaction", + ) + + comp = resolve_plugin_compactor(plugins, "alpha-plugin") + assert comp is not None + assert getattr(type(comp), "PLUGIN_MARK", None) == "alpha" + + llm = _make_test_llm("main-chat") + result = await comp.compact([Message(role="user", content=[TextPart(text="hi")])], llm) + assert result.messages[0].extract_text("\n") == "plugin compacted via main-chat" + + +def test_resolve_plugin_compactor_requires_explicit_selected_plugin(tmp_path: Path) -> None: + plugins = tmp_path / "plugins" + alpha = plugins / "alpha-plugin" + beta = plugins / "beta-plugin" + _write_alpha_compactor(alpha) + _write_beta_compactor(beta) + _write_plugin(alpha, name="alpha-plugin", entrypoint="alpha_compaction.AlphaCompaction") + _write_plugin(beta, name="beta-plugin", entrypoint="beta_compaction.BetaCompaction") + + comp = resolve_plugin_compactor(plugins, "beta-plugin") + assert comp is not None + assert getattr(type(comp), "PLUGIN_MARK", None) == "beta" + + +def test_resolve_plugin_compactor_raises_for_missing_plugin(tmp_path: Path) -> None: + (tmp_path / "plugins").mkdir() + with pytest.raises(PluginError, match="missing-plugin"): + resolve_plugin_compactor(tmp_path / "plugins", "missing-plugin") + + +@pytest.mark.asyncio +async def test_resolve_plugin_compactor_allows_stdlib_colliding_module_names( + tmp_path: Path, +) -> None: + plugins = tmp_path / "plugins" + pdir = plugins / "collision-plugin" + pdir.mkdir(parents=True, exist_ok=True) + (pdir / "json.py").write_text( + "from kosong.message import Message\nfrom kimi_cli.soul.compaction import CompactionResult\nfrom kimi_cli.wire.types import TextPart\n\nclass MyCompactor:\n async def compact(self, messages, llm, *, custom_instruction=''):\n return CompactionResult(messages=[Message(role='user', content=[TextPart(text='json plugin ok')])], usage=None)\n", + encoding="utf-8", + ) + _write_plugin(pdir, name="collision-plugin", entrypoint="json.MyCompactor") + + comp = resolve_plugin_compactor(plugins, "collision-plugin") + assert comp is not None + result = await comp.compact( + [Message(role="user", content=[TextPart(text="hi")])], _make_test_llm("main-chat") + ) + assert result.messages[0].extract_text("\n") == "json plugin ok" + + assert json.dumps({"ok": True}) == '{"ok": true}' + + +@pytest.mark.asyncio +async def test_resolve_plugin_compactor_supports_lazy_sibling_imports(tmp_path: Path) -> None: + plugins = tmp_path / "plugins" + pdir = plugins / "lazy-plugin" + pdir.mkdir(parents=True, exist_ok=True) + (pdir / "helper_mod.py").write_text( + "def build_message(model_name):\n return f'lazy via {model_name}'\n", + encoding="utf-8", + ) + (pdir / "lazy_compaction.py").write_text( + "from kosong.message import Message\nfrom kimi_cli.soul.compaction import CompactionResult\nfrom kimi_cli.wire.types import TextPart\n\nclass LazyCompaction:\n async def compact(self, messages, llm, *, custom_instruction=''):\n import helper_mod\n model_name = llm.model_config.model if llm.model_config is not None else 'unknown'\n return CompactionResult(messages=[Message(role='user', content=[TextPart(text=helper_mod.build_message(model_name))])], usage=None)\n", + encoding="utf-8", + ) + _write_plugin(pdir, name="lazy-plugin", entrypoint="lazy_compaction.LazyCompaction") + + comp = resolve_plugin_compactor(plugins, "lazy-plugin") + assert comp is not None + result = await comp.compact( + [Message(role="user", content=[TextPart(text="hi")])], _make_test_llm("main-chat") + ) + assert result.messages[0].extract_text("\n") == "lazy via main-chat" + + +@pytest.mark.asyncio +async def test_resolve_plugin_compactor_supports_eager_sibling_imports(tmp_path: Path) -> None: + plugins = tmp_path / "plugins" + pdir = plugins / "eager-plugin" + pdir.mkdir(parents=True, exist_ok=True) + (pdir / "helper_mod.py").write_text( + "def build_message(model_name):\n return f'eager via {model_name}'\n", + encoding="utf-8", + ) + (pdir / "eager_compaction.py").write_text( + "import helper_mod\nfrom kosong.message import Message\nfrom kimi_cli.soul.compaction import CompactionResult\nfrom kimi_cli.wire.types import TextPart\n\nclass EagerCompaction:\n async def compact(self, messages, llm, *, custom_instruction=''):\n model_name = llm.model_config.model if llm.model_config is not None else 'unknown'\n return CompactionResult(messages=[Message(role='user', content=[TextPart(text=helper_mod.build_message(model_name))])], usage=None)\n", + encoding="utf-8", + ) + _write_plugin(pdir, name="eager-plugin", entrypoint="eager_compaction.EagerCompaction") + + comp = resolve_plugin_compactor(plugins, "eager-plugin") + assert comp is not None + result = await comp.compact( + [Message(role="user", content=[TextPart(text="hi")])], _make_test_llm("main-chat") + ) + assert result.messages[0].extract_text("\n") == "eager via main-chat" + + +def test_resolve_plugin_compactor_wraps_import_failures(tmp_path: Path) -> None: + plugins = tmp_path / "plugins" + pdir = plugins / "broken-plugin" + pdir.mkdir(parents=True, exist_ok=True) + (pdir / "broken_compaction.py").write_text( + "import missing_helper\n\nclass BrokenCompaction:\n async def compact(self, messages, llm, *, custom_instruction=''):\n return messages\n", + encoding="utf-8", + ) + _write_plugin(pdir, name="broken-plugin", entrypoint="broken_compaction.BrokenCompaction") + + with pytest.raises(PluginError, match="Failed to import compaction module 'broken_compaction'"): + resolve_plugin_compactor(plugins, "broken-plugin") + + +def test_subagent_copies_get_fresh_compaction_instances(runtime: Runtime, tmp_path: Path) -> None: + plugins = tmp_path / "plugins" + pdir = plugins / "alpha-plugin" + _write_alpha_compactor(pdir) + _write_plugin(pdir, name="alpha-plugin", entrypoint="alpha_compaction.AlphaCompaction") + + runtime.compaction = resolve_plugin_compactor(plugins, "alpha-plugin") + assert runtime.compaction is not None + + subagent_a = runtime.copy_for_subagent(agent_id="alpha-one", subagent_type="coder") + subagent_b = runtime.copy_for_subagent(agent_id="alpha-two", subagent_type="plan") + + assert subagent_a.compaction is not runtime.compaction + assert subagent_b.compaction is not runtime.compaction + assert subagent_a.compaction is not subagent_b.compaction + assert getattr(type(subagent_a.compaction), "PLUGIN_MARK", None) == "alpha" + assert getattr(type(subagent_b.compaction), "PLUGIN_MARK", None) == "alpha" + + +@pytest.mark.asyncio +async def test_compact_context_uses_runtime_plugin_compactor( + runtime: Runtime, tmp_path: Path +) -> None: + plugins = tmp_path / "plugins" + pdir = plugins / "alpha-plugin" + _write_alpha_compactor(pdir) + _write_plugin( + pdir, + name="alpha-plugin", + entrypoint="alpha_compaction.AlphaCompaction", + ) + runtime.llm = _make_test_llm("main-chat") + runtime.compaction_llm = _make_test_llm("compact-chat") + runtime.compaction = resolve_plugin_compactor(plugins, "alpha-plugin") + assert runtime.compaction is not None + + soul, context = _make_soul(runtime, tmp_path) + await context.append_message( + [ + Message(role="user", content=[TextPart(text="message 1")]), + Message(role="assistant", content=[TextPart(text="message 2")]), + Message(role="user", content=[TextPart(text="message 3")]), + Message(role="assistant", content=[TextPart(text="message 4")]), + ] + ) + + wire = Wire() + wire_ui = wire.ui_side(merge=False) + token = _current_wire.set(wire) + try: + await soul.compact_context() + finally: + _current_wire.reset(token) + + begin = await wire_ui.receive() + end = await wire_ui.receive() + assert isinstance(begin, CompactionBegin) + assert isinstance(end, CompactionEnd) + assert getattr(runtime.compaction, "calls", 0) == 1 + assert getattr(runtime.compaction, "last_model", None) == "compact-chat" + assert [message.extract_text("\n") for message in context.history] == [ + "plugin compacted via compact-chat" + ] + + +@pytest.mark.asyncio +async def test_compact_context_falls_back_to_active_llm_when_compaction_llm_is_too_small( + runtime: Runtime, tmp_path: Path +) -> None: + plugins = tmp_path / "plugins" + pdir = plugins / "alpha-plugin" + _write_alpha_compactor(pdir) + _write_plugin( + pdir, + name="alpha-plugin", + entrypoint="alpha_compaction.AlphaCompaction", + ) + runtime.llm = _make_test_llm("main-chat", max_context_size=200_000) + runtime.compaction_llm = _make_test_llm("compact-chat", max_context_size=50_000) + runtime.compaction = resolve_plugin_compactor(plugins, "alpha-plugin") + assert runtime.compaction is not None + + soul, context = _make_soul(runtime, tmp_path) + await context.append_message( + [ + Message(role="user", content=[TextPart(text="message 1")]), + Message(role="assistant", content=[TextPart(text="message 2")]), + Message(role="user", content=[TextPart(text="message 3")]), + Message(role="assistant", content=[TextPart(text="message 4")]), + ] + ) + + token = _current_wire.set(Wire()) + try: + await soul.compact_context() + finally: + _current_wire.reset(token) + + assert getattr(runtime.compaction, "last_model", None) == "main-chat" diff --git a/tests/core/test_runtime_roles.py b/tests/core/test_runtime_roles.py index 33a169959..830379011 100644 --- a/tests/core/test_runtime_roles.py +++ b/tests/core/test_runtime_roles.py @@ -1,5 +1,23 @@ from __future__ import annotations +from kosong.chat_provider.echo import EchoChatProvider + +from kimi_cli.config import LLMModel +from kimi_cli.llm import LLM + + +def _make_test_llm(model_name: str, *, max_context_size: int) -> LLM: + return LLM( + chat_provider=EchoChatProvider(), + max_context_size=max_context_size, + capabilities=set(), + model_config=LLMModel( + provider="_echo", + model=model_name, + max_context_size=max_context_size, + ), + ) + def test_runtime_roles_are_root_and_subagent_only(runtime): assert runtime.role == "root" @@ -10,3 +28,18 @@ def test_runtime_roles_are_root_and_subagent_only(runtime): ) assert subagent_runtime.role == "subagent" + + +def test_subagent_runtime_reuses_active_llm_when_root_compaction_llm_is_too_small(runtime): + runtime.llm = _make_test_llm("root-main", max_context_size=100_000) + runtime.compaction_llm = _make_test_llm("root-compact", max_context_size=50_000) + subagent_llm = _make_test_llm("subagent-main", max_context_size=200_000) + + subagent_runtime = runtime.copy_for_subagent( + agent_id="atestcompact", + subagent_type="coder", + llm_override=subagent_llm, + ) + + assert subagent_runtime.llm is subagent_llm + assert subagent_runtime.compaction_llm is subagent_llm diff --git a/tests/core/test_startup_progress.py b/tests/core/test_startup_progress.py index cd1f180c2..0a3f500ee 100644 --- a/tests/core/test_startup_progress.py +++ b/tests/core/test_startup_progress.py @@ -1,13 +1,18 @@ from __future__ import annotations +import json from types import SimpleNamespace from unittest.mock import AsyncMock, Mock import pytest +from pydantic import SecretStr import kimi_cli.app as app_module import kimi_cli.ui.shell.startup as startup_module from kimi_cli.app import KimiCLI +from kimi_cli.config import LLMModel, LLMProvider +from kimi_cli.exception import ConfigError +from kimi_cli.plugin import PluginError from kimi_cli.ui.shell.startup import ShellStartupProgress @@ -89,6 +94,9 @@ async def fake_restore() -> None: monkeypatch.setattr(app_module, "load_config", lambda conf: conf) monkeypatch.setattr(app_module, "augment_provider_with_env_vars", lambda provider, model: {}) + monkeypatch.setattr( + app_module, "augment_provider_credentials_with_env_vars", lambda provider: {} + ) monkeypatch.setattr(app_module, "create_llm", lambda *args, **kwargs: None) monkeypatch.setattr(app_module.Runtime, "create", fake_runtime_create) monkeypatch.setattr(app_module, "load_agent", fake_load_agent) @@ -107,6 +115,93 @@ async def fake_restore() -> None: write_system_prompt.assert_awaited_once_with("Test system prompt") +@pytest.mark.asyncio +async def test_kimi_cli_create_passes_compaction_llm(session, config, monkeypatch) -> None: + config.default_model = "main" + config.models = { + "main": LLMModel(provider="main-provider", model="main-model", max_context_size=4096), + "compact": LLMModel( + provider="compact-provider", + model="compact-model", + max_context_size=8192, + ), + } + config.providers = { + "main-provider": LLMProvider( + type="_echo", + base_url="", + api_key=SecretStr(""), + ), + "compact-provider": LLMProvider( + type="_echo", + base_url="", + api_key=SecretStr(""), + ), + } + config.loop_control.compaction_model = "compact" + + main_llm = SimpleNamespace(name="main-llm", max_context_size=4096) + compact_llm = SimpleNamespace(name="compact-llm", max_context_size=8192) + captured: dict[str, object] = {} + + fake_runtime = SimpleNamespace( + session=session, + config=config, + llm=main_llm, + compaction_llm=compact_llm, + notifications=SimpleNamespace(recover=lambda: None), + background_tasks=SimpleNamespace(reconcile=lambda: None), + ) + fake_agent = SimpleNamespace(name="Test Agent", system_prompt="Test system prompt") + fake_context = SimpleNamespace(system_prompt=None) + write_system_prompt = AsyncMock() + + def fake_create_llm(provider, model, **kwargs): + if model.model == "compact-model": + return compact_llm + if model.model == "main-model": + return main_llm + raise AssertionError(f"Unexpected model {model.model!r}") + + async def fake_runtime_create( + config_arg, oauth, llm, compaction_llm, session_arg, yolo, skills_dir + ): + captured["llm"] = llm + captured["compaction_llm"] = compaction_llm + captured["session"] = session_arg + return fake_runtime + + async def fake_load_agent(*args, **kwargs): + return fake_agent + + async def fake_restore() -> None: + return None + + fake_context.restore = fake_restore + fake_context.write_system_prompt = write_system_prompt + + monkeypatch.setattr(app_module, "load_config", lambda conf: conf) + monkeypatch.setattr(app_module, "augment_provider_with_env_vars", lambda provider, model: {}) + monkeypatch.setattr( + app_module, "augment_provider_credentials_with_env_vars", lambda provider: {} + ) + monkeypatch.setattr(app_module, "create_llm", fake_create_llm) + monkeypatch.setattr(app_module.Runtime, "create", fake_runtime_create) + monkeypatch.setattr(app_module, "load_agent", fake_load_agent) + monkeypatch.setattr(app_module, "Context", lambda _path: fake_context) + monkeypatch.setattr(app_module, "KimiSoul", lambda agent, context: (agent, context)) + + cli = await KimiCLI.create(session, config=config) + + assert isinstance(cli, KimiCLI) + assert captured == { + "llm": main_llm, + "compaction_llm": compact_llm, + "session": session, + } + write_system_prompt.assert_awaited_once_with("Test system prompt") + + @pytest.mark.asyncio async def test_kimi_cli_create_cleans_stale_running_foreground_subagents( session, config, monkeypatch @@ -145,6 +240,9 @@ async def fake_restore() -> None: monkeypatch.setattr(app_module, "load_config", lambda conf: conf) monkeypatch.setattr(app_module, "augment_provider_with_env_vars", lambda provider, model: {}) + monkeypatch.setattr( + app_module, "augment_provider_credentials_with_env_vars", lambda provider: {} + ) monkeypatch.setattr(app_module, "create_llm", lambda *args, **kwargs: None) monkeypatch.setattr(app_module.Runtime, "create", fake_runtime_create) monkeypatch.setattr(app_module, "load_agent", fake_load_agent) @@ -154,3 +252,471 @@ async def fake_restore() -> None: await KimiCLI.create(session, config=config) update_instance.assert_called_once_with("afg1", status="failed") + + +@pytest.mark.asyncio +async def test_kimi_cli_create_warns_for_unknown_model_name(session, config, monkeypatch) -> None: + warnings: list[str] = [] + fake_runtime = SimpleNamespace( + session=session, + config=config, + llm=None, + notifications=SimpleNamespace(recover=lambda: None), + background_tasks=SimpleNamespace(reconcile=lambda: None), + ) + fake_agent = SimpleNamespace(name="Test Agent", system_prompt="Test system prompt") + fake_context = SimpleNamespace(system_prompt=None) + write_system_prompt = AsyncMock() + + async def fake_runtime_create(*args, **kwargs): + return fake_runtime + + async def fake_load_agent(*args, **kwargs): + return fake_agent + + async def fake_restore() -> None: + return None + + def fake_warning(message: str, **kwargs) -> None: + warnings.append(message.format(**kwargs)) + + fake_context.restore = fake_restore + fake_context.write_system_prompt = write_system_prompt + + monkeypatch.setattr(app_module, "load_config", lambda conf: conf) + monkeypatch.setattr(app_module, "augment_provider_with_env_vars", lambda provider, model: {}) + monkeypatch.setattr( + app_module, "augment_provider_credentials_with_env_vars", lambda provider: {} + ) + monkeypatch.setattr(app_module, "create_llm", lambda *args, **kwargs: None) + monkeypatch.setattr(app_module.Runtime, "create", fake_runtime_create) + monkeypatch.setattr(app_module, "load_agent", fake_load_agent) + monkeypatch.setattr(app_module, "Context", lambda _path: fake_context) + monkeypatch.setattr(app_module, "KimiSoul", lambda agent, context: (agent, context)) + monkeypatch.setattr(app_module.logger, "warning", fake_warning) + + cli = await KimiCLI.create(session, config=config, model_name="missing") + + assert isinstance(cli, KimiCLI) + assert warnings == ["Model 'missing' not found in config, using placeholder"] + write_system_prompt.assert_awaited_once_with("Test system prompt") + + +@pytest.mark.asyncio +async def test_kimi_cli_create_warns_for_missing_model_provider( + session, config, monkeypatch +) -> None: + config.default_model = "main" + config.models = { + "main": LLMModel(provider="missing-provider", model="main-model", max_context_size=4096), + } + config.providers = {} + + warnings: list[str] = [] + fake_runtime = SimpleNamespace( + session=session, + config=config, + llm=None, + notifications=SimpleNamespace(recover=lambda: None), + background_tasks=SimpleNamespace(reconcile=lambda: None), + ) + fake_agent = SimpleNamespace(name="Test Agent", system_prompt="Test system prompt") + fake_context = SimpleNamespace(system_prompt=None) + write_system_prompt = AsyncMock() + + async def fake_runtime_create(*args, **kwargs): + return fake_runtime + + async def fake_load_agent(*args, **kwargs): + return fake_agent + + async def fake_restore() -> None: + return None + + def fake_warning(message: str, **kwargs) -> None: + warnings.append(message.format(**kwargs)) + + fake_context.restore = fake_restore + fake_context.write_system_prompt = write_system_prompt + + monkeypatch.setattr(app_module, "load_config", lambda conf: conf) + monkeypatch.setattr(app_module, "augment_provider_with_env_vars", lambda provider, model: {}) + monkeypatch.setattr(app_module, "create_llm", lambda *args, **kwargs: None) + monkeypatch.setattr(app_module.Runtime, "create", fake_runtime_create) + monkeypatch.setattr(app_module, "load_agent", fake_load_agent) + monkeypatch.setattr(app_module, "Context", lambda _path: fake_context) + monkeypatch.setattr(app_module, "KimiSoul", lambda agent, context: (agent, context)) + monkeypatch.setattr(app_module.logger, "warning", fake_warning) + + cli = await KimiCLI.create(session, config=config) + + assert isinstance(cli, KimiCLI) + assert warnings == ["Provider 'missing-provider' for model 'main' missing; using placeholder"] + write_system_prompt.assert_awaited_once_with("Test system prompt") + + +@pytest.mark.asyncio +async def test_kimi_cli_create_skips_compaction_llm_when_provider_is_missing( + session, config, monkeypatch +) -> None: + config.default_model = "main" + config.models = { + "main": LLMModel(provider="main-provider", model="main-model", max_context_size=4096), + "compact": LLMModel( + provider="missing-provider", + model="compact-model", + max_context_size=8192, + ), + } + config.providers = { + "main-provider": LLMProvider( + type="_echo", + base_url="", + api_key=SecretStr(""), + ), + } + config.loop_control.compaction_model = "compact" + + warnings: list[str] = [] + main_llm = SimpleNamespace(name="main-llm", max_context_size=4096) + captured: dict[str, object] = {} + fake_runtime = SimpleNamespace( + session=session, + config=config, + llm=main_llm, + compaction_llm=None, + notifications=SimpleNamespace(recover=lambda: None), + background_tasks=SimpleNamespace(reconcile=lambda: None), + ) + fake_agent = SimpleNamespace(name="Test Agent", system_prompt="Test system prompt") + fake_context = SimpleNamespace(system_prompt=None) + write_system_prompt = AsyncMock() + + def fake_create_llm(provider, model, **kwargs): + assert model.model == "main-model" + return main_llm + + async def fake_runtime_create( + config_arg, oauth, llm, compaction_llm, session_arg, yolo, skills_dir + ): + captured["llm"] = llm + captured["compaction_llm"] = compaction_llm + captured["session"] = session_arg + return fake_runtime + + async def fake_load_agent(*args, **kwargs): + return fake_agent + + async def fake_restore() -> None: + return None + + def fake_warning(message: str, **kwargs) -> None: + warnings.append(message.format(**kwargs)) + + fake_context.restore = fake_restore + fake_context.write_system_prompt = write_system_prompt + + monkeypatch.setattr(app_module, "load_config", lambda conf: conf) + monkeypatch.setattr(app_module, "augment_provider_with_env_vars", lambda provider, model: {}) + monkeypatch.setattr( + app_module, "augment_provider_credentials_with_env_vars", lambda provider: {} + ) + monkeypatch.setattr(app_module, "create_llm", fake_create_llm) + monkeypatch.setattr(app_module.Runtime, "create", fake_runtime_create) + monkeypatch.setattr(app_module, "load_agent", fake_load_agent) + monkeypatch.setattr(app_module, "Context", lambda _path: fake_context) + monkeypatch.setattr(app_module, "KimiSoul", lambda agent, context: (agent, context)) + monkeypatch.setattr(app_module.logger, "warning", fake_warning) + + cli = await KimiCLI.create(session, config=config) + + assert isinstance(cli, KimiCLI) + assert captured == { + "llm": main_llm, + "compaction_llm": None, + "session": session, + } + assert warnings == ["Compaction provider 'missing-provider' not found in config, skipping"] + write_system_prompt.assert_awaited_once_with("Test system prompt") + + +@pytest.mark.asyncio +async def test_kimi_cli_create_rejects_smaller_compaction_model( + session, config, monkeypatch +) -> None: + config.default_model = "main" + config.models = { + "main": LLMModel(provider="main-provider", model="main-model", max_context_size=8192), + "compact": LLMModel( + provider="compact-provider", + model="compact-model", + max_context_size=4096, + ), + } + config.providers = { + "main-provider": LLMProvider(type="_echo", base_url="", api_key=SecretStr("")), + "compact-provider": LLMProvider(type="_echo", base_url="", api_key=SecretStr("")), + } + config.loop_control.compaction_model = "compact" + + runtime_create_called = False + + def fake_create_llm(provider, model, **kwargs): + return SimpleNamespace(name=model.model, max_context_size=model.max_context_size) + + async def fake_runtime_create(*args, **kwargs): + nonlocal runtime_create_called + runtime_create_called = True + raise AssertionError("Runtime.create should not be called for invalid compaction sizing") + + monkeypatch.setattr(app_module, "load_config", lambda conf: conf) + monkeypatch.setattr(app_module, "augment_provider_with_env_vars", lambda provider, model: {}) + monkeypatch.setattr( + app_module, "augment_provider_credentials_with_env_vars", lambda provider: {} + ) + monkeypatch.setattr(app_module, "create_llm", fake_create_llm) + monkeypatch.setattr(app_module.Runtime, "create", fake_runtime_create) + + with pytest.raises(ConfigError, match="smaller than active model"): + await KimiCLI.create(session, config=config) + + assert runtime_create_called is False + + +@pytest.mark.asyncio +async def test_kimi_cli_create_allows_smaller_compaction_model_without_active_llm( + session, config, monkeypatch +) -> None: + config.models = { + "compact": LLMModel( + provider="compact-provider", + model="compact-model", + max_context_size=4096, + ), + } + config.providers = { + "compact-provider": LLMProvider(type="_echo", base_url="", api_key=SecretStr("")), + } + config.loop_control.compaction_model = "compact" + + captured: dict[str, object] = {} + fake_runtime = SimpleNamespace( + session=session, + config=config, + llm=None, + compaction_llm=SimpleNamespace(name="compact-model", max_context_size=4096), + notifications=SimpleNamespace(recover=lambda: None), + background_tasks=SimpleNamespace(reconcile=lambda: None), + ) + fake_agent = SimpleNamespace(name="Test Agent", system_prompt="Test system prompt") + fake_context = SimpleNamespace(system_prompt=None) + write_system_prompt = AsyncMock() + + def fake_create_llm(provider, model, **kwargs): + if model.model == "": + return None + return SimpleNamespace(name=model.model, max_context_size=model.max_context_size) + + async def fake_runtime_create( + config_arg, oauth, llm, compaction_llm, session_arg, yolo, skills_dir + ): + captured["llm"] = llm + captured["compaction_llm"] = compaction_llm + return fake_runtime + + async def fake_load_agent(*args, **kwargs): + return fake_agent + + async def fake_restore() -> None: + return None + + fake_context.restore = fake_restore + fake_context.write_system_prompt = write_system_prompt + + monkeypatch.setattr(app_module, "load_config", lambda conf: conf) + monkeypatch.setattr(app_module, "augment_provider_with_env_vars", lambda provider, model: {}) + monkeypatch.setattr( + app_module, "augment_provider_credentials_with_env_vars", lambda provider: {} + ) + monkeypatch.setattr(app_module, "create_llm", fake_create_llm) + monkeypatch.setattr(app_module.Runtime, "create", fake_runtime_create) + monkeypatch.setattr(app_module, "load_agent", fake_load_agent) + monkeypatch.setattr(app_module, "Context", lambda _path: fake_context) + monkeypatch.setattr(app_module, "KimiSoul", lambda agent, context: (agent, context)) + + cli = await KimiCLI.create(session, config=config) + + assert isinstance(cli, KimiCLI) + assert captured["llm"] is None + assert getattr(captured["compaction_llm"], "name", None) == "compact-model" + write_system_prompt.assert_awaited_once_with("Test system prompt") + + +@pytest.mark.asyncio +async def test_kimi_cli_create_keeps_explicit_compaction_model_under_env_override( + session, config, monkeypatch +) -> None: + config.default_model = "main" + config.models = { + "main": LLMModel(provider="main-provider", model="main-model", max_context_size=4096), + "compact": LLMModel( + provider="compact-provider", + model="compact-model", + max_context_size=8192, + ), + } + config.providers = { + "main-provider": LLMProvider( + type="kimi", + base_url="https://config.test/v1", + api_key=SecretStr("config-main"), + ), + "compact-provider": LLMProvider( + type="kimi", + base_url="https://config.test/v1", + api_key=SecretStr("config-compact"), + ), + } + config.loop_control.compaction_model = "compact" + + captured_calls: list[tuple[str, int, str, str]] = [] + fake_runtime = SimpleNamespace( + session=session, + config=config, + llm=SimpleNamespace(name="main-llm"), + compaction_llm=SimpleNamespace(name="compact-llm"), + notifications=SimpleNamespace(recover=lambda: None), + background_tasks=SimpleNamespace(reconcile=lambda: None), + ) + fake_agent = SimpleNamespace(name="Test Agent", system_prompt="Test system prompt") + fake_context = SimpleNamespace(system_prompt=None) + write_system_prompt = AsyncMock() + + def fake_create_llm(provider, model, **kwargs): + captured_calls.append( + ( + model.model, + model.max_context_size, + provider.base_url, + provider.api_key.get_secret_value(), + ) + ) + return SimpleNamespace(name=model.model, max_context_size=model.max_context_size) + + async def fake_runtime_create(*args, **kwargs): + return fake_runtime + + async def fake_load_agent(*args, **kwargs): + return fake_agent + + async def fake_restore() -> None: + return None + + fake_context.restore = fake_restore + fake_context.write_system_prompt = write_system_prompt + + monkeypatch.setenv("KIMI_MODEL_NAME", "env-main-model") + monkeypatch.setenv("KIMI_BASE_URL", "https://env.test/v1") + monkeypatch.setenv("KIMI_API_KEY", "env-key") + monkeypatch.setattr(app_module, "load_config", lambda conf: conf) + monkeypatch.setattr(app_module, "create_llm", fake_create_llm) + monkeypatch.setattr(app_module.Runtime, "create", fake_runtime_create) + monkeypatch.setattr(app_module, "load_agent", fake_load_agent) + monkeypatch.setattr(app_module, "Context", lambda _path: fake_context) + monkeypatch.setattr(app_module, "KimiSoul", lambda agent, context: (agent, context)) + + cli = await KimiCLI.create(session, config=config) + + assert isinstance(cli, KimiCLI) + assert captured_calls == [ + ("env-main-model", 4096, "https://env.test/v1", "env-key"), + ("compact-model", 8192, "https://env.test/v1", "env-key"), + ] + write_system_prompt.assert_awaited_once_with("Test system prompt") + + +@pytest.mark.asyncio +async def test_kimi_cli_create_surfaces_invalid_compaction_plugin( + session, config, monkeypatch +) -> None: + fake_runtime = SimpleNamespace( + session=session, + config=config, + llm=None, + notifications=SimpleNamespace(recover=lambda: None), + background_tasks=SimpleNamespace(reconcile=lambda: None), + ) + + async def fake_runtime_create(*args, **kwargs): + return fake_runtime + + monkeypatch.setattr(app_module, "load_config", lambda conf: conf) + monkeypatch.setattr(app_module, "augment_provider_with_env_vars", lambda provider, model: {}) + monkeypatch.setattr( + app_module, "augment_provider_credentials_with_env_vars", lambda provider: {} + ) + monkeypatch.setattr(app_module, "create_llm", lambda *args, **kwargs: None) + monkeypatch.setattr(app_module.Runtime, "create", fake_runtime_create) + monkeypatch.setattr( + "kimi_cli.plugin.compaction.resolve_plugin_compactor", + lambda *args, **kwargs: (_ for _ in ()).throw(PluginError("broken entrypoint")), + ) + monkeypatch.setattr( + "kimi_cli.plugin.manager.get_plugins_dir", lambda: session.context_file.parent + ) + + config.loop_control.compaction_plugin = "broken-plugin" + + with pytest.raises(ConfigError, match="Invalid compaction plugin 'broken-plugin'"): + await KimiCLI.create(session, config=config) + + +@pytest.mark.asyncio +async def test_kimi_cli_create_wraps_plugin_import_failures_as_config_error( + session, config, monkeypatch, tmp_path +) -> None: + plugins = tmp_path / "plugins" + pdir = plugins / "broken-plugin" + pdir.mkdir(parents=True, exist_ok=True) + (pdir / "broken_compaction.py").write_text( + "import missing_helper\n\nclass BrokenCompaction:\n async def compact(self, messages, llm, *, custom_instruction=''):\n return messages\n", + encoding="utf-8", + ) + (pdir / "plugin.json").write_text( + json.dumps( + { + "name": "broken-plugin", + "version": "1.0.0", + "compaction": {"entrypoint": "broken_compaction.BrokenCompaction"}, + } + ), + encoding="utf-8", + ) + + fake_runtime = SimpleNamespace( + session=session, + config=config, + llm=None, + notifications=SimpleNamespace(recover=lambda: None), + background_tasks=SimpleNamespace(reconcile=lambda: None), + ) + + async def fake_runtime_create(*args, **kwargs): + return fake_runtime + + monkeypatch.setattr(app_module, "load_config", lambda conf: conf) + monkeypatch.setattr(app_module, "augment_provider_with_env_vars", lambda provider, model: {}) + monkeypatch.setattr( + app_module, "augment_provider_credentials_with_env_vars", lambda provider: {} + ) + monkeypatch.setattr(app_module, "create_llm", lambda *args, **kwargs: None) + monkeypatch.setattr(app_module.Runtime, "create", fake_runtime_create) + monkeypatch.setattr("kimi_cli.plugin.manager.get_plugins_dir", lambda: plugins) + + config.loop_control.compaction_plugin = "broken-plugin" + + with pytest.raises( + ConfigError, + match="Invalid compaction plugin 'broken-plugin': Failed to import compaction module 'broken_compaction'", + ): + await KimiCLI.create(session, config=config)