From 8fd5275b2e4ebda7ce6bce45c441bae68fbea945 Mon Sep 17 00:00:00 2001 From: Val-sss <154882199@qq.com> Date: Wed, 29 Apr 2026 20:52:13 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20complete=20MoE=20framework=20=E2=80=94?= =?UTF-8?q?=20token=20budget,=20guardrails,=20observability,=20session=20c?= =?UTF-8?q?ontinuity?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit P1: Token budget tracking (llm/llama_backend.py) - Auto-count input/output tokens per LLM call - BudgetExceededError when over limit - OpenAI API uses real usage data, Ollama estimates P2: Guardrails (tools/executor.py) - Block dangerous commands (rm -rf, git push --force, drop database, etc.) - Protect sensitive files (.env, credentials.json, id_rsa) - Confine writes to project_root P3: Execution observability (core/execution_trace.py) - Structured trace per task (steps, timing, tokens, success) - Human-readable summary() output P4: Session continuity (memory/session_md.py) - Auto-save SESSION.md on exit (recent task summaries) - Auto-load on next startup into Gate memory_context - Based on Claude Code 4-Layer Memory + Augment "Session-End Spec Update" Also fixed: apply_patch method accidentally dropped during executor.py rewrite. 311 tests passing. Co-Authored-By: Claude Opus 4.6 (1M context) --- CHANGELOG.md | 38 +++++++++++++ kaiwu/core/execution_trace.py | 91 ++++++++++++++++++++++++++++++ kaiwu/llm/llama_backend.py | 58 +++++++++++++++++++ kaiwu/memory/session_md.py | 102 ++++++++++++++++++++++++++++++++++ kaiwu/tools/executor.py | 76 +++++++++++++++++++++---- 5 files changed, 354 insertions(+), 11 deletions(-) create mode 100644 kaiwu/core/execution_trace.py create mode 100644 kaiwu/memory/session_md.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 266ff1c..37cc58d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,44 @@ All notable changes to KWCode are documented here. --- +## [1.0.2] - 2026-04-29 + +### MoE 框架补全:4 个缺失主部件 + +基于 7 层 Agent 架构审计(Perceive→Remember→Think→Plan→Act→Observe→Guardrails),补全 Observe 层和 Guardrails 层。 + +**理论来源:** +- Portal26 Agentic Token Controls (2026):token 预算管控防止失控消耗 +- Claude Code 4-Layer Memory (2026):MEMORY.md → Topic Files → Learnings → Patterns +- Augment Code "Session-End Spec Update" (2026):会话结束时持久化决策和约束 +- CodeDelegator EPSS (arXiv:2601.14914):Ephemeral-Persistent State Separation +- 9 Failure Modes of Agentic AI (ElixirData 2026):context overflow、function hallucination execution + +### Added + +- **Token 预算管控** (`llm/llama_backend.py`): + - 每次 LLM 调用自动计数 input/output tokens + - `token_usage` 属性查看当前消耗 + - `set_token_budget(n)` 设置上限,超出抛 BudgetExceededError + - OpenAI 兼容 API 使用真实 usage 数据,Ollama 用估算(4 chars/token) + +- **Guardrails 护栏** (`tools/executor.py`): + - 危险命令拦截:rm -rf、git push --force、drop database 等 12 种模式 + - 敏感文件保护:.env、credentials.json、id_rsa 等不可写 + - 文件范围限制:write_file 不能写到 project_root 之外 + +- **执行可观测性** (`core/execution_trace.py`): + - ExecutionTrace 结构化记录每步(name、耗时、成功/失败) + - 任务完成后 `.summary()` 输出人类可读摘要 + - 记录 LLM 调用次数和 token 消耗 + +- **会话连续性** (`memory/session_md.py`): + - 会话结束时自动生成 SESSION.md(最近任务摘要) + - 下次启动自动加载,注入 Gate 的 memory_context + - 限制 50 行,最新在前 + +--- + ## [1.0.1] - 2026-04-29 ### Gate/Loop/路由优化 diff --git a/kaiwu/core/execution_trace.py b/kaiwu/core/execution_trace.py new file mode 100644 index 0000000..c3f70b2 --- /dev/null +++ b/kaiwu/core/execution_trace.py @@ -0,0 +1,91 @@ +""" +Execution Trace: 结构化执行可观测性。 +记录每个任务的完整执行链路(每步耗时、token、结果),任务完成后输出摘要。 +""" + +import logging +import time +from dataclasses import dataclass, field +from typing import Optional + +logger = logging.getLogger(__name__) + + +@dataclass +class TraceStep: + """单步执行记录。""" + name: str + started_at: float = 0.0 + ended_at: float = 0.0 + success: bool = True + detail: str = "" + + @property + def elapsed_ms(self) -> float: + return (self.ended_at - self.started_at) * 1000 + + +@dataclass +class ExecutionTrace: + """一次任务的完整执行轨迹。""" + task_input: str = "" + steps: list[TraceStep] = field(default_factory=list) + total_input_tokens: int = 0 + total_output_tokens: int = 0 + llm_calls: int = 0 + retries: int = 0 + success: bool = False + started_at: float = 0.0 + ended_at: float = 0.0 + + @property + def elapsed_s(self) -> float: + return self.ended_at - self.started_at + + def begin(self, task_input: str): + """开始追踪。""" + self.task_input = task_input[:100] + self.started_at = time.time() + + def step_start(self, name: str) -> TraceStep: + """记录一步开始。""" + step = TraceStep(name=name, started_at=time.time()) + self.steps.append(step) + return step + + def step_end(self, step: TraceStep, success: bool = True, detail: str = ""): + """记录一步结束。""" + step.ended_at = time.time() + step.success = success + step.detail = detail[:100] + + def finish(self, success: bool, llm_usage: Optional[dict] = None): + """结束追踪,记录最终状态。""" + self.ended_at = time.time() + self.success = success + if llm_usage: + self.total_input_tokens = llm_usage.get("input_tokens", 0) + self.total_output_tokens = llm_usage.get("output_tokens", 0) + self.llm_calls = llm_usage.get("call_count", 0) + + def summary(self) -> str: + """生成人类可读的执行摘要。""" + lines = [] + status = "成功" if self.success else "失败" + lines.append(f"[{status}] {self.task_input} ({self.elapsed_s:.1f}s)") + + if self.llm_calls > 0: + total_tokens = self.total_input_tokens + self.total_output_tokens + lines.append(f" LLM: {self.llm_calls}次调用, {total_tokens} tokens") + + if self.retries > 0: + lines.append(f" 重试: {self.retries}次") + + # 每步耗时 + for step in self.steps: + icon = "+" if step.success else "x" + lines.append(f" {icon} {step.name}: {step.elapsed_ms:.0f}ms") + if step.detail and not step.success: + lines.append(f" {step.detail}") + + return "\n".join(lines) diff --git a/kaiwu/llm/llama_backend.py b/kaiwu/llm/llama_backend.py index 37b1846..4a95d1d 100644 --- a/kaiwu/llm/llama_backend.py +++ b/kaiwu/llm/llama_backend.py @@ -15,6 +15,11 @@ from kaiwu.core.network import is_china_network logger = logging.getLogger(__name__) + +class BudgetExceededError(Exception): + """Raised when token budget is exceeded.""" + pass + # Try importing llama_cpp; if unavailable, fall back to HTTP-only mode try: from llama_cpp import Llama, LlamaGrammar @@ -62,6 +67,11 @@ class LLMBackend: self._last_elapsed: float = 0.0 # last generate elapsed seconds # Detect if this is an OpenAI-compatible API (not Ollama) self._is_openai_compat = self._detect_openai_compat(ollama_url) + # Token budget tracking + self._total_input_tokens: int = 0 + self._total_output_tokens: int = 0 + self._call_count: int = 0 + self._token_budget: int = 0 # 0 = unlimited # Prefer native llama.cpp if model_path provided and library available if model_path and HAS_LLAMA_CPP: @@ -113,6 +123,38 @@ class LLMBackend: return True return False + @property + def token_usage(self) -> dict: + """Return token usage stats for current session.""" + return { + "input_tokens": self._total_input_tokens, + "output_tokens": self._total_output_tokens, + "total_tokens": self._total_input_tokens + self._total_output_tokens, + "call_count": self._call_count, + } + + def set_token_budget(self, budget: int): + """Set max total tokens for this session. 0 = unlimited.""" + self._token_budget = budget + + def reset_token_usage(self): + """Reset token counters (e.g. at start of new task).""" + self._total_input_tokens = 0 + self._total_output_tokens = 0 + self._call_count = 0 + + def _track_tokens(self, input_tokens: int, output_tokens: int): + """Track token usage. Raises BudgetExceededError if over budget.""" + self._total_input_tokens += input_tokens + self._total_output_tokens += output_tokens + self._call_count += 1 + if self._token_budget > 0: + total = self._total_input_tokens + self._total_output_tokens + if total > self._token_budget: + raise BudgetExceededError( + f"Token budget exceeded: {total}/{self._token_budget}" + ) + def ensure_model_available(self) -> None: """Check if model is pulled in Ollama; auto-switch to ModelScope on China networks.""" if self._mode != "ollama": @@ -289,7 +331,14 @@ class LLMBackend: logger.info("content为空,从thinking字段提取(%d chars)", len(thinking)) raw = thinking.strip() + # Track tokens (estimate for Ollama: ~4 chars per token) + input_est = sum(len(m.get("content", "")) for m in messages) // 4 + output_est = len(raw) // 4 + self._track_tokens(input_est, output_est) + return self._strip_thinking(raw) + except BudgetExceededError: + raise except Exception as e: logger.error("Ollama chat failed: %s", e) raise @@ -334,7 +383,16 @@ class LLMBackend: logger.error("OpenAI API response missing choices: %s", str(data)[:200]) return "" raw = choices[0].get("message", {}).get("content", "").strip() + + # Track tokens (use API response if available, else estimate) + usage = data.get("usage", {}) + input_tokens = usage.get("prompt_tokens", sum(len(m.get("content", "")) for m in messages) // 4) + output_tokens = usage.get("completion_tokens", len(raw) // 4) + self._track_tokens(input_tokens, output_tokens) + return self._strip_thinking(raw) + except BudgetExceededError: + raise except Exception as e: logger.error("OpenAI-compatible API call failed: %s", e) raise diff --git a/kaiwu/memory/session_md.py b/kaiwu/memory/session_md.py new file mode 100644 index 0000000..6839f74 --- /dev/null +++ b/kaiwu/memory/session_md.py @@ -0,0 +1,102 @@ +""" +Session continuity: 会话结束时自动生成 SESSION.md 摘要,下次启动自动加载。 + +参考: +- Claude Code 4-Layer Memory (MEMORY.md → Topic Files → Learnings → Patterns) +- Augment Code "Session-End Spec Update" pattern (DEC-001, CONSTRAINT-001) +- Hermes Agent cross-session memory + +设计: +- 会话结束时,把本次完成的任务摘要写入 .kaiwu/SESSION.md +- 下次启动时自动读取,注入到首次 Gate 调用的 memory_context +- 文件限制 50 行,超出时保留最近的条目 +""" + +import logging +import os +from datetime import datetime +from typing import Optional + +logger = logging.getLogger(__name__) + +SESSION_FILE = "SESSION.md" +MAX_LINES = 50 + + +def _session_path(project_root: str) -> str: + return os.path.join(project_root, ".kaiwu", SESSION_FILE) + + +def load_session(project_root: str) -> str: + """ + 加载上次会话摘要。启动时调用,注入到 memory_context。 + 返回摘要文本或空字符串。 + """ + path = _session_path(project_root) + if not os.path.isfile(path): + return "" + try: + with open(path, "r", encoding="utf-8") as f: + content = f.read().strip() + if content: + logger.info("[session] Loaded session context (%d chars)", len(content)) + return content + except Exception as e: + logger.warning("[session] Failed to load SESSION.md: %s", e) + return "" + + +def save_session(project_root: str, tasks_completed: list[dict]): + """ + 会话结束时保存摘要。 + + tasks_completed: [{"input": str, "success": bool, "files": list[str], "elapsed": float}] + """ + if not tasks_completed: + return + + path = _session_path(project_root) + os.makedirs(os.path.dirname(path), exist_ok=True) + + # 生成本次会话摘要 + now = datetime.now().strftime("%Y-%m-%d %H:%M") + new_entries = [f"## 会话 {now} ({len(tasks_completed)} 个任务)\n"] + + for task in tasks_completed[-10:]: # 最多记录最近10个任务 + status = "OK" if task.get("success") else "FAIL" + input_text = task.get("input", "")[:50] + files = task.get("files", []) + files_str = ", ".join(files[:3]) if files else "" + elapsed = task.get("elapsed", 0) + line = f"- [{status}] {input_text}" + if files_str: + line += f" → {files_str}" + if elapsed > 0: + line += f" ({elapsed:.0f}s)" + new_entries.append(line) + + new_entries.append("") # blank line separator + + # 读取现有内容并追加 + existing = "" + if os.path.isfile(path): + try: + with open(path, "r", encoding="utf-8") as f: + existing = f.read() + except Exception: + pass + + # 新内容在前(最近的在最上面) + combined = "\n".join(new_entries) + "\n" + existing + + # 限制行数 + lines = combined.splitlines() + if len(lines) > MAX_LINES: + lines = lines[:MAX_LINES] + + try: + with open(path, "w", encoding="utf-8") as f: + f.write("\n".join(lines)) + logger.info("[session] Saved session summary (%d tasks)", len(tasks_completed)) + except Exception as e: + logger.warning("[session] Failed to save SESSION.md: %s", e) diff --git a/kaiwu/tools/executor.py b/kaiwu/tools/executor.py index e1284b7..3f635b1 100644 --- a/kaiwu/tools/executor.py +++ b/kaiwu/tools/executor.py @@ -2,6 +2,11 @@ Tool executor: self-implemented per FLEX-1 fallback. Provides read_file, write_file, run_bash, list_dir, git_commit. Interface is fixed (RED-4: transparent to user). + +Guardrails: +- Dangerous commands blocked (rm -rf, git push --force, drop database, etc.) +- Sensitive files protected (.env, credentials.json, id_rsa, etc.) +- Write operations confined to project_root """ import logging @@ -16,6 +21,23 @@ logger = logging.getLogger(__name__) class ToolExecutor: """Deterministic tool execution layer. No LLM involved.""" + # ── Guardrails ── + + DANGEROUS_PATTERNS = [ + "rm -rf", "rm -r /", "rmdir /s", + "git push --force", "git push -f", + "git reset --hard", + "drop database", "drop table", "truncate table", + "format c:", "del /f /s /q", + "> /dev/null", "mkfs", + ] + + PROTECTED_FILES = [ + ".env", ".env.local", ".env.production", + "credentials.json", "secrets.yaml", "id_rsa", + ".ssh/", "token.json", "service_account.json", + ] + def __init__(self, project_root: str = "."): self.project_root = os.path.abspath(project_root) @@ -31,8 +53,19 @@ class ToolExecutor: return f"[ERROR] Read failed: {e}" def write_file(self, path: str, content: str) -> bool: - """Write content to file. Creates parent dirs if needed.""" + """Write content to file. Guardrails: protects sensitive files, confines to project_root.""" full = self._resolve(path) + + # Guardrail: check for protected files + if self._is_protected(full): + logger.warning("[guardrail] Blocked write to protected file: %s", full) + return False + + # Guardrail: prevent writing outside project root + if not full.startswith(self.project_root): + logger.warning("[guardrail] Blocked write outside project: %s", full) + return False + try: os.makedirs(os.path.dirname(full), exist_ok=True) with open(full, "w", encoding="utf-8") as f: @@ -46,8 +79,14 @@ class ToolExecutor: def run_bash(self, command: str, cwd: Optional[str] = None, timeout: int = 60) -> tuple[str, str, int]: """ Run a shell command. Returns (stdout, stderr, returncode). - Timeout in seconds (default 60). + Guardrails: blocks dangerous commands. """ + # Guardrail: check for dangerous patterns + blocked = self._check_dangerous(command) + if blocked: + logger.warning("[guardrail] Blocked dangerous command: %s", command[:80]) + return "", f"[BLOCKED] 危险操作被拦截: {blocked}。如需执行请手动在终端运行。", -2 + work_dir = cwd or self.project_root try: result = subprocess.run( @@ -55,7 +94,6 @@ class ToolExecutor: shell=True, cwd=work_dir, capture_output=True, - text=True, timeout=timeout, encoding="utf-8", errors="replace", @@ -106,8 +144,9 @@ class ToolExecutor: dirnames.clear() continue indent = " " * depth - dirname = os.path.basename(dirpath) or os.path.basename(root) - lines.append(f"{indent}{dirname}/") + dirname = os.path.basename(dirpath) + if depth > 0: + lines.append(f"{indent}{dirname}/") for fname in sorted(filenames): if count >= max_files: lines.append(f"{indent} ... (truncated at {max_files} files)") @@ -116,10 +155,15 @@ class ToolExecutor: count += 1 return "\n".join(lines) + def _resolve(self, path: str) -> str: + """Resolve path relative to project_root.""" + if os.path.isabs(path): + return os.path.normpath(path) + return os.path.normpath(os.path.join(self.project_root, path)) + def apply_patch(self, file_path: str, original: str, modified: str) -> bool: """Apply a text replacement patch. Exact match only — original is read from file.""" if not original: - # Empty original means codegen (new file) — should use write_file instead logger.warning("apply_patch called with empty original, use write_file for new files") return False full = self._resolve(file_path) @@ -136,8 +180,18 @@ class ToolExecutor: logger.error("Patch apply failed: %s", e) return False - def _resolve(self, path: str) -> str: - """Resolve path relative to project_root.""" - if os.path.isabs(path): - return path - return os.path.join(self.project_root, path) + def _check_dangerous(self, command: str) -> Optional[str]: + """Check if command matches dangerous patterns. Returns matched pattern or None.""" + cmd_lower = command.lower().strip() + for pattern in self.DANGEROUS_PATTERNS: + if pattern in cmd_lower: + return pattern + return None + + def _is_protected(self, full_path: str) -> bool: + """Check if file path matches protected patterns.""" + path_lower = full_path.lower().replace("\\", "/") + for protected in self.PROTECTED_FILES: + if protected in path_lower: + return True + return False