nju-coding-agent-harness 0.1.0__py3-none-any.whl

harness/agent.py

@@ -0,0 +1,334 @@
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass, field
+from datetime import datetime
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Callable
+
+from harness.guardrails import evaluate
+from harness.registry import (
+    Context,
+    ToolResult,
+    build_request_tools,
+    validate_args,
+)
+
+if TYPE_CHECKING:
+    from harness.config import Config
+    from harness.fake_llm import FakeLLM
+    from harness.hooks import HookBus
+    from harness.llm import LLM
+    from harness.memory import MemoryStore
+    from harness.policy import Policy
+    from harness.sandbox import Sandbox
+    from harness.state import StateMachine
+
+SYSTEM_PROMPT = (
+    "你是编码代理助手。可调用工具完成任务："
+    "先规划，再执行，最后给出最终答案。"
+)
+
+_ASK_OPTIONS = ["y", "n", "always_allow", "never_allow"]
+
+
+@dataclass
+class AgentResult:
+    text: str = ""
+    steps_used: int = 0
+    tool_results: list[dict] = field(default_factory=list)
+    policy_changes: list[dict] = field(default_factory=list)
+    messages: list[dict] = field(default_factory=list)
+    failed_sequence: int = 0
+    transcript_path: str | None = None
+
+
+class Agent:
+    def __init__(
+        self,
+        llm: LLM,
+        registry: dict[str, Any],
+        sandbox: Sandbox,
+        hooks: HookBus,
+        policy: Policy,
+        state: StateMachine,
+        memory: MemoryStore | None,
+        config: Config,
+        ask_callback: Callable | None = None,
+        on_text: Callable[[str], None] | None = None,
+    ):
+        self.llm = llm
+        self.registry = registry
+        self.sandbox = sandbox
+        self.hooks = hooks
+        self.policy = policy
+        self.state = state
+        self.memory = memory
+        self.config = config
+        self.ask_callback = ask_callback
+        self.on_text = on_text
+        self._compress_calls = 0
+        self.warnings: list[str] = []
+        self._tool_calls: list[dict] = []
+
+    def context_for_tool(self) -> Context:
+        return Context(
+            workspace=self.config.workspace,
+            sandbox=self.sandbox,
+            hooks=self.hooks,
+            policy=self.policy,
+            state=self.state,
+            memory=self.memory,
+            config=self.config,
+            ask_callback=self.ask_callback,
+            llm=self.llm,
+            registry=self.registry,
+        )
+
+    def pipeline(self, call: dict, ctx: Context) -> ToolResult:
+        name = call["name"]
+        args = call["arguments"]
+        verdict = evaluate(self.policy.rules, name, args)
+        if verdict.action == "deny":
+            return ToolResult(status="error", error=f"guardrail denied: {verdict.reason}")
+        if verdict.action == "ask":
+            self.state.fire("approval_needed", "guardrail")
+            answer = self._ask(verdict.matched_rule, verdict.reason)
+            self.policy.apply_answer(verdict.matched_rule, answer)
+            if self.state.state == "awaiting_user":
+                self.state.fire("user_answered", "user")
+            if answer in ("n", "never_allow"):
+                return ToolResult(status="error", error=f"guardrail denied: {verdict.reason}")
+        if self.state.state == "awaiting_user":
+            self.state.fire("user_answered", "user")
+        args, ok = self.hooks.pre_tool_use(name, args)
+        if not ok:
+            return ToolResult(status="error", error=f"pre_tool_use hook rejected {name}")
+        self.state.fire("tool_requested", "loop")
+        tool = self.registry.get(name)
+        if tool is None:
+            result = ToolResult(status="error", error=f"unknown tool: {name}")
+        else:
+            err = validate_args(tool.parameters, args, self.config.workspace)
+            if err is not None:
+                result = ToolResult(status="error", error=f"参数错误: {err}")
+            else:
+                try:
+                    result = tool.handler(args, ctx)
+                except Exception as exc:
+                    result = ToolResult(
+                        status="error",
+                        error=f"工具 {name} 异常：{type(exc).__name__}: {exc}",
+                    )
+        self.state.fire("tool_finished", "loop")
+        self.hooks.post_tool_use(name, args, result)
+        return result
+
+    def _emit_text(self, text: str) -> None:
+        if text and self.on_text is not None:
+            self.on_text(text)
+
+    def _ask(self, rule, reason: str) -> str:
+        question = f"是否允许执行该操作？\n规则: {rule.pattern}\n原因: {reason}"
+        if self.ask_callback is None:
+            return "n"
+        try:
+            answer = self.ask_callback(question, list(_ASK_OPTIONS))
+        except Exception:
+            return "n"
+        if answer not in _ASK_OPTIONS:
+            return "n"
+        return answer
+
+    def _check_budget(self, messages: list[dict]) -> bool:
+        total = sum(len(m.get("content", "")) / 4 for m in messages)
+        return total > self.config.max_budget_tokens
+
+    def _compress(self, messages: list[dict]) -> list[dict]:
+        keep = self.config.compression_keep_turns
+        oldest = messages[:-keep]
+        if not oldest:
+            return messages
+        self._compress_calls += 1
+        if self._compress_calls > self.config.compression_max_rounds:
+            return self._drop_oldest(messages, keep)
+        try:
+            summary = self.llm.complete(
+                [
+                    {
+                        "role": "system",
+                        "content": (
+                            "请将以下较早回合总结为简洁摘要，"
+                            "保留关键事实、决定与结果："
+                        ),
+                    },
+                    {
+                        "role": "user",
+                        "content": json.dumps(oldest, ensure_ascii=False),
+                    },
+                ],
+                tools=[],
+            )
+        except Exception:
+            return self._drop_oldest(messages, keep)
+        text = (summary.text or "").strip()
+        if not text:
+            return self._drop_oldest(messages, keep)
+        return [{"role": "system", "content": f"[summary] {text}"}] + messages[-keep:]
+
+    def _drop_oldest(self, messages: list[dict], keep: int) -> list[dict]:
+        window = messages[-keep:]
+        if window and window[0]["role"] not in ("user", "system"):
+            for m in reversed(messages[:-keep]):
+                if m["role"] in ("user", "system"):
+                    return [m] + window
+        return window
+
+    def run(self, task: str) -> AgentResult:
+        result = AgentResult()
+        messages = [
+            {"role": "system", "content": SYSTEM_PROMPT},
+            {"role": "user", "content": task},
+        ]
+        if self.memory is not None:
+            for chunk in self.memory.top_k_chunks(task):
+                messages.append(
+                    {"role": "system", "content": f"[memory] {chunk['chunk']}"}
+                )
+        call_uid = 0
+        fail_seq = 0
+        fail_tool: str | None = None
+        max_fail_seq = 0
+        self._compress_calls = 0
+        self._tool_calls = []
+        self.state.fire("task_submitted", "loop")
+        while result.steps_used < self.config.max_steps:
+            if self._check_budget(messages):
+                messages = self._compress(messages)
+            response = self.llm.complete(messages, build_request_tools(self.registry))
+            result.steps_used += 1
+            if not response.tool_calls:
+                final = response.text or "任务完成"
+                self._emit_text(final)
+                messages.append({"role": "assistant", "content": final})
+                result.text = final
+                return self._finish(result, messages, max_fail_seq)
+            self._emit_text(response.text)
+            assistant_call = []
+            for i, call in enumerate(response.tool_calls):
+                assistant_call.append({
+                    "id": f"call_{call_uid}",
+                    "type": "function",
+                    "function": {
+                        "name": call["name"],
+                        "arguments": json.dumps(call["arguments"], ensure_ascii=False),
+                    },
+                })
+                call_uid += 1
+            messages.append({
+                "role": "assistant",
+                "content": response.text,
+                "tool_calls": assistant_call,
+            })
+            for i, call in enumerate(response.tool_calls):
+                tool_id = f"call_{call_uid - len(response.tool_calls) + i}"
+                tool_result = self.pipeline(call, self.context_for_tool())
+                result.tool_results.append(tool_result)
+                self._tool_calls.append({"name": call["name"], "arguments": call["arguments"]})
+                messages.append({
+                    "role": "tool",
+                    "tool_call_id": tool_id,
+                    "name": call["name"],
+                    "content": json.dumps(
+                        self._result_to_dict(tool_result), ensure_ascii=False
+                    ),
+                })
+                norm = self._result_to_dict(tool_result)
+                failed = norm.get("status") != "success" or bool(norm.get("error"))
+                if failed:
+                    if call["name"] == fail_tool:
+                        fail_seq += 1
+                    else:
+                        fail_seq = 1
+                        fail_tool = call["name"]
+                    max_fail_seq = max(max_fail_seq, fail_seq)
+                    if fail_seq >= self.config.failure_budget:
+                        final = (
+                            f"连续失败 {fail_seq} 次（工具 {fail_tool}），"
+                            f"超过失败预算 {self.config.failure_budget}，停止重试。"
+                        )
+                        self._emit_text(final)
+                        messages.append({"role": "assistant", "content": final})
+                        result.text = final
+                        return self._finish(result, messages, max_fail_seq)
+                else:
+                    fail_seq = 0
+                    fail_tool = None
+        final = f"达到步数上限 {self.config.max_steps}，任务终止，未挂死。"
+        self._emit_text(final)
+        messages.append({"role": "assistant", "content": final})
+        result.text = final
+        return self._finish(result, messages, max_fail_seq)
+
+    def _finish(self, result: AgentResult, messages: list[dict], max_fail_seq: int) -> AgentResult:
+        result.messages = messages
+        self.messages = messages
+        result.failed_sequence = max_fail_seq
+        self.state.fire("final_answer", "loop")
+        self._finalize(result)
+        return result
+
+    def _finalize(self, result: AgentResult) -> None:
+        self.hooks.session_data["tool_calls"] = list(self._tool_calls)
+        self.hooks.session_data["policy_changes"] = self.policy.changes()
+        self.hooks.session_end(self.messages)
+        if self.hooks.transcript_dir is not None:
+            files = [p for p in Path(self.hooks.transcript_dir).glob("*.json")]
+            if files:
+                result.transcript_path = str(
+                    max(files, key=lambda p: p.stat().st_mtime)
+                )
+        if self.memory is not None and self.llm is not None:
+            self._consolidate()
+
+    def _consolidate(self) -> None:
+        try:
+            response = self.llm.complete(
+                [
+                    {
+                        "role": "system",
+                        "content": (
+                            "请总结本次会话的关键事实、决策与工具结果，"
+                            "作为长期记忆保存："
+                        ),
+                    }
+                ]
+                + self.messages,
+                tools=[],
+            )
+        except Exception as exc:
+            self.warnings.append(f"memory consolidation failed: {exc}")
+            return
+        summary = (response.text or "").strip()
+        if not summary:
+            self.warnings.append("memory consolidation skipped: empty summary")
+            return
+        try:
+            self.memory.save(
+                f"session-summary-{datetime.now():%Y%m%d-%H%M%S}", summary
+            )
+        except Exception as exc:
+            self.warnings.append(f"memory save failed: {exc}")
+
+    @staticmethod
+    def _result_to_dict(tool_result) -> dict:
+        if isinstance(tool_result, ToolResult):
+            return {
+                "status": tool_result.status,
+                "output": tool_result.output,
+                "error": tool_result.error,
+                "exit_code": tool_result.exit_code,
+            }
+        if isinstance(tool_result, dict):
+            return tool_result
+        return {"status": "success", "output": str(tool_result)}

harness/config.py

@@ -0,0 +1,45 @@
+from __future__ import annotations
+
+import tomllib
+import warnings
+from dataclasses import dataclass, field, fields
+from pathlib import Path
+
+
+@dataclass
+class Config:
+    model: str = "deepseek-chat"
+    base_url: str = "https://api.deepseek.com"
+    max_steps: int = 50
+    failure_budget: int = 3
+    tool_timeout: int = 30
+    memory_top_k: int = 2
+    max_budget_tokens: int = 6000
+    compression_keep_turns: int = 10
+    compression_max_rounds: int = 3
+    workspace: Path = field(default_factory=Path.cwd)
+    max_output_bytes: int = 51200
+    mcp_servers: list[dict] = field(default_factory=list)
+
+    @classmethod
+    def load(cls, path: Path | None = None) -> Config:
+        cfg = cls()
+        if path is None:
+            return cfg
+        path = Path(path)
+        if not path.exists():
+            return cfg
+        try:
+            with path.open("rb") as f:
+                data = tomllib.load(f)
+        except Exception as exc:
+            warnings.warn(f"config 解析失败，使用默认值: {exc}")
+            return cfg
+        known = {f.name for f in fields(cls)}
+        for key, value in data.items():
+            if key not in known:
+                continue
+            if key == "workspace":
+                value = Path(value)
+            setattr(cfg, key, value)
+        return cfg

harness/credentials.py

@@ -0,0 +1,86 @@
+from __future__ import annotations
+
+import getpass
+from pathlib import Path
+
+_ENV_KEY = "DEEPSEEK_API_KEY"
+
+
+class CredentialStore:
+    def __init__(self, service: str = "coding-agent-harness", env_file: str = ".env", keyring_backend=None):
+        self.service = service
+        self.env_file = Path(env_file)
+        self._keyring = keyring_backend
+        if self._keyring is None:
+            try:
+                import keyring
+                self._keyring = keyring.get_keyring()
+            except Exception:
+                self._keyring = None
+
+    def _keyring_get(self, user: str) -> str | None:
+        if self._keyring is None:
+            return None
+        try:
+            return self._keyring.get_password(self.service, user)
+        except Exception:
+            return None
+
+    def _env_value(self) -> str | None:
+        try:
+            if not self.env_file.exists():
+                return None
+            for line in self.env_file.read_text(encoding="utf-8").splitlines():
+                line = line.strip()
+                if line.startswith(_ENV_KEY + "="):
+                    value = line[len(_ENV_KEY) + 1:].strip()
+                    if value:
+                        return value
+        except Exception:
+            return None
+        return None
+
+    def get(self) -> str | None:
+        key = self._keyring_get("api_key")
+        if key:
+            return key
+        return self._env_value()
+
+    def set(self, key: str) -> None:
+        if self._keyring is None:
+            raise RuntimeError("keyring 不可用（无凭据服务），无法保存 API Key")
+        try:
+            self._keyring.set_password(self.service, "api_key", key)
+        except Exception:
+            raise RuntimeError("keyring 写入失败，无法保存 API Key") from None
+
+    def clear(self) -> None:
+        if self._keyring is None:
+            raise RuntimeError("keyring 不可用（无凭据服务），无法删除 API Key")
+        try:
+            self._keyring.delete_password(self.service, "api_key")
+        except Exception:
+            raise RuntimeError("keyring 删除失败，无法删除 API Key") from None
+
+    def verified_at(self) -> str | None:
+        return self._keyring_get("verified_at")
+
+    def status(self) -> dict:
+        if self._keyring_get("api_key"):
+            source = "keyring"
+        elif self._env_value():
+            source = "env"
+        else:
+            source = None
+        return {
+            "configured": source is not None,
+            "source": source,
+            "verified_at": self.verified_at(),
+        }
+
+
+def wizard_enter_key() -> str:
+    key = getpass.getpass("请粘贴 API Key（输入不可见）: ")
+    if not key.strip():
+        raise ValueError("API Key 不能为空")
+    return key

harness/fake_llm.py

@@ -0,0 +1,34 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from harness.llm import LLM, LLMResult
+
+
+@dataclass
+class FakeTurn:
+    text: str = ""
+    tool_calls: list[dict] | None = None
+    usage_approx: int = 10
+
+
+class FakeLLM(LLM):
+    """脚本化 LLM：按序重放 turns，耗尽后重放最后一个；绝无网络。"""
+
+    def __init__(self, turns: list[FakeTurn]):
+        self.turns = list(turns)
+        self.calls = 0
+        self.turn_index = 0
+
+    def complete(self, messages: list[dict], tools: list[dict]) -> LLMResult:
+        self.calls += 1
+        if not self.turns:
+            return LLMResult(text="", tool_calls=[], usage={"approx_tokens": 0})
+        index = min(self.turn_index, len(self.turns) - 1)
+        self.turn_index += 1
+        turn = self.turns[index]
+        return LLMResult(
+            text=turn.text,
+            tool_calls=turn.tool_calls or [],
+            usage={"approx_tokens": turn.usage_approx},
+        )

harness/guardrails.py

@@ -0,0 +1,46 @@
+import json
+import re
+from dataclasses import dataclass
+
+
+@dataclass
+class Rule:
+    pattern: str
+    action: str
+    source: str
+
+
+@dataclass
+class Verdict:
+    action: str
+    matched_rule: Rule | None
+    reason: str
+
+
+def pattern_matches(pattern: str, tool_name: str, args: dict) -> bool:
+    if ":" in pattern:
+        tool, _, regex = pattern.partition(":")
+        if tool != tool_name:
+            return False
+        return re.search(regex, json.dumps(args, ensure_ascii=False)) is not None
+    return pattern == tool_name
+
+
+def default_rules() -> list[Rule]:
+    return [
+        Rule(r"bash:rm -rf\s+[\\/]?(/|[A-Z]:[\\/]|[A-Z](\b|:)|etc\b|boot\b|bin\b)",
+             "deny", "builtin"),
+        Rule(r"bash:.*:\(\)\s*\{.*:.*\};", "deny", "builtin"),
+        Rule(r"bash:format.*", "deny", "builtin"),
+        Rule(r"bash:del /f.*", "deny", "builtin"),
+    ]
+
+
+def evaluate(rules: list[Rule], tool_name: str, args: dict) -> Verdict:
+    matched = None
+    for rule in rules:
+        if pattern_matches(rule.pattern, tool_name, args):
+            matched = rule
+    if matched is None:
+        return Verdict("allow", None, "no rule matched")
+    return Verdict(matched.action, matched, f"matched {matched.pattern}")

harness/hooks.py

@@ -0,0 +1,57 @@
+from datetime import datetime
+from typing import Callable
+
+from harness.transcript import default_session_end_hook
+
+
+class HookBus:
+    def __init__(self, transcript_dir=None):
+        self.transcript_dir = transcript_dir
+        self.session_data: dict = {"tool_calls": [], "policy_changes": []}
+        self._hooks: dict[str, list[Callable]] = {}
+        self._records: list[dict] = []
+        self.errors: list[str] = []
+        if transcript_dir is not None:
+            self.register("session_end", default_session_end_hook(transcript_dir, self.session_data))
+
+    def register(self, name: str, fn: Callable) -> None:
+        self._hooks.setdefault(name, []).append(fn)
+
+    def _record(self, hook_name: str, tool_name, args, result) -> None:
+        self._records.append({
+            "hook_name": hook_name,
+            "tool_name": tool_name,
+            "args": args,
+            "result": result,
+            "timestamp": datetime.now().isoformat(),
+        })
+
+    def pre_tool_use(self, tool_name: str, args: dict) -> tuple[dict, bool]:
+        ok = True
+        for hook in self._hooks.get("pre", []):
+            try:
+                args, flag = hook(tool_name, args)
+                ok = ok and flag
+                self._record("pre", tool_name, args, None)
+            except Exception as exc:
+                self.errors.append(f"pre_tool_use({tool_name}): {exc}")
+        return args, ok
+
+    def post_tool_use(self, tool_name: str, args: dict, result) -> None:
+        for hook in self._hooks.get("post", []):
+            try:
+                hook(tool_name, args, result)
+                self._record("post", tool_name, args, result)
+            except Exception as exc:
+                self.errors.append(f"post_tool_use({tool_name}): {exc}")
+
+    def session_end(self, messages: list[dict]) -> None:
+        for hook in self._hooks.get("session_end", []):
+            try:
+                hook(messages)
+                self._record("session_end", None, None, None)
+            except Exception as exc:
+                self.errors.append(f"session_end: {exc}")
+
+    def records(self) -> list[dict]:
+        return self._records