from __future__ import annotations import ast import hashlib import json import os import re import shutil import time import zipfile from dataclasses import dataclass, asdict from pathlib import Path, PurePosixPath from typing import Any, Callable try: from jsonschema import Draft202012Validator except Exception: # pragma: no cover - target installer provides jsonschema Draft202012Validator = None class ToolRuntimeError(RuntimeError): def __init__(self, code: str, message: str, details: dict[str, Any] | None = None): super().__init__(message) self.code = code self.details = details or {} @dataclass(frozen=True) class ToolDefinition: id: str version: int name: str description: str risk_class: str approval_required: bool input_schema: dict[str, Any] output_limit: int = 262_144 timeout_seconds: int = 30 offline: bool = True compatible_models: tuple[str, ...] = ("k2.6", "k2.7-code") def public(self) -> dict[str, Any]: value = asdict(self) value["compatible_models"] = list(self.compatible_models) return value SAFE_SEGMENT = re.compile(r"^[^\x00-\x1f<>:\"|?*]+$") def safe_relative_path(value: str, *, allow_root: bool = False) -> PurePosixPath: value = str(value or "").replace("\\", "/").strip() if value in {"", "."} and allow_root: return PurePosixPath(".") path = PurePosixPath(value) if path.is_absolute() or not path.parts or any(part in {"", ".", ".."} for part in path.parts): raise ToolRuntimeError("KIMU_TOOL_PATH_INVALID", "The requested path is outside the workspace.") if any(not SAFE_SEGMENT.match(part) for part in path.parts): raise ToolRuntimeError("KIMU_TOOL_PATH_INVALID", "The requested path contains unsupported characters.") return path def resolve_workspace_path(root: Path, value: str, *, allow_root: bool = False) -> Path: relative = safe_relative_path(value, allow_root=allow_root) root = root.resolve() target = (root / relative).resolve() try: target.relative_to(root) except ValueError as exc: raise ToolRuntimeError("KIMU_TOOL_PATH_ESCAPE", "The requested path escapes the workspace.") from exc return target def file_sha256(path: Path) -> str: digest = hashlib.sha256() with path.open("rb") as handle: for chunk in iter(lambda: handle.read(1024 * 1024), b""): digest.update(chunk) return digest.hexdigest() def _schema_object(properties: dict[str, Any], required: list[str] | None = None) -> dict[str, Any]: return { "type": "object", "properties": properties, "required": required or [], "additionalProperties": False, } TOOL_DEFINITIONS: dict[str, ToolDefinition] = { "workspace.list": ToolDefinition( "workspace.list", 1, "List workspace files", "List authorised files under a workspace path.", "read", False, _schema_object({"path": {"type": "string", "default": "."}, "recursive": {"type": "boolean", "default": True}, "limit": {"type": "integer", "minimum": 1, "maximum": 2000, "default": 500}}), ), "workspace.read": ToolDefinition( "workspace.read", 1, "Read workspace file", "Read a UTF-8 text file with a strict output limit.", "read", False, _schema_object({"path": {"type": "string"}, "max_chars": {"type": "integer", "minimum": 1, "maximum": 250000, "default": 50000}}, ["path"]), ), "workspace.search": ToolDefinition( "workspace.search", 1, "Search workspace", "Search text files for a literal or regular expression.", "read", False, _schema_object({"query": {"type": "string", "minLength": 1, "maxLength": 500}, "path": {"type": "string", "default": "."}, "regex": {"type": "boolean", "default": False}, "case_sensitive": {"type": "boolean", "default": False}, "limit": {"type": "integer", "minimum": 1, "maximum": 500, "default": 100}}, ["query"]), ), "workspace.write": ToolDefinition( "workspace.write", 1, "Create or replace file", "Write one UTF-8 file inside the approved workspace.", "write", True, _schema_object({"path": {"type": "string"}, "content": {"type": "string", "maxLength": 2_000_000}, "expected_sha256": {"type": ["string", "null"], "pattern": "^[a-f0-9]{64}$"}}, ["path", "content"]), output_limit=65_536, ), "workspace.patch": ToolDefinition( "workspace.patch", 1, "Apply exact text patch", "Replace an exact text block once with a version check.", "write", True, _schema_object({"path": {"type": "string"}, "old_text": {"type": "string", "minLength": 1, "maxLength": 1_000_000}, "new_text": {"type": "string", "maxLength": 1_000_000}, "expected_sha256": {"type": ["string", "null"], "pattern": "^[a-f0-9]{64}$"}}, ["path", "old_text", "new_text"]), ), "workspace.move": ToolDefinition( "workspace.move", 1, "Move file", "Move or rename one file inside the workspace.", "write", True, _schema_object({"source": {"type": "string"}, "destination": {"type": "string"}, "overwrite": {"type": "boolean", "default": False}}, ["source", "destination"]), ), "workspace.delete": ToolDefinition( "workspace.delete", 1, "Delete file", "Delete one file or an empty directory.", "destructive", True, _schema_object({"path": {"type": "string"}, "expected_sha256": {"type": ["string", "null"], "pattern": "^[a-f0-9]{64}$"}}, ["path"]), ), "archive.inspect": ToolDefinition( "archive.inspect", 1, "Inspect ZIP", "Inspect a ZIP inventory without extraction.", "read", False, _schema_object({"path": {"type": "string"}, "limit": {"type": "integer", "minimum": 1, "maximum": 5000, "default": 1000}}, ["path"]), ), "archive.create": ToolDefinition( "archive.create", 1, "Create ZIP", "Create a ZIP from selected workspace paths.", "write", True, _schema_object({"output_path": {"type": "string"}, "paths": {"type": "array", "items": {"type": "string"}, "minItems": 1, "maxItems": 500}}, ["output_path", "paths"]), ), "checks.static": ToolDefinition( "checks.static", 1, "Run static checks", "Parse Python, JSON, JavaScript balance and HTML link references without arbitrary command execution.", "execute", True, _schema_object({"path": {"type": "string", "default": "."}, "limit": {"type": "integer", "minimum": 1, "maximum": 2000, "default": 1000}}), timeout_seconds=60, ), } def validate_arguments(definition: ToolDefinition, arguments: dict[str, Any]) -> dict[str, Any]: if not isinstance(arguments, dict): raise ToolRuntimeError("KIMU_TOOL_ARGUMENTS_INVALID", "Tool arguments must be a JSON object.") if Draft202012Validator is not None: errors = sorted(Draft202012Validator(definition.input_schema).iter_errors(arguments), key=lambda item: list(item.path)) if errors: details = [{"path": ".".join(str(part) for part in item.path), "message": item.message} for item in errors[:20]] raise ToolRuntimeError("KIMU_TOOL_ARGUMENTS_INVALID", "Tool arguments failed schema validation.", {"errors": details}) else: allowed = set(definition.input_schema.get("properties", {})) unknown = sorted(set(arguments) - allowed) missing = sorted(set(definition.input_schema.get("required", [])) - set(arguments)) if unknown or missing: raise ToolRuntimeError("KIMU_TOOL_ARGUMENTS_INVALID", "Tool arguments failed schema validation.", {"unknown": unknown, "missing": missing}) result = dict(arguments) for key, schema in definition.input_schema.get("properties", {}).items(): if key not in result and "default" in schema: result[key] = schema["default"] return result def _text_files(root: Path, limit: int = 1000): count = 0 for path in sorted(root.rglob("*")): if path.is_symlink() or not path.is_file(): continue if path.stat().st_size > 4 * 1024 * 1024: continue if path.suffix.lower() in {".png", ".jpg", ".jpeg", ".gif", ".webp", ".mp3", ".wav", ".mp4", ".zip", ".pdf", ".woff", ".woff2"}: continue yield path count += 1 if count >= limit: break def _relative(root: Path, path: Path) -> str: return path.resolve().relative_to(root.resolve()).as_posix() def _check_expected(path: Path, expected: str | None) -> None: if expected is None: return if not path.is_file() or file_sha256(path) != expected: raise ToolRuntimeError("KIMU_TOOL_VERSION_CONFLICT", "The file changed before the operation could be applied.") def execute_tool(definition: ToolDefinition, root: Path, arguments: dict[str, Any], cancel_check: Callable[[], bool] | None = None) -> dict[str, Any]: arguments = validate_arguments(definition, arguments) root = root.resolve() root.mkdir(parents=True, exist_ok=True) started = time.monotonic() def cancelled() -> bool: return bool(cancel_check and cancel_check()) if cancelled(): raise ToolRuntimeError("KIMU_TOOL_CANCELLED", "The operation was cancelled before execution.") if definition.id == "workspace.list": base = resolve_workspace_path(root, arguments["path"], allow_root=True) if not base.exists(): raise ToolRuntimeError("KIMU_TOOL_NOT_FOUND", "The requested workspace path does not exist.") candidates = base.rglob("*") if arguments["recursive"] else base.glob("*") entries = [] for path in sorted(candidates): if path.is_symlink(): continue entries.append({"path": _relative(root, path), "type": "directory" if path.is_dir() else "file", "size": 0 if path.is_dir() else path.stat().st_size, "sha256": file_sha256(path) if path.is_file() and path.stat().st_size <= 16 * 1024 * 1024 else None}) if len(entries) >= arguments["limit"]: break result = {"entries": entries, "truncated": len(entries) >= arguments["limit"]} elif definition.id == "workspace.read": path = resolve_workspace_path(root, arguments["path"]) if not path.is_file(): raise ToolRuntimeError("KIMU_TOOL_NOT_FOUND", "The requested file does not exist.") data = path.read_bytes() if b"\x00" in data[:8192]: raise ToolRuntimeError("KIMU_TOOL_BINARY_FILE", "The requested file is not a supported text file.") text = data.decode("utf-8", errors="replace") maximum = arguments["max_chars"] result = {"path": _relative(root, path), "content": text[:maximum], "truncated": len(text) > maximum, "sha256": hashlib.sha256(data).hexdigest(), "size": len(data)} elif definition.id == "workspace.search": base = resolve_workspace_path(root, arguments["path"], allow_root=True) flags = 0 if arguments["case_sensitive"] else re.IGNORECASE try: pattern = re.compile(arguments["query"] if arguments["regex"] else re.escape(arguments["query"]), flags) except re.error as exc: raise ToolRuntimeError("KIMU_TOOL_REGEX_INVALID", str(exc)) from exc matches = [] for path in _text_files(base, 2000): if cancelled(): raise ToolRuntimeError("KIMU_TOOL_CANCELLED", "The search was cancelled.") try: lines = path.read_text(encoding="utf-8", errors="replace").splitlines() except OSError: continue for number, line in enumerate(lines, 1): if pattern.search(line): matches.append({"path": _relative(root, path), "line": number, "text": line[:1000]}) if len(matches) >= arguments["limit"]: break if len(matches) >= arguments["limit"]: break result = {"matches": matches, "truncated": len(matches) >= arguments["limit"]} elif definition.id == "workspace.write": path = resolve_workspace_path(root, arguments["path"]) _check_expected(path, arguments.get("expected_sha256")) path.parent.mkdir(parents=True, exist_ok=True) temp = path.with_name(path.name + ".kimu-tmp") temp.write_text(arguments["content"], encoding="utf-8") os.replace(temp, path) result = {"path": _relative(root, path), "size": path.stat().st_size, "sha256": file_sha256(path)} elif definition.id == "workspace.patch": path = resolve_workspace_path(root, arguments["path"]) if not path.is_file(): raise ToolRuntimeError("KIMU_TOOL_NOT_FOUND", "The requested file does not exist.") _check_expected(path, arguments.get("expected_sha256")) text = path.read_text(encoding="utf-8") count = text.count(arguments["old_text"]) if count != 1: raise ToolRuntimeError("KIMU_TOOL_PATCH_AMBIGUOUS", "The old text must occur exactly once.", {"occurrences": count}) updated = text.replace(arguments["old_text"], arguments["new_text"], 1) path.write_text(updated, encoding="utf-8") result = {"path": _relative(root, path), "sha256": file_sha256(path), "changed_chars": len(updated) - len(text)} elif definition.id == "workspace.move": source = resolve_workspace_path(root, arguments["source"]) destination = resolve_workspace_path(root, arguments["destination"]) if not source.exists() or source.is_symlink(): raise ToolRuntimeError("KIMU_TOOL_NOT_FOUND", "The source path does not exist.") if destination.exists() and not arguments["overwrite"]: raise ToolRuntimeError("KIMU_TOOL_DESTINATION_EXISTS", "The destination already exists.") destination.parent.mkdir(parents=True, exist_ok=True) if destination.exists(): if destination.is_dir(): shutil.rmtree(destination) else: destination.unlink() shutil.move(str(source), str(destination)) result = {"source": arguments["source"], "destination": _relative(root, destination)} elif definition.id == "workspace.delete": path = resolve_workspace_path(root, arguments["path"]) if not path.exists() or path.is_symlink(): raise ToolRuntimeError("KIMU_TOOL_NOT_FOUND", "The requested path does not exist.") if path.is_file(): _check_expected(path, arguments.get("expected_sha256")) size = path.stat().st_size path.unlink() else: try: path.rmdir() except OSError as exc: raise ToolRuntimeError("KIMU_TOOL_DIRECTORY_NOT_EMPTY", "Only empty directories may be deleted.") from exc size = 0 result = {"path": arguments["path"], "deleted": True, "size": size} elif definition.id == "archive.inspect": path = resolve_workspace_path(root, arguments["path"]) if not path.is_file() or not zipfile.is_zipfile(path): raise ToolRuntimeError("KIMU_TOOL_ZIP_INVALID", "The requested file is not a valid ZIP archive.") entries = [] total_uncompressed = 0 with zipfile.ZipFile(path) as archive: for info in archive.infolist()[: arguments["limit"]]: safe_relative_path(info.filename.rstrip("/")) if info.filename.rstrip("/") else None total_uncompressed += info.file_size entries.append({"name": info.filename, "size": info.file_size, "compressed_size": info.compress_size, "directory": info.is_dir()}) result = {"path": _relative(root, path), "entries": entries, "total_uncompressed": total_uncompressed, "truncated": len(entries) >= arguments["limit"]} elif definition.id == "archive.create": output = resolve_workspace_path(root, arguments["output_path"]) output.parent.mkdir(parents=True, exist_ok=True) included = [] with zipfile.ZipFile(output, "w", compression=zipfile.ZIP_DEFLATED, compresslevel=6) as archive: for value in arguments["paths"]: source = resolve_workspace_path(root, value) if not source.exists() or source.is_symlink(): raise ToolRuntimeError("KIMU_TOOL_NOT_FOUND", f"Archive input does not exist: {value}") candidates = [source] if source.is_file() else [p for p in source.rglob("*") if p.is_file() and not p.is_symlink()] for path in candidates: if path.resolve() == output.resolve(): continue arcname = _relative(root, path) archive.write(path, arcname) included.append(arcname) result = {"path": _relative(root, output), "entries": included, "size": output.stat().st_size, "sha256": file_sha256(output)} elif definition.id == "checks.static": base = resolve_workspace_path(root, arguments["path"], allow_root=True) checked = [] failures = [] for path in _text_files(base, arguments["limit"]): if cancelled(): raise ToolRuntimeError("KIMU_TOOL_CANCELLED", "Static checks were cancelled.") relative = _relative(root, path) suffix = path.suffix.lower() try: text = path.read_text(encoding="utf-8", errors="strict") if suffix == ".py": ast.parse(text, filename=relative) elif suffix == ".json": json.loads(text) elif suffix in {".js", ".jsx", ".mjs", ".cjs"}: pairs = {"(": ")", "[": "]", "{": "}"} stack = [] quote = None escaped = False for char in text: if quote: if escaped: escaped = False elif char == "\\": escaped = True elif char == quote: quote = None continue if char in {"'", '"', "`"}: quote = char elif char in pairs: stack.append(pairs[char]) elif char in pairs.values(): if not stack or stack.pop() != char: raise SyntaxError("Unbalanced delimiter") if stack or quote: raise SyntaxError("Unbalanced delimiter or string") checked.append(relative) except Exception as exc: failures.append({"path": relative, "message": str(exc)[:1000]}) result = {"checked": checked, "failures": failures, "passed": not failures} else: # pragma: no cover raise ToolRuntimeError("KIMU_TOOL_UNKNOWN", "Unknown tool.") elapsed_ms = int((time.monotonic() - started) * 1000) encoded = json.dumps(result, ensure_ascii=False).encode("utf-8") if len(encoded) > definition.output_limit: raise ToolRuntimeError("KIMU_TOOL_OUTPUT_LIMIT", "Tool output exceeded the configured limit.", {"bytes": len(encoded), "limit": definition.output_limit}) result["elapsed_ms"] = elapsed_ms return result def registry_payload() -> list[dict[str, Any]]: return [TOOL_DEFINITIONS[key].public() for key in sorted(TOOL_DEFINITIONS)] def provider_tools() -> list[dict[str, Any]]: """Return the registry in the OpenAI-compatible function-calling shape.""" return [ { "type": "function", "function": { "name": definition.id, "description": definition.description, "parameters": definition.input_schema, }, } for definition in (TOOL_DEFINITIONS[key] for key in sorted(TOOL_DEFINITIONS)) ]